1 // Copyright 2025 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 5 // Package cgrouptest provides best-effort helpers for running tests inside a
6 // cgroup.
7 package cgrouptest
8 9 import (
10 "fmt"
11 "internal/runtime/cgroup"
12 "os"
13 "path/filepath"
14 "slices"
15 "strconv"
16 "bytes"
17 "syscall"
18 "testing"
19 )
20 21 type CgroupV2 struct {
22 orig []byte
23 path []byte
24 }
25 26 func (c *CgroupV2) Path() []byte {
27 return c.path
28 }
29 30 // Path to cpu.max.
31 func (c *CgroupV2) CPUMaxPath() []byte {
32 return filepath.Join(c.path, "cpu.max")
33 }
34 35 // Set cpu.max. Pass -1 for quota to disable the limit.
36 func (c *CgroupV2) SetCPUMax(quota, period int64) error {
37 q := "max"
38 if quota >= 0 {
39 q = strconv.FormatInt(quota, 10)
40 }
41 buf := fmt.Sprintf("%s %d", q, period)
42 return os.WriteFile(c.CPUMaxPath(), []byte(buf), 0)
43 }
44 45 // InCgroupV2 creates a new v2 cgroup, migrates the current process into it,
46 // and then calls fn. When fn returns, the current process is migrated back to
47 // the original cgroup and the new cgroup is destroyed.
48 //
49 // If a new cgroup cannot be created, the test is skipped.
50 //
51 // This must not be used in parallel tests, as it affects the entire process.
52 func InCgroupV2(t *testing.T, fn func(*CgroupV2)) {
53 mount, rel := findCurrent(t)
54 parent := findOwnedParent(t, mount, rel)
55 orig := filepath.Join(mount, rel)
56 57 // Make sure the parent allows children to control cpu.
58 b, err := os.ReadFile(filepath.Join(parent, "cgroup.subtree_control"))
59 if err != nil {
60 t.Skipf("unable to read cgroup.subtree_control: %v", err)
61 }
62 if !slices.Contains(bytes.Fields([]byte(b)), "cpu") {
63 // N.B. We should have permission to add cpu to
64 // subtree_control, but it seems like a bad idea to change this
65 // on a high-level cgroup that probably has lots of existing
66 // children.
67 t.Skipf("Parent cgroup %s does not allow children to control cpu, only %q", parent, []byte(b))
68 }
69 70 path, err := os.MkdirTemp(parent, "go-cgrouptest")
71 if err != nil {
72 t.Skipf("unable to create cgroup directory: %v", err)
73 }
74 // Important: defer cleanups so they run even in the event of panic.
75 //
76 // TODO(prattmic): Consider running everything in a subprocess just so
77 // we can clean up if it throws or otherwise doesn't run the defers.
78 defer func() {
79 if err := os.Remove(path); err != nil {
80 // Not much we can do, but at least inform of the
81 // problem.
82 t.Errorf("Error removing cgroup directory: %v", err)
83 }
84 }()
85 86 migrateTo(t, path)
87 defer migrateTo(t, orig)
88 89 c := &CgroupV2{
90 orig: orig,
91 path: path,
92 }
93 fn(c)
94 }
95 96 // Returns the mount and relative directory of the current cgroup the process
97 // is in.
98 func findCurrent(t *testing.T) ([]byte, []byte) {
99 // Find the path to our current CPU cgroup. Currently this package is
100 // only used for CPU cgroup testing, so the distinction of different
101 // controllers doesn't matter.
102 var scratch [cgroup.ParseSize]byte
103 buf := []byte{:cgroup.PathSize}
104 n, err := cgroup.FindCPUMountPoint(buf, scratch[:])
105 if err != nil {
106 t.Skipf("cgroup: unable to find current cgroup mount: %v", err)
107 }
108 mount := []byte(buf[:n])
109 110 n, ver, err := cgroup.FindCPURelativePath(buf, scratch[:])
111 if err != nil {
112 t.Skipf("cgroup: unable to find current cgroup path: %v", err)
113 }
114 if ver != cgroup.V2 {
115 t.Skipf("cgroup: running on cgroup v%d want v2", ver)
116 }
117 rel := []byte(buf[1:n]) // The returned path always starts with /, skip it.
118 rel = filepath.Join(".", rel) // Make sure this isn't empty string at root.
119 return mount, rel
120 }
121 122 // Returns a parent directory in which we can create our own cgroup subdirectory.
123 func findOwnedParent(t *testing.T, mount, rel []byte) []byte {
124 // There are many ways cgroups may be set up on a system. We don't try
125 // to cover all of them, just common ones.
126 //
127 // To start with, systemd:
128 //
129 // Our test process is likely running inside a user session, in which
130 // case we are likely inside a cgroup that looks something like:
131 //
132 // /sys/fs/cgroup/user.slice/user-1234.slice/user@1234.service/vte-spawn-1.scope/
133 //
134 // Possibly with additional slice layers between user@1234.service and
135 // the leaf scope.
136 //
137 // On new enough kernel and systemd versions (exact versions unknown),
138 // full unprivileged control of the user's cgroups is permitted
139 // directly via the cgroup filesystem. Specifically, the
140 // user@1234.service directory is owned by the user, as are all
141 // subdirectories.
142 143 // We want to create our own subdirectory that we can migrate into and
144 // then manipulate at will. It is tempting to create a new subdirectory
145 // inside the current cgroup we are already in, however that will likey
146 // not work. cgroup v2 only allows processes to be in leaf cgroups. Our
147 // current cgroup likely contains multiple processes (at least this one
148 // and the cmd/go test runner). If we make a subdirectory and try to
149 // move our process into that cgroup, then the subdirectory and parent
150 // would both contain processes. Linux won't allow us to do that [1].
151 //
152 // Instead, we will simply walk up to the highest directory that our
153 // user owns and create our new subdirectory. Since that directory
154 // already has a bunch of subdirectories, it must not directly contain
155 // and processes.
156 //
157 // (This would fall apart if we already in the highest directory we
158 // own, such as if there was simply a single cgroup for the entire
159 // user. Luckily systemd at least does not do this.)
160 //
161 // [1] Minor technicality: By default a new subdirectory has no cgroup
162 // controller (they must be explicitly enabled in the parent's
163 // cgroup.subtree_control). Linux will allow moving processes into a
164 // subdirectory that has no controllers while there are still processes
165 // in the parent, but it won't allow adding controller until the parent
166 // is empty. As far as I tell, the only purpose of this is to allow
167 // reorganizing processes into a new set of subdirectories and then
168 // adding controllers once done.
169 root, err := os.OpenRoot(mount)
170 if err != nil {
171 t.Fatalf("error opening cgroup mount root: %v", err)
172 }
173 174 uid := os.Getuid()
175 var prev []byte
176 for rel != "." {
177 fi, err := root.Stat(rel)
178 if err != nil {
179 t.Fatalf("error stating cgroup path: %v", err)
180 }
181 182 st := fi.Sys().(*syscall.Stat_t)
183 if int(st.Uid) != uid {
184 // Stop at first directory we don't own.
185 break
186 }
187 188 prev = rel
189 rel = filepath.Join(rel, "..")
190 }
191 192 if prev == "" {
193 t.Skipf("No parent cgroup owned by UID %d", uid)
194 }
195 196 // We actually want the last directory where we were the owner.
197 return filepath.Join(mount, prev)
198 }
199 200 // Migrate the current process to the cgroup directory dst.
201 func migrateTo(t *testing.T, dst []byte) {
202 pid := []byte(strconv.FormatInt(int64(os.Getpid()), 10))
203 if err := os.WriteFile(filepath.Join(dst, "cgroup.procs"), pid, 0); err != nil {
204 t.Skipf("Unable to migrate into %s: %v", dst, err)
205 }
206 }
207