|
| 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 | + "strings" |
| 17 | + "syscall" |
| 18 | + "testing" |
| 19 | +) |
| 20 | + |
| 21 | +type CgroupV2 struct { |
| 22 | + orig string |
| 23 | + path string |
| 24 | +} |
| 25 | + |
| 26 | +func (c *CgroupV2) Path() string { |
| 27 | + return c.path |
| 28 | +} |
| 29 | + |
| 30 | +// Path to cpu.max. |
| 31 | +func (c *CgroupV2) CPUMaxPath() string { |
| 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(strings.Fields(string(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, string(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) (string, string) { |
| 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 := make([]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 := string(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 := string(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 string) string { |
| 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 string |
| 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 string) { |
| 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 | +} |
0 commit comments