cgroup: Enhance cgroup mock setup

Cgroup testing suite has an utility to mock underlying cgroupfs files.
This mocking utility hard-codes the values or accepts a very limited set
of configs. This commits adds a map that allows this utility to mock
custom values of files.

2 files changed, 78 insertions, 45 deletions

diff --git a/internal/cgroups/mock_linux_test.go b/internal/cgroups/mock_linux_test.go
index 135dca76f..72e040ca7 100644
--- a/internal/cgroups/mock_linux_test.go
+++ b/internal/cgroups/mock_linux_test.go
@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"strconv"
 	"testing"
 
 	cgrps "github.com/containerd/cgroups/v3"
@@ -59,50 +58,73 @@ func newMock(t *testing.T) *mockCgroup {
 	}
 }
 
+type mockCgroupFile struct {
+	name    string
+	content string
+}
+
+var defaultV1MockMemoryFiles = map[string]string{
+	"cgroup.procs":                       "",
+	"memory.stat":                        "",
+	"memory.oom_control":                 "",
+	"memory.usage_in_bytes":              "0",
+	"memory.max_usage_in_bytes":          "0",
+	"memory.limit_in_bytes":              "0",
+	"memory.failcnt":                     "0",
+	"memory.memsw.failcnt":               "0",
+	"memory.memsw.usage_in_bytes":        "0",
+	"memory.memsw.max_usage_in_bytes":    "0",
+	"memory.memsw.limit_in_bytes":        "0",
+	"memory.kmem.usage_in_bytes":         "0",
+	"memory.kmem.max_usage_in_bytes":     "0",
+	"memory.kmem.failcnt":                "0",
+	"memory.kmem.limit_in_bytes":         "0",
+	"memory.kmem.tcp.usage_in_bytes":     "0",
+	"memory.kmem.tcp.max_usage_in_bytes": "0",
+	"memory.kmem.tcp.failcnt":            "0",
+	"memory.kmem.tcp.limit_in_bytes":     "0",
+}
+
+var defaultV1MockCPUFiles = map[string]string{
+	"cgroup.procs": "",
+	"cpu.shares":   "0",
+	"cpu.stat": `nr_periods 10
+nr_throttled 20
+throttled_time 1000000`,
+}
+
 func (m *mockCgroup) setupMockCgroupFiles(
 	t *testing.T,
 	manager *CGroupManager,
-	memFailCount int,
+	inputContent ...mockCgroupFile,
 ) {
 	for _, s := range m.subsystems {
 		cgroupPath := filepath.Join(m.root, string(s.Name()), manager.currentProcessCgroup())
 		require.NoError(t, os.MkdirAll(cgroupPath, perm.SharedDir))
 
-		contentByFilename := map[string]string{
-			"cgroup.procs": "",
+		content := map[string]string{}
+		for _, ic := range inputContent {
+			content[ic.name] = ic.content
 		}
 
 		switch s.Name() {
 		case "memory":
-			contentByFilename["memory.stat"] = ""
-			contentByFilename["memory.oom_control"] = ""
-			contentByFilename["memory.usage_in_bytes"] = "0"
-			contentByFilename["memory.max_usage_in_bytes"] = "0"
-			contentByFilename["memory.limit_in_bytes"] = "0"
-			contentByFilename["memory.failcnt"] = "0"
-			contentByFilename["memory.memsw.failcnt"] = "0"
-			contentByFilename["memory.memsw.usage_in_bytes"] = "0"
-			contentByFilename["memory.memsw.max_usage_in_bytes"] = "0"
-			contentByFilename["memory.memsw.limit_in_bytes"] = "0"
-			contentByFilename["memory.kmem.usage_in_bytes"] = "0"
-			contentByFilename["memory.kmem.max_usage_in_bytes"] = "0"
-			contentByFilename["memory.kmem.failcnt"] = "0"
-			contentByFilename["memory.kmem.limit_in_bytes"] = "0"
-			contentByFilename["memory.kmem.tcp.usage_in_bytes"] = "0"
-			contentByFilename["memory.kmem.tcp.max_usage_in_bytes"] = "0"
-			contentByFilename["memory.kmem.tcp.failcnt"] = "0"
-			contentByFilename["memory.kmem.tcp.limit_in_bytes"] = "0"
-			contentByFilename["memory.failcnt"] = strconv.Itoa(memFailCount)
+			for key, value := range defaultV1MockMemoryFiles {
+				if _, exist := content[key]; !exist {
+					content[key] = value
+				}
+			}
 		case "cpu":
-			contentByFilename["cpu.shares"] = "0"
-			contentByFilename["cpu.stat"] = `nr_periods 10
-nr_throttled 20
-throttled_time 1000000`
+			for key, value := range defaultV1MockCPUFiles {
+				if _, exist := content[key]; !exist {
+					content[key] = value
+				}
+			}
 		default:
 			require.FailNow(t, "cannot set up subsystem", "unknown subsystem %q", s.Name())
 		}
 
-		for filename, content := range contentByFilename {
+		for filename, content := range content {
 			controlFilePath := filepath.Join(cgroupPath, filename)
 			require.NoError(t, os.WriteFile(controlFilePath, []byte(content), perm.SharedFile))
 		}
@@ -111,7 +133,7 @@ throttled_time 1000000`
 			shardPath := filepath.Join(cgroupPath, fmt.Sprintf("repos-%d", shard))
 			require.NoError(t, os.MkdirAll(shardPath, perm.SharedDir))
 
-			for filename, content := range contentByFilename {
+			for filename, content := range content {
 				shardControlFilePath := filepath.Join(shardPath, filename)
 				require.NoError(t, os.WriteFile(shardControlFilePath, []byte(content), perm.SharedFile))
 			}
@@ -139,31 +161,42 @@ func newMockV2(t *testing.T) *mockCgroupV2 {
 	}
 }
 
+var defaultV2MockFiles = map[string]string{
+	"cgroup.procs":           "",
+	"cgroup.subtree_control": "cpu cpuset memory",
+	"cgroup.controllers":     "cpu cpuset memory",
+	"cpu.max":                "max 100000",
+	"cpu.weight":             "10",
+	"memory.max":             "max",
+	"cpu.stat": `nr_periods 10
+	nr_throttled 20
+	throttled_usec 1000000`,
+}
+
 func (m *mockCgroupV2) setupMockCgroupFiles(
 	t *testing.T,
 	manager *CGroupManager,
+	inputContent ...mockCgroupFile,
 ) {
+	content := map[string]string{}
+	for _, ic := range inputContent {
+		content[ic.name] = ic.content
+	}
+	for key, value := range defaultV2MockFiles {
+		if _, exist := content[key]; !exist {
+			content[key] = value
+		}
+	}
+
 	cgroupPath := filepath.Join(m.root, manager.currentProcessCgroup())
 	require.NoError(t, os.MkdirAll(cgroupPath, perm.SharedDir))
 
-	contentByFilename := map[string]string{
-		"cgroup.procs":           "",
-		"cgroup.subtree_control": "cpu cpuset memory",
-		"cgroup.controllers":     "cpu cpuset memory",
-		"cpu.max":                "max 100000",
-		"cpu.weight":             "10",
-		"memory.max":             "max",
-		"cpu.stat": `nr_periods 10
-		nr_throttled 20
-		throttled_usec 1000000`,
-	}
-
-	for filename, content := range contentByFilename {
+	for filename, content := range content {
 		controlFilePath := filepath.Join(m.root, manager.cfg.HierarchyRoot, filename)
 		require.NoError(t, os.WriteFile(controlFilePath, []byte(content), perm.SharedFile))
 	}
 
-	for filename, content := range contentByFilename {
+	for filename, content := range content {
 		controlFilePath := filepath.Join(cgroupPath, filename)
 		require.NoError(t, os.WriteFile(controlFilePath, []byte(content), perm.SharedFile))
 	}
@@ -172,7 +205,7 @@ func (m *mockCgroupV2) setupMockCgroupFiles(
 		shardPath := filepath.Join(cgroupPath, fmt.Sprintf("repos-%d", shard))
 		require.NoError(t, os.MkdirAll(shardPath, perm.SharedDir))
 
-		for filename, content := range contentByFilename {
+		for filename, content := range content {
 			shardControlFilePath := filepath.Join(shardPath, filename)
 			require.NoError(t, os.WriteFile(shardControlFilePath, []byte(content), perm.SharedFile))
 		}
diff --git a/internal/cgroups/v1_linux_test.go b/internal/cgroups/v1_linux_test.go
index a68ebed4d..234020468 100644
--- a/internal/cgroups/v1_linux_test.go
+++ b/internal/cgroups/v1_linux_test.go
@@ -344,7 +344,7 @@ gitaly_cgroup_cpu_cfs_throttled_seconds_total{path="%s"} 0.001
 
 			v1Manager1 := mock.newCgroupManager(config, tt.pid)
 
-			mock.setupMockCgroupFiles(t, v1Manager1, 2)
+			mock.setupMockCgroupFiles(t, v1Manager1, mockCgroupFile{"memory.failcnt", "2"})
 			require.NoError(t, v1Manager1.Setup())
 
 			ctx := testhelper.Context(t)