pdh/registry: reuse the raw data buffer across Collect calls (#2567)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

internal/pdh/registry/collector.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 
 	"github.com/prometheus-community/windows_exporter/internal/pdh"
 )
@@ -32,6 +33,12 @@ type Collector[T any] struct {
 
 	counters       map[string]Counter
 	nameIndexValue int
+
+	// bufferMu guards buffer, the raw data buffer of the last query. A Collect
+	// takes it out for the duration of the query, so concurrent calls never
+	// share it; they allocate their own buffer instead.
+	bufferMu sync.Mutex
+	buffer   []byte
 }
 
 type Counter struct {
@@ -126,7 +133,20 @@ func (c *Collector[T]) Collect(dst *[]T) error {
 		return errors.New("dst must not be nil")
 	}
 
-	perfObjects, err := QueryPerformanceData(c.query, c.object)
+	c.bufferMu.Lock()
+	buffer := c.buffer
+	c.buffer = nil
+	c.bufferMu.Unlock()
+
+	perfObjects, buffer, err := queryPerformanceData(c.query, c.object, buffer)
+
+	// Keep the larger buffer if a concurrent call put one back in the meantime.
+	c.bufferMu.Lock()
+	if cap(buffer) > cap(c.buffer) {
+		c.buffer = buffer
+	}
+	c.bufferMu.Unlock()
+
 	if err != nil {
 		return fmt.Errorf("QueryPerformanceData: %w", err)
 	}

internal/pdh/registry/nametable.go

@@ -86,7 +86,7 @@ func (t *NameTable) load() error {
 		return nil
 	}
 
-	buffer, err := queryRawData(t.name)
+	buffer, err := queryRawData(t.name, nil)
 	if err != nil {
 		return fmt.Errorf("failed to query perflib name table %q: %w", t.name, err)
 	}

internal/pdh/registry/perflib.go

@@ -204,10 +204,12 @@ var (
 
 // queryRawData Queries the performance counter buffer using RegQueryValueEx, returning raw bytes. See:
 // https://msdn.microsoft.com/de-de/library/windows/desktop/aa373219(v=vs.85).aspx
-func queryRawData(query string) ([]byte, error) {
+//
+// If buffer has enough capacity for the initial size guess, it is reused instead
+// of allocating a new one. The returned slice may share memory with buffer.
+func queryRawData(query string, buffer []byte) ([]byte, error) {
 	var (
 		valType uint32
-		buffer  []byte
 		bufLen  uint32
 	)
 
@@ -223,7 +225,13 @@ func queryRawData(query string) ([]byte, error) {
 		bufLen = uint32(150000 * numCounters)
 	}
 
-	buffer = make([]byte, bufLen)
+	// A buffer from a previous call keeps any growth from ERROR_MORE_DATA, so
+	// using its full capacity avoids growing it again.
+	if uint32(cap(buffer)) >= bufLen {
+		buffer = buffer[:cap(buffer)]
+	} else {
+		buffer = make([]byte, bufLen)
+	}
 
 	name, err := windows.UTF16PtrFromString(query)
 	if err != nil {
@@ -293,22 +301,31 @@ Many objects have dependencies - if you query one of them, you often get back
 more than you asked for.
 */
 func QueryPerformanceData(query string, counterName string) ([]*PerfObject, error) {
+	objects, _, err := queryPerformanceData(query, counterName, nil)
+
+	return objects, err
+}
+
+// queryPerformanceData is QueryPerformanceData with a reusable raw data buffer.
+// It returns the buffer for the next call. The parsed objects copy everything
+// they need out of the buffer, so the buffer can be reused once this returns.
+func queryPerformanceData(query string, counterName string, buffer []byte) ([]*PerfObject, []byte, error) {
 	// Object and counter names are resolved through the name table.
 	if err := CounterNameTable.load(); err != nil {
-		return nil, err
+		return nil, buffer, err
 	}
 
-	buffer, err := queryRawData(query)
+	buffer, err := queryRawData(query, buffer)
 	if err != nil {
-		return nil, err
+		return nil, buffer, err
 	}
 
 	objects, err := parsePerformanceData(buffer, counterName)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse performance data for %q: %w", query, err)
+		return nil, buffer, fmt.Errorf("failed to parse performance data for %q: %w", query, err)
 	}
 
-	return objects, nil
+	return objects, buffer, nil
 }
 
 //nolint:gochecknoglobals

internal/pdh/registry/perflib_parse_test.go

@@ -21,6 +21,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"testing"
+	"unsafe"
 
 	"github.com/prometheus-community/windows_exporter/internal/pdh"
 	"github.com/stretchr/testify/require"
@@ -202,12 +203,57 @@ func TestParsePerformanceDataLive(t *testing.T) {
 	require.NotEmpty(t, objects[0].CounterDefs)
 }
 
+// TestParsePerformanceDataNoAlias checks that the parsed objects do not
+// reference the buffer, which Collector.Collect reuses for the next query.
+func TestParsePerformanceDataNoAlias(t *testing.T) {
+	t.Parallel()
+
+	for _, buffer := range [][]byte{
+		newSyntheticPerfData(true).bytes(t),
+		livePerformanceData(t, "Process"),
+	} {
+		objects, err := parsePerformanceData(buffer, "")
+		require.NoError(t, err)
+
+		expected, err := parsePerformanceData(bytes.Clone(buffer), "")
+		require.NoError(t, err)
+
+		for i := range buffer {
+			buffer[i] = 0xFF
+		}
+
+		require.Equal(t, expected, objects)
+	}
+}
+
+func TestQueryRawDataReusesBuffer(t *testing.T) {
+	t.Parallel()
+
+	require.NoError(t, CounterNameTable.load())
+
+	query := MapCounterToIndex("Process")
+
+	buffer, err := queryRawData(query, nil)
+	require.NoError(t, err)
+
+	reused, err := queryRawData(query, buffer)
+	require.NoError(t, err)
+	require.Same(t, unsafe.SliceData(buffer), unsafe.SliceData(reused))
+
+	// A buffer below the initial size guess is not used.
+	small := make([]byte, 16)
+
+	buffer, err = queryRawData(query, small)
+	require.NoError(t, err)
+	require.NotSame(t, unsafe.SliceData(small), unsafe.SliceData(buffer))
+}
+
 func livePerformanceData(tb testing.TB, object string) []byte {
 	tb.Helper()
 
 	require.NoError(tb, CounterNameTable.load())
 
-	buffer, err := queryRawData(MapCounterToIndex(object))
+	buffer, err := queryRawData(MapCounterToIndex(object), nil)
 	require.NoError(tb, err)
 
 	return buffer

internal/pdh/registry/perflib_test.go

@@ -18,6 +18,7 @@
 package registry
 
 import (
+	"sync"
 	"testing"
 )
 
@@ -59,8 +60,77 @@ func TestCollectNilDestination(t *testing.T) {
 	}
 }
 
+// TestCollectConcurrent runs Collect concurrently on one collector, as
+// overlapping scrapes do, to let the race detector check the buffer reuse.
+func TestCollectConcurrent(t *testing.T) {
+	t.Parallel()
+
+	type systemCounterValues struct {
+		Name string
+
+		ProcessorQueueLength float64 `perfdata:"Processor Queue Length"`
+	}
+
+	collector, err := NewCollector[systemCounterValues]("System", nil)
+	if err != nil {
+		t.Skipf("skipping: failed to create collector: %v", err)
+	}
+
+	var wg sync.WaitGroup
+
+	for range 8 {
+		wg.Go(func() {
+			for range 10 {
+				var dst []systemCounterValues
+
+				if err := collector.Collect(&dst); err != nil {
+					t.Error(err)
+
+					return
+				}
+
+				if len(dst) != 1 {
+					t.Errorf("expected 1 row, got %d", len(dst))
+
+					return
+				}
+			}
+		})
+	}
+
+	wg.Wait()
+}
+
 func BenchmarkQueryPerformanceData(b *testing.B) {
 	for b.Loop() {
 		_, _ = QueryPerformanceData("Global", "")
 	}
 }
+
+func BenchmarkCollectorCollect(b *testing.B) {
+	type processCounterValues struct {
+		Name string
+
+		ProcessID      float64 `perfdata:"ID Process"`
+		ProcessorTime  float64 `perfdata:"% Processor Time"`
+		WorkingSet     float64 `perfdata:"Working Set"`
+		HandleCount    float64 `perfdata:"Handle Count"`
+		ThreadCount    float64 `perfdata:"Thread Count"`
+		IOReadBytesSec float64 `perfdata:"IO Read Bytes/sec"`
+	}
+
+	collector, err := NewCollector[processCounterValues]("Process", nil)
+	if err != nil {
+		b.Skipf("skipping: failed to create collector: %v", err)
+	}
+
+	var dst []processCounterValues
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		if err := collector.Collect(&dst); err != nil {
+			b.Fatal(err)
+		}
+	}
+}