command changes

1 files changed, 14 insertions, 10 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index a131f88c1..0d45e14e3 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -155,7 +155,7 @@ type Command struct {
 	metricsCmd    string
 	metricsSubCmd string
 	cgroupPath    string
-	finishers     []func()
+	releaseCgroup func()
 }
 
 type stdinSentinel struct{}
@@ -204,8 +204,8 @@ func (c *Command) SetCgroupPath(path string) {
 	c.cgroupPath = path
 }
 
-func (c *Command) AddFinisher(f func()) {
-	c.finishers = append(c.finishers, f)
+func (c *Command) SetCgroupReleaseFunc(f func()) {
+	c.releaseCgroup = f
 }
 
 // SetMetricsCmd overrides the "cmd" label used in metrics
@@ -257,10 +257,11 @@ func New(ctx context.Context, cmd *exec.Cmd, stdin io.Reader, stdout, stderr io.
 	}()
 
 	command := &Command{
-		cmd:       cmd,
-		startTime: time.Now(),
-		context:   ctx,
-		span:      span,
+		cmd:           cmd,
+		startTime:     time.Now(),
+		context:       ctx,
+		span:          span,
+		releaseCgroup: func() {},
 	}
 
 	// Explicitly set the environment for the command
@@ -370,9 +371,7 @@ func (c *Command) wait() {
 
 	inFlightCommandGauge.Dec()
 
-	for _, finisher := range c.finishers {
-		finisher()
-	}
+	c.ReleaseCgroupResources()
 
 	c.logProcessComplete()
 
@@ -386,6 +385,11 @@ func (c *Command) wait() {
 	commandcounter.Decrement()
 }
 
+func (c *Command) ReleaseCgroupResources() {
+	c.releaseCgroup()
+	c.releaseCgroup = func() {}
+}
+
 // ExitStatus will return the exit-code from an error returned by Wait().
 func ExitStatus(err error) (int, bool) {
 	exitError, ok := err.(*exec.ExitError)