cgrouptest_linux.go

     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  }
   207
View as plain text