From 9e565c7c6e830119146285b971a5fedfff655196 Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Thu, 12 Oct 2023 11:08:31 -0700
Subject: [PATCH 1/6] add option to round up CPU quota

---
 internal/runtime/cpu_quota_linux.go       |  5 +++-
 internal/runtime/cpu_quota_unsupported.go |  2 +-
 maxprocs/maxprocs.go                      | 12 ++++++--
 maxprocs/maxprocs_test.go                 | 36 ++++++++++++++++++-----
 4 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go
index 3b97475..4dd54cb 100644
--- a/internal/runtime/cpu_quota_linux.go
+++ b/internal/runtime/cpu_quota_linux.go
@@ -32,7 +32,7 @@ import (
 
 // CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
 // to a valid GOMAXPROCS value.
-func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
+func CPUQuotaToGOMAXPROCS(minValue int, roundUpQuota bool) (int, CPUQuotaStatus, error) {
 	cgroups, err := newQueryer()
 	if err != nil {
 		return -1, CPUQuotaUndefined, err
@@ -44,6 +44,9 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
 	}
 
 	maxProcs := int(math.Floor(quota))
+	if roundUpQuota {
+		maxProcs = int(math.Ceil(quota))
+	}
 	if minValue > 0 && maxProcs < minValue {
 		return minValue, CPUQuotaMinUsed, nil
 	}
diff --git a/internal/runtime/cpu_quota_unsupported.go b/internal/runtime/cpu_quota_unsupported.go
index 6922554..2984603 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, _ bool) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }
diff --git a/maxprocs/maxprocs.go b/maxprocs/maxprocs.go
index 98176d6..e7dff7f 100644
--- a/maxprocs/maxprocs.go
+++ b/maxprocs/maxprocs.go
@@ -38,8 +38,9 @@ func currentMaxProcs() int {
 
 type config struct {
 	printf        func(string, ...interface{})
-	procs         func(int) (int, iruntime.CPUQuotaStatus, error)
+	procs         func(int, bool) (int, iruntime.CPUQuotaStatus, error)
 	minGOMAXPROCS int
+	roundUpQuota  bool
 }
 
 func (c *config) log(fmt string, args ...interface{}) {
@@ -71,6 +72,13 @@ func Min(n int) Option {
 	})
 }
 
+// RoundUpQuota controls whether the CPU quota should be rounded up instead of down.
+func RoundUpQuota(v bool) Option {
+	return optionFunc(func(cfg *config) {
+		cfg.roundUpQuota = v
+	})
+}
+
 type optionFunc func(*config)
 
 func (of optionFunc) apply(cfg *config) { of(cfg) }
@@ -102,7 +110,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.roundUpQuota)
 	if err != nil {
 		return undoNoop, err
 	}
diff --git a/maxprocs/maxprocs_test.go b/maxprocs/maxprocs_test.go
index a2aa7cb..914a301 100644
--- a/maxprocs/maxprocs_test.go
+++ b/maxprocs/maxprocs_test.go
@@ -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, bool) (int, iruntime.CPUQuotaStatus, error)) Option {
 	return optionFunc(func(cfg *config) {
 		cfg.procs = f
 	})
@@ -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, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
 		})
 		prev := currentMaxProcs()
@@ -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, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -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, bool) (int, iruntime.CPUQuotaStatus, error) {
 			return 7, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -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, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		undo, err := Set(logOpt, quotaOpt, Min(5))
@@ -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, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			return min, iruntime.CPUQuotaMinUsed, nil
 		})
 		// Min(-1) should be ignored.
@@ -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, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
 			assert.Equal(t, 1, min, "Default minimum value should be 1")
 			return 42, iruntime.CPUQuotaUsed, nil
 		})
@@ -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("RoundUpQuotaSetToTrue", func(t *testing.T) {
+		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, true, roundUpQuota, "roundUpQuota should be true")
+			return 43, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundUpQuota(true))
+		defer undo()
+		require.NoError(t, err, "Set failed")
+		assert.Equal(t, 43, currentMaxProcs(), "should change GOMAXPROCS to match rounded up quota")
+	})
+
+	t.Run("RoundUpQuotaSetToFalse", func(t *testing.T) {
+		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, false, roundUpQuota, "roundUpQuota should be false")
+			return 42, iruntime.CPUQuotaUsed, nil
+		})
+		undo, err := Set(opt, RoundUpQuota(false))
+		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) {

From c022213085f344dbfb6cc1c3b6eaf67b5c6e3422 Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Mon, 20 Nov 2023 19:39:16 -0800
Subject: [PATCH 2/6] add Rounding type to control rounding opt

---
 internal/runtime/cpu_quota_linux.go       |  7 +++---
 internal/runtime/cpu_quota_unsupported.go |  2 +-
 internal/runtime/runtime.go               | 10 ++++++++
 maxprocs/maxprocs.go                      | 12 ++++-----
 maxprocs/maxprocs_test.go                 | 30 +++++++++++------------
 5 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go
index 4dd54cb..fec01c3 100644
--- a/internal/runtime/cpu_quota_linux.go
+++ b/internal/runtime/cpu_quota_linux.go
@@ -31,8 +31,9 @@ import (
 )
 
 // CPUQuotaToGOMAXPROCS converts the CPU quota applied to the calling process
-// to a valid GOMAXPROCS value.
-func CPUQuotaToGOMAXPROCS(minValue int, roundUpQuota bool) (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
@@ -44,7 +45,7 @@ func CPUQuotaToGOMAXPROCS(minValue int, roundUpQuota bool) (int, CPUQuotaStatus,
 	}
 
 	maxProcs := int(math.Floor(quota))
-	if roundUpQuota {
+	if roundOpt == Ceil {
 		maxProcs = int(math.Ceil(quota))
 	}
 	if minValue > 0 && maxProcs < minValue {
diff --git a/internal/runtime/cpu_quota_unsupported.go b/internal/runtime/cpu_quota_unsupported.go
index 2984603..bfff88b 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, _ bool) (int, CPUQuotaStatus, error) {
+func CPUQuotaToGOMAXPROCS(_ int, _ Rounding) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }
diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go
index df6eacf..b27526d 100644
--- a/internal/runtime/runtime.go
+++ b/internal/runtime/runtime.go
@@ -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
+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
+	Floor
+)
diff --git a/maxprocs/maxprocs.go b/maxprocs/maxprocs.go
index e7dff7f..a5d63a0 100644
--- a/maxprocs/maxprocs.go
+++ b/maxprocs/maxprocs.go
@@ -38,9 +38,9 @@ func currentMaxProcs() int {
 
 type config struct {
 	printf        func(string, ...interface{})
-	procs         func(int, bool) (int, iruntime.CPUQuotaStatus, error)
+	procs         func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error)
 	minGOMAXPROCS int
-	roundUpQuota  bool
+	roundQuota    iruntime.Rounding
 }
 
 func (c *config) log(fmt string, args ...interface{}) {
@@ -72,10 +72,10 @@ func Min(n int) Option {
 	})
 }
 
-// RoundUpQuota controls whether the CPU quota should be rounded up instead of down.
-func RoundUpQuota(v bool) Option {
+// RoundQuota controls whether the CPU quota should be rounded using ceil or floor.
+func RoundQuota(v iruntime.Rounding) Option {
 	return optionFunc(func(cfg *config) {
-		cfg.roundUpQuota = v
+		cfg.roundQuota = v
 	})
 }
 
@@ -110,7 +110,7 @@ func Set(opts ...Option) (func(), error) {
 		return undoNoop, nil
 	}
 
-	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS, cfg.roundUpQuota)
+	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS, cfg.roundQuota)
 	if err != nil {
 		return undoNoop, err
 	}
diff --git a/maxprocs/maxprocs_test.go b/maxprocs/maxprocs_test.go
index 914a301..d988fc6 100644
--- a/maxprocs/maxprocs_test.go
+++ b/maxprocs/maxprocs_test.go
@@ -55,7 +55,7 @@ func testLogger() (*bytes.Buffer, Option) {
 	return buf, Logger(printf)
 }
 
-func stubProcs(f func(int, bool) (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
 	})
@@ -96,7 +96,7 @@ func TestSet(t *testing.T) {
 	})
 
 	t.Run("ErrorReadingQuota", func(t *testing.T) {
-		opt := stubProcs(func(int, bool) (int, iruntime.CPUQuotaStatus, error) {
+		opt := stubProcs(func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
 		})
 		prev := currentMaxProcs()
@@ -109,7 +109,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("QuotaUndefined", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(int, bool) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -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, bool) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubProcs(func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
 			return 7, iruntime.CPUQuotaUndefined, nil
 		})
 		prev := currentMaxProcs()
@@ -135,7 +135,7 @@ func TestSet(t *testing.T) {
 
 	t.Run("QuotaTooSmall", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(min int, roundUpQuota bool) (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))
@@ -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, roundUpQuota bool) (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.
@@ -159,7 +159,7 @@ func TestSet(t *testing.T) {
 	})
 
 	t.Run("QuotaUsed", func(t *testing.T) {
-		opt := stubProcs(func(min int, roundUpQuota bool) (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
 		})
@@ -169,23 +169,23 @@ func TestSet(t *testing.T) {
 		assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match quota")
 	})
 
-	t.Run("RoundUpQuotaSetToTrue", func(t *testing.T) {
-		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
-			assert.Equal(t, true, roundUpQuota, "roundUpQuota should be true")
+	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, RoundUpQuota(true))
+		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("RoundUpQuotaSetToFalse", func(t *testing.T) {
-		opt := stubProcs(func(min int, roundUpQuota bool) (int, iruntime.CPUQuotaStatus, error) {
-			assert.Equal(t, false, roundUpQuota, "roundUpQuota should be false")
+	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, RoundUpQuota(false))
+		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")

From 57c88e3d496259dc5136816f87bcf8ebc109a6a9 Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Mon, 20 Nov 2023 19:39:31 -0800
Subject: [PATCH 3/6] add tests for ceil and floor rounding opts

---
 internal/runtime/cpu_quota_linux_test.go | 30 ++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/internal/runtime/cpu_quota_linux_test.go b/internal/runtime/cpu_quota_linux_test.go
index c896a19..8050be2 100644
--- a/internal/runtime/cpu_quota_linux_test.go
+++ b/internal/runtime/cpu_quota_linux_test.go
@@ -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 {

From b80bb9d8a909af2027f103446b3e3b87b089d4c5 Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Mon, 20 Nov 2023 19:44:33 -0800
Subject: [PATCH 4/6] set config.roundQuota init value to Floor

---
 maxprocs/maxprocs.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/maxprocs/maxprocs.go b/maxprocs/maxprocs.go
index a5d63a0..254b89f 100644
--- a/maxprocs/maxprocs.go
+++ b/maxprocs/maxprocs.go
@@ -92,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)

From 5a55c64c8c0633620c6a4973c4d55f1a0c4e9bbf Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Wed, 6 Dec 2023 19:37:26 -0800
Subject: [PATCH 5/6] update CPUQuotaToGOMAXPROCS to pass a round function as
 arg

---
 internal/runtime/cpu_quota_linux.go       | 18 +++++++--------
 internal/runtime/cpu_quota_linux_test.go  | 13 ++++++-----
 internal/runtime/cpu_quota_unsupported.go |  2 +-
 internal/runtime/runtime.go               | 15 +++++--------
 maxprocs/maxprocs.go                      | 22 +++++++++---------
 maxprocs/maxprocs_test.go                 | 27 ++++++++++++-----------
 6 files changed, 48 insertions(+), 49 deletions(-)

diff --git a/internal/runtime/cpu_quota_linux.go b/internal/runtime/cpu_quota_linux.go
index fec01c3..f9057fd 100644
--- a/internal/runtime/cpu_quota_linux.go
+++ b/internal/runtime/cpu_quota_linux.go
@@ -25,16 +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. 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()
+// 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
 	}
@@ -44,10 +46,7 @@ func CPUQuotaToGOMAXPROCS(minValue int, roundOpt Rounding) (int, CPUQuotaStatus,
 		return -1, CPUQuotaUndefined, err
 	}
 
-	maxProcs := int(math.Floor(quota))
-	if roundOpt == Ceil {
-		maxProcs = int(math.Ceil(quota))
-	}
+	maxProcs := round(quota)
 	if minValue > 0 && maxProcs < minValue {
 		return minValue, CPUQuotaMinUsed, nil
 	}
@@ -61,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 8050be2..e85074b 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"
@@ -86,22 +87,22 @@ func TestNewQueryer(t *testing.T) {
 		stubs := newStubs(t)
 
 		q := testQueryer{v: 2.7}
-		stubs.StubFunc(&_newCgroups2, q, nil)
+		stubs.StubFunc(&_newQueryer, q, nil)
 
-		got, _, err := CPUQuotaToGOMAXPROCS(0, Ceil)
+		got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Ceil(v)) })
 		require.NoError(t, err)
-		assert.Same(t, 3, got)
+		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(&_newCgroups2, q, nil)
+		stubs.StubFunc(&_newQueryer, q, nil)
 
-		got, _, err := CPUQuotaToGOMAXPROCS(0, Floor)
+		got, _, err := CPUQuotaToGOMAXPROCS(0, func(v float64) int { return int(math.Floor(v)) })
 		require.NoError(t, err)
-		assert.Same(t, 2, got)
+		assert.Equal(t, 2, got)
 	})
 }
 
diff --git a/internal/runtime/cpu_quota_unsupported.go b/internal/runtime/cpu_quota_unsupported.go
index bfff88b..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, _ Rounding) (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 b27526d..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
 
@@ -32,12 +34,7 @@ const (
 	CPUQuotaMinUsed
 )
 
-// Rounding controls how the CPU quota value should be rounded to an int
-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
-	Floor
-)
+// 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 254b89f..e561fe6 100644
--- a/maxprocs/maxprocs.go
+++ b/maxprocs/maxprocs.go
@@ -37,10 +37,10 @@ func currentMaxProcs() int {
 }
 
 type config struct {
-	printf        func(string, ...interface{})
-	procs         func(int, iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error)
-	minGOMAXPROCS int
-	roundQuota    iruntime.Rounding
+	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{}) {
@@ -72,10 +72,10 @@ func Min(n int) Option {
 	})
 }
 
-// RoundQuota controls whether the CPU quota should be rounded using ceil or floor.
-func RoundQuota(v iruntime.Rounding) 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.roundQuota = v
+		cfg.roundQuotaFunc = rf
 	})
 }
 
@@ -90,9 +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,
-		roundQuota:    iruntime.Floor,
+		procs:          iruntime.CPUQuotaToGOMAXPROCS,
+		roundQuotaFunc: iruntime.DefaultRoundFunc,
+		minGOMAXPROCS:  1,
 	}
 	for _, o := range opts {
 		o.apply(cfg)
@@ -111,7 +111,7 @@ func Set(opts ...Option) (func(), error) {
 		return undoNoop, nil
 	}
 
-	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS, cfg.roundQuota)
+	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 d988fc6..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, iruntime.Rounding) (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, iruntime.Rounding) (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, iruntime.Rounding) (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, iruntime.Rounding) (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, round iruntime.Rounding) (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, round iruntime.Rounding) (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, round iruntime.Rounding) (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
 		})
@@ -170,22 +171,22 @@ func TestSet(t *testing.T) {
 	})
 
 	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")
+		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, RoundQuota(iruntime.Ceil))
+		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 iruntime.Rounding) (int, iruntime.CPUQuotaStatus, error) {
-			assert.Equal(t, iruntime.Floor, round, "round should be Floor")
+		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, RoundQuota(iruntime.Floor))
+		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")

From b686bfabbcc9ce67d9181c5ce614bc18296f1d61 Mon Sep 17 00:00:00 2001
From: Walther Lee <walthere.lee@gmail.com>
Date: Wed, 13 Dec 2023 22:56:19 -0800
Subject: [PATCH 6/6] add test for rounding quota with a nil round function

Signed-off-by: Walther Lee <walthere.lee@gmail.com>
---
 internal/runtime/cpu_quota_linux_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/internal/runtime/cpu_quota_linux_test.go b/internal/runtime/cpu_quota_linux_test.go
index e85074b..213c165 100644
--- a/internal/runtime/cpu_quota_linux_test.go
+++ b/internal/runtime/cpu_quota_linux_test.go
@@ -83,6 +83,18 @@ func TestNewQueryer(t *testing.T) {
 		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)