Merge branch 'wc/cgroup-refactor' into 'master'

cgroups: Deduplicate repeated code

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/6500

Merged-by: Will Chandler <wchandler@gitlab.com>
Approved-by: Quang-Minh Nguyen <qmnguyen@gitlab.com>
Approved-by: Eric Ju <eju@gitlab.com>

6 files changed, 63 insertions, 51 deletions

diff --git a/internal/cgroups/manager_linux.go b/internal/cgroups/manager_linux.go
index 37b719e80..f2cf383dd 100644
--- a/internal/cgroups/manager_linux.go
+++ b/internal/cgroups/manager_linux.go
@@ -24,7 +24,7 @@ import (
 const cfsPeriodUs uint64 = 100000
 
 type cgroupHandler interface {
-	setupParent(reposResources *specs.LinuxResources) error
+	setupParent(parentResources *specs.LinuxResources) error
 	setupRepository(reposResources *specs.LinuxResources) error
 	addToCgroup(pid int, cgroupPath string) error
 	collect(ch chan<- prometheus.Metric)
diff --git a/internal/cgroups/testhelper_test.go b/internal/cgroups/testhelper_test.go
index 5f594195a..25b7e008e 100644
--- a/internal/cgroups/testhelper_test.go
+++ b/internal/cgroups/testhelper_test.go
@@ -3,11 +3,22 @@
 package cgroups
 
 import (
+	"hash/crc32"
+	"strings"
 	"testing"
 
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 )
 
+// cmdArgs are Command arguments used processes to be added to a cgroup.
+var cmdArgs = []string{"ls", "-hal", "."}
+
 func TestMain(m *testing.M) {
 	testhelper.Run(m)
 }
+
+// calcGroupID calculates the repository cgroup ID for the key provided.
+func calcGroupID(key []string, ct uint) uint {
+	checksum := crc32.ChecksumIEEE([]byte(strings.Join(key, "/")))
+	return uint(checksum) % ct
+}
diff --git a/internal/cgroups/v1_linux.go b/internal/cgroups/v1_linux.go
index 52b1106f2..92f2b6031 100644
--- a/internal/cgroups/v1_linux.go
+++ b/internal/cgroups/v1_linux.go
@@ -62,12 +62,9 @@ func (cvh *cgroupV1Handler) setupRepository(reposResources *specs.LinuxResources
 }
 
 func (cvh *cgroupV1Handler) addToCgroup(pid int, cgroupPath string) error {
-	control, err := cgroup1.Load(
-		cgroup1.StaticPath(cgroupPath),
-		cgroup1.WithHiearchy(cvh.hierarchy),
-	)
+	control, err := cvh.loadCgroup(cgroupPath)
 	if err != nil {
-		return fmt.Errorf("failed loading %s cgroup: %w", cgroupPath, err)
+		return err
 	}
 
 	if err := control.Add(cgroup1.Process{Pid: pid}); err != nil {
@@ -82,6 +79,17 @@ func (cvh *cgroupV1Handler) addToCgroup(pid int, cgroupPath string) error {
 	return nil
 }
 
+func (cvh *cgroupV1Handler) loadCgroup(cgroupPath string) (cgroup1.Cgroup, error) {
+	control, err := cgroup1.Load(
+		cgroup1.StaticPath(cgroupPath),
+		cgroup1.WithHiearchy(cvh.hierarchy),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed loading %s cgroup: %w", cgroupPath, err)
+	}
+	return control, nil
+}
+
 func (cvh *cgroupV1Handler) collect(ch chan<- prometheus.Metric) {
 	if !cvh.cfg.MetricsEnabled {
 		return
@@ -90,10 +98,7 @@ func (cvh *cgroupV1Handler) collect(ch chan<- prometheus.Metric) {
 	for i := 0; i < int(cvh.cfg.Repositories.Count); i++ {
 		repoPath := cvh.repoPath(i)
 		logger := cvh.logger.WithField("cgroup_path", repoPath)
-		control, err := cgroup1.Load(
-			cgroup1.StaticPath(repoPath),
-			cgroup1.WithHiearchy(cvh.hierarchy),
-		)
+		control, err := cvh.loadCgroup(repoPath)
 		if err != nil {
 			logger.WithError(err).Warn("unable to load cgroup controller")
 			return
@@ -159,12 +164,9 @@ func (cvh *cgroupV1Handler) collect(ch chan<- prometheus.Metric) {
 func (cvh *cgroupV1Handler) cleanup() error {
 	processCgroupPath := cvh.currentProcessCgroup()
 
-	control, err := cgroup1.Load(
-		cgroup1.StaticPath(processCgroupPath),
-		cgroup1.WithHiearchy(cvh.hierarchy),
-	)
+	control, err := cvh.loadCgroup(processCgroupPath)
 	if err != nil {
-		return fmt.Errorf("failed loading cgroup %s: %w", processCgroupPath, err)
+		return err
 	}
 
 	if err := control.Delete(); err != nil {
@@ -185,12 +187,9 @@ func (cvh *cgroupV1Handler) currentProcessCgroup() string {
 func (cvh *cgroupV1Handler) stats() (Stats, error) {
 	processCgroupPath := cvh.currentProcessCgroup()
 
-	control, err := cgroup1.Load(
-		cgroup1.StaticPath(processCgroupPath),
-		cgroup1.WithHiearchy(cvh.hierarchy),
-	)
+	control, err := cvh.loadCgroup(processCgroupPath)
 	if err != nil {
-		return Stats{}, fmt.Errorf("failed loading cgroup %s: %w", processCgroupPath, err)
+		return Stats{}, err
 	}
 
 	metrics, err := control.Stat()
diff --git a/internal/cgroups/v1_linux_test.go b/internal/cgroups/v1_linux_test.go
index ba8a55b29..cb28d40c0 100644
--- a/internal/cgroups/v1_linux_test.go
+++ b/internal/cgroups/v1_linux_test.go
@@ -4,7 +4,6 @@ package cgroups
 
 import (
 	"fmt"
-	"hash/crc32"
 	"io/fs"
 	"os"
 	"os/exec"
@@ -229,18 +228,17 @@ func TestAddCommand(t *testing.T) {
 	require.NoError(t, v1Manager1.Setup())
 	ctx := testhelper.Context(t)
 
-	cmd2 := exec.CommandContext(ctx, "ls", "-hal", ".")
+	cmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 	require.NoError(t, cmd2.Run())
 
 	v1Manager2 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
 
 	t.Run("without overridden key", func(t *testing.T) {
+		groupID := calcGroupID(cmd2.Args, config.Repositories.Count)
+
 		_, err := v1Manager2.AddCommand(cmd2)
 		require.NoError(t, err)
 
-		checksum := crc32.ChecksumIEEE([]byte(strings.Join(cmd2.Args, "/")))
-		groupID := uint(checksum) % config.Repositories.Count
-
 		for _, s := range mock.subsystems {
 			path := filepath.Join(mock.root, string(s.Name()), "gitaly",
 				fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
@@ -254,15 +252,14 @@ func TestAddCommand(t *testing.T) {
 	})
 
 	t.Run("with overridden key", func(t *testing.T) {
+		overridenGroupID := calcGroupID([]string{"foobar"}, config.Repositories.Count)
+
 		_, err := v1Manager2.AddCommand(cmd2, WithCgroupKey("foobar"))
 		require.NoError(t, err)
 
-		checksum := crc32.ChecksumIEEE([]byte("foobar"))
-		groupID := uint(checksum) % config.Repositories.Count
-
 		for _, s := range mock.subsystems {
 			path := filepath.Join(mock.root, string(s.Name()), "gitaly",
-				fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
+				fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", overridenGroupID), "cgroup.procs")
 			content := readCgroupFile(t, path)
 
 			cmdPid, err := strconv.Atoi(string(content))
@@ -354,17 +351,17 @@ gitaly_cgroup_cpu_cfs_throttled_seconds_total{path="%s"} 0.001
 
 			ctx := testhelper.Context(t)
 
-			cmd := exec.CommandContext(ctx, "ls", "-hal", ".")
+			cmd := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, cmd.Start())
 			_, err := v1Manager1.AddCommand(cmd)
 			require.NoError(t, err)
 
-			gitCmd1 := exec.CommandContext(ctx, "ls", "-hal", ".")
+			gitCmd1 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, gitCmd1.Start())
 			_, err = v1Manager1.AddCommand(gitCmd1)
 			require.NoError(t, err)
 
-			gitCmd2 := exec.CommandContext(ctx, "ls", "-hal", ".")
+			gitCmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, gitCmd2.Start())
 			_, err = v1Manager1.AddCommand(gitCmd2)
 			require.NoError(t, err)
diff --git a/internal/cgroups/v2_linux.go b/internal/cgroups/v2_linux.go
index c0487aade..969c94e1f 100644
--- a/internal/cgroups/v2_linux.go
+++ b/internal/cgroups/v2_linux.go
@@ -66,9 +66,9 @@ func (cvh *cgroupV2Handler) setupRepository(reposResources *specs.LinuxResources
 }
 
 func (cvh *cgroupV2Handler) addToCgroup(pid int, cgroupPath string) error {
-	control, err := cgroup2.Load("/"+cgroupPath, cgroup2.WithMountpoint(cvh.cfg.Mountpoint))
+	control, err := cvh.loadCgroup(cgroupPath)
 	if err != nil {
-		return fmt.Errorf("failed loading %s cgroup: %w", cgroupPath, err)
+		return err
 	}
 
 	if err := control.AddProc(uint64(pid)); err != nil {
@@ -83,6 +83,14 @@ func (cvh *cgroupV2Handler) addToCgroup(pid int, cgroupPath string) error {
 	return nil
 }
 
+func (cvh *cgroupV2Handler) loadCgroup(cgroupPath string) (*cgroup2.Manager, error) {
+	control, err := cgroup2.Load("/"+cgroupPath, cgroup2.WithMountpoint(cvh.cfg.Mountpoint))
+	if err != nil {
+		return nil, fmt.Errorf("failed loading %s cgroup: %w", cgroupPath, err)
+	}
+	return control, nil
+}
+
 func (cvh *cgroupV2Handler) collect(ch chan<- prometheus.Metric) {
 	if !cvh.cfg.MetricsEnabled {
 		return
@@ -91,7 +99,7 @@ func (cvh *cgroupV2Handler) collect(ch chan<- prometheus.Metric) {
 	for i := 0; i < int(cvh.cfg.Repositories.Count); i++ {
 		repoPath := cvh.repoPath(i)
 		logger := cvh.logger.WithField("cgroup_path", repoPath)
-		control, err := cgroup2.Load("/"+repoPath, cgroup2.WithMountpoint(cvh.cfg.Mountpoint))
+		control, err := cvh.loadCgroup(repoPath)
 		if err != nil {
 			logger.WithError(err).Warn("unable to load cgroup controller")
 			return
@@ -152,9 +160,9 @@ func (cvh *cgroupV2Handler) collect(ch chan<- prometheus.Metric) {
 func (cvh *cgroupV2Handler) cleanup() error {
 	processCgroupPath := cvh.currentProcessCgroup()
 
-	control, err := cgroup2.Load("/"+processCgroupPath, cgroup2.WithMountpoint(cvh.cfg.Mountpoint))
+	control, err := cvh.loadCgroup(processCgroupPath)
 	if err != nil {
-		return fmt.Errorf("failed loading cgroup %s: %w", processCgroupPath, err)
+		return err
 	}
 
 	if err := control.Delete(); err != nil {
@@ -175,9 +183,9 @@ func (cvh *cgroupV2Handler) currentProcessCgroup() string {
 func (cvh *cgroupV2Handler) stats() (Stats, error) {
 	processCgroupPath := cvh.currentProcessCgroup()
 
-	control, err := cgroup2.Load("/"+processCgroupPath, cgroup2.WithMountpoint(cvh.cfg.Mountpoint))
+	control, err := cvh.loadCgroup(processCgroupPath)
 	if err != nil {
-		return Stats{}, fmt.Errorf("failed loading cgroup %s: %w", processCgroupPath, err)
+		return Stats{}, err
 	}
 
 	metrics, err := control.Stat()
diff --git a/internal/cgroups/v2_linux_test.go b/internal/cgroups/v2_linux_test.go
index 96c5a3fdb..67b5bc2dc 100644
--- a/internal/cgroups/v2_linux_test.go
+++ b/internal/cgroups/v2_linux_test.go
@@ -4,7 +4,6 @@ package cgroups
 
 import (
 	"fmt"
-	"hash/crc32"
 	"io/fs"
 	"os"
 	"os/exec"
@@ -217,18 +216,17 @@ func TestAddCommandV2(t *testing.T) {
 	require.NoError(t, v2Manager1.Setup())
 	ctx := testhelper.Context(t)
 
-	cmd2 := exec.CommandContext(ctx, "ls", "-hal", ".")
+	cmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 	require.NoError(t, cmd2.Run())
 
 	v2Manager2 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
 
 	t.Run("without overridden key", func(t *testing.T) {
+		groupID := calcGroupID(cmd2.Args, config.Repositories.Count)
+
 		_, err := v2Manager2.AddCommand(cmd2)
 		require.NoError(t, err)
 
-		checksum := crc32.ChecksumIEEE([]byte(strings.Join(cmd2.Args, "/")))
-		groupID := uint(checksum) % config.Repositories.Count
-
 		path := filepath.Join(mock.root, "gitaly",
 			fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
 		content := readCgroupFile(t, path)
@@ -240,14 +238,13 @@ func TestAddCommandV2(t *testing.T) {
 	})
 
 	t.Run("with overridden key", func(t *testing.T) {
+		overriddenGroupID := calcGroupID([]string{"foobar"}, config.Repositories.Count)
+
 		_, err := v2Manager2.AddCommand(cmd2, WithCgroupKey("foobar"))
 		require.NoError(t, err)
 
-		checksum := crc32.ChecksumIEEE([]byte("foobar"))
-		groupID := uint(checksum) % config.Repositories.Count
-
 		path := filepath.Join(mock.root, "gitaly",
-			fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
+			fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", overriddenGroupID), "cgroup.procs")
 		content := readCgroupFile(t, path)
 
 		cmdPid, err := strconv.Atoi(string(content))
@@ -333,17 +330,17 @@ gitaly_cgroup_procs_total{path="%s",subsystem="memory"} 1
 
 			ctx := testhelper.Context(t)
 
-			cmd := exec.CommandContext(ctx, "ls", "-hal", ".")
+			cmd := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, cmd.Start())
 			_, err := v2Manager1.AddCommand(cmd)
 			require.NoError(t, err)
 
-			gitCmd1 := exec.CommandContext(ctx, "ls", "-hal", ".")
+			gitCmd1 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, gitCmd1.Start())
 			_, err = v2Manager1.AddCommand(gitCmd1)
 			require.NoError(t, err)
 
-			gitCmd2 := exec.CommandContext(ctx, "ls", "-hal", ".")
+			gitCmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
 			require.NoError(t, gitCmd2.Start())
 			_, err = v2Manager1.AddCommand(gitCmd2)
 			require.NoError(t, err)