From eeaefba0d7ff3d3641643a8aae8f7376e5426e27 Mon Sep 17 00:00:00 2001
From: Ben Reedy <breed808@users.noreply.github.com>
Date: Sat, 18 May 2024 06:51:33 +1000
Subject: [PATCH] feat: allow setting of exporter process priority (#1488)

---
 exporter.go | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/exporter.go b/exporter.go
index 38e48f90b..f97c0964d 100644
--- a/exporter.go
+++ b/exporter.go
@@ -31,8 +31,12 @@ import (
 	"github.com/prometheus/common/version"
 	"github.com/prometheus/exporter-toolkit/web"
 	webflag "github.com/prometheus/exporter-toolkit/web/kingpinflag"
+	"golang.org/x/sys/windows"
 )
 
+// https://learn.microsoft.com/en-us/windows/win32/procthread/process-security-and-access-rights
+const PROCESS_ALL_ACCESS = windows.STANDARD_RIGHTS_REQUIRED | windows.SYNCHRONIZE | windows.SPECIFIC_RIGHTS_ALL
+
 // Same struct prometheus uses for their /version endpoint.
 // Separate copy to avoid pulling all of prometheus as a dependency
 type prometheusVersion struct {
@@ -44,6 +48,32 @@ type prometheusVersion struct {
 	GoVersion string `json:"goVersion"`
 }
 
+// Mapping of priority names to uin32 values required by windows.SetPriorityClass
+var priorityStringToInt = map[string]uint32{
+	"realtime":    windows.REALTIME_PRIORITY_CLASS,
+	"high":        windows.HIGH_PRIORITY_CLASS,
+	"abovenormal": windows.ABOVE_NORMAL_PRIORITY_CLASS,
+	"normal":      windows.NORMAL_PRIORITY_CLASS,
+	"belownormal": windows.BELOW_NORMAL_PRIORITY_CLASS,
+	"low":         windows.IDLE_PRIORITY_CLASS,
+}
+
+func setPriorityWindows(pid int, priority uint32) error {
+	handle, err := windows.OpenProcess(PROCESS_ALL_ACCESS, false, uint32(pid))
+	if err != nil {
+		return err
+	}
+	//nolint:errcheck
+	defer windows.CloseHandle(handle) // Technically this can fail, but we ignore it
+
+	err = windows.SetPriorityClass(handle, priority)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func main() {
 	app := kingpin.New("windows_exporter", "A metrics collector for Windows.")
 	var (
@@ -84,6 +114,10 @@ func main() {
 			"debug.enabled",
 			"If true, windows_exporter will expose debug endpoints under /debug/pprof.",
 		).Default("false").Bool()
+		processPriority = app.Flag(
+			"process.priority",
+			"Priority of the exporter process. Higher priorities may improve exporter responsiveness during periods of system load. Can be one of [\"realtime\", \"high\", \"abovenormal\", \"normal\", \"belownormal\", \"low\"]",
+		).Default("normal").String()
 	)
 
 	winlogConfig := &winlog.Config{}
@@ -145,6 +179,16 @@ func main() {
 		return
 	}
 
+	// Only set process priority if a non-default and valid value has been set
+	if *processPriority != "normal" && priorityStringToInt[*processPriority] != 0 {
+		_ = level.Debug(logger).Log("msg", "setting process priority to "+*processPriority)
+		err = setPriorityWindows(os.Getpid(), priorityStringToInt[*processPriority])
+		if err != nil {
+			_ = level.Error(logger).Log("msg", "failed to set process priority", "err", err)
+			os.Exit(1)
+		}
+	}
+
 	if err = wmi.InitWbem(logger); err != nil {
 		_ = level.Error(logger).Log("err", err)
 		os.Exit(1)