Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add option to round up CPU quota #79

Merged
merged 6 commits into from
Feb 13, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions internal/runtime/cpu_quota_linux.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,9 @@ import (
)

// CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
// to a valid GOMAXPROCS value.
func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
// to a valid GOMAXPROCS value. The quota is rounded to an int using roundOpt.
// The default is rounding down (Floor).
func CPUQuotaToGOMAXPROCS(minValue int, roundOpt Rounding) (int, CPUQuotaStatus, error) {
cgroups, err := newQueryer()
if err != nil {
return -1, CPUQuotaUndefined, err
Expand All @@ -44,6 +45,9 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
}

maxProcs := int(math.Floor(quota))
if roundOpt == Ceil {
maxProcs = int(math.Ceil(quota))
}
wallee94 marked this conversation as resolved.
Show resolved Hide resolved
if minValue > 0 && maxProcs < minValue {
return minValue, CPUQuotaMinUsed, nil
}
Expand Down
30 changes: 30 additions & 0 deletions internal/runtime/cpu_quota_linux_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,36 @@ func TestNewQueryer(t *testing.T) {
_, err := newQueryer()
assert.ErrorIs(t, err, giveErr)
})

t.Run("round quota with ceil", func(t *testing.T) {
stubs := newStubs(t)

q := testQueryer{v: 2.7}
stubs.StubFunc(&_newCgroups2, q, nil)

got, _, err := CPUQuotaToGOMAXPROCS(0, Ceil)
require.NoError(t, err)
assert.Same(t, 3, got)
})

t.Run("round quota with floor", func(t *testing.T) {
stubs := newStubs(t)

q := testQueryer{v: 2.7}
stubs.StubFunc(&_newCgroups2, q, nil)

got, _, err := CPUQuotaToGOMAXPROCS(0, Floor)
require.NoError(t, err)
assert.Same(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 {
Expand Down
2 changes: 1 addition & 1 deletion internal/runtime/cpu_quota_unsupported.go
Original file line number Diff line number Diff line change
Expand Up @@ -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, _ Rounding) (int, CPUQuotaStatus, error) {
return -1, CPUQuotaUndefined, nil
}
10 changes: 10 additions & 0 deletions internal/runtime/runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,13 @@ const (
// CPUQuotaMinUsed is returned when CPU quota is smaller than the min value
CPUQuotaMinUsed
)

// Rounding controls how the CPU quota value should be rounded to an int
wallee94 marked this conversation as resolved.
Show resolved Hide resolved
type Rounding int

const (
// Ceil is used to return a CPU quota rounded up
Ceil Rounding = iota
// Floor is used to return a CPU quota rounded down
wallee94 marked this conversation as resolved.
Show resolved Hide resolved
Floor
)
13 changes: 11 additions & 2 deletions maxprocs/maxprocs.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,9 @@ func currentMaxProcs() int {

type config struct {
printf func(string, ...interface{})
procs func(int) (int, iruntime.CPUQuotaStatus, error)
procs func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error)
minGOMAXPROCS int
roundQuota iruntime.Rounding
}

func (c *config) log(fmt string, args ...interface{}) {
Expand Down Expand Up @@ -71,6 +72,13 @@ func Min(n int) Option {
})
}

// RoundQuota controls whether the CPU quota should be rounded using ceil or floor.
func RoundQuota(v iruntime.Rounding) Option {
wallee94 marked this conversation as resolved.
Show resolved Hide resolved
return optionFunc(func(cfg *config) {
cfg.roundQuota = v
})
}

type optionFunc func(*config)

func (of optionFunc) apply(cfg *config) { of(cfg) }
Expand All @@ -84,6 +92,7 @@ func Set(opts ...Option) (func(), error) {
cfg := &config{
procs: iruntime.CPUQuotaToGOMAXPROCS,
minGOMAXPROCS: 1,
roundQuota: iruntime.Floor,
}
for _, o := range opts {
o.apply(cfg)
Expand All @@ -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.roundQuota)
if err != nil {
return undoNoop, err
}
Expand Down
36 changes: 29 additions & 7 deletions maxprocs/maxprocs_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error)) Option {
return optionFunc(func(cfg *config) {
cfg.procs = f
})
Expand Down Expand Up @@ -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, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
})
prev := currentMaxProcs()
Expand All @@ -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, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
return 0, iruntime.CPUQuotaUndefined, nil
})
prev := currentMaxProcs()
Expand All @@ -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, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
return 7, iruntime.CPUQuotaUndefined, nil
})
prev := currentMaxProcs()
Expand All @@ -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, round iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
return min, iruntime.CPUQuotaMinUsed, nil
})
undo, err := Set(logOpt, quotaOpt, Min(5))
Expand All @@ -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, round iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
return min, iruntime.CPUQuotaMinUsed, nil
})
// Min(-1) should be ignored.
Expand All @@ -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, round iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
assert.Equal(t, 1, min, "Default minimum value should be 1")
return 42, iruntime.CPUQuotaUsed, nil
})
Expand All @@ -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("RoundQuotaSetToCeil", func(t *testing.T) {
opt := stubProcs(func(min int, round iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
assert.Equal(t, iruntime.Ceil, round, "round should be Ceil")
return 43, iruntime.CPUQuotaUsed, nil
})
undo, err := Set(opt, RoundQuota(iruntime.Ceil))
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 iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
assert.Equal(t, iruntime.Floor, round, "round should be Floor")
return 42, iruntime.CPUQuotaUsed, nil
})
undo, err := Set(opt, RoundQuota(iruntime.Floor))
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) {
Expand Down
Loading