Add option to round up CPU quota (#79)

* add option to round up CPU quota * add Rounding type to control rounding opt * add tests for ceil and floor rounding opts * set config.roundQuota init value to Floor * update CPUQuotaToGOMAXPROCS to pass a round function as arg * add test for rounding quota with a nil round function Signed-off-by: Walther Lee <[email protected]> --------- Signed-off-by: Walther Lee <[email protected]>
uber-go · Feb 13, 2024 · 8553d3b · 8553d3b
1 parent c9adbb9
commit 8553d3b
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 19 deletions.
diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go
@@ -25,15 +25,18 @@ package runtime
 
 import (
 	"errors"
-	"math"
 
 	cg "go.uber.org/automaxprocs/internal/cgroups"
 )
 
 // CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
-// to a valid GOMAXPROCS value.
-func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
-	cgroups, err := newQueryer()
+// to a valid GOMAXPROCS value. The quota is converted from float to int using round.
+// If round == nil, DefaultRoundFunc is used.
+func CPUQuotaToGOMAXPROCS(minValue int, round func(v float64) int) (int, CPUQuotaStatus, error) {
+	if round == nil {
+		round = DefaultRoundFunc
+	}
+	cgroups, err := _newQueryer()
 	if err != nil {
 		return -1, CPUQuotaUndefined, err
 	}
@@ -43,7 +46,7 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
 		return -1, CPUQuotaUndefined, err
 	}
 
-	maxProcs := int(math.Floor(quota))
+	maxProcs := round(quota)
 	if minValue > 0 && maxProcs < minValue {
 		return minValue, CPUQuotaMinUsed, nil
 	}
@@ -57,6 +60,7 @@ type queryer interface {
 var (
 	_newCgroups2 = cg.NewCGroups2ForCurrentProcess
 	_newCgroups  = cg.NewCGroupsForCurrentProcess
+	_newQueryer  = newQueryer
 )
 
 func newQueryer() (queryer, error) {

diff --git a/internal/runtime/cpu_quota_linux_test.go b/internal/runtime/cpu_quota_linux_test.go
@@ -26,6 +26,7 @@ package runtime
 import (
 	"errors"
 	"fmt"
+	"math"
 	"testing"
 
 	"github.com/prashantv/gostub"
@@ -81,6 +82,48 @@ func TestNewQueryer(t *testing.T) {
 		_, err := newQueryer()
 		assert.ErrorIs(t, err, giveErr)
 	})
+
+	t.Run("round quota with a nil round function", func(t *testing.T) {
+		stubs := newStubs(t)
+
+		q := testQueryer{v: 2.7}
+		stubs.StubFunc(&_newQueryer, q, nil)
+
+		// If round function is nil, CPUQuotaToGOMAXPROCS uses DefaultRoundFunc, which rounds down the value
+		got, _, err := CPUQuotaToGOMAXPROCS(0, nil)
+		require.NoError(t, err)
+		assert.Equal(t, 2, got)
+	})
+
+	t.Run("round quota with ceil", func(t *testing.T) {
+		stubs := newStubs(t)
+
+		q := testQueryer{v: 2.7}
+		stubs.StubFunc(&_newQueryer, q, nil)
+
+		got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Ceil(v)) })
+		require.NoError(t, err)
+		assert.Equal(t, 3, got)
+	})
+
+	t.Run("round quota with floor", func(t *testing.T) {
+		stubs := newStubs(t)
+
+		q := testQueryer{v: 2.7}
+		stubs.StubFunc(&_newQueryer, q, nil)
+
+		got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Floor(v)) })
+		require.NoError(t, err)
+		assert.Equal(t, 2, got)
+	})
+}
+
+type testQueryer struct {
+	v float64
+}
+
+func (tq testQueryer) CPUQuota() (float64, bool, error) {
+	return tq.v, true, nil
 }
 
 func newStubs(t *testing.T) *gostub.Stubs {

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, _ func(v float64) int) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }
diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go
@@ -20,6 +20,8 @@
 
 package runtime
 
+import "math"
+
 // CPUQuotaStatus presents the status of how CPU quota is used
 type CPUQuotaStatus int
 
@@ -31,3 +33,8 @@ const (
 	// CPUQuotaMinUsed is returned when CPU quota is smaller than the min value
 	CPUQuotaMinUsed
 )
+
+// DefaultRoundFunc is the default function to convert CPU quota from float to int. It rounds the value down (floor).
+func DefaultRoundFunc(v float64) int {
+	return int(math.Floor(v))
+}
diff --git a/maxprocs/maxprocs.go b/maxprocs/maxprocs.go
@@ -37,9 +37,10 @@ func currentMaxProcs() int {
 }
 
 type config struct {
-	printf        func(string, ...interface{})
-	procs         func(int) (int, iruntime.CPUQuotaStatus, error)
-	minGOMAXPROCS int
+	printf         func(string, ...interface{})
+	procs          func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error)
+	minGOMAXPROCS  int
+	roundQuotaFunc func(v float64) int
 }
 
 func (c *config) log(fmt string, args ...interface{}) {
@@ -71,6 +72,13 @@ func Min(n int) Option {
 	})
 }
 
+// RoundQuotaFunc sets the function that will be used to covert the CPU quota from float to int.
+func RoundQuotaFunc(rf func(v float64) int) Option {
+	return optionFunc(func(cfg *config) {
+		cfg.roundQuotaFunc = rf
+	})
+}
+
 type optionFunc func(*config)
 
 func (of optionFunc) apply(cfg *config) { of(cfg) }
@@ -82,8 +90,9 @@ func (of optionFunc) apply(cfg *config) { of(cfg) }
 // configured CPU quota.
 func Set(opts ...Option) (func(), error) {
 	cfg := &config{
-		procs:         iruntime.CPUQuotaToGOMAXPROCS,
-		minGOMAXPROCS: 1,
+		procs:          iruntime.CPUQuotaToGOMAXPROCS,
+		roundQuotaFunc: iruntime.DefaultRoundFunc,
+		minGOMAXPROCS:  1,
 	}
 	for _, o := range opts {
 		o.apply(cfg)
@@ -102,7 +111,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.roundQuotaFunc)
 	if err != nil {
 		return undoNoop, err
 	}

diff --git a/maxprocs/maxprocs_test.go b/maxprocs/maxprocs_test.go
@@ -25,6 +25,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"math"
 	"os"
 	"strconv"
 	"testing"
@@ -55,7 +56,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, func(v float64) int) (int, iruntime.CPUQuotaStatus, error)) Option {
 	return optionFunc(func(cfg *config) {
 		cfg.procs = f
 	})
@@ -96,7 +97,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, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
 		})
 		prev := currentMaxProcs()
@@ -109,7 +110,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, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -122,7 +123,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, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			return 7, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -135,7 +136,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, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		undo, err := Set(logOpt, quotaOpt, Min(5))
@@ -147,7 +148,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, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		// Min(-1) should be ignored.
@@ -159,7 +160,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, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
 			assert.Equal(t, 1, min, "Default minimum value should be 1")
 			return 42, iruntime.CPUQuotaUsed, nil
 		})
@@ -168,6 +169,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("RoundQuotaSetToCeil", func(t *testing.T) {
+		opt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, round(2.4), 3, "round should be math.Ceil")
+			return 43, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundQuotaFunc(func(v float64) int { return int(math.Ceil(v)) }))
+		defer undo()
+		require.NoError(t, err, "Set failed")
+		assert.Equal(t, 43, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
+	})
+
+	t.Run("RoundQuotaSetToFloor", func(t *testing.T) {
+		opt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, round(2.6), 2, "round should be math.Floor")
+			return 42, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundQuotaFunc(func(v float64) int { return int(math.Floor(v)) }))
+		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) {