add option to round up CPU quota

uber-go · Oct 12, 2023 · f5ff8d2 · f5ff8d2
1 parent e83e959
commit f5ff8d2
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 11 deletions.
diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go
@@ -32,7 +32,7 @@ import (
 
 // CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
 // to a valid GOMAXPROCS value.
-func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
+func CPUQuotaToGOMAXPROCS(minValue int, roundUpQuota bool) (int, CPUQuotaStatus, error) {
 	cgroups, err := newQueryer()
 	if err != nil {
 		return -1, CPUQuotaUndefined, err
@@ -44,6 +44,9 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
 	}
 
 	maxProcs := int(math.Floor(quota))
+	if roundUpQuota {
+		maxProcs = int(math.Ceil(quota))
+	}
 	if minValue > 0 && maxProcs < minValue {
 		return minValue, CPUQuotaMinUsed, nil
 	}

diff --git a/internal/runtime/cpu_quota_unsupported.go b/internal/runtime/cpu_quota_unsupported.go
@@ -26,6 +26,6 @@ package runtime
 // CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
 // to a valid GOMAXPROCS value. This is Linux-specific and not supported in the
 // current OS.
-func CPUQuotaToGOMAXPROCS(_ int) (int, CPUQuotaStatus, error) {
+func CPUQuotaToGOMAXPROCS(_ int, _ bool) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }
diff --git a/maxprocs/maxprocs.go b/maxprocs/maxprocs.go
@@ -38,8 +38,9 @@ func currentMaxProcs() int {
 
 type config struct {
 	printf        func(string, ...interface{})
-	procs         func(int) (int, iruntime.CPUQuotaStatus, error)
+	procs         func(int, bool) (int, iruntime.CPUQuotaStatus, error)
 	minGOMAXPROCS int
+	roundUpQuota  bool
 }
 
 func (c *config) log(fmt string, args ...interface{}) {
@@ -71,6 +72,13 @@ func Min(n int) Option {
 	})
 }
 
+// RoundUpQuota controls whether the CPU quota should be rounded up instead of down.
+func RoundUpQuota(v bool) Option {
+	return optionFunc(func(cfg *config) {
+		cfg.roundUpQuota = v
+	})
+}
+
 type optionFunc func(*config)
 
 func (of optionFunc) apply(cfg *config) { of(cfg) }
@@ -102,7 +110,7 @@ func Set(opts ...Option) (func(), error) {
 		return undoNoop, nil
 	}
 
-	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS)
+	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS, cfg.roundUpQuota)
 	if err != nil {
 		return undoNoop, err
 	}

diff --git a/maxprocs/maxprocs_test.go b/maxprocs/maxprocs_test.go
@@ -55,7 +55,7 @@ func testLogger() (*bytes.Buffer, Option) {
 	return buf, Logger(printf)
 }
 
-func stubProcs(f func(int) (int, iruntime.CPUQuotaStatus, error)) Option {
+func stubProcs(f func(int, bool) (int, iruntime.CPUQuotaStatus, error)) Option {
 	return optionFunc(func(cfg *config) {
 		cfg.procs = f
 	})
@@ -96,7 +96,7 @@ func TestSet(t *testing.T) {
 	})
 
 	t.Run("ErrorReadingQuota", func(t *testing.T) {
-		opt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		opt := stubProcs(func(int, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
 		})
 		prev := currentMaxProcs()
@@ -109,7 +109,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("QuotaUndefined", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(int, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -122,7 +122,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("QuotaUndefined return maxProcs=7", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(int, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 7, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -135,7 +135,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("QuotaTooSmall", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		undo, err := Set(logOpt, quotaOpt, Min(5))
@@ -147,7 +147,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("Min unused", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		// Min(-1) should be ignored.
@@ -159,7 +159,7 @@ func TestSet(t *testing.T) {
 	})
 
 	t.Run("QuotaUsed", func(t *testing.T) {
-		opt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
+		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			assert.Equal(t, 1, min, "Default minimum value should be 1")
 			return 42, iruntime.CPUQuotaUsed, nil
 		})
@@ -168,6 +168,28 @@ func TestSet(t *testing.T) {
 		require.NoError(t, err, "Set failed")
 		assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match quota")
 	})
+
+	t.Run("RoundUpQuotaSetToTrue", func(t *testing.T) {
+		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, true, roundUpQuota, "roundUpQuota should be true")
+			return 43, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundUpQuota(true))
+		defer undo()
+		require.NoError(t, err, "Set failed")
+		assert.Equal(t, 43, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
+	})
+
+	t.Run("RoundUpQuotaSetToFalse", func(t *testing.T) {
+		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, false, roundUpQuota, "roundUpQuota should be false")
+			return 42, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundUpQuota(false))
+		defer undo()
+		require.NoError(t, err, "Set failed")
+		assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
+	})
 }
 
 func TestMain(m *testing.M) {