From e2887ffef447f5098f1a52a1564fe1e5b9a47adc Mon Sep 17 00:00:00 2001
From: Haoyu Sun <hasun@redhat.com>
Date: Mon, 23 Oct 2023 17:43:13 +0200
Subject: [PATCH] Automatically set GOMAXPROCS according to available CPU quota
 if --runtime.gomaxprocs < 1.

---
 go.mod           |  1 +
 go.sum           |  3 +++
 node_exporter.go | 15 ++++++++++++++-
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 4686e5edc1..d0ae58976f 100644
--- a/go.mod
+++ b/go.mod
@@ -28,6 +28,7 @@ require (
 	github.com/prometheus/exporter-toolkit v0.10.0
 	github.com/prometheus/procfs v0.12.0
 	github.com/safchain/ethtool v0.3.0
+	go.uber.org/automaxprocs v1.5.3
 	golang.org/x/exp v0.0.0-20230522175609-2e198f4a06a1
 	golang.org/x/sys v0.13.0
 	howett.net/plist v1.0.0
diff --git a/go.sum b/go.sum
index e87556951d..c3f922c611 100644
--- a/go.sum
+++ b/go.sum
@@ -74,6 +74,7 @@ github.com/opencontainers/selinux v1.11.0 h1:+5Zbo97w3Lbmb3PeqQtpmTkMwsW5nRI3YaL
 github.com/opencontainers/selinux v1.11.0/go.mod h1:E5dMC3VPuVvVHDYmi78qvhJp8+M586T4DlDRYpFkyec=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
 github.com/prometheus-community/go-runit v0.1.0 h1:uTWEj/Fn2RoLdfg/etSqwzgYNOYPrARx1BHUN052tGA=
 github.com/prometheus-community/go-runit v0.1.0/go.mod h1:AvJ9Jo3gAFu2lbM4+qfjdpq30FfiLDJZKbQ015u08IQ=
 github.com/prometheus/client_golang v1.17.0 h1:rl2sfwZMtSthVU752MqfjQozy7blglC+1SOtjMAMh+Q=
@@ -105,6 +106,8 @@ github.com/xhit/go-str2duration/v2 v2.1.0/go.mod h1:ohY8p+0f07DiV6Em5LKB0s2YpLtX
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/automaxprocs v1.5.3 h1:kWazyxZUrS3Gs4qUpbwo5kEIMGe/DAvi5Z4tl2NW4j8=
+go.uber.org/automaxprocs v1.5.3/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0=
 go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
diff --git a/node_exporter.go b/node_exporter.go
index cc24f7dfba..5a4e6518ce 100644
--- a/node_exporter.go
+++ b/node_exporter.go
@@ -25,6 +25,7 @@ import (
 
 	"github.com/prometheus/common/promlog"
 	"github.com/prometheus/common/promlog/flag"
+	"go.uber.org/automaxprocs/maxprocs"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/go-kit/log"
@@ -173,7 +174,7 @@ func main() {
 			"Set all collectors to disabled by default.",
 		).Default("false").Bool()
 		maxProcs = kingpin.Flag(
-			"runtime.gomaxprocs", "The target number of CPUs Go will run on (GOMAXPROCS)",
+			"runtime.gomaxprocs", "The target number of CPUs Go will run on (GOMAXPROCS). If < 1, it will use all available CPU quota.",
 		).Envar("GOMAXPROCS").Default("1").Int()
 		toolkitFlags = kingpinflag.AddFlags(kingpin.CommandLine, ":9100")
 	)
@@ -195,6 +196,18 @@ func main() {
 		level.Warn(logger).Log("msg", "Node Exporter is running as root user. This exporter is designed to run as unprivileged user, root is not required.")
 	}
 	runtime.GOMAXPROCS(*maxProcs)
+
+	if *maxProcs < 1 {
+		debugPrintf := func(format string, v ...interface{}) {
+			level.Debug(logger).Log("msg", fmt.Sprintf(format, v...))
+		}
+		undo, err := maxprocs.Set(maxprocs.Logger(debugPrintf))
+		defer undo()
+		if err != nil {
+			level.Warn(logger).Log("failed to set GOMAXPROCS automatically: %v", err)
+		}
+	}
+
 	level.Debug(logger).Log("msg", "Go MAXPROCS", "procs", runtime.GOMAXPROCS(0))
 
 	http.Handle(*metricsPath, newHandler(!*disableExporterMetrics, *maxRequests, logger))