Add a Reserved option

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## v1.2.0 (unreleased)
+
+- Add a `Reserved(quota float64)` option to allow users to reserve part of
+  their CPU quota things outside the Go runtime.
+
 ## v1.1.0 (2017-11-10)
 
 - Log the new value of `GOMAXPROCS` rather than the current value.

internal/runtime/cpu_quota_linux.go

@@ -30,7 +30,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(cfg CPUQuotaConfig) (int, CPUQuotaStatus, error) {
 	cgroups, err := cg.NewCGroupsForCurrentProcess()
 	if err != nil {
 		return -1, CPUQuotaUndefined, err
@@ -41,9 +41,14 @@ func CPUQuotaToGOMAXPROCS(minValue int) (int, CPUQuotaStatus, error) {
 		return -1, CPUQuotaUndefined, err
 	}
 
+	if cfg.Reserved > 0 {
+		quota -= cfg.Reserved
+	}
+
 	maxProcs := int(math.Ceil(quota))
-	if minValue > 0 && maxProcs < minValue {
-		return minValue, CPUQuotaMinUsed, nil
+	if cfg.MinValue > 0 && maxProcs < cfg.MinValue {
+		return cfg.MinValue, CPUQuotaMinUsed, nil
 	}
+
 	return maxProcs, CPUQuotaUsed, nil
 }

internal/runtime/cpu_quota_unsupported.go

@@ -25,6 +25,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(_ CPUQuotaConfig) (int, CPUQuotaStatus, error) {
 	return -1, CPUQuotaUndefined, nil
 }

internal/runtime/runtime.go

@@ -31,3 +31,18 @@ const (
 	// CPUQuotaMinUsed is return when CPU quota is larger than the min value
 	CPUQuotaMinUsed
 )
+
+// CPUQuotaConfig specifies configured constraints for automatic GOMAXPROCS
+// configuration.
+type CPUQuotaConfig struct {
+	MinValue int
+	Reserved float64
+}
+
+// CPUQuotaFunc represents a strategy for auto-configuring MAXPROCS based on
+// CPU quota based on configured constraints and environmental data.
+type CPUQuotaFunc func(CPUQuotaConfig) (
+	maxprocs int,
+	status CPUQuotaStatus,
+	err error,
+)

maxprocs/maxprocs.go

@@ -32,20 +32,13 @@ import (
 
 const _maxProcsKey = "GOMAXPROCS"
 
-func currentMaxProcs() int {
-	return runtime.GOMAXPROCS(0)
-}
+func noopLog(fmt string, args ...interface{}) {}
 
 type config struct {
-	printf        func(string, ...interface{})
-	procs         func(int) (int, iruntime.CPUQuotaStatus, error)
-	minGOMAXPROCS int
-}
-
-func (c *config) log(fmt string, args ...interface{}) {
-	if c.printf != nil {
-		c.printf(fmt, args...)
-	}
+	iruntime.CPUQuotaConfig
+	log       func(string, ...interface{})
+	quotaFunc iruntime.CPUQuotaFunc
+	origValue int
 }
 
 // An Option alters the behavior of Set.
@@ -57,7 +50,7 @@ type Option interface {
 // Set doesn't log anything.
 func Logger(printf func(string, ...interface{})) Option {
 	return optionFunc(func(cfg *config) {
-		cfg.printf = printf
+		cfg.log = printf
 	})
 }
 
@@ -66,7 +59,23 @@ func Logger(printf func(string, ...interface{})) Option {
 func Min(n int) Option {
 	return optionFunc(func(cfg *config) {
 		if n >= 1 {
-			cfg.minGOMAXPROCS = n
+			cfg.MinValue = n
+		}
+	})
+}
+
+// ReservedCPUQuota specifies a CPU quota amount to consider reserved for use
+// outside of the Go runtime; e.g. by a cgo extension, or another process
+// within the container.
+//
+// GOMAXPROCS will be set to the remaining (rounded down) quota, clamped to the
+// minimum limit (default 1).
+//
+// Any non-positive value is ignored.
+func ReservedCPUQuota(reserved float64) Option {
+	return optionFunc(func(cfg *config) {
+		if reserved >= 0 {
+			cfg.Reserved = reserved
 		}
 	})
 }
@@ -75,46 +84,50 @@ type optionFunc func(*config)
 
 func (of optionFunc) apply(cfg *config) { of(cfg) }
 
+func (c config) undoNoop() {
+	c.log("maxprocs: No GOMAXPROCS change to reset")
+}
+
+func (c config) undo() {
+	c.log("maxprocs: Resetting GOMAXPROCS to %d", c.origValue)
+	runtime.GOMAXPROCS(c.origValue)
+}
+
 // Set GOMAXPROCS to match the Linux container CPU quota (if any), returning
 // any error encountered and an undo function.
 //
 // Set is a no-op on non-Linux systems and in Linux environments without a
 // configured CPU quota.
 func Set(opts ...Option) (func(), error) {
-	cfg := &config{
-		procs:         iruntime.CPUQuotaToGOMAXPROCS,
-		minGOMAXPROCS: 1,
+	cfg := config{
+		log:       noopLog,
+		quotaFunc: iruntime.CPUQuotaToGOMAXPROCS,
+		CPUQuotaConfig: iruntime.CPUQuotaConfig{
+			MinValue: 1,
+		},
 	}
 	for _, o := range opts {
-		o.apply(cfg)
-	}
-
-	undoNoop := func() {
-		cfg.log("maxprocs: No GOMAXPROCS change to reset")
+		o.apply(&cfg)
 	}
 
 	// Honor the GOMAXPROCS environment variable if present. Otherwise, amend
 	// `runtime.GOMAXPROCS()` with the current process' CPU quota if the OS is
 	// Linux, and guarantee a minimum value of 2 to ensure efficiency.
 	if max, exists := os.LookupEnv(_maxProcsKey); exists {
 		cfg.log("maxprocs: Honoring GOMAXPROCS=%d as set in environment", max)
-		return undoNoop, nil
+		return cfg.undoNoop, nil
 	}
 
-	maxProcs, status, err := cfg.procs(cfg.minGOMAXPROCS)
+	maxProcs, status, err := cfg.quotaFunc(cfg.CPUQuotaConfig)
 	if err != nil {
-		return undoNoop, err
+		return cfg.undoNoop, err
 	}
 
-	if status == iruntime.CPUQuotaUndefined {
-		cfg.log("maxprocs: Leaving GOMAXPROCS=%d: CPU quota undefined", currentMaxProcs())
-		return undoNoop, nil
-	}
+	cfg.origValue = runtime.GOMAXPROCS(0)
 
-	prev := currentMaxProcs()
-	undo := func() {
-		cfg.log("maxprocs: Resetting GOMAXPROCS to %d", prev)
-		runtime.GOMAXPROCS(prev)
+	if status == iruntime.CPUQuotaUndefined {
+		cfg.log("maxprocs: Leaving GOMAXPROCS=%d: CPU quota undefined", cfg.origValue)
+		return cfg.undoNoop, nil
 	}
 
 	switch status {
@@ -125,5 +138,5 @@ func Set(opts ...Option) (func(), error) {
 	}
 
 	runtime.GOMAXPROCS(maxProcs)
-	return undo, nil
+	return cfg.undo, nil
 }

maxprocs/maxprocs_test.go

@@ -26,6 +26,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"runtime"
 	"strconv"
 	"testing"
 
@@ -55,9 +56,9 @@ func testLogger() (*bytes.Buffer, Option) {
 	return buf, Logger(printf)
 }
 
-func stubProcs(f func(int) (int, iruntime.CPUQuotaStatus, error)) Option {
+func stubQuotaFunc(f iruntime.CPUQuotaFunc) Option {
 	return optionFunc(func(cfg *config) {
-		cfg.procs = f
+		cfg.quotaFunc = f
 	})
 }
 
@@ -80,93 +81,93 @@ func TestLogger(t *testing.T) {
 
 func TestSet(t *testing.T) {
 	// Ensure that we've undone any modifications correctly.
-	prev := currentMaxProcs()
+	prev := runtime.GOMAXPROCS(0)
 	defer func() {
-		require.Equal(t, prev, currentMaxProcs(), "didn't undo GOMAXPROCS changes")
+		require.Equal(t, prev, runtime.GOMAXPROCS(0), "didn't undo GOMAXPROCS changes")
 	}()
 
 	t.Run("EnvVarPresent", func(t *testing.T) {
 		withMax(t, 42, func() {
-			prev := currentMaxProcs()
+			prev := runtime.GOMAXPROCS(0)
 			undo, err := Set()
 			defer undo()
 			require.NoError(t, err, "Set failed")
-			assert.Equal(t, prev, currentMaxProcs(), "shouldn't alter GOMAXPROCS")
+			assert.Equal(t, prev, runtime.GOMAXPROCS(0), "shouldn't alter GOMAXPROCS")
 		})
 	})
 
 	t.Run("ErrorReadingQuota", func(t *testing.T) {
-		opt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		opt := stubQuotaFunc(func(iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, errors.New("failed")
 		})
-		prev := currentMaxProcs()
+		prev := runtime.GOMAXPROCS(0)
 		undo, err := Set(opt)
 		defer undo()
 		require.Error(t, err, "Set should have failed")
 		assert.Equal(t, "failed", err.Error(), "should pass errors up the stack")
-		assert.Equal(t, prev, currentMaxProcs(), "shouldn't alter GOMAXPROCS")
+		assert.Equal(t, prev, runtime.GOMAXPROCS(0), "shouldn't alter GOMAXPROCS")
 	})
 
 	t.Run("QuotaUndefined", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubQuotaFunc(func(iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
 			return 0, iruntime.CPUQuotaUndefined, nil
 		})
-		prev := currentMaxProcs()
+		prev := runtime.GOMAXPROCS(0)
 		undo, err := Set(logOpt, quotaOpt)
 		defer undo()
 		require.NoError(t, err, "Set failed")
-		assert.Equal(t, prev, currentMaxProcs(), "shouldn't alter GOMAXPROCS")
+		assert.Equal(t, prev, runtime.GOMAXPROCS(0), "shouldn't alter GOMAXPROCS")
 		assert.Contains(t, buf.String(), "quota undefined", "unexpected log output")
 	})
 
 	t.Run("QuotaUndefined return maxProcs=7", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(int) (int, iruntime.CPUQuotaStatus, error) {
+		quotaOpt := stubQuotaFunc(func(iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
 			return 7, iruntime.CPUQuotaUndefined, nil
 		})
-		prev := currentMaxProcs()
+		prev := runtime.GOMAXPROCS(0)
 		undo, err := Set(logOpt, quotaOpt)
 		defer undo()
 		require.NoError(t, err, "Set failed")
-		assert.Equal(t, prev, currentMaxProcs(), "shouldn't alter GOMAXPROCS")
+		assert.Equal(t, prev, runtime.GOMAXPROCS(0), "shouldn't alter GOMAXPROCS")
 		assert.Contains(t, buf.String(), "quota undefined", "unexpected log output")
 	})
 
 	t.Run("QuotaTooSmall", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
-			return min, iruntime.CPUQuotaMinUsed, nil
+		quotaOpt := stubQuotaFunc(func(cfg iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
+			return cfg.MinValue, iruntime.CPUQuotaMinUsed, nil
 		})
 		undo, err := Set(logOpt, quotaOpt, Min(5))
 		defer undo()
 		require.NoError(t, err, "Set failed")
-		assert.Equal(t, 5, currentMaxProcs(), "should use min allowed GOMAXPROCS")
+		assert.Equal(t, 5, runtime.GOMAXPROCS(0), "should use min allowed GOMAXPROCS")
 		assert.Contains(t, buf.String(), "using minimum allowed", "unexpected log output")
 	})
 
 	t.Run("Min unused", func(t *testing.T) {
 		buf, logOpt := testLogger()
-		quotaOpt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
-			return min, iruntime.CPUQuotaMinUsed, nil
+		quotaOpt := stubQuotaFunc(func(cfg iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
+			return cfg.MinValue, iruntime.CPUQuotaMinUsed, nil
 		})
 		// Min(-1) should be ignored.
 		undo, err := Set(logOpt, quotaOpt, Min(5), Min(-1))
 		defer undo()
 		require.NoError(t, err, "Set failed")
-		assert.Equal(t, 5, currentMaxProcs(), "should use min allowed GOMAXPROCS")
+		assert.Equal(t, 5, runtime.GOMAXPROCS(0), "should use min allowed GOMAXPROCS")
 		assert.Contains(t, buf.String(), "using minimum allowed", "unexpected log output")
 	})
 
 	t.Run("QuotaUsed", func(t *testing.T) {
-		opt := stubProcs(func(min int) (int, iruntime.CPUQuotaStatus, error) {
-			assert.Equal(t, 1, min, "Default minimum value should be 1")
+		opt := stubQuotaFunc(func(cfg iruntime.CPUQuotaConfig) (int, iruntime.CPUQuotaStatus, error) {
+			assert.Equal(t, 1, cfg.MinValue, "Default minimum value should be 1")
 			return 42, iruntime.CPUQuotaUsed, nil
 		})
 		undo, err := Set(opt)
 		defer undo()
 		require.NoError(t, err, "Set failed")
-		assert.Equal(t, 42, currentMaxProcs(), "should change GOMAXPROCS to match quota")
+		assert.Equal(t, 42, runtime.GOMAXPROCS(0), "should change GOMAXPROCS to match quota")
 	})
 }