command: Fix race with setting Cgroup path and context cancellation

When spawning commands, callers can afterwards add these commands to a
Cgroup and then call `SetCgroupPath()` to add the path to log messages.
But because we spawn a Goroutine that creates these logs as soon as the
context is done, this creates a race between setting the Cgroup path and
context cancellation.

Fix this race by handling addition to the Cgroups manager in the command
package itself. Like this, we can set up the Cgroup path before starting
the Goroutine and thus fix the race.

Changelog: fixed

9 files changed, 87 insertions, 49 deletions

diff --git a/internal/cgroups/cgroups.go b/internal/cgroups/cgroups.go
index c3e148bba..b2209b509 100644
--- a/internal/cgroups/cgroups.go
+++ b/internal/cgroups/cgroups.go
@@ -13,8 +13,8 @@ type Manager interface {
 	// It is expected to be called once at Gitaly startup from any
 	// instance of the Manager.
 	Setup() error
-	// AddCommand adds a Command to a cgroup
-	AddCommand(*command.Command, repository.GitRepo) error
+	// AddCommand adds a Command to a cgroup.
+	AddCommand(*command.Command, repository.GitRepo) (string, error)
 	// Cleanup cleans up cgroups created in Setup.
 	// It is expected to be called once at Gitaly shutdown from any
 	// instance of the Manager.
diff --git a/internal/cgroups/noop.go b/internal/cgroups/noop.go
index 399bd4931..feaa0d6ef 100644
--- a/internal/cgroups/noop.go
+++ b/internal/cgroups/noop.go
@@ -15,8 +15,8 @@ func (cg *NoopManager) Setup() error {
 }
 
 //nolint: revive,stylecheck // This is unintentionally missing documentation.
-func (cg *NoopManager) AddCommand(cmd *command.Command, repo repository.GitRepo) error {
-	return nil
+func (cg *NoopManager) AddCommand(cmd *command.Command, repo repository.GitRepo) (string, error) {
+	return "", nil
 }
 
 //nolint: revive,stylecheck // This is unintentionally missing documentation.
diff --git a/internal/cgroups/v1_linux.go b/internal/cgroups/v1_linux.go
index e4b72f28b..3fa5da4c3 100644
--- a/internal/cgroups/v1_linux.go
+++ b/internal/cgroups/v1_linux.go
@@ -102,7 +102,7 @@ func (cg *CGroupV1Manager) Setup() error {
 func (cg *CGroupV1Manager) AddCommand(
 	cmd *command.Command,
 	repo repository.GitRepo,
-) error {
+) (string, error) {
 	var key string
 	if repo == nil {
 		key = strings.Join(cmd.Args(), "/")
@@ -117,9 +117,7 @@ func (cg *CGroupV1Manager) AddCommand(
 	groupID := uint(checksum) % cg.cfg.Repositories.Count
 	cgroupPath := cg.repoPath(int(groupID))
 
-	cmd.SetCgroupPath(cgroupPath)
-
-	return cg.addToCgroup(cmd.Pid(), cgroupPath)
+	return cgroupPath, cg.addToCgroup(cmd.Pid(), cgroupPath)
 }
 
 func (cg *CGroupV1Manager) addToCgroup(pid int, cgroupPath string) error {
diff --git a/internal/cgroups/v1_linux_test.go b/internal/cgroups/v1_linux_test.go
index 73bb7df6c..aaa095db0 100644
--- a/internal/cgroups/v1_linux_test.go
+++ b/internal/cgroups/v1_linux_test.go
@@ -91,7 +91,8 @@ func TestAddCommand(t *testing.T) {
 	}
 
 	t.Run("without a repository", func(t *testing.T) {
-		require.NoError(t, v1Manager2.AddCommand(cmd2, nil))
+		_, err := v1Manager2.AddCommand(cmd2, nil)
+		require.NoError(t, err)
 
 		checksum := crc32.ChecksumIEEE([]byte(strings.Join(cmd2.Args(), "/")))
 		groupID := uint(checksum) % config.Repositories.Count
@@ -109,7 +110,8 @@ func TestAddCommand(t *testing.T) {
 	})
 
 	t.Run("with a repository", func(t *testing.T) {
-		require.NoError(t, v1Manager2.AddCommand(cmd2, repo))
+		_, err := v1Manager2.AddCommand(cmd2, repo)
+		require.NoError(t, err)
 
 		checksum := crc32.ChecksumIEEE([]byte(strings.Join([]string{
 			"default",
@@ -150,6 +152,8 @@ func TestCleanup(t *testing.T) {
 }
 
 func TestMetrics(t *testing.T) {
+	t.Parallel()
+
 	mock := newMock(t)
 	repo := &gitalypb.Repository{
 		StorageName:  "default",
@@ -167,22 +171,23 @@ func TestMetrics(t *testing.T) {
 	mock.setupMockCgroupFiles(t, v1Manager1, 2)
 
 	require.NoError(t, v1Manager1.Setup())
-	ctx := testhelper.Context(t)
 
+	ctx := testhelper.Context(t)
 	logger, hook := test.NewNullLogger()
 	logger.SetLevel(logrus.DebugLevel)
 	ctx = ctxlogrus.ToContext(ctx, logrus.NewEntry(logger))
 
-	cmd, err := command.New(ctx, exec.Command("ls", "-hal", "."))
+	cmd, err := command.New(ctx, exec.Command("ls", "-hal", "."), command.WithCgroup(v1Manager1, repo))
 	require.NoError(t, err)
-	gitCmd1, err := command.New(ctx, exec.Command("ls", "-hal", "."))
+	gitCmd1, err := command.New(ctx, exec.Command("ls", "-hal", "."), command.WithCgroup(v1Manager1, repo))
 	require.NoError(t, err)
-	gitCmd2, err := command.New(ctx, exec.Command("ls", "-hal", "."))
+	gitCmd2, err := command.New(ctx, exec.Command("ls", "-hal", "."), command.WithCgroup(v1Manager1, repo))
 	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, gitCmd2.Wait())
+	}()
 
-	require.NoError(t, v1Manager1.AddCommand(cmd, repo))
-	require.NoError(t, v1Manager1.AddCommand(gitCmd1, repo))
-	require.NoError(t, v1Manager1.AddCommand(gitCmd2, repo))
+	require.NoError(t, err)
 	require.NoError(t, cmd.Wait())
 	require.NoError(t, gitCmd1.Wait())
 
diff --git a/internal/command/command.go b/internal/command/command.go
index e1f3f0458..64f29673e 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -20,6 +20,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/command/commandcounter"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/labkit/tracing"
 )
 
@@ -242,25 +243,41 @@ func New(ctx context.Context, cmd *exec.Cmd, opts ...Option) (*Command, error) {
 	if err := cmd.Start(); err != nil {
 		return nil, fmt.Errorf("GitCommand: start %v: %v", cmd.Args, err)
 	}
-	inFlightCommandGauge.Inc()
 
-	// The goroutine below is responsible for terminating and reaping the
-	// process when ctx is canceled.
+	inFlightCommandGauge.Inc()
 	commandcounter.Increment()
-	go func() {
-		<-ctx.Done()
 
-		if process := cmd.Process; process != nil && process.Pid > 0 {
-			//nolint:errcheck // TODO: do we want to report errors?
-			// Send SIGTERM to the process group of cmd
-			syscall.Kill(-process.Pid, syscall.SIGTERM)
-		}
+	// The goroutine below is responsible for terminating and reaping the process when ctx is
+	// canceled. While we must ensure that it does run when `cmd.Start()` was successful, it
+	// must not run before have fully set up the command. Otherwise, we may end up with racy
+	// access patterns when the context gets terminated early.
+	//
+	// We thus defer spawning the Goroutine.
+	defer func() {
+		go func() {
+			<-ctx.Done()
 
-		// We do not care for any potential erorr code, but just want to make sure that the
-		// subprocess gets properly killed and processed.
-		_ = command.Wait()
+			if process := cmd.Process; process != nil && process.Pid > 0 {
+				//nolint:errcheck // TODO: do we want to report errors?
+				// Send SIGTERM to the process group of cmd
+				syscall.Kill(-process.Pid, syscall.SIGTERM)
+			}
+
+			// We do not care for any potential erorr code, but just want to make sure that the
+			// subprocess gets properly killed and processed.
+			_ = command.Wait()
+		}()
 	}()
 
+	if featureflag.RunCommandsInCGroup.IsEnabled(ctx) && cfg.cgroupsManager != nil {
+		cgroupPath, err := cfg.cgroupsManager.AddCommand(command, cfg.cgroupsRepo)
+		if err != nil {
+			return nil, err
+		}
+
+		command.cgroupPath = cgroupPath
+	}
+
 	logPid = cmd.Process.Pid
 
 	return command, nil
@@ -294,11 +311,6 @@ func (c *Command) Wait() error {
 	return c.waitError
 }
 
-// SetCgroupPath sets the cgroup path for logging
-func (c *Command) SetCgroupPath(path string) {
-	c.cgroupPath = path
-}
-
 // This function should never be called directly, use Wait().
 func (c *Command) wait() {
 	if c.writer != nil {
diff --git a/internal/command/command_test.go b/internal/command/command_test.go
index dbe7d04dd..5ac1aed85 100644
--- a/internal/command/command_test.go
+++ b/internal/command/command_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/repository"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 )
 
@@ -396,6 +397,12 @@ func TestCommand_stderrLoggingMaxBytes(t *testing.T) {
 	require.Len(t, hook.LastEntry().Message, maxStderrBytes)
 }
 
+type mockCgroupManager string
+
+func (m mockCgroupManager) AddCommand(*Command, repository.GitRepo) (string, error) {
+	return string(m), nil
+}
+
 func TestCommand_logMessage(t *testing.T) {
 	t.Parallel()
 
@@ -404,17 +411,17 @@ func TestCommand_logMessage(t *testing.T) {
 
 	ctx := ctxlogrus.ToContext(testhelper.Context(t), logrus.NewEntry(logger))
 
-	cmd, err := New(ctx, exec.Command("echo", "hello world"))
+	cmd, err := New(ctx, exec.Command("echo", "hello world"),
+		WithCgroup(mockCgroupManager("/sys/fs/cgroup/1"), nil),
+	)
 	require.NoError(t, err)
-	cgroupPath := "/sys/fs/cgroup/1"
-	cmd.SetCgroupPath(cgroupPath)
 
 	require.NoError(t, cmd.Wait())
 	logEntry := hook.LastEntry()
 	assert.Equal(t, cmd.Pid(), logEntry.Data["pid"])
 	assert.Equal(t, []string{"echo", "hello world"}, logEntry.Data["args"])
 	assert.Equal(t, 0, logEntry.Data["command.exitCode"])
-	assert.Equal(t, cgroupPath, logEntry.Data["command.cgroup_path"])
+	assert.Equal(t, "/sys/fs/cgroup/1", logEntry.Data["command.cgroup_path"])
 }
 
 func TestNew_commandSpawnTokenMetrics(t *testing.T) {
diff --git a/internal/command/option.go b/internal/command/option.go
index ade36dbd5..21d2d9c33 100644
--- a/internal/command/option.go
+++ b/internal/command/option.go
@@ -1,6 +1,10 @@
 package command
 
-import "io"
+import (
+	"io"
+
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/repository"
+)
 
 type config struct {
 	stdin       io.Reader
@@ -10,6 +14,9 @@ type config struct {
 
 	commandName    string
 	subcommandName string
+
+	cgroupsManager CgroupsManager
+	cgroupsRepo    repository.GitRepo
 }
 
 // Option is an option that can be passed to `New()` for controlling how the command is being
@@ -60,3 +67,18 @@ func WithCommandName(commandName, subcommandName string) Option {
 		cfg.subcommandName = subcommandName
 	}
 }
+
+// CgroupsManager is a subset of the `cgroups.Manager` interface. We need to replicate it here to
+// avoid a cyclic dependency between both packages.
+type CgroupsManager interface {
+	AddCommand(*Command, repository.GitRepo) (string, error)
+}
+
+// WithCgroup adds the spawned command to a Cgroup. The bucket used will be derived from the
+// command's arguments and/or from the repository.
+func WithCgroup(cgroupsManager CgroupsManager, repo repository.GitRepo) Option {
+	return func(cfg *config) {
+		cfg.cgroupsManager = cgroupsManager
+		cfg.cgroupsRepo = repo
+	}
+}
diff --git a/internal/git/command_factory.go b/internal/git/command_factory.go
index 93242d1f3..86918985f 100644
--- a/internal/git/command_factory.go
+++ b/internal/git/command_factory.go
@@ -17,7 +17,6 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/log"
-	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 )
 
 // CommandFactory is designed to create and run git commands in a protected and fully managed manner.
@@ -401,17 +400,12 @@ func (cf *ExecCommandFactory) newCommand(ctx context.Context, repo repository.Gi
 		config.commandOpts,
 		command.WithEnvironment(env),
 		command.WithCommandName("git", sc.Subcommand()),
+		command.WithCgroup(cf.cgroupsManager, repo),
 	)...)
 	if err != nil {
 		return nil, err
 	}
 
-	if featureflag.RunCommandsInCGroup.IsEnabled(ctx) {
-		if err := cf.cgroupsManager.AddCommand(command, repo); err != nil {
-			return nil, err
-		}
-	}
-
 	return command, nil
 }
 
diff --git a/internal/git/command_factory_cgroup_test.go b/internal/git/command_factory_cgroup_test.go
index 69c3126bf..e504880b6 100644
--- a/internal/git/command_factory_cgroup_test.go
+++ b/internal/git/command_factory_cgroup_test.go
@@ -24,9 +24,9 @@ func (m *mockCgroupsManager) Setup() error {
 	return nil
 }
 
-func (m *mockCgroupsManager) AddCommand(c *command.Command, repo repository.GitRepo) error {
+func (m *mockCgroupsManager) AddCommand(c *command.Command, repo repository.GitRepo) (string, error) {
 	m.commands = append(m.commands, c)
-	return nil
+	return "", nil
 }
 
 func (m *mockCgroupsManager) Cleanup() error {