Skip to content

Commit

Permalink
Add option to round up CPU quota (#79)
Browse files Browse the repository at this point in the history 
* 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>
  • Loading branch information
@wallee94
wallee94 authored Feb 13, 2024
1 parent c9adbb9 commit 8553d3b
Show file tree
Hide file tree
Showing 6 changed files with 105 additions and 19 deletions.
14 changes: 9 additions & 5 deletions internal/runtime/cpu_quota_linux.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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
}
Expand All @@ -57,6 +60,7 @@ type queryer interface {
var (
_newCgroups2 = cg.NewCGroups2ForCurrentProcess
_newCgroups = cg.NewCGroupsForCurrentProcess
_newQueryer = newQueryer
)

func newQueryer() (queryer, error) {
Expand Down
43 changes: 43 additions & 0 deletions internal/runtime/cpu_quota_linux_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ package runtime
import (
"errors"
"fmt"
"math"
"testing"

"github.com/prashantv/gostub"
Expand Down Expand Up @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion internal/runtime/cpu_quota_unsupported.go
View file
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, _ func(v float64) int) (int, CPUQuotaStatus, error) {
return -1, CPUQuotaUndefined, nil
}
7 changes: 7 additions & 0 deletions internal/runtime/runtime.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

package runtime

import "math"

// CPUQuotaStatus presents the status of how CPU quota is used
type CPUQuotaStatus int

Expand All @@ -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))
}
21 changes: 15 additions & 6 deletions maxprocs/maxprocs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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{}) {
Expand Down Expand Up @@ -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) }
Expand All @@ -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)
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.roundQuotaFunc)
if err != nil {
return undoNoop, err
}
Expand Down
37 changes: 30 additions & 7 deletions maxprocs/maxprocs_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import (
"errors"
"fmt"
"log"
"math"
"os"
"strconv"
"testing"
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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))
Expand All @@ -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.
Expand All @@ -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
})
Expand All @@ -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) {
Expand Down

0 comments on commit 8553d3b

Please sign in to comment.