Stub cpuset cgroup control files.

PiperOrigin-RevId: 404382475

7 files changed, 446 insertions, 5 deletions

diff --git a/pkg/sentry/fsimpl/cgroupfs/BUILD b/pkg/sentry/fsimpl/cgroupfs/BUILD
index e5fdcc776..60ee5ede2 100644
--- a/pkg/sentry/fsimpl/cgroupfs/BUILD
+++ b/pkg/sentry/fsimpl/cgroupfs/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library")
+load("//tools:defs.bzl", "go_library", "go_test")
 load("//tools/go_generics:defs.bzl", "go_template_instance")
 
 licenses(["notice"])
@@ -18,6 +18,7 @@ go_library(
     name = "cgroupfs",
     srcs = [
         "base.go",
+        "bitmap.go",
         "cgroupfs.go",
         "cpu.go",
         "cpuacct.go",
@@ -29,10 +30,12 @@ go_library(
     visibility = ["//pkg/sentry:internal"],
     deps = [
         "//pkg/abi/linux",
+        "//pkg/bitmap",
         "//pkg/context",
         "//pkg/coverage",
         "//pkg/errors/linuxerr",
         "//pkg/fspath",
+        "//pkg/hostarch",
         "//pkg/log",
         "//pkg/refs",
         "//pkg/refsvfs2",
@@ -47,3 +50,11 @@ go_library(
         "//pkg/usermem",
     ],
 )
+
+go_test(
+    name = "cgroupfs_test",
+    size = "small",
+    srcs = ["bitmap_test.go"],
+    library = ":cgroupfs",
+    deps = ["//pkg/bitmap"],
+)
diff --git a/pkg/sentry/fsimpl/cgroupfs/bitmap.go b/pkg/sentry/fsimpl/cgroupfs/bitmap.go
new file mode 100644
index 000000000..8074641db
--- /dev/null
+++ b/pkg/sentry/fsimpl/cgroupfs/bitmap.go
@@ -0,0 +1,139 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cgroupfs
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"gvisor.dev/gvisor/pkg/bitmap"
+)
+
+// formatBitmap produces a string representation of b, which lists the indicies
+// of set bits in the bitmap. Indicies are separated by commas and ranges of
+// set bits are abbreviated. Example outputs: "0,2,4", "0,3-7,10", "0-10".
+//
+// Inverse of parseBitmap.
+func formatBitmap(b *bitmap.Bitmap) string {
+	ones := b.ToSlice()
+	if len(ones) == 0 {
+		return ""
+	}
+
+	elems := make([]string, 0, len(ones))
+	runStart := ones[0]
+	lastVal := ones[0]
+	inRun := false
+
+	for _, v := range ones[1:] {
+		last := lastVal
+		lastVal = v
+
+		if last+1 == v {
+			// In a contiguous block of ones.
+			if !inRun {
+				runStart = last
+				inRun = true
+			}
+
+			continue
+		}
+
+		// Non-contiguous bit.
+		if inRun {
+			// Render a run
+			elems = append(elems, fmt.Sprintf("%d-%d", runStart, last))
+			inRun = false
+			continue
+		}
+
+		// Lone non-contiguous bit.
+		elems = append(elems, fmt.Sprintf("%d", last))
+
+	}
+
+	// Process potential final run
+	if inRun {
+		elems = append(elems, fmt.Sprintf("%d-%d", runStart, lastVal))
+	} else {
+		elems = append(elems, fmt.Sprintf("%d", lastVal))
+	}
+
+	return strings.Join(elems, ",")
+}
+
+func parseToken(token string) (start, end uint32, err error) {
+	ts := strings.SplitN(token, "-", 2)
+	switch len(ts) {
+	case 0:
+		return 0, 0, fmt.Errorf("invalid token %q", token)
+	case 1:
+		val, err := strconv.ParseUint(ts[0], 10, 32)
+		if err != nil {
+			return 0, 0, err
+		}
+		return uint32(val), uint32(val), nil
+	case 2:
+		val1, err := strconv.ParseUint(ts[0], 10, 32)
+		if err != nil {
+			return 0, 0, err
+		}
+		val2, err := strconv.ParseUint(ts[1], 10, 32)
+		if err != nil {
+			return 0, 0, err
+		}
+		if val1 >= val2 {
+			return 0, 0, fmt.Errorf("start (%v) must be less than end (%v)", val1, val2)
+		}
+		return uint32(val1), uint32(val2), nil
+	default:
+		panic(fmt.Sprintf("Unreachable: got %d substrs", len(ts)))
+	}
+}
+
+// parseBitmap parses input as a bitmap. input should be a comma separated list
+// of indices, and ranges of set bits may be abbreviated. Examples: "0,2,4",
+// "0,3-7,10", "0-10". Input after the first newline or null byte is discarded.
+//
+// sizeHint sets the initial size of the bitmap, which may prevent reallocation
+// when growing the bitmap during parsing. Ideally sizeHint should be at least
+// as large as the bitmap represented by input, but this is not required.
+//
+// Inverse of formatBitmap.
+func parseBitmap(input string, sizeHint uint32) (*bitmap.Bitmap, error) {
+	b := bitmap.New(sizeHint)
+
+	if termIdx := strings.IndexAny(input, "\n\000"); termIdx != -1 {
+		input = input[:termIdx]
+	}
+	input = strings.TrimSpace(input)
+
+	if len(input) == 0 {
+		return &b, nil
+	}
+	tokens := strings.Split(input, ",")
+
+	for _, t := range tokens {
+		start, end, err := parseToken(strings.TrimSpace(t))
+		if err != nil {
+			return nil, err
+		}
+		for i := start; i <= end; i++ {
+			b.Add(i)
+		}
+	}
+	return &b, nil
+}
diff --git a/pkg/sentry/fsimpl/cgroupfs/bitmap_test.go b/pkg/sentry/fsimpl/cgroupfs/bitmap_test.go
new file mode 100644
index 000000000..5cc56de3b
--- /dev/null
+++ b/pkg/sentry/fsimpl/cgroupfs/bitmap_test.go
@@ -0,0 +1,99 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cgroupfs
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/bitmap"
+)
+
+func TestFormat(t *testing.T) {
+	tests := []struct {
+		input  []uint32
+		output string
+	}{
+		{[]uint32{1, 2, 3, 4, 7}, "1-4,7"},
+		{[]uint32{2}, "2"},
+		{[]uint32{0, 1, 2}, "0-2"},
+		{[]uint32{}, ""},
+		{[]uint32{1, 3, 4, 5, 6, 9, 11, 13, 14, 15, 16, 17}, "1,3-6,9,11,13-17"},
+		{[]uint32{2, 3, 10, 12, 13, 14, 15, 16, 20, 21, 33, 34, 47}, "2-3,10,12-16,20-21,33-34,47"},
+	}
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("case-%d", i), func(t *testing.T) {
+			b := bitmap.New(64)
+			for _, v := range tt.input {
+				b.Add(v)
+			}
+			s := formatBitmap(&b)
+			if s != tt.output {
+				t.Errorf("Expected %q, got %q", tt.output, s)
+			}
+			b1, err := parseBitmap(s, 64)
+			if err != nil {
+				t.Fatalf("Failed to parse formatted bitmap: %v", err)
+			}
+			if got, want := b1.ToSlice(), b.ToSlice(); !reflect.DeepEqual(got, want) {
+				t.Errorf("Parsing formatted output doesn't result in the original bitmap. Got %v, want %v", got, want)
+			}
+		})
+	}
+}
+
+func TestParse(t *testing.T) {
+	tests := []struct {
+		input      string
+		output     []uint32
+		shouldFail bool
+	}{
+		{"1", []uint32{1}, false},
+		{"", []uint32{}, false},
+		{"1,2,3,4", []uint32{1, 2, 3, 4}, false},
+		{"1-4", []uint32{1, 2, 3, 4}, false},
+		{"1,2-4", []uint32{1, 2, 3, 4}, false},
+		{"1,2-3,4", []uint32{1, 2, 3, 4}, false},
+		{"1-2,3,4,10,11", []uint32{1, 2, 3, 4, 10, 11}, false},
+		{"1,2-4,5,16", []uint32{1, 2, 3, 4, 5, 16}, false},
+		{"abc", []uint32{}, true},
+		{"1,3-2,4", []uint32{}, true},
+		{"1,3-3,4", []uint32{}, true},
+		{"1,2,3\000,4", []uint32{1, 2, 3}, false},
+		{"1,2,3\n,4", []uint32{1, 2, 3}, false},
+	}
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("case-%d", i), func(t *testing.T) {
+			b, err := parseBitmap(tt.input, 64)
+			if tt.shouldFail {
+				if err == nil {
+					t.Fatalf("Expected parsing of %q to fail, but it didn't", tt.input)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("Failed to parse bitmap: %v", err)
+				return
+			}
+
+			got := b.ToSlice()
+			if !reflect.DeepEqual(got, tt.output) {
+				t.Errorf("Parsed bitmap doesn't match what we expected. Got %v, want %v", got, tt.output)
+			}
+
+		})
+	}
+}
diff --git a/pkg/sentry/fsimpl/cgroupfs/cgroupfs.go b/pkg/sentry/fsimpl/cgroupfs/cgroupfs.go
index edc3b50b9..e089b2c28 100644
--- a/pkg/sentry/fsimpl/cgroupfs/cgroupfs.go
+++ b/pkg/sentry/fsimpl/cgroupfs/cgroupfs.go
@@ -269,7 +269,7 @@ func (fsType FilesystemType) GetFilesystem(ctx context.Context, vfsObj *vfs.Virt
 		case controllerCPUAcct:
 			c = newCPUAcctController(fs)
 		case controllerCPUSet:
-			c = newCPUSetController(fs)
+			c = newCPUSetController(k, fs)
 		case controllerJob:
 			c = newJobController(fs)
 		case controllerMemory:
diff --git a/pkg/sentry/fsimpl/cgroupfs/cpuset.go b/pkg/sentry/fsimpl/cgroupfs/cpuset.go
index ac547f8e2..62e7029da 100644
--- a/pkg/sentry/fsimpl/cgroupfs/cpuset.go
+++ b/pkg/sentry/fsimpl/cgroupfs/cpuset.go
@@ -15,25 +15,133 @@
 package cgroupfs
 
 import (
+	"bytes"
+	"fmt"
+
+	"gvisor.dev/gvisor/pkg/bitmap"
 	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
+	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/kernfs"
+	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
+	"gvisor.dev/gvisor/pkg/usermem"
 )
 
 // +stateify savable
 type cpusetController struct {
 	controllerCommon
+
+	maxCpus uint32
+	maxMems uint32
+
+	cpus *bitmap.Bitmap
+	mems *bitmap.Bitmap
 }
 
 var _ controller = (*cpusetController)(nil)
 
-func newCPUSetController(fs *filesystem) *cpusetController {
-	c := &cpusetController{}
+func newCPUSetController(k *kernel.Kernel, fs *filesystem) *cpusetController {
+	cores := uint32(k.ApplicationCores())
+	cpus := bitmap.New(cores)
+	cpus.FlipRange(0, cores)
+	mems := bitmap.New(1)
+	mems.FlipRange(0, 1)
+	c := &cpusetController{
+		cpus:    &cpus,
+		mems:    &mems,
+		maxCpus: uint32(k.ApplicationCores()),
+		maxMems: 1, // We always report a single NUMA node.
+	}
 	c.controllerCommon.init(controllerCPUSet, fs)
 	return c
 }
 
 // AddControlFiles implements controller.AddControlFiles.
 func (c *cpusetController) AddControlFiles(ctx context.Context, creds *auth.Credentials, _ *cgroupInode, contents map[string]kernfs.Inode) {
-	// This controller is currently intentionally empty.
+	contents["cpuset.cpus"] = c.fs.newControllerWritableFile(ctx, creds, &cpusData{c: c})
+	contents["cpuset.mems"] = c.fs.newControllerWritableFile(ctx, creds, &memsData{c: c})
+}
+
+// +stateify savable
+type cpusData struct {
+	c *cpusetController
+}
+
+// Generate implements vfs.DynamicBytesSource.Generate.
+func (d *cpusData) Generate(ctx context.Context, buf *bytes.Buffer) error {
+	fmt.Fprintf(buf, "%s\n", formatBitmap(d.c.cpus))
+	return nil
+}
+
+// Write implements vfs.WritableDynamicBytesSource.Write.
+func (d *cpusData) Write(ctx context.Context, src usermem.IOSequence, offset int64) (int64, error) {
+	src = src.DropFirst64(offset)
+	if src.NumBytes() > hostarch.PageSize {
+		return 0, linuxerr.EINVAL
+	}
+
+	t := kernel.TaskFromContext(ctx)
+	buf := t.CopyScratchBuffer(hostarch.PageSize)
+	n, err := src.CopyIn(ctx, buf)
+	if err != nil {
+		return 0, err
+	}
+	buf = buf[:n]
+
+	b, err := parseBitmap(string(buf), d.c.maxCpus)
+	if err != nil {
+		log.Warningf("cgroupfs cpuset controller: Failed to parse bitmap: %v", err)
+		return 0, linuxerr.EINVAL
+	}
+
+	if got, want := b.Maximum(), d.c.maxCpus; got > want {
+		log.Warningf("cgroupfs cpuset controller: Attempted to specify cpuset.cpus beyond highest available cpu: got %d, want %d", got, want)
+		return 0, linuxerr.EINVAL
+	}
+
+	d.c.cpus = b
+	return int64(n), nil
+}
+
+// +stateify savable
+type memsData struct {
+	c *cpusetController
+}
+
+// Generate implements vfs.DynamicBytesSource.Generate.
+func (d *memsData) Generate(ctx context.Context, buf *bytes.Buffer) error {
+	fmt.Fprintf(buf, "%s\n", formatBitmap(d.c.mems))
+	return nil
+}
+
+// Write implements vfs.WritableDynamicBytesSource.Write.
+func (d *memsData) Write(ctx context.Context, src usermem.IOSequence, offset int64) (int64, error) {
+	src = src.DropFirst64(offset)
+	if src.NumBytes() > hostarch.PageSize {
+		return 0, linuxerr.EINVAL
+	}
+
+	t := kernel.TaskFromContext(ctx)
+	buf := t.CopyScratchBuffer(hostarch.PageSize)
+	n, err := src.CopyIn(ctx, buf)
+	if err != nil {
+		return 0, err
+	}
+	buf = buf[:n]
+
+	b, err := parseBitmap(string(buf), d.c.maxMems)
+	if err != nil {
+		log.Warningf("cgroupfs cpuset controller: Failed to parse bitmap: %v", err)
+		return 0, linuxerr.EINVAL
+	}
+
+	if got, want := b.Maximum(), d.c.maxMems; got > want {
+		log.Warningf("cgroupfs cpuset controller: Attempted to specify cpuset.mems beyond highest available node: got %d, want %d", got, want)
+		return 0, linuxerr.EINVAL
+	}
+
+	d.c.mems = b
+	return int64(n), nil
 }
diff --git a/test/syscalls/linux/cgroup.cc b/test/syscalls/linux/cgroup.cc
index ca23dfeee..50705a86e 100644
--- a/test/syscalls/linux/cgroup.cc
+++ b/test/syscalls/linux/cgroup.cc
@@ -37,6 +37,8 @@ namespace {
 
 using ::testing::_;
 using ::testing::Contains;
+using ::testing::Each;
+using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::Gt;
 using ::testing::Key;
@@ -383,6 +385,32 @@ PosixError WriteAndVerifyControlValue(const Cgroup& c, std::string_view path,
   return NoError();
 }
 
+PosixErrorOr<std::vector<bool>> ParseBitmap(std::string s) {
+  std::vector<bool> bitmap;
+  bitmap.reserve(64);
+  for (const std::string_view& t : absl::StrSplit(s, ',')) {
+    std::vector<std::string> parts = absl::StrSplit(t, absl::MaxSplits('-', 2));
+    if (parts.size() == 2) {
+      ASSIGN_OR_RETURN_ERRNO(int64_t start, Atoi<int64_t>(parts[0]));
+      ASSIGN_OR_RETURN_ERRNO(int64_t end, Atoi<int64_t>(parts[1]));
+      // Note: start and end are indices into bitmap.
+      if (end >= bitmap.size()) {
+        bitmap.resize(end + 1, false);
+      }
+      for (int i = start; i <= end; ++i) {
+        bitmap[i] = true;
+      }
+    } else {  // parts.size() == 1, 0 not possible.
+      ASSIGN_OR_RETURN_ERRNO(int64_t i, Atoi<int64_t>(parts[0]));
+      if (i >= bitmap.size()) {
+        bitmap.resize(i + 1, false);
+      }
+      bitmap[i] = true;
+    }
+  }
+  return bitmap;
+}
+
 TEST(JobCgroup, ReadWriteRead) {
   SKIP_IF(!CgroupsAvailable());
 
@@ -396,6 +424,59 @@ TEST(JobCgroup, ReadWriteRead) {
   EXPECT_NO_ERRNO(WriteAndVerifyControlValue(c, "job.id", LLONG_MAX));
 }
 
+TEST(CpusetCgroup, Defaults) {
+  SKIP_IF(!CgroupsAvailable());
+  Mounter m(ASSERT_NO_ERRNO_AND_VALUE(TempPath::CreateDir()));
+  Cgroup c = ASSERT_NO_ERRNO_AND_VALUE(m.MountCgroupfs("cpuset"));
+
+  std::string cpus =
+      ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.cpus"));
+  std::vector<bool> cpus_bitmap = ASSERT_NO_ERRNO_AND_VALUE(ParseBitmap(cpus));
+  EXPECT_GT(cpus_bitmap.size(), 0);
+  EXPECT_THAT(cpus_bitmap, Each(Eq(true)));
+
+  std::string mems =
+      ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.mems"));
+  std::vector<bool> mems_bitmap = ASSERT_NO_ERRNO_AND_VALUE(ParseBitmap(mems));
+  EXPECT_GT(mems_bitmap.size(), 0);
+  EXPECT_THAT(mems_bitmap, Each(Eq(true)));
+}
+
+TEST(CpusetCgroup, SetMask) {
+  SKIP_IF(!CgroupsAvailable());
+  Mounter m(ASSERT_NO_ERRNO_AND_VALUE(TempPath::CreateDir()));
+  Cgroup c = ASSERT_NO_ERRNO_AND_VALUE(m.MountCgroupfs("cpuset"));
+
+  std::string cpus =
+      ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.cpus"));
+  std::vector<bool> cpus_bitmap = ASSERT_NO_ERRNO_AND_VALUE(ParseBitmap(cpus));
+
+  SKIP_IF(cpus_bitmap.size() <= 1);  // "Not enough CPUs"
+
+  int max_cpu = cpus_bitmap.size() - 1;
+  ASSERT_NO_ERRNO(
+      c.WriteControlFile("cpuset.cpus", absl::StrCat("1-", max_cpu)));
+  cpus_bitmap[0] = false;
+  cpus = ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.cpus"));
+  std::vector<bool> cpus_bitmap_after =
+      ASSERT_NO_ERRNO_AND_VALUE(ParseBitmap(cpus));
+  EXPECT_EQ(cpus_bitmap_after, cpus_bitmap);
+}
+
+TEST(CpusetCgroup, SetEmptyMask) {
+  SKIP_IF(!CgroupsAvailable());
+  Mounter m(ASSERT_NO_ERRNO_AND_VALUE(TempPath::CreateDir()));
+  Cgroup c = ASSERT_NO_ERRNO_AND_VALUE(m.MountCgroupfs("cpuset"));
+  ASSERT_NO_ERRNO(c.WriteControlFile("cpuset.cpus", ""));
+  std::string_view cpus = absl::StripAsciiWhitespace(
+      ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.cpus")));
+  EXPECT_EQ(cpus, "");
+  ASSERT_NO_ERRNO(c.WriteControlFile("cpuset.mems", ""));
+  std::string_view mems = absl::StripAsciiWhitespace(
+      ASSERT_NO_ERRNO_AND_VALUE(c.ReadControlFile("cpuset.mems")));
+  EXPECT_EQ(mems, "");
+}
+
 TEST(ProcCgroups, Empty) {
   SKIP_IF(!CgroupsAvailable());
 
diff --git a/test/util/cgroup_util.cc b/test/util/cgroup_util.cc
index df3c57b87..0308c2153 100644
--- a/test/util/cgroup_util.cc
+++ b/test/util/cgroup_util.cc
@@ -69,6 +69,9 @@ PosixError Cgroup::WriteControlFile(absl::string_view name,
                                     const std::string& value) const {
   ASSIGN_OR_RETURN_ERRNO(FileDescriptor fd, Open(Relpath(name), O_WRONLY));
   RETURN_ERROR_IF_SYSCALL_FAIL(WriteFd(fd.get(), value.c_str(), value.size()));
+  const std::string alias_path = absl::StrFormat("[cg#%d]/%s", id_, name);
+  std::cerr << absl::StreamFormat("echo '%s' > %s", value, alias_path)
+            << std::endl;
   return NoError();
 }