diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go index 3b97475..f9057fd 100644 --- 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 index c896a19..213c165 100644 --- 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 index 6922554..e747015 100644 --- 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 index df6eacf..f8a2834 100644 --- 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 index 98176d6..e561fe6 100644 --- 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 index a2aa7cb..310764c 100644 --- 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) {