command: Add option to set up finalizer

We have cases where we need to clean up resources as soon as a command
has finished. Add the ability to set up a finalizer that runs exactly
once after the command has been reaped to make it easier to implement
this in a race-free manner.

3 files changed, 26 insertions, 0 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index 477b58de1..61d7fb6ef 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -134,6 +134,8 @@ type Command struct {
 	waitError error
 	waitOnce  sync.Once
 
+	finalizer func(*Command)
+
 	span opentracing.Span
 
 	metricsCmd    string
@@ -190,6 +192,7 @@ func New(ctx context.Context, cmd *exec.Cmd, opts ...Option) (*Command, error) {
 		startTime:     time.Now(),
 		context:       ctx,
 		span:          span,
+		finalizer:     cfg.finalizer,
 		metricsCmd:    cfg.commandName,
 		metricsSubCmd: cfg.subcommandName,
 	}
@@ -337,6 +340,10 @@ func (c *Command) wait() {
 	// counter again. So we instead do it here to accelerate the process, even though it's less
 	// idiomatic.
 	commandcounter.Decrement()
+
+	if c.finalizer != nil {
+		c.finalizer(c)
+	}
 }
 
 func (c *Command) logProcessComplete() {
diff --git a/internal/command/option.go b/internal/command/option.go
index 21d2d9c33..08acf0f33 100644
--- a/internal/command/option.go
+++ b/internal/command/option.go
@@ -12,6 +12,8 @@ type config struct {
 	stderr      io.Writer
 	environment []string
 
+	finalizer func(*Command)
+
 	commandName    string
 	subcommandName string
 
@@ -82,3 +84,11 @@ func WithCgroup(cgroupsManager CgroupsManager, repo repository.GitRepo) Option {
 		cfg.cgroupsRepo = repo
 	}
 }
+
+// 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 {
+	return func(cfg *config) {
+		cfg.finalizer = finalizer
+	}
+}
diff --git a/internal/git/command_options.go b/internal/git/command_options.go
index f2ed3bb04..044ae754d 100644
--- a/internal/git/command_options.go
+++ b/internal/git/command_options.go
@@ -309,3 +309,12 @@ func withInternalFetch(req repoScopedRequest, withSidechannel bool) func(ctx con
 		return nil
 	}
 }
+
+// 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 {
+	return func(_ context.Context, _ config.Cfg, _ CommandFactory, c *cmdCfg) error {
+		c.commandOpts = append(c.commandOpts, command.WithFinalizer(finalizer))
+		return nil
+	}
+}