Merge pull request #193 from nixpanic/automaxprocs

Add package standardflags for -automaxprocs commandline argument

Author: k8s-ci-robot

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/stretchr/testify v1.10.0
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.58.0
 	go.opentelemetry.io/otel/trace v1.33.0
+	go.uber.org/automaxprocs v1.6.0
 	google.golang.org/grpc v1.69.0
 	google.golang.org/protobuf v1.36.0
 	k8s.io/api v0.32.0

go.sum

@@ -75,6 +75,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
+github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
 github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y=
 github.com/prometheus/client_golang v1.20.5/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE=
 github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E=
@@ -109,6 +111,8 @@ go.opentelemetry.io/otel/sdk/metric v1.31.0 h1:i9hxxLJF/9kkvfHppyLL55aW7iIJz4Jjx
 go.opentelemetry.io/otel/sdk/metric v1.31.0/go.mod h1:CRInTMVvNhUKgSAMbKyTMxqOBC0zgyxzW55lZzX43Y8=
 go.opentelemetry.io/otel/trace v1.33.0 h1:cCJuF7LRjUFso9LPnEAHJDB2pqzp+hbO8eu1qqW2d/s=
 go.opentelemetry.io/otel/trace v1.33.0/go.mod h1:uIcdVUZMpTAmz0tI1z04GoVSezK37CbGV4fr1f2nBck=
+go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
+go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

standardflags/automaxprocs.go

@@ -0,0 +1,105 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package standardflags contains flags that multiple CSI sidecars and
+// drivers may want to support.
+package standardflags
+
+import (
+	"flag"
+	"fmt"
+	"strconv"
+
+	"go.uber.org/automaxprocs/maxprocs"
+)
+
+var (
+	logFunc          func(format string, args ...interface{}) = nil
+	undoAutomaxprocs func()                                   = nil
+)
+
+// AddAutomaxprocs adds the -automaxprocs boolean flag to the commandline options.
+// By default the flag is disabled, use [EnableAutomaxprocs] to enable it during
+// startup of an application.
+//
+// The printf function that is passed as an argument, will be used by the
+// [maxprocs.Logger] option when the GOMAXPROCS runtime configuration is adjusted.
+func AddAutomaxprocs(printf func(format string, args ...interface{})) {
+	flag.BoolFunc("automaxprocs",
+		"automatically set GOMAXPROCS to match Linux container CPU quota",
+		handleAutomaxprocs,
+	)
+
+	if printf != nil {
+		// maxprocs.Logger expects a Printf like function.
+		// klog.Info() isn't one, so wrap the contents in a
+		// fmt.Sprintf() for %-formatting substitution.
+		logFunc = func(f string, a ...interface{}) {
+			printf(fmt.Sprintf(f, a...))
+		}
+	}
+}
+
+// EnableAutomaxprocs can be used as an equivalent of -automaxprocs=true on the
+// commandline.
+func EnableAutomaxprocs() {
+	if automaxprocsIsEnabled() {
+		// enabled already, don't enable again
+		return
+	}
+
+	flag.Set("automaxprocs", "true")
+}
+
+// automaxprocsIsEnabled returns true if maxprocs.Set() was successfully
+// executed.
+func automaxprocsIsEnabled() bool {
+	return undoAutomaxprocs != nil
+}
+
+// handleAutomaxprocs parses the passed string into a bool, and enables
+// automaxprocs according to it. If the passed string is empty, automaxprocs
+// is enabled as well.
+func handleAutomaxprocs(s string) error {
+	var err error
+	enabled := true
+
+	if s == "" {
+		EnableAutomaxprocs()
+		return nil
+	}
+
+	enabled, err = strconv.ParseBool(s)
+	if err != nil {
+		return err
+	}
+
+	switch enabled {
+	case true:
+		opts := make([]maxprocs.Option, 0)
+		if logFunc != nil {
+			opts = append(opts, maxprocs.Logger(logFunc))
+		}
+		undoAutomaxprocs, err = maxprocs.Set(opts...)
+	case false:
+		if undoAutomaxprocs != nil {
+			undoAutomaxprocs()
+			undoAutomaxprocs = nil
+		}
+	}
+
+	return err
+}

standardflags/automaxprocs_test.go

@@ -0,0 +1,86 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package standardflags
+
+import (
+	"flag"
+	"testing"
+)
+
+func TestAutomaxprocsArgument(t *testing.T) {
+	tests := []struct {
+		name        string
+		value       string
+		enabled     bool
+		expectError bool
+	}{
+		{
+			name:    "with value as true",
+			value:   "true",
+			enabled: true,
+		},
+		{
+			name:    "with value as false",
+			value:   "false",
+			enabled: false,
+		},
+		{
+			name:    "without value",
+			enabled: true,
+		},
+		{
+			name:        "with invalid value",
+			value:       "error",
+			expectError: true,
+		},
+	}
+
+	if flag.Lookup("automaxprocs") == nil {
+		AddAutomaxprocs(nil) // pass t.Logf to see the logs
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			err := flag.Set("automaxprocs", test.value)
+			if !test.expectError && err != nil {
+				t.Errorf("test %s: failed to set value to %q", test.name, test.value)
+			}
+		})
+	}
+}
+
+func TestEnableDisableAutomaxprocs(t *testing.T) {
+	// make sure the flags is not enabled yet
+	f := flag.Lookup("automaxprocs")
+	if f == nil {
+		AddAutomaxprocs(nil)
+	}
+	if automaxprocsIsEnabled() {
+		handleAutomaxprocs("false")
+	}
+
+	EnableAutomaxprocs()
+	if !automaxprocsIsEnabled() {
+		t.Errorf("failed to enable automaxprocs")
+	}
+
+	// disable again
+	handleAutomaxprocs("false")
+	if automaxprocsIsEnabled() {
+		t.Errorf("failed to disable automaxprocs")
+	}
+}