fix: tcp, udp sample the IPv4 and IPv6 counters in one PDH query (#2684)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/collector/tcp/tcp.go

@@ -60,6 +60,9 @@ var ConfigDefaults = Config{
 type Collector struct {
 	config Config
 
+	// perfDataQuery samples the IPv4 and IPv6 objects together, so the
+	// collector waits only once for the other performance counter queries.
+	perfDataQuery      *pdh.Query
 	perfDataCollector4 *pdh.Collector[perfDataCounterValues]
 	perfDataCollector6 *pdh.Collector[perfDataCounterValues]
 	perfDataObject4    []perfDataCounterValues
@@ -123,6 +126,7 @@ func (c *Collector) Close() error {
 	if slices.Contains(c.config.CollectorsEnabled, subCollectorMetrics) {
 		c.perfDataCollector4.Close()
 		c.perfDataCollector6.Close()
+		c.perfDataQuery.Close()
 	}
 
 	return nil
@@ -207,12 +211,17 @@ func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error {
 	if slices.Contains(c.config.CollectorsEnabled, subCollectorMetrics) {
 		var err error
 
-		c.perfDataCollector4, err = pdh.NewCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), pdh.CounterTypeRaw, "TCPv4", nil)
+		c.perfDataQuery, err = pdh.NewQuery()
+		if err != nil {
+			return fmt.Errorf("failed to open TCP performance counter query: %w", err)
+		}
+
+		c.perfDataCollector4, err = pdh.NewQueryCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), c.perfDataQuery, pdh.CounterTypeRaw, "TCPv4", nil)
 		if err != nil {
 			errs = append(errs, fmt.Errorf("failed to create TCPv4 collector: %w", err))
 		}
 
-		c.perfDataCollector6, err = pdh.NewCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), pdh.CounterTypeRaw, "TCPv6", nil)
+		c.perfDataCollector6, err = pdh.NewQueryCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), c.perfDataQuery, pdh.CounterTypeRaw, "TCPv6", nil)
 		if err != nil {
 			errs = append(errs, fmt.Errorf("failed to create TCPv6 collector: %w", err))
 		}
@@ -244,6 +253,10 @@ func (c *Collector) Collect(ch chan<- prometheus.Metric, _ time.Duration) error
 func (c *Collector) collect(ch chan<- prometheus.Metric) error {
 	errs := make([]error, 0)
 
+	if err := c.perfDataQuery.Collect(); err != nil {
+		return fmt.Errorf("failed to collect TCP metrics: %w", err)
+	}
+
 	if err := c.perfDataCollector4.Collect(&c.perfDataObject4); err != nil {
 		errs = append(errs, fmt.Errorf("failed to collect TCPv4 metrics. %w", err))
 	} else if len(c.perfDataObject4) == 0 {

internal/collector/udp/udp.go

@@ -41,6 +41,9 @@ var ConfigDefaults = Config{}
 type Collector struct {
 	config Config
 
+	// perfDataQuery samples the IPv4 and IPv6 objects together, so the
+	// collector waits only once for the other performance counter queries.
+	perfDataQuery      *pdh.Query
 	perfDataCollector4 *pdh.Collector[perfDataCounterValues]
 	perfDataCollector6 *pdh.Collector[perfDataCounterValues]
 	perfDataObject4    []perfDataCounterValues
@@ -79,6 +82,7 @@ func (c *Collector) GetName() string {
 func (c *Collector) Close() error {
 	c.perfDataCollector4.Close()
 	c.perfDataCollector6.Close()
+	c.perfDataQuery.Close()
 
 	return nil
 }
@@ -113,12 +117,17 @@ func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error {
 
 	var err error
 
-	c.perfDataCollector4, err = pdh.NewCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), pdh.CounterTypeRaw, "UDPv4", nil)
+	c.perfDataQuery, err = pdh.NewQuery()
+	if err != nil {
+		return fmt.Errorf("failed to open UDP performance counter query: %w", err)
+	}
+
+	c.perfDataCollector4, err = pdh.NewQueryCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), c.perfDataQuery, pdh.CounterTypeRaw, "UDPv4", nil)
 	if err != nil {
 		errs = append(errs, fmt.Errorf("failed to create UDPv4 collector: %w", err))
 	}
 
-	c.perfDataCollector6, err = pdh.NewCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), pdh.CounterTypeRaw, "UDPv6", nil)
+	c.perfDataCollector6, err = pdh.NewQueryCollector[perfDataCounterValues](logger.With(slog.String("collector", Name)), c.perfDataQuery, pdh.CounterTypeRaw, "UDPv6", nil)
 	if err != nil {
 		errs = append(errs, fmt.Errorf("failed to create UDPv6 collector: %w", err))
 	}
@@ -133,6 +142,10 @@ func (c *Collector) Collect(ch chan<- prometheus.Metric, _ time.Duration) error
 }
 
 func (c *Collector) collect(ch chan<- prometheus.Metric) error {
+	if err := c.perfDataQuery.Collect(); err != nil {
+		return fmt.Errorf("failed to collect UDP metrics: %w", err)
+	}
+
 	errs := make([]error, 0)
 
 	if err := c.perfDataCollector4.Collect(&c.perfDataObject4); err != nil {

internal/pdh/collector.go

@@ -44,6 +44,9 @@ var (
 type CounterValues = map[string]map[string]CounterValue
 
 type Collector[T any] struct {
+	// query is the shared query of a collector from [NewQueryCollector], or
+	// nil if the collector owns handle.
+	query                 *Query
 	object                string
 	resultType            CounterType
 	counters              []Counter
@@ -113,6 +116,22 @@ type Counter struct {
 // could not be added or the initial collection failed. The caller owns that
 // Collector and must Close it. If the Collector is nil, no resources are held.
 func NewCollector[T any](logger *slog.Logger, resultType CounterType, object string, instances []string) (*Collector[T], error) {
+	return newStructCollector[T](logger, nil, resultType, object, instances)
+}
+
+// NewQueryCollector creates a collector like [NewCollector] whose counters are
+// part of query. [Query.Collect] samples them together with the other
+// collectors of the query; [Collector.Collect] reads the latest sample. Close
+// the collector before the query.
+func NewQueryCollector[T any](logger *slog.Logger, query *Query, resultType CounterType, object string, instances []string) (*Collector[T], error) {
+	if query == nil {
+		return nil, ErrPerformanceCounterNotInitialized
+	}
+
+	return newStructCollector[T](logger, query, resultType, object, instances)
+}
+
+func newStructCollector[T any](logger *slog.Logger, query *Query, resultType CounterType, object string, instances []string) (*Collector[T], error) {
 	valueType := reflect.TypeFor[T]()
 	if valueType.Kind() != reflect.Struct {
 		return nil, fmt.Errorf("expected a struct, got %s", valueType)
@@ -167,7 +186,7 @@ func NewCollector[T any](logger *slog.Logger, resultType CounterType, object str
 		},
 	}
 
-	return newCollector(logger, resultType, object, instances, fields, rows, false, errs)
+	return newCollector(logger, query, resultType, object, instances, fields, rows, false, errs)
 }
 
 // NewDynamicCollector creates a collector for counters that are only known at runtime.
@@ -192,10 +211,10 @@ func NewDynamicCollector(logger *slog.Logger, resultType CounterType, object str
 		},
 	}
 
-	return newCollector(logger, resultType, object, instances, fields, rows, true, nil)
+	return newCollector(logger, nil, resultType, object, instances, fields, rows, true, nil)
 }
 
-func newCollector[T any](logger *slog.Logger, resultType CounterType, object string, instances []string,
+func newCollector[T any](logger *slog.Logger, query *Query, resultType CounterType, object string, instances []string,
 	fields []counterField, rows rowAccessor[T], partialRows bool, errs []error,
 ) (*Collector[T], error) {
 	if resultType != CounterTypeRaw && resultType != CounterTypeFormatted {
@@ -204,7 +223,17 @@ func newCollector[T any](logger *slog.Logger, resultType CounterType, object str
 
 	var handle pdhQueryHandle
 
-	if ret := OpenQuery(0, 0, &handle); ret != ErrorSuccess {
+	if query != nil {
+		// Adding counters while the query samples is not safe.
+		query.mu.Lock()
+		defer query.mu.Unlock()
+
+		if query.handle == 0 {
+			return nil, ErrPerformanceCounterNotInitialized
+		}
+
+		handle = query.handle
+	} else if ret := OpenQuery(0, 0, &handle); ret != ErrorSuccess {
 		return nil, NewPdhError(ret)
 	}
 
@@ -213,6 +242,7 @@ func newCollector[T any](logger *slog.Logger, resultType CounterType, object str
 	}
 
 	collector := &Collector[T]{
+		query:                 query,
 		object:                object,
 		resultType:            resultType,
 		counters:              make([]Counter, 0, len(fields)),
@@ -326,7 +356,11 @@ func newCollector[T any](logger *slog.Logger, resultType CounterType, object str
 	}
 
 	if len(collector.counters) == 0 {
-		collector.Close()
+		// A shared query stays open for its other collectors, and its lock
+		// is held here, so only close a query of our own.
+		if query == nil {
+			collector.Close()
+		}
 
 		errs = append(errs, errors.New("no counters configured"))
 
@@ -340,6 +374,14 @@ func newCollector[T any](logger *slog.Logger, resultType CounterType, object str
 	collector.state = &collectState{}
 
 	// Collect initial data because some counters need to be read twice to get the correct value.
+	if query != nil {
+		if ret := CollectQueryData(handle); ret != ErrorSuccess && ret != NoData {
+			return collector, fmt.Errorf("failed to collect initial data: %w", NewPdhError(ret))
+		}
+
+		return collector, nil
+	}
+
 	var collectValues []T
 	if err := collector.Collect(&collectValues); err != nil && !errors.Is(err, ErrNoData) {
 		return collector, fmt.Errorf("failed to collect initial data: %w", err)
@@ -408,6 +450,12 @@ func (c *Collector[T]) Collect(dst *[]T) error {
 		return errors.New("dst must not be nil")
 	}
 
+	if c.query != nil {
+		// Reading the counter arrays while the query samples is not safe.
+		c.query.mu.Lock()
+		defer c.query.mu.Unlock()
+	}
+
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -445,8 +493,11 @@ func (c *Collector[T]) collect(dst *[]T, state *collectState) (err error) {
 		}
 	}()
 
-	if ret := CollectQueryData(c.handle); ret != ErrorSuccess {
-		return fmt.Errorf("failed to collect query data: %w", NewPdhError(ret))
+	// A collector in a shared query reads the sample taken by Query.Collect.
+	if c.query == nil {
+		if ret := CollectQueryData(c.handle); ret != ErrorSuccess {
+			return fmt.Errorf("failed to collect query data: %w", NewPdhError(ret))
+		}
 	}
 
 	*dst = (*dst)[:0:0]
@@ -787,11 +838,26 @@ func (c *Collector[T]) Close() {
 		return
 	}
 
+	if c.query != nil {
+		c.query.mu.Lock()
+		defer c.query.mu.Unlock()
+	}
+
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	if c.handle != 0 {
+	switch {
+	case c.handle == 0:
+	case c.query == nil:
 		CloseQuery(c.handle)
+	case c.query.handle != 0:
+		// The query belongs to its other collectors; only remove the
+		// counters of this one.
+		for _, counter := range c.counters {
+			for _, instance := range counter.Instances {
+				RemoveCounter(instance)
+			}
+		}
 	}
 
 	c.handle = 0

internal/pdh/pdh.go

@@ -290,6 +290,7 @@ var (
 	pdhGetRawCounterValue        = libPdhDll.NewProc("PdhGetRawCounterValue")
 	pdhGetRawCounterArrayW       = libPdhDll.NewProc("PdhGetRawCounterArrayW")
 	pdhPdhGetCounterTimeBase     = libPdhDll.NewProc("PdhGetCounterTimeBase")
+	pdhRemoveCounter             = libPdhDll.NewProc("PdhRemoveCounter")
 )
 
 // AddCounter adds the specified counter to the query. This is the internationalized version. Preferably, use the
@@ -358,6 +359,13 @@ func AddEnglishCounter(hQuery pdhQueryHandle, szFullCounterPath string, dwUserDa
 	return uint32(ret)
 }
 
+// RemoveCounter removes a counter from its query.
+func RemoveCounter(hCounter pdhCounterHandle) uint32 {
+	ret, _, _ := pdhRemoveCounter.Call(uintptr(hCounter))
+
+	return uint32(ret)
+}
+
 // CloseQuery closes all counters contained in the specified query, closes all handles related to the query,
 // and frees all memory associated with the query.
 func CloseQuery(hQuery pdhQueryHandle) uint32 {

internal/pdh/query.go

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright The Prometheus 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.
+
+//go:build windows
+
+package pdh
+
+import (
+	"fmt"
+	"sync"
+)
+
+// Query is a PDH query that several collectors share, so one
+// PdhCollectQueryData samples all of them. Windows serves counter samples of
+// a process one at a time, so every query waits for the others during a
+// scrape; collectors that are always read together, such as the IPv4 and IPv6
+// objects of a protocol, wait only once in one query.
+//
+// Create the collectors with [NewQueryCollector]. [Query.Collect] takes the
+// sample; [Collector.Collect] of a collector in the query reads the latest
+// sample without taking a new one.
+type Query struct {
+	mu     sync.Mutex
+	handle pdhQueryHandle
+}
+
+// NewQuery opens a query for [NewQueryCollector]. The caller must Close it
+// after the collectors in it.
+func NewQuery() (*Query, error) {
+	var handle pdhQueryHandle
+
+	if ret := OpenQuery(0, 0, &handle); ret != ErrorSuccess {
+		return nil, NewPdhError(ret)
+	}
+
+	return &Query{handle: handle}, nil
+}
+
+// Collect takes one sample of all counters in the query.
+func (q *Query) Collect() error {
+	if q == nil {
+		return ErrPerformanceCounterNotInitialized
+	}
+
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if q.handle == 0 {
+		return ErrPerformanceCounterNotInitialized
+	}
+
+	if ret := CollectQueryData(q.handle); ret != ErrorSuccess {
+		return fmt.Errorf("failed to collect query data: %w", NewPdhError(ret))
+	}
+
+	return nil
+}
+
+// Close closes the query and all counters in it.
+func (q *Query) Close() {
+	if q == nil {
+		return
+	}
+
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if q.handle != 0 {
+		CloseQuery(q.handle)
+	}
+
+	q.handle = 0
+}

internal/pdh/query_test.go

@@ -0,0 +1,155 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Copyright The Prometheus 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.
+
+//go:build windows
+
+package pdh_test
+
+import (
+	"log/slog"
+	"sync"
+	"testing"
+
+	"github.com/prometheus-community/windows_exporter/internal/pdh"
+	"github.com/stretchr/testify/require"
+)
+
+type systemValues struct {
+	Processes float64 `perfdata:"Processes"`
+	Threads   float64 `perfdata:"Threads"`
+}
+
+type memoryValues struct {
+	AvailableBytes float64 `perfdata:"Available Bytes"`
+}
+
+func TestQuery(t *testing.T) {
+	t.Parallel()
+
+	query, err := pdh.NewQuery()
+	require.NoError(t, err)
+
+	system, err := pdh.NewQueryCollector[systemValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "System", nil)
+	require.NoError(t, err)
+
+	memory, err := pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "Memory", nil)
+	require.NoError(t, err)
+
+	var (
+		systemData []systemValues
+		memoryData []memoryValues
+	)
+
+	// One sample serves both collectors.
+	require.NoError(t, query.Collect())
+	require.NoError(t, system.Collect(&systemData))
+	require.NoError(t, memory.Collect(&memoryData))
+	require.Len(t, systemData, 1)
+	require.Len(t, memoryData, 1)
+	require.NotZero(t, systemData[0].Processes)
+	require.NotZero(t, memoryData[0].AvailableBytes)
+
+	// Closing one collector removes only its counters.
+	system.Close()
+	require.ErrorIs(t, system.Collect(&systemData), pdh.ErrPerformanceCounterNotInitialized)
+	require.NoError(t, query.Collect())
+	require.NoError(t, memory.Collect(&memoryData))
+	require.NotZero(t, memoryData[0].AvailableBytes)
+
+	memory.Close()
+	query.Close()
+	query.Close()
+
+	require.ErrorIs(t, query.Collect(), pdh.ErrPerformanceCounterNotInitialized)
+
+	_, err = pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "Memory", nil)
+	require.ErrorIs(t, err, pdh.ErrPerformanceCounterNotInitialized)
+
+	_, err = pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), nil, pdh.CounterTypeRaw, "Memory", nil)
+	require.ErrorIs(t, err, pdh.ErrPerformanceCounterNotInitialized)
+}
+
+// A collector whose object is missing reports the error, and the other
+// collector of the query keeps working.
+func TestQueryMissingObject(t *testing.T) {
+	t.Parallel()
+
+	query, err := pdh.NewQuery()
+	require.NoError(t, err)
+
+	defer query.Close()
+
+	missing, err := pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "Does Not Exist", nil)
+	require.Error(t, err)
+
+	defer missing.Close()
+
+	memory, err := pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "Memory", nil)
+	require.NoError(t, err)
+
+	defer memory.Close()
+
+	var data []memoryValues
+
+	require.NoError(t, query.Collect())
+	require.NoError(t, memory.Collect(&data))
+	require.NotEmpty(t, data)
+}
+
+func TestQueryConcurrentCollect(t *testing.T) {
+	t.Parallel()
+
+	query, err := pdh.NewQuery()
+	require.NoError(t, err)
+
+	system, err := pdh.NewQueryCollector[systemValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "System", nil)
+	require.NoError(t, err)
+
+	memory, err := pdh.NewQueryCollector[memoryValues](slog.New(slog.DiscardHandler), query, pdh.CounterTypeRaw, "Memory", nil)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+
+	for range 8 {
+		wg.Go(func() {
+			var (
+				systemData []systemValues
+				memoryData []memoryValues
+			)
+
+			for range 20 {
+				if err := query.Collect(); err != nil {
+					t.Error(err)
+
+					return
+				}
+
+				if err := system.Collect(&systemData); err != nil {
+					t.Error(err)
+				}
+
+				if err := memory.Collect(&memoryData); err != nil {
+					t.Error(err)
+				}
+			}
+		})
+	}
+
+	wg.Wait()
+
+	system.Close()
+	memory.Close()
+	query.Close()
+}