process: add collector.process.counter-version CLI parameter (#2064)

internal/collector/process/process.go

@@ -18,6 +18,7 @@
 package process
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -31,6 +32,7 @@ import (
 	"github.com/prometheus-community/windows_exporter/internal/mi"
 	"github.com/prometheus-community/windows_exporter/internal/pdh"
 	"github.com/prometheus-community/windows_exporter/internal/pdh/registry"
+	pdhtypes "github.com/prometheus-community/windows_exporter/internal/pdh/types"
 	"github.com/prometheus-community/windows_exporter/internal/types"
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/sys/windows"
@@ -42,6 +44,7 @@ type Config struct {
 	ProcessInclude      *regexp.Regexp `yaml:"include"`
 	ProcessExclude      *regexp.Regexp `yaml:"exclude"`
 	EnableWorkerProcess bool           `yaml:"iis"`
+	CounterVersion      uint8          `yaml:"counter-version"`
 }
 
 //nolint:gochecknoglobals
@@ -49,6 +52,7 @@ var ConfigDefaults = Config{
 	ProcessInclude:      types.RegExpAny,
 	ProcessExclude:      types.RegExpEmpty,
 	EnableWorkerProcess: false,
+	CounterVersion:      0,
 }
 
 type Collector struct {
@@ -59,10 +63,9 @@ type Collector struct {
 	miSession                 *mi.Session
 	workerProcessMIQueryQuery mi.Query
 
-	collectorVersion int
-
-	collectorV1
-	collectorV2
+	perfDataCollector pdhtypes.Collector
+	perfDataObject    []perfDataCounterValues
+	workerCh          chan processWorkerRequest
 
 	lookupCache sync.Map
 
@@ -130,6 +133,11 @@ func NewWithFlags(app *kingpin.Application) *Collector {
 		"Enable IIS collectWorker process name queries. May cause the collector to leak memory.",
 	).Default(strconv.FormatBool(c.config.EnableWorkerProcess)).BoolVar(&c.config.EnableWorkerProcess)
 
+	app.Flag(
+		"collector.process.counter-version",
+		"Version of the process collector to use. 1 for Process V1, 2 for Process V2. Defaults to 0 which will use the latest version available.",
+	).Default(strconv.FormatUint(uint64(c.config.CounterVersion), 10)).Uint8Var(&c.config.CounterVersion)
+
 	app.Action(func(*kingpin.ParseContext) error {
 		var err error
 
@@ -157,8 +165,12 @@ func (c *Collector) Close() error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	c.closeV1()
-	c.closeV2()
+	c.perfDataCollector.Close()
+
+	if c.workerCh != nil {
+		close(c.workerCh)
+		c.workerCh = nil
+	}
 
 	return nil
 }
@@ -166,42 +178,47 @@ func (c *Collector) Close() error {
 func (c *Collector) Build(logger *slog.Logger, miSession *mi.Session) error {
 	c.logger = logger.With(slog.String("collector", Name))
 
-	if miSession == nil {
-		return errors.New("miSession is nil")
+	var err error
+
+	if c.config.EnableWorkerProcess {
+		if miSession == nil {
+			return errors.New("miSession is nil")
+		}
+
+		miQuery, err := mi.NewQuery("SELECT AppPoolName, ProcessId FROM WorkerProcess")
+		if err != nil {
+			return fmt.Errorf("failed to create WMI query: %w", err)
+		}
+
+		c.workerProcessMIQueryQuery = miQuery
+		c.miSession = miSession
 	}
 
-	miQuery, err := mi.NewQuery("SELECT AppPoolName, ProcessId FROM WorkerProcess")
-	if err != nil {
-		return fmt.Errorf("failed to create WMI query: %w", err)
-	}
+	switch c.config.CounterVersion {
+	case 2:
+		c.perfDataCollector, err = pdh.NewCollector[perfDataCounterValues](pdh.CounterTypeRaw, "Process V2", pdh.InstancesAll)
+	case 1:
+		c.perfDataCollector, err = registry.NewCollector[perfDataCounterValues]("Process", pdh.InstancesAll)
+	default:
+		c.perfDataCollector, err = pdh.NewCollector[perfDataCounterValues](pdh.CounterTypeRaw, "Process V2", pdh.InstancesAll)
+		c.config.CounterVersion = 2
 
-	c.workerProcessMIQueryQuery = miQuery
-	c.miSession = miSession
+		if errors.Is(err, pdh.NewPdhError(pdh.CstatusNoObject)) {
+			c.perfDataCollector, err = registry.NewCollector[perfDataCounterValues]("Process", pdh.InstancesAll)
+			c.config.CounterVersion = 1
+		}
 
-	c.collectorVersion = 2
-	c.perfDataCollectorV2, err = pdh.NewCollector[perfDataCounterValuesV2](pdh.CounterTypeRaw, "Process V2", pdh.InstancesAll)
-
-	if errors.Is(err, pdh.NewPdhError(pdh.CstatusNoObject)) {
-		c.collectorVersion = 1
-		c.perfDataCollectorV1, err = registry.NewCollector[perfDataCounterValuesV1]("Process", pdh.InstancesAll)
+		c.logger.LogAttrs(context.Background(), slog.LevelDebug, fmt.Sprintf("Using process collector V%d", c.config.CounterVersion))
 	}
 
 	if err != nil {
-		return fmt.Errorf("failed to create Process collector: %w", err)
+		return fmt.Errorf("failed to create Process V%d collector: %w", c.config.CounterVersion, err)
 	}
 
-	if c.collectorVersion == 1 {
-		c.workerChV1 = make(chan processWorkerRequestV1, 32)
+	c.workerCh = make(chan processWorkerRequest, 32)
 
-		for range 4 {
-			go c.collectWorkerV1()
-		}
-	} else {
-		c.workerChV2 = make(chan processWorkerRequestV2, 32)
-
-		for range 4 {
-			go c.collectWorkerV2()
-		}
+	for range 4 {
+		go c.collectWorker()
 	}
 
 	c.mu = sync.RWMutex{}
@@ -320,18 +337,7 @@ func (c *Collector) Build(logger *slog.Logger, miSession *mi.Session) error {
 }
 
 func (c *Collector) Collect(ch chan<- prometheus.Metric) error {
-	var workerProcesses []WorkerProcess
-	if c.config.EnableWorkerProcess {
-		if err := c.miSession.Query(&workerProcesses, mi.NamespaceRootWebAdministration, c.workerProcessMIQueryQuery); err != nil {
-			return fmt.Errorf("WMI query failed: %w", err)
-		}
-	}
-
-	if c.collectorVersion == 1 {
-		return c.collectV1(ch, workerProcesses)
-	}
-
-	return c.collectV2(ch, workerProcesses)
+	return c.collect(ch)
 }
 
 // ref: https://github.com/microsoft/hcsshim/blob/8beabacfc2d21767a07c20f8dd5f9f3932dbf305/internal/uvm/stats.go#L25