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  	orig := findCurrent(t)
    54  	parent := findOwnedParent(t, orig)
    55  
    56  	// Make sure the parent allows children to control cpu.
    57  	b, err := os.ReadFile(filepath.Join(parent, "cgroup.subtree_control"))
    58  	if err != nil {
    59  		t.Skipf("unable to read cgroup.subtree_control: %v", err)
    60  	}
    61  	if !slices.Contains(strings.Fields(string(b)), "cpu") {
    62  		// N.B. We should have permission to add cpu to
    63  		// subtree_control, but it seems like a bad idea to change this
    64  		// on a high-level cgroup that probably has lots of existing
    65  		// children.
    66  		t.Skipf("Parent cgroup %s does not allow children to control cpu, only %q", parent, string(b))
    67  	}
    68  
    69  	path, err := os.MkdirTemp(parent, "go-cgrouptest")
    70  	if err != nil {
    71  		t.Skipf("unable to create cgroup directory: %v", err)
    72  	}
    73  	// Important: defer cleanups so they run even in the event of panic.
    74  	//
    75  	// TODO(prattmic): Consider running everything in a subprocess just so
    76  	// we can clean up if it throws or otherwise doesn't run the defers.
    77  	defer func() {
    78  		if err := os.Remove(path); err != nil {
    79  			// Not much we can do, but at least inform of the
    80  			// problem.
    81  			t.Errorf("Error removing cgroup directory: %v", err)
    82  		}
    83  	}()
    84  
    85  	migrateTo(t, path)
    86  	defer migrateTo(t, orig)
    87  
    88  	c := &CgroupV2{
    89  		orig: orig,
    90  		path: path,
    91  	}
    92  	fn(c)
    93  }
    94  
    95  // Returns the filesystem path to the current cgroup the process is in.
    96  func findCurrent(t *testing.T) string {
    97  	// Find the path to our current CPU cgroup. Currently this package is
    98  	// only used for CPU cgroup testing, so the distinction of different
    99  	// controllers doesn't matter.
   100  	var scratch [cgroup.ParseSize]byte
   101  	buf := make([]byte, cgroup.PathSize)
   102  	n, ver, err := cgroup.FindCPU(buf, scratch[:])
   103  	if err != nil {
   104  		t.Skipf("cgroup: unable to find current cgroup mount: %v", err)
   105  	}
   106  	if ver != cgroup.V2 {
   107  		t.Skipf("cgroup: running on cgroup v%d want v2", ver)
   108  	}
   109  	return string(buf[:n])
   110  }
   111  
   112  // Returns a parent directory in which we can create our own cgroup subdirectory.
   113  func findOwnedParent(t *testing.T, orig string) string {
   114  	// There are many ways cgroups may be set up on a system. We don't try
   115  	// to cover all of them, just common ones.
   116  	//
   117  	// To start with, systemd:
   118  	//
   119  	// Our test process is likely running inside a user session, in which
   120  	// case we are likely inside a cgroup that looks something like:
   121  	//
   122  	//   /sys/fs/cgroup/user.slice/user-1234.slice/user@1234.service/vte-spawn-1.scope/
   123  	//
   124  	// Possibly with additional slice layers between user@1234.service and
   125  	// the leaf scope.
   126  	//
   127  	// On new enough kernel and systemd versions (exact versions unknown),
   128  	// full unprivileged control of the user's cgroups is permitted
   129  	// directly via the cgroup filesystem. Specifically, the
   130  	// user@1234.service directory is owned by the user, as are all
   131  	// subdirectories.
   132  
   133  	// We want to create our own subdirectory that we can migrate into and
   134  	// then manipulate at will. It is tempting to create a new subdirectory
   135  	// inside the current cgroup we are already in, however that will likely
   136  	// not work. cgroup v2 only allows processes to be in leaf cgroups. Our
   137  	// current cgroup likely contains multiple processes (at least this one
   138  	// and the cmd/go test runner). If we make a subdirectory and try to
   139  	// move our process into that cgroup, then the subdirectory and parent
   140  	// would both contain processes. Linux won't allow us to do that [1].
   141  	//
   142  	// Instead, we will simply walk up to the highest directory that our
   143  	// user owns and create our new subdirectory. Since that directory
   144  	// already has a bunch of subdirectories, it must not directly contain
   145  	// and processes.
   146  	//
   147  	// (This would fall apart if we already in the highest directory we
   148  	// own, such as if there was simply a single cgroup for the entire
   149  	// user. Luckily systemd at least does not do this.)
   150  	//
   151  	// [1] Minor technicality: By default a new subdirectory has no cgroup
   152  	// controller (they must be explicitly enabled in the parent's
   153  	// cgroup.subtree_control). Linux will allow moving processes into a
   154  	// subdirectory that has no controllers while there are still processes
   155  	// in the parent, but it won't allow adding controller until the parent
   156  	// is empty. As far as I tell, the only purpose of this is to allow
   157  	// reorganizing processes into a new set of subdirectories and then
   158  	// adding controllers once done.
   159  	var stat syscall.Stat_t
   160  	err := syscall.Stat(orig, &stat)
   161  	if err != nil {
   162  		t.Fatalf("error stating orig cgroup: %v", err)
   163  	}
   164  
   165  	uid := os.Getuid()
   166  	var prev string
   167  	cur := filepath.Dir(orig)
   168  	for cur != "/" {
   169  		var curStat syscall.Stat_t
   170  		err = syscall.Stat(cur, &curStat)
   171  		if err != nil {
   172  			t.Fatalf("error stating cgroup path: %v", err)
   173  		}
   174  
   175  		if int(curStat.Uid) != uid || curStat.Dev != stat.Dev {
   176  			// Stop at first directory we don't own or filesystem boundary.
   177  			break
   178  		}
   179  
   180  		prev = cur
   181  		cur = filepath.Dir(cur)
   182  	}
   183  
   184  	if prev == "" {
   185  		t.Skipf("No parent cgroup owned by UID %d", uid)
   186  	}
   187  
   188  	// We actually want the last directory where we were the owner.
   189  	return prev
   190  }
   191  
   192  // Migrate the current process to the cgroup directory dst.
   193  func migrateTo(t *testing.T, dst string) {
   194  	pid := []byte(strconv.FormatInt(int64(os.Getpid()), 10))
   195  	if err := os.WriteFile(filepath.Join(dst, "cgroup.procs"), pid, 0); err != nil {
   196  		t.Skipf("Unable to migrate into %s: %v", dst, err)
   197  	}
   198  }
   199
View as plain text