Skip to content

Commit 8553d3b

Browse files
wallee94wallee94
authored
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]>
1 parent c9adbb9 commit 8553d3b

File tree

6 files changed

+105
-19
lines changed

6 files changed

+105
-19
lines changed

‎internal/runtime/cpu_quota_linux.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,18 @@ package runtime
2525

2626
import (
2727
"errors"
28-
"math"
2928

3029
cg "go.uber.org/automaxprocs/internal/cgroups"
3130
)
3231

3332
// CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
34-
// to a valid GOMAXPROCS value.
35-
func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
36-
cgroups, err := newQueryer()
33+
// to a valid GOMAXPROCS value. The quota is converted from float to int using round.
34+
// If round == nil, DefaultRoundFunc is used.
35+
func CPUQuotaToGOMAXPROCS(minValue int, round func(v float64) int) (int, CPUQuotaStatus, error) {
36+
if round == nil {
37+
round = DefaultRoundFunc
38+
}
39+
cgroups, err := _newQueryer()
3740
if err != nil {
3841
return -1, CPUQuotaUndefined, err
3942
}
@@ -43,7 +46,7 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
4346
return -1, CPUQuotaUndefined, err
4447
}
4548

46-
maxProcs := int(math.Floor(quota))
49+
maxProcs := round(quota)
4750
if minValue > 0 && maxProcs < minValue {
4851
return minValue, CPUQuotaMinUsed, nil
4952
}
@@ -57,6 +60,7 @@ type queryer interface {
5760
var (
5861
_newCgroups2 = cg.NewCGroups2ForCurrentProcess
5962
_newCgroups = cg.NewCGroupsForCurrentProcess
63+
_newQueryer = newQueryer
6064
)
6165

6266
func newQueryer() (queryer, error) {

‎internal/runtime/cpu_quota_linux_test.go

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ package runtime
2626
import (
2727
"errors"
2828
"fmt"
29+
"math"
2930
"testing"
3031

3132
"github.com/prashantv/gostub"
@@ -81,6 +82,48 @@ func TestNewQueryer(t *testing.T) {
8182
_, err := newQueryer()
8283
assert.ErrorIs(t, err, giveErr)
8384
})
85+
86+
t.Run("round quota with a nil round function", func(t *testing.T) {
87+
stubs := newStubs(t)
88+
89+
q := testQueryer{v: 2.7}
90+
stubs.StubFunc(&_newQueryer, q, nil)
91+
92+
// If round function is nil, CPUQuotaToGOMAXPROCS uses DefaultRoundFunc, which rounds down the value
93+
got, _, err := CPUQuotaToGOMAXPROCS(0, nil)
94+
require.NoError(t, err)
95+
assert.Equal(t, 2, got)
96+
})
97+
98+
t.Run("round quota with ceil", func(t *testing.T) {
99+
stubs := newStubs(t)
100+
101+
q := testQueryer{v: 2.7}
102+
stubs.StubFunc(&_newQueryer, q, nil)
103+
104+
got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Ceil(v)) })
105+
require.NoError(t, err)
106+
assert.Equal(t, 3, got)
107+
})
108+
109+
t.Run("round quota with floor", func(t *testing.T) {
110+
stubs := newStubs(t)
111+
112+
q := testQueryer{v: 2.7}
113+
stubs.StubFunc(&_newQueryer, q, nil)
114+
115+
got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Floor(v)) })
116+
require.NoError(t, err)
117+
assert.Equal(t, 2, got)
118+
})
119+
}
120+
121+
type testQueryer struct {
122+
v float64
123+
}
124+
125+
func (tq testQueryer) CPUQuota() (float64, bool, error) {
126+
return tq.v, true, nil
84127
}
85128

86129
func newStubs(t *testing.T) *gostub.Stubs {

‎internal/runtime/cpu_quota_unsupported.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,6 @@ package runtime
2626
// CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
2727
// to a valid GOMAXPROCS value. This is Linux-specific and not supported in the
2828
// current OS.
29-
func CPUQuotaToGOMAXPROCS(_ int) (int, CPUQuotaStatus, error) {
29+
func CPUQuotaToGOMAXPROCS(_ int, _ func(v float64) int) (int, CPUQuotaStatus, error) {
3030
return -1, CPUQuotaUndefined, nil
3131
}

‎internal/runtime/runtime.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020

2121
package runtime
2222

23+
import "math"
24+
2325
// CPUQuotaStatus presents the status of how CPU quota is used
2426
type CPUQuotaStatus int
2527

@@ -31,3 +33,8 @@ const (
3133
// CPUQuotaMinUsed is returned when CPU quota is smaller than the min value
3234
CPUQuotaMinUsed
3335
)
36+
37+
// DefaultRoundFunc is the default function to convert CPU quota from float to int. It rounds the value down (floor).
38+
func DefaultRoundFunc(v float64) int {
39+
return int(math.Floor(v))
40+
}

‎maxprocs/maxprocs.go

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,10 @@ func currentMaxProcs() int {
3737
}
3838

3939
type config struct {
40-
printf func(string, ...interface{})
41-
procs func(int) (int, iruntime.CPUQuotaStatus, error)
42-
minGOMAXPROCS int
40+
printf func(string, ...interface{})
41+
procs func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error)
42+
minGOMAXPROCS int
43+
roundQuotaFunc func(v float64) int
4344
}
4445

4546
func (c *config) log(fmt string, args ...interface{}) {
@@ -71,6 +72,13 @@ func Min(n int) Option {
7172
})
7273
}
7374

75+
// RoundQuotaFunc sets the function that will be used to covert the CPU quota from float to int.
76+
func RoundQuotaFunc(rf func(v float64) int) Option {
77+
return optionFunc(func(cfg *config) {
78+
cfg.roundQuotaFunc = rf
79+
})
80+
}
81+
7482
type optionFunc func(*config)
7583

7684
func (of optionFunc) apply(cfg *config) { of(cfg) }
@@ -82,8 +90,9 @@ func (of optionFunc) apply(cfg *config) { of(cfg) }
8290
// configured CPU quota.
8391
func Set(opts ...Option) (func(), error) {
8492
cfg := &config{
85-
procs: iruntime.CPUQuotaToGOMAXPROCS,
86-
minGOMAXPROCS: 1,
93+
procs: iruntime.CPUQuotaToGOMAXPROCS,
94+
roundQuotaFunc: iruntime.DefaultRoundFunc,
95+
minGOMAXPROCS: 1,
8796
}
8897
for _, o := range opts {
8998
o.apply(cfg)
@@ -102,7 +111,7 @@ func Set(opts ...Option) (func(), error) {
102111
return undoNoop, nil
103112
}
104113

105-
maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS)
114+
maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS, cfg.roundQuotaFunc)
106115
if err != nil {
107116
return undoNoop, err
108117
}

‎maxprocs/maxprocs_test.go

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"errors"
2626
"fmt"
2727
"log"
28+
"math"
2829
"os"
2930
"strconv"
3031
"testing"
@@ -55,7 +56,7 @@ func testLogger() (*bytes.Buffer, Option) {
5556
return buf, Logger(printf)
5657
}
5758

58-
func stubProcs(f func(int) (int, iruntime.CPUQuotaStatus, error)) Option {
59+
func stubProcs(f func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error)) Option {
5960
return optionFunc(func(cfg *config) {
6061
cfg.procs = f
6162
})
@@ -96,7 +97,7 @@ func TestSet(t *testing.T) {
9697
})
9798

9899
t.Run("ErrorReadingQuota", func(t *testing.T) {
99-
opt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
100+
opt := stubProcs(func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
100101
return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
101102
})
102103
prev := currentMaxProcs()
@@ -109,7 +110,7 @@ func TestSet(t *testing.T) {
109110

110111
t.Run("QuotaUndefined", func(t *testing.T) {
111112
buf, logOpt := testLogger()
112-
quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
113+
quotaOpt := stubProcs(func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
113114
return 0, iruntime.CPUQuotaUndefined, nil
114115
})
115116
prev := currentMaxProcs()
@@ -122,7 +123,7 @@ func TestSet(t *testing.T) {
122123

123124
t.Run("QuotaUndefined return maxProcs=7", func(t *testing.T) {
124125
buf, logOpt := testLogger()
125-
quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
126+
quotaOpt := stubProcs(func(int, func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
126127
return 7, iruntime.CPUQuotaUndefined, nil
127128
})
128129
prev := currentMaxProcs()
@@ -135,7 +136,7 @@ func TestSet(t *testing.T) {
135136

136137
t.Run("QuotaTooSmall", func(t *testing.T) {
137138
buf, logOpt := testLogger()
138-
quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
139+
quotaOpt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
139140
return min, iruntime.CPUQuotaMinUsed, nil
140141
})
141142
undo, err := Set(logOpt, quotaOpt, Min(5))
@@ -147,7 +148,7 @@ func TestSet(t *testing.T) {
147148

148149
t.Run("Min unused", func(t *testing.T) {
149150
buf, logOpt := testLogger()
150-
quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
151+
quotaOpt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
151152
return min, iruntime.CPUQuotaMinUsed, nil
152153
})
153154
// Min(-1) should be ignored.
@@ -159,7 +160,7 @@ func TestSet(t *testing.T) {
159160
})
160161

161162
t.Run("QuotaUsed", func(t *testing.T) {
162-
opt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
163+
opt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
163164
assert.Equal(t, 1, min, "Default minimum value should be 1")
164165
return 42, iruntime.CPUQuotaUsed, nil
165166
})
@@ -168,6 +169,28 @@ func TestSet(t *testing.T) {
168169
require.NoError(t, err, "Set failed")
169170
assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match quota")
170171
})
172+
173+
t.Run("RoundQuotaSetToCeil", func(t *testing.T) {
174+
opt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
175+
assert.Equal(t, round(2.4), 3, "round should be math.Ceil")
176+
return 43, iruntime.CPUQuotaUsed, nil
177+
})
178+
undo, err := Set(opt, RoundQuotaFunc(func(v float64) int { return int(math.Ceil(v)) }))
179+
defer undo()
180+
require.NoError(t, err, "Set failed")
181+
assert.Equal(t, 43, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
182+
})
183+
184+
t.Run("RoundQuotaSetToFloor", func(t *testing.T) {
185+
opt := stubProcs(func(min int, round func(v float64) int) (int, iruntime.CPUQuotaStatus, error) {
186+
assert.Equal(t, round(2.6), 2, "round should be math.Floor")
187+
return 42, iruntime.CPUQuotaUsed, nil
188+
})
189+
undo, err := Set(opt, RoundQuotaFunc(func(v float64) int { return int(math.Floor(v)) }))
190+
defer undo()
191+
require.NoError(t, err, "Set failed")
192+
assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
193+
})
171194
}
172195

173196
func TestMain(m *testing.M) {

0 commit comments

Comments
 (0)