Expose a method to configure a command directly into a cgroup

This commit exposes a method from the cgroup manager to configure a
command directly into a cgroup. This is done by configuring the cgroup
in the SysProcAttr so the command gets spawned directly in the cgroup.
The functionality is only available with cgroups version 2 and kernel
version 5.7 or later.

6 files changed, 182 insertions, 12 deletions

diff --git a/internal/cgroups/cgroups.go b/internal/cgroups/cgroups.go
index 2d77c732f..9a8b96492 100644
--- a/internal/cgroups/cgroups.go
+++ b/internal/cgroups/cgroups.go
@@ -1,6 +1,7 @@
 package cgroups
 
 import (
+	"io"
 	"os/exec"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -65,6 +66,14 @@ type Manager interface {
 	Ready() bool
 	// AddCommand adds a Cmd to a cgroup.
 	AddCommand(*exec.Cmd, ...AddCommandOption) (string, error)
+	// SupportsCloneIntoCgroup returns whether this Manager supports the CloneIntoCgroup method.
+	// CloneIntoCgroup requires CLONE_INTO_CGROUP which is only supported with cgroups version 2
+	// with Linux 5.7 or newer.
+	SupportsCloneIntoCgroup() bool
+	// CloneIntoCgroup configures the cgroup parameters UseCgroupFD and CgroupFD in SysProcAttr
+	// to start the command directly in the correct cgroup. On success, the function returns an io.Closer
+	// that must be closed after the command has been started to close the cgroup's file descriptor.
+	CloneIntoCgroup(*exec.Cmd, ...AddCommandOption) (string, io.Closer, 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/manager_linux.go b/internal/cgroups/manager_linux.go
index 2680b2862..37b719e80 100644
--- a/internal/cgroups/manager_linux.go
+++ b/internal/cgroups/manager_linux.go
@@ -3,10 +3,15 @@
 package cgroups
 
 import (
+	"errors"
 	"fmt"
 	"hash/crc32"
+	"io"
+	"os"
 	"os/exec"
+	"path/filepath"
 	"strings"
+	"syscall"
 
 	cgrps "github.com/containerd/cgroups/v3"
 	"github.com/opencontainers/runtime-spec/specs-go"
@@ -27,6 +32,7 @@ type cgroupHandler interface {
 	currentProcessCgroup() string
 	repoPath(groupID int) string
 	stats() (Stats, error)
+	supportsCloneIntoCgroup() bool
 }
 
 // CGroupManager is a manager class that implements specific methods related to cgroups
@@ -82,6 +88,45 @@ func (cgm *CGroupManager) Ready() bool {
 
 // AddCommand adds a Cmd to a cgroup
 func (cgm *CGroupManager) AddCommand(cmd *exec.Cmd, opts ...AddCommandOption) (string, error) {
+	if cmd.Process == nil {
+		return "", errors.New("cannot add command that has not yet been started")
+	}
+
+	cgroupPath := cgm.cgroupPathForCommand(cmd, opts)
+
+	return cgroupPath, cgm.handler.addToCgroup(cmd.Process.Pid, cgroupPath)
+}
+
+// SupportsCloneIntoCgroup returns whether this Manager supports the CloneIntoCgroup method.
+// CloneIntoCgroup requires CLONE_INTO_CGROUP which is only supported with cgroup version 2
+// with Linux 5.7 or newer.
+func (cgm *CGroupManager) SupportsCloneIntoCgroup() bool {
+	return cgm.handler.supportsCloneIntoCgroup()
+}
+
+// CloneIntoCgroup configures the cgroup parameters UseCgroupFD and CgroupFD in SysProcAttr
+// to start the command directly in the correct cgroup. On success, the function returns an io.Closer
+// that must be closed after the command has been started to close the cgroup's file descriptor.
+func (cgm *CGroupManager) CloneIntoCgroup(cmd *exec.Cmd, opts ...AddCommandOption) (string, io.Closer, error) {
+	cgroupPath := filepath.Join(cgm.cfg.Mountpoint, cgm.cgroupPathForCommand(cmd, opts))
+
+	file, err := os.Open(cgroupPath)
+	if err != nil {
+		return "", nil, fmt.Errorf("open file: %w", err)
+	}
+
+	if cmd.SysProcAttr == nil {
+		cmd.SysProcAttr = &syscall.SysProcAttr{}
+	}
+
+	cmd.SysProcAttr.UseCgroupFD = true
+	cmd.SysProcAttr.CgroupFD = int(file.Fd())
+
+	return cgroupPath, file, nil
+}
+
+// cgroupPathForCommand returns the path of the cgroup a given command should go in.
+func (cgm *CGroupManager) cgroupPathForCommand(cmd *exec.Cmd, opts []AddCommandOption) string {
 	var cfg addCommandCfg
 	for _, opt := range opts {
 		opt(&cfg)
@@ -96,14 +141,8 @@ func (cgm *CGroupManager) AddCommand(cmd *exec.Cmd, opts ...AddCommandOption) (s
 		[]byte(key),
 	)
 
-	if cmd.Process == nil {
-		return "", fmt.Errorf("cannot add command that has not yet been started")
-	}
-
 	groupID := uint(checksum) % cgm.cfg.Repositories.Count
-	cgroupPath := cgm.handler.repoPath(int(groupID))
-
-	return cgroupPath, cgm.handler.addToCgroup(cmd.Process.Pid, cgroupPath)
+	return cgm.handler.repoPath(int(groupID))
 }
 
 // Cleanup cleans up cgroups created in Setup.
diff --git a/internal/cgroups/manager_linux_test.go b/internal/cgroups/manager_linux_test.go
new file mode 100644
index 000000000..c1563e7bf
--- /dev/null
+++ b/internal/cgroups/manager_linux_test.go
@@ -0,0 +1,94 @@
+package cgroups
+
+import (
+	"io/fs"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"syscall"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config/cgroups"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
+)
+
+func TestCloneIntoCgroup(t *testing.T) {
+	mountPoint := t.TempDir()
+	hierarchyRoot := filepath.Join(mountPoint, "gitaly")
+
+	// Create the files we expect the manager to open.
+	require.NoError(t, os.MkdirAll(filepath.Join(hierarchyRoot, "gitaly-1"), fs.ModePerm))
+	require.NoError(t, os.WriteFile(filepath.Join(hierarchyRoot, "gitaly-1", "repos-3"), nil, fs.ModePerm))
+	require.NoError(t, os.WriteFile(filepath.Join(hierarchyRoot, "gitaly-1", "repos-5"), nil, fs.ModePerm))
+
+	mgr := NewManager(cgroups.Config{
+		Mountpoint:    mountPoint,
+		HierarchyRoot: "gitaly",
+		Repositories: cgroups.Repositories{
+			Count: 10,
+		},
+	}, testhelper.NewLogger(t), 1)
+
+	t.Run("command args used as key", func(t *testing.T) {
+		for _, tc := range []struct {
+			desc                string
+			existingSysProcAttr *syscall.SysProcAttr
+			expectedSysProcAttr *syscall.SysProcAttr
+		}{
+			{
+				desc: "nil SysProcAttr",
+				expectedSysProcAttr: &syscall.SysProcAttr{
+					UseCgroupFD: true,
+				},
+			},
+			{
+				desc: "existing SysProcAttr",
+				existingSysProcAttr: &syscall.SysProcAttr{
+					Chroot: "preserved",
+				},
+				expectedSysProcAttr: &syscall.SysProcAttr{
+					Chroot:      "preserved",
+					UseCgroupFD: true,
+				},
+			},
+		} {
+			t.Run(tc.desc, func(t *testing.T) {
+				command := exec.Command("command", "arg")
+				command.SysProcAttr = tc.existingSysProcAttr
+
+				path, closer, err := mgr.CloneIntoCgroup(command)
+				require.NoError(t, err)
+				defer testhelper.MustClose(t, closer)
+				require.Equal(t, filepath.Join(hierarchyRoot, "gitaly-1/repos-5"), path)
+
+				// CgroupFD may vary so just set it in the assertion after verifying it is set.
+				require.NotEqual(t, 0, command.SysProcAttr.UseCgroupFD)
+				tc.expectedSysProcAttr.CgroupFD = command.SysProcAttr.CgroupFD
+
+				require.Equal(t, tc.expectedSysProcAttr, command.SysProcAttr)
+			})
+		}
+	})
+
+	t.Run("key provided", func(t *testing.T) {
+		commandWithKey := exec.Command("command", "arg")
+		pathWithKey, closeWithKey, err := mgr.CloneIntoCgroup(commandWithKey, WithCgroupKey("some-key"))
+		require.NoError(t, err)
+		defer testhelper.MustClose(t, closeWithKey)
+		require.Equal(t, filepath.Join(hierarchyRoot, "gitaly-1/repos-3"), pathWithKey)
+		require.True(t, commandWithKey.SysProcAttr.UseCgroupFD)
+		require.NotEqual(t, 0, commandWithKey.SysProcAttr.UseCgroupFD)
+
+		commandWithoutKey := exec.Command("command", "arg")
+		pathWithoutKey, closeWithoutKey, err := mgr.CloneIntoCgroup(commandWithoutKey)
+		require.NoError(t, err)
+		defer testhelper.MustClose(t, closeWithoutKey)
+		require.Equal(t, filepath.Join(hierarchyRoot, "gitaly-1/repos-5"), pathWithoutKey)
+		require.True(t, commandWithoutKey.SysProcAttr.UseCgroupFD)
+		require.NotEqual(t, 0, commandWithoutKey.SysProcAttr.UseCgroupFD)
+
+		require.NotEqual(t, pathWithKey, pathWithoutKey, "commands should be placed in different groups")
+		require.NotEqual(t, commandWithKey.SysProcAttr.CgroupFD, commandWithoutKey.SysProcAttr.CgroupFD)
+	})
+}
diff --git a/internal/cgroups/noop.go b/internal/cgroups/noop.go
index 69407e413..535c7df4b 100644
--- a/internal/cgroups/noop.go
+++ b/internal/cgroups/noop.go
@@ -1,6 +1,7 @@
 package cgroups
 
 import (
+	"io"
 	"os/exec"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -29,6 +30,16 @@ func (cg *NoopManager) AddCommand(*exec.Cmd, ...AddCommandOption) (string, error
 	return "", nil
 }
 
+// SupportsCloneIntoCgroup returns false.
+func (cg *NoopManager) SupportsCloneIntoCgroup() bool {
+	return false
+}
+
+// CloneIntoCgroup does nothing.
+func (cg *NoopManager) CloneIntoCgroup(*exec.Cmd, ...AddCommandOption) (string, io.Closer, error) {
+	return "", io.NopCloser(nil), nil
+}
+
 //nolint:revive // This is unintentionally missing documentation.
 func (cg *NoopManager) Cleanup() error {
 	return nil
diff --git a/internal/cgroups/v1_linux.go b/internal/cgroups/v1_linux.go
index 8de8d283f..52b1106f2 100644
--- a/internal/cgroups/v1_linux.go
+++ b/internal/cgroups/v1_linux.go
@@ -210,6 +210,10 @@ func (cvh *cgroupV1Handler) stats() (Stats, error) {
 	}, nil
 }
 
+func (cvh *cgroupV1Handler) supportsCloneIntoCgroup() bool {
+	return false
+}
+
 func defaultSubsystems(root string) ([]cgroup1.Subsystem, error) {
 	subsystems := []cgroup1.Subsystem{
 		cgroup1.NewMemory(root, cgroup1.OptionalSwap()),
diff --git a/internal/cgroups/v2_linux.go b/internal/cgroups/v2_linux.go
index f8970dbf7..c0487aade 100644
--- a/internal/cgroups/v2_linux.go
+++ b/internal/cgroups/v2_linux.go
@@ -15,6 +15,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config"
 	cgroupscfg "gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config/cgroups"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/kernel"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/log"
 )
 
@@ -23,15 +24,23 @@ type cgroupV2Handler struct {
 	logger log.Logger
 
 	*cgroupsMetrics
-	pid int
+	pid             int
+	cloneIntoCgroup bool
 }
 
 func newV2Handler(cfg cgroupscfg.Config, logger log.Logger, pid int) *cgroupV2Handler {
+	cloneIntoCgroup, err := kernel.IsAtLeast(kernel.Version{Major: 5, Minor: 7})
+	if err != nil {
+		// Log the error for now as we're only rolling out functionality behind feature flag.
+		logger.WithError(err).Error("failed detecting kernel version, CLONE_INTO_CGROUP support disabled")
+	}
+
 	return &cgroupV2Handler{
-		cfg:            cfg,
-		logger:         logger,
-		pid:            pid,
-		cgroupsMetrics: newV2CgroupsMetrics(),
+		cfg:             cfg,
+		logger:          logger,
+		pid:             pid,
+		cgroupsMetrics:  newV2CgroupsMetrics(),
+		cloneIntoCgroup: cloneIntoCgroup,
 	}
 }
 
@@ -190,6 +199,10 @@ func (cvh *cgroupV2Handler) stats() (Stats, error) {
 	return stats, nil
 }
 
+func (cvh *cgroupV2Handler) supportsCloneIntoCgroup() bool {
+	return cvh.cloneIntoCgroup
+}
+
 func pruneOldCgroupsV2(cfg cgroupscfg.Config, logger log.Logger) {
 	if err := config.PruneOldGitalyProcessDirectories(
 		logger,