src/internal/cgrouptest/cgrouptest_linux.go - go - Git at Google

 // Copyright 2025 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 // Package cgrouptest provides best-effort helpers for running tests inside a
 // cgroup.
 package cgrouptest

 import (
 	"fmt"
 	"internal/runtime/cgroup"
 	"os"
 	"path/filepath"
 	"slices"
 	"strconv"
 	"strings"
 	"syscall"
 	"testing"
 )

 type CgroupV2 struct {
 	orig string
 	path string
 }

 func (c *CgroupV2) Path() string {
 	return c.path
 }

 // Path to cpu.max.
 func (c *CgroupV2) CPUMaxPath() string {
 	return filepath.Join(c.path, "cpu.max")
 }

 // Set cpu.max. Pass -1 for quota to disable the limit.
 func (c *CgroupV2) SetCPUMax(quota, period int64) error {
 	q := "max"
 	if quota >= 0 {
 		q = strconv.FormatInt(quota, 10)
 	}
 	buf := fmt.Sprintf("%s %d", q, period)
 	return os.WriteFile(c.CPUMaxPath(), []byte(buf), 0)
 }

 // InCgroupV2 creates a new v2 cgroup, migrates the current process into it,
 // and then calls fn. When fn returns, the current process is migrated back to
 // the original cgroup and the new cgroup is destroyed.
 //
 // If a new cgroup cannot be created, the test is skipped.
 //
 // This must not be used in parallel tests, as it affects the entire process.
 func InCgroupV2(t *testing.T, fn func(*CgroupV2)) {
 	mount, rel := findCurrent(t)
 	parent := findOwnedParent(t, mount, rel)
 	orig := filepath.Join(mount, rel)

 	// Make sure the parent allows children to control cpu.
 	b, err := os.ReadFile(filepath.Join(parent, "cgroup.subtree_control"))
 	if err != nil {
 		t.Skipf("unable to read cgroup.subtree_control: %v", err)
 	}
 	if !slices.Contains(strings.Fields(string(b)), "cpu") {
 		// N.B. We should have permission to add cpu to
 		// subtree_control, but it seems like a bad idea to change this
 		// on a high-level cgroup that probably has lots of existing
 		// children.
 		t.Skipf("Parent cgroup %s does not allow children to control cpu, only %q", parent, string(b))
 	}

 	path, err := os.MkdirTemp(parent, "go-cgrouptest")
 	if err != nil {
 		t.Skipf("unable to create cgroup directory: %v", err)
 	}
 	// Important: defer cleanups so they run even in the event of panic.
 	//
 	// TODO(prattmic): Consider running everything in a subprocess just so
 	// we can clean up if it throws or otherwise doesn't run the defers.
 	defer func() {
 		if err := os.Remove(path); err != nil {
 			// Not much we can do, but at least inform of the
 			// problem.
 			t.Errorf("Error removing cgroup directory: %v", err)
 		}
 	}()

 	migrateTo(t, path)
 	defer migrateTo(t, orig)

 	c := &CgroupV2{
 		orig: orig,
 		path: path,
 	}
 	fn(c)
 }

 // Returns the mount and relative directory of the current cgroup the process
 // is in.
 func findCurrent(t *testing.T) (string, string) {
 	// Find the path to our current CPU cgroup. Currently this package is
 	// only used for CPU cgroup testing, so the distinction of different
 	// controllers doesn't matter.
 	var scratch [cgroup.ParseSize]byte
 	buf := make([]byte, cgroup.PathSize)
 	n, err := cgroup.FindCPUMountPoint(buf, scratch[:])
 	if err != nil {
 		t.Skipf("cgroup: unable to find current cgroup mount: %v", err)
 	}
 	mount := string(buf[:n])

 	n, ver, err := cgroup.FindCPURelativePath(buf, scratch[:])
 	if err != nil {
 		t.Skipf("cgroup: unable to find current cgroup path: %v", err)
 	}
 	if ver != cgroup.V2 {
 		t.Skipf("cgroup: running on cgroup v%d want v2", ver)
 	}
 	rel := string(buf[1:n])       // The returned path always starts with /, skip it.
 	rel = filepath.Join(".", rel) // Make sure this isn't empty string at root.
 	return mount, rel
 }

 // Returns a parent directory in which we can create our own cgroup subdirectory.
 func findOwnedParent(t *testing.T, mount, rel string) string {
 	// There are many ways cgroups may be set up on a system. We don't try
 	// to cover all of them, just common ones.
 	//
 	// To start with, systemd:
 	//
 	// Our test process is likely running inside a user session, in which
 	// case we are likely inside a cgroup that looks something like:
 	//
 	//   /sys/fs/cgroup/user.slice/user-1234.slice/user@1234.service/vte-spawn-1.scope/
 	//
 	// Possibly with additional slice layers between user@1234.service and
 	// the leaf scope.
 	//
 	// On new enough kernel and systemd versions (exact versions unknown),
 	// full unprivileged control of the user's cgroups is permitted
 	// directly via the cgroup filesystem. Specifically, the
 	// user@1234.service directory is owned by the user, as are all
 	// subdirectories.

 	// We want to create our own subdirectory that we can migrate into and
 	// then manipulate at will. It is tempting to create a new subdirectory
 	// inside the current cgroup we are already in, however that will likey
 	// not work. cgroup v2 only allows processes to be in leaf cgroups. Our
 	// current cgroup likely contains multiple processes (at least this one
 	// and the cmd/go test runner). If we make a subdirectory and try to
 	// move our process into that cgroup, then the subdirectory and parent
 	// would both contain processes. Linux won't allow us to do that [1].
 	//
 	// Instead, we will simply walk up to the highest directory that our
 	// user owns and create our new subdirectory. Since that directory
 	// already has a bunch of subdirectories, it must not directly contain
 	// and processes.
 	//
 	// (This would fall apart if we already in the highest directory we
 	// own, such as if there was simply a single cgroup for the entire
 	// user. Luckily systemd at least does not do this.)
 	//
 	// [1] Minor technicality: By default a new subdirectory has no cgroup
 	// controller (they must be explicitly enabled in the parent's
 	// cgroup.subtree_control). Linux will allow moving processes into a
 	// subdirectory that has no controllers while there are still processes
 	// in the parent, but it won't allow adding controller until the parent
 	// is empty. As far as I tell, the only purpose of this is to allow
 	// reorganizing processes into a new set of subdirectories and then
 	// adding controllers once done.
 	root, err := os.OpenRoot(mount)
 	if err != nil {
 		t.Fatalf("error opening cgroup mount root: %v", err)
 	}

 	uid := os.Getuid()
 	var prev string
 	for rel != "." {
 		fi, err := root.Stat(rel)
 		if err != nil {
 			t.Fatalf("error stating cgroup path: %v", err)
 		}

 		st := fi.Sys().(*syscall.Stat_t)
 		if int(st.Uid) != uid {
 			// Stop at first directory we don't own.
 			break
 		}

 		prev = rel
 		rel = filepath.Join(rel, "..")
 	}

 	if prev == "" {
 		t.Skipf("No parent cgroup owned by UID %d", uid)
 	}

 	// We actually want the last directory where we were the owner.
 	return filepath.Join(mount, prev)
 }

 // Migrate the current process to the cgroup directory dst.
 func migrateTo(t *testing.T, dst string) {
 	pid := []byte(strconv.FormatInt(int64(os.Getpid()), 10))
 	if err := os.WriteFile(filepath.Join(dst, "cgroup.procs"), pid, 0); err != nil {
 		t.Skipf("Unable to migrate into %s: %v", dst, err)
 	}
 }
	// Copyright 2025 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	// Package cgrouptest provides best-effort helpers for running tests inside a
	// cgroup.
	package cgrouptest

	import (
	"fmt"
	"internal/runtime/cgroup"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"syscall"
	"testing"
	)

	type CgroupV2 struct {
	orig string
	path string
	}

	func (c *CgroupV2) Path() string {
	return c.path
	}

	// Path to cpu.max.
	func (c *CgroupV2) CPUMaxPath() string {
	return filepath.Join(c.path, "cpu.max")
	}

	// Set cpu.max. Pass -1 for quota to disable the limit.
	func (c *CgroupV2) SetCPUMax(quota, period int64) error {
	q := "max"
	if quota >= 0 {
	q = strconv.FormatInt(quota, 10)
	}
	buf := fmt.Sprintf("%s %d", q, period)
	return os.WriteFile(c.CPUMaxPath(), []byte(buf), 0)
	}

	// InCgroupV2 creates a new v2 cgroup, migrates the current process into it,
	// and then calls fn. When fn returns, the current process is migrated back to
	// the original cgroup and the new cgroup is destroyed.
	//
	// If a new cgroup cannot be created, the test is skipped.
	//
	// This must not be used in parallel tests, as it affects the entire process.
	func InCgroupV2(t testing.T, fn func(CgroupV2)) {
	mount, rel := findCurrent(t)
	parent := findOwnedParent(t, mount, rel)
	orig := filepath.Join(mount, rel)

	// Make sure the parent allows children to control cpu.
	b, err := os.ReadFile(filepath.Join(parent, "cgroup.subtree_control"))
	if err != nil {
	t.Skipf("unable to read cgroup.subtree_control: %v", err)
	}
	if !slices.Contains(strings.Fields(string(b)), "cpu") {
	// N.B. We should have permission to add cpu to
	// subtree_control, but it seems like a bad idea to change this
	// on a high-level cgroup that probably has lots of existing
	// children.
	t.Skipf("Parent cgroup %s does not allow children to control cpu, only %q", parent, string(b))
	}

	path, err := os.MkdirTemp(parent, "go-cgrouptest")
	if err != nil {
	t.Skipf("unable to create cgroup directory: %v", err)
	}
	// Important: defer cleanups so they run even in the event of panic.
	//
	// TODO(prattmic): Consider running everything in a subprocess just so
	// we can clean up if it throws or otherwise doesn't run the defers.
	defer func() {
	if err := os.Remove(path); err != nil {
	// Not much we can do, but at least inform of the
	// problem.
	t.Errorf("Error removing cgroup directory: %v", err)
	}
	}()

	migrateTo(t, path)
	defer migrateTo(t, orig)

	c := &CgroupV2{
	orig: orig,
	path: path,
	}
	fn(c)
	}

	// Returns the mount and relative directory of the current cgroup the process
	// is in.
	func findCurrent(t *testing.T) (string, string) {
	// Find the path to our current CPU cgroup. Currently this package is
	// only used for CPU cgroup testing, so the distinction of different
	// controllers doesn't matter.
	var scratch [cgroup.ParseSize]byte
	buf := make([]byte, cgroup.PathSize)
	n, err := cgroup.FindCPUMountPoint(buf, scratch[:])
	if err != nil {
	t.Skipf("cgroup: unable to find current cgroup mount: %v", err)
	}
	mount := string(buf[:n])

	n, ver, err := cgroup.FindCPURelativePath(buf, scratch[:])
	if err != nil {
	t.Skipf("cgroup: unable to find current cgroup path: %v", err)
	}
	if ver != cgroup.V2 {
	t.Skipf("cgroup: running on cgroup v%d want v2", ver)
	}
	rel := string(buf[1:n]) // The returned path always starts with /, skip it.
	rel = filepath.Join(".", rel) // Make sure this isn't empty string at root.
	return mount, rel
	}

	// Returns a parent directory in which we can create our own cgroup subdirectory.
	func findOwnedParent(t *testing.T, mount, rel string) string {
	// There are many ways cgroups may be set up on a system. We don't try
	// to cover all of them, just common ones.
	//
	// To start with, systemd:
	//
	// Our test process is likely running inside a user session, in which
	// case we are likely inside a cgroup that looks something like:
	//
	// /sys/fs/cgroup/user.slice/user-1234.slice/user@1234.service/vte-spawn-1.scope/
	//
	// Possibly with additional slice layers between user@1234.service and
	// the leaf scope.
	//
	// On new enough kernel and systemd versions (exact versions unknown),
	// full unprivileged control of the user's cgroups is permitted
	// directly via the cgroup filesystem. Specifically, the
	// user@1234.service directory is owned by the user, as are all
	// subdirectories.

	// We want to create our own subdirectory that we can migrate into and
	// then manipulate at will. It is tempting to create a new subdirectory
	// inside the current cgroup we are already in, however that will likey
	// not work. cgroup v2 only allows processes to be in leaf cgroups. Our
	// current cgroup likely contains multiple processes (at least this one
	// and the cmd/go test runner). If we make a subdirectory and try to
	// move our process into that cgroup, then the subdirectory and parent
	// would both contain processes. Linux won't allow us to do that [1].
	//
	// Instead, we will simply walk up to the highest directory that our
	// user owns and create our new subdirectory. Since that directory
	// already has a bunch of subdirectories, it must not directly contain
	// and processes.
	//
	// (This would fall apart if we already in the highest directory we
	// own, such as if there was simply a single cgroup for the entire
	// user. Luckily systemd at least does not do this.)
	//
	// [1] Minor technicality: By default a new subdirectory has no cgroup
	// controller (they must be explicitly enabled in the parent's
	// cgroup.subtree_control). Linux will allow moving processes into a
	// subdirectory that has no controllers while there are still processes
	// in the parent, but it won't allow adding controller until the parent
	// is empty. As far as I tell, the only purpose of this is to allow
	// reorganizing processes into a new set of subdirectories and then
	// adding controllers once done.
	root, err := os.OpenRoot(mount)
	if err != nil {
	t.Fatalf("error opening cgroup mount root: %v", err)
	}

	uid := os.Getuid()
	var prev string
	for rel != "." {
	fi, err := root.Stat(rel)
	if err != nil {
	t.Fatalf("error stating cgroup path: %v", err)
	}

	st := fi.Sys().(*syscall.Stat_t)
	if int(st.Uid) != uid {
	// Stop at first directory we don't own.
	break
	}

	prev = rel
	rel = filepath.Join(rel, "..")
	}

	if prev == "" {
	t.Skipf("No parent cgroup owned by UID %d", uid)
	}

	// We actually want the last directory where we were the owner.
	return filepath.Join(mount, prev)
	}

	// Migrate the current process to the cgroup directory dst.
	func migrateTo(t *testing.T, dst string) {
	pid := []byte(strconv.FormatInt(int64(os.Getpid()), 10))
	if err := os.WriteFile(filepath.Join(dst, "cgroup.procs"), pid, 0); err != nil {
	t.Skipf("Unable to migrate into %s: %v", dst, err)
	}
	}