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 <walthere.lee@gmail.com> --------- Signed-off-by: Walther Lee <walthere.lee@gmail.com>
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) {