From 1b7787e9967951aa86db507cab60bc96d0953250 Mon Sep 17 00:00:00 2001 From: Quang-Minh Nguyen Date: Wed, 8 Mar 2023 21:23:11 +0700 Subject: 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. --- internal/command/command.go | 8 ++--- internal/command/command_test.go | 37 ++++++++++++++++++++-- internal/command/option.go | 7 ++-- internal/git/command_options.go | 2 +- .../gitaly/service/ssh/monitor_stdin_command.go | 2 +- 5 files changed, 44 insertions(+), 12 deletions(-) (limited to 'internal') 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) } } diff --git a/internal/git/command_options.go b/internal/git/command_options.go index 7b1c3ddb2..1dee4d82f 100644 --- a/internal/git/command_options.go +++ b/internal/git/command_options.go @@ -278,7 +278,7 @@ func withInternalFetch(req repoScopedRequest, withSidechannel bool) func(ctx con // 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.Command)) CmdOpt { +func WithFinalizer(finalizer func(context.Context, *command.Command)) CmdOpt { return func(_ context.Context, _ config.Cfg, _ CommandFactory, c *cmdCfg) error { c.commandOpts = append(c.commandOpts, command.WithFinalizer(finalizer)) return nil diff --git a/internal/gitaly/service/ssh/monitor_stdin_command.go b/internal/gitaly/service/ssh/monitor_stdin_command.go index 6e04bbdb1..fe2446a52 100644 --- a/internal/gitaly/service/ssh/monitor_stdin_command.go +++ b/internal/gitaly/service/ssh/monitor_stdin_command.go @@ -29,7 +29,7 @@ func monitorStdinCommand( git.WithStdin(stdinPipe), git.WithStdout(stdout), git.WithStderr(stderr), - git.WithFinalizer(func(*command.Command) { cleanup() }), + git.WithFinalizer(func(context.Context, *command.Command) { cleanup() }), }, opts...)...) stdinPipe.Close() // this now belongs to cmd if err != nil { -- cgit v1.2.3