Add function to apply cgroup settings

buildbuddy-io · Nov 18, 2024 · 38e22bf · 38e22bf
1 parent fb37b57
commit 38e22bf
Show file tree

Hide file tree

Showing 4 changed files with 175 additions and 0 deletions.
diff --git a/enterprise/server/remote_execution/cgroup/BUILD b/enterprise/server/remote_execution/cgroup/BUILD
@@ -12,6 +12,7 @@ go_library(
         "@io_bazel_rules_go//go/platform:linux": [
             "//enterprise/server/remote_execution/block_io",
             "//proto:remote_execution_go_proto",
+            "//proto:scheduler_go_proto",
             "//server/util/log",
             "//server/util/status",
         ],
@@ -25,7 +26,10 @@ go_test(
     embed = [":cgroup"],
     target_compatible_with = ["@platforms//os:linux"],
     deps = [
+        "//enterprise/server/remote_execution/block_io",
         "//proto:remote_execution_go_proto",
+        "//proto:scheduler_go_proto",
+        "//server/util/proto",
         "@com_github_google_go_cmp//cmp",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_protobuf//testing/protocmp",

diff --git a/enterprise/server/remote_execution/cgroup/cgroup.go b/enterprise/server/remote_execution/cgroup/cgroup.go
@@ -21,6 +21,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
 
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
+	scpb "github.com/buildbuddy-io/buildbuddy/proto/scheduler"
 )
 
 const (
@@ -31,6 +32,105 @@ const (
 	cidPlaceholder = "{{.ContainerID}}"
 )
 
+// min/max weight values for cgroup2 weight-based control
+const (
+	minWeight = 1
+	maxWeight = 10_000
+)
+
+// Setup configures the cgroup at the given path with the given settings.
+// IO limits are applied to the given block device, if specified.
+func Setup(ctx context.Context, path string, s *scpb.CgroupSettings, blockDevice *block_io.Device) error {
+	m, err := settingsMap(s, blockDevice)
+	if err != nil {
+		return err
+	}
+	for name, value := range m {
+		settingFilePath := filepath.Join(path, name)
+		if err := os.WriteFile(settingFilePath, []byte(value), 0); err != nil {
+			if os.IsNotExist(err) {
+				log.CtxWarningf(ctx, "Failed to set up cgroup file %q: %s", settingFilePath, err)
+			} else {
+				return fmt.Errorf("write %q: %w", settingFilePath, err)
+			}
+		}
+	}
+	return nil
+}
+
+func settingsMap(s *scpb.CgroupSettings, blockDevice *block_io.Device) (map[string]string, error) {
+	m := map[string]string{}
+	if s.CpuWeight != nil {
+		m["cpu.weight"] = fmt.Sprintf("%d", s.GetCpuWeight())
+	}
+	if s.CpuQuotaPeriodUsec != nil {
+		m["cpu.max"] = fmt.Sprintf("%d %s", s.GetCpuQuotaLimitUsec(), strconv.Itoa(int(s.GetCpuQuotaPeriodUsec())))
+	}
+	if s.CpuMaxBurstUsec != nil {
+		m["cpu.max.burst"] = fmt.Sprintf("%d", s.GetCpuMaxBurstUsec())
+	}
+	if s.CpuUclampMin != nil {
+		m["cpu.uclamp.min"] = fmtPercent(s.GetCpuUclampMin())
+	}
+	if s.CpuUclampMax != nil {
+		m["cpu.uclamp.max"] = fmtPercent(s.GetCpuUclampMax())
+	}
+	if s.MemoryThrottleLimitBytes != nil {
+		m["memory.high"] = strconv.Itoa(int(s.GetMemoryThrottleLimitBytes()))
+	}
+	if s.MemoryLimitBytes != nil {
+		m["memory.max"] = strconv.Itoa(int(s.GetMemoryLimitBytes()))
+	}
+	if s.MemorySoftGuaranteeBytes != nil {
+		m["memory.low"] = strconv.Itoa(int(s.GetMemorySoftGuaranteeBytes()))
+	}
+	if s.MemoryMinimumBytes != nil {
+		m["memory.min"] = strconv.Itoa(int(s.GetMemoryMinimumBytes()))
+	}
+	if s.SwapThrottleLimitBytes != nil {
+		m["memory.swap.high"] = strconv.Itoa(int(s.GetSwapThrottleLimitBytes()))
+	}
+	if s.SwapLimitBytes != nil {
+		m["memory.swap.max"] = strconv.Itoa(int(s.GetSwapLimitBytes()))
+	}
+	if blockDevice != nil {
+		if s.BlockIoLatencyTargetUsec != nil {
+			m["io.latency"] = fmt.Sprintf("%d:%d target=%d", blockDevice.Maj, blockDevice.Min, s.GetBlockIoLatencyTargetUsec())
+		}
+		if s.BlockIoWeight != nil {
+			m["io.weight"] = fmt.Sprintf("%d:%d %d", blockDevice.Maj, blockDevice.Min, clampWeight(s.GetBlockIoWeight()))
+		}
+		var limitFields []string
+		limits := s.GetBlockIoLimit()
+		if limits != nil {
+			if limits.Riops != nil {
+				limitFields = append(limitFields, fmt.Sprintf("riops=%d", limits.GetRiops()))
+			}
+			if limits.Wiops != nil {
+				limitFields = append(limitFields, fmt.Sprintf("wiops=%d", limits.GetWiops()))
+			}
+			if limits.Rbps != nil {
+				limitFields = append(limitFields, fmt.Sprintf("rbps=%d", limits.GetRbps()))
+			}
+			if limits.Wbps != nil {
+				limitFields = append(limitFields, fmt.Sprintf("wbps=%d", limits.GetWbps()))
+			}
+		}
+		if len(limitFields) > 0 {
+			m["io.max"] = fmt.Sprintf("%d:%d %s", blockDevice.Maj, blockDevice.Min, strings.Join(limitFields, " "))
+		}
+	}
+	return m, nil
+}
+
+func clampWeight(v int64) int64 {
+	return min(maxWeight, max(minWeight, v))
+}
+
+func fmtPercent(v float32) string {
+	return fmt.Sprintf("%.2f", v)
+}
+
 // Paths holds cgroup path templates that map a container ID to their cgroupfs
 // file paths.
 //

diff --git a/enterprise/server/remote_execution/cgroup/cgroup_test.go b/enterprise/server/remote_execution/cgroup/cgroup_test.go
@@ -4,13 +4,83 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/block_io"
+	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/testing/protocmp"
 
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
+	scpb "github.com/buildbuddy-io/buildbuddy/proto/scheduler"
 )
 
+func TestSettingsMap(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		settings *scpb.CgroupSettings
+		expected map[string]string
+	}{
+		{
+			name:     "nil",
+			settings: nil,
+			expected: map[string]string{},
+		},
+		{
+			name:     "all fields unset",
+			settings: nil,
+			expected: map[string]string{},
+		},
+		{
+			name: "all values set",
+			settings: &scpb.CgroupSettings{
+				CpuWeight:                proto.Int64(200),
+				CpuQuotaLimitUsec:        proto.Int64(400e3),
+				CpuQuotaPeriodUsec:       proto.Int64(100e3),
+				CpuMaxBurstUsec:          proto.Int64(50e3),
+				CpuUclampMin:             proto.Float32(12.34),
+				CpuUclampMax:             proto.Float32(98.76),
+				MemoryThrottleLimitBytes: proto.Int64(777e6),
+				MemoryLimitBytes:         proto.Int64(800e6),
+				MemorySoftGuaranteeBytes: proto.Int64(100e6),
+				MemoryMinimumBytes:       proto.Int64(50e6),
+				SwapThrottleLimitBytes:   proto.Int64(800e6),
+				SwapLimitBytes:           proto.Int64(1e9),
+				BlockIoLatencyTargetUsec: proto.Int64(100e3),
+				BlockIoWeight:            proto.Int64(300),
+				BlockIoLimit: &scpb.CgroupSettings_BlockIOLimits{
+					Riops: proto.Int64(1000),
+					Wiops: proto.Int64(500),
+					Rbps:  proto.Int64(4096e3),
+					Wbps:  proto.Int64(1024e3),
+				},
+			},
+			expected: map[string]string{
+				"cpu.weight":       "200",
+				"cpu.max":          "400000 100000",
+				"cpu.max.burst":    "50000",
+				"cpu.uclamp.min":   "12.34",
+				"cpu.uclamp.max":   "98.76",
+				"memory.high":      "777000000",
+				"memory.max":       "800000000",
+				"memory.low":       "100000000",
+				"memory.min":       "50000000",
+				"memory.swap.high": "800000000",
+				"memory.swap.max":  "1000000000",
+				"io.latency":       "279:8 target=100000",
+				"io.weight":        "279:8 300",
+				"io.max":           "279:8 riops=1000 wiops=500 rbps=4096000 wbps=1024000",
+			},
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			device := &block_io.Device{Maj: 279, Min: 8}
+			m, err := settingsMap(test.settings, device)
+			require.NoError(t, err)
+			require.Equal(t, test.expected, m)
+		})
+	}
+}
+
 func TestParsePSI(t *testing.T) {
 	r := strings.NewReader(`some avg10=0.00 avg60=1.00 avg300=4.11 total=123456
 full avg10=0.01 avg60=0.50 avg300=1.23 total=23456

diff --git a/server/util/proto/proto.go b/server/util/proto/proto.go
@@ -10,6 +10,7 @@ var Equal = gproto.Equal
 var MarshalOld = gproto.Marshal
 
 var String = gproto.String
+var Float32 = gproto.Float32
 var Float64 = gproto.Float64
 var Uint64 = gproto.Uint64
 var Int32 = gproto.Int32