Add multiple finalizers support to internal/command

When spawning a command, we support adding a finalizer option via
`command.WithFinalizer`.  Originally, there is only one finalizer. To
prepare for trace2 integration, Git command factory needs to find a way
to hook to when after a command finishes. Finalizer is a good way to
achieve this purpose. Unfortunately, it's not easy to create a wrapper
function then trigger the pre-existing finalizer if any. So, expanding
this option to support multiple finalizers is a good approach.

This commit also adds context to the function signature of finalizer.

3 files changed, 42 insertions, 10 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index 3bae16807..0d5d30b43 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -141,7 +141,7 @@ type Command struct {
 	waitOnce        sync.Once
 	processExitedCh chan struct{}
 
-	finalizer func(*Command)
+	finalizers []func(context.Context, *Command)
 
 	span opentracing.Span
 
@@ -215,7 +215,7 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 		startTime:       time.Now(),
 		context:         ctx,
 		span:            span,
-		finalizer:       cfg.finalizer,
+		finalizers:      cfg.finalizers,
 		metricsCmd:      cfg.commandName,
 		metricsSubCmd:   cfg.subcommandName,
 		cmdGitVersion:   cfg.gitVersion,
@@ -393,8 +393,8 @@ func (c *Command) wait() {
 	// idiomatic.
 	commandcounter.Decrement()
 
-	if c.finalizer != nil {
-		c.finalizer(c)
+	for _, finalizer := range c.finalizers {
+		finalizer(c.context, c)
 	}
 }
 
diff --git a/internal/command/command_test.go b/internal/command/command_test.go
index c82cdbd0f..6fd0026ea 100644
--- a/internal/command/command_test.go
+++ b/internal/command/command_test.go
@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -473,7 +474,7 @@ func TestCommand_withFinalizer(t *testing.T) {
 		ctx, cancel := context.WithCancel(testhelper.Context(t))
 
 		finalizerCh := make(chan struct{})
-		_, err := New(ctx, []string{"echo"}, WithFinalizer(func(*Command) {
+		_, err := New(ctx, []string{"echo"}, WithFinalizer(func(context.Context, *Command) {
 			close(finalizerCh)
 		}))
 		require.NoError(t, err)
@@ -487,7 +488,7 @@ func TestCommand_withFinalizer(t *testing.T) {
 		ctx := testhelper.Context(t)
 
 		finalizerCh := make(chan struct{})
-		cmd, err := New(ctx, []string{"echo"}, WithFinalizer(func(*Command) {
+		cmd, err := New(ctx, []string{"echo"}, WithFinalizer(func(context.Context, *Command) {
 			close(finalizerCh)
 		}))
 		require.NoError(t, err)
@@ -497,11 +498,41 @@ func TestCommand_withFinalizer(t *testing.T) {
 		<-finalizerCh
 	})
 
+	t.Run("Wait runs multiple finalizers", func(t *testing.T) {
+		ctx := testhelper.Context(t)
+
+		wg := sync.WaitGroup{}
+		wg.Add(2)
+		cmd, err := New(
+			ctx,
+			[]string{"echo"},
+			WithFinalizer(func(context.Context, *Command) { wg.Done() }),
+			WithFinalizer(func(context.Context, *Command) { wg.Done() }),
+		)
+		require.NoError(t, err)
+		require.NoError(t, cmd.Wait())
+
+		wg.Wait()
+	})
+
+	t.Run("Wait runs finalizer with the latest context", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(testhelper.Context(t))
+		//nolint:staticcheck
+		ctx = context.WithValue(ctx, "hello", "world")
+
+		_, err := New(ctx, []string{"echo"}, WithFinalizer(func(ctx context.Context, _ *Command) {
+			require.Equal(t, "world", ctx.Value("hello"))
+		}))
+		require.NoError(t, err)
+
+		cancel()
+	})
+
 	t.Run("process exit does not run finalizer", func(t *testing.T) {
 		ctx := testhelper.Context(t)
 
 		finalizerCh := make(chan struct{})
-		_, err := New(ctx, []string{"echo"}, WithFinalizer(func(*Command) {
+		_, err := New(ctx, []string{"echo"}, WithFinalizer(func(context.Context, *Command) {
 			close(finalizerCh)
 		}))
 		require.NoError(t, err)
diff --git a/internal/command/option.go b/internal/command/option.go
index 8f1d64f1d..862dc8d83 100644
--- a/internal/command/option.go
+++ b/internal/command/option.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"context"
 	"io"
 
 	"gitlab.com/gitlab-org/gitaly/v15/internal/cgroups"
@@ -13,7 +14,7 @@ type config struct {
 	dir         string
 	environment []string
 
-	finalizer func(*Command)
+	finalizers []func(context.Context, *Command)
 
 	commandName    string
 	subcommandName string
@@ -97,8 +98,8 @@ func WithCgroup(cgroupsManager cgroups.Manager, opts ...cgroups.AddCommandOption
 
 // WithFinalizer sets up the finalizer to be run when the command is being wrapped up. It will be
 // called after `Wait()` has returned.
-func WithFinalizer(finalizer func(*Command)) Option {
+func WithFinalizer(finalizer func(context.Context, *Command)) Option {
 	return func(cfg *config) {
-		cfg.finalizer = finalizer
+		cfg.finalizers = append(cfg.finalizers, finalizer)
 	}
 }