pdh/registry: decode performance data without per-counter allocations (#2572)

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

internal/pdh/registry/perflib.go

@@ -128,11 +128,9 @@ Data for some of the objects is also available through WMI:
 */
 
 import (
-	"bytes"
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"io"
 	"strings"
 	"time"
 	"unsafe"
@@ -151,23 +149,20 @@ type PerfObject struct {
 	Name string
 	// NameIndex Same index you pass to QueryPerformanceData
 	NameIndex   uint
-	Instances   []*PerfInstance
-	CounterDefs []*PerfCounterDef
+	Instances   []PerfInstance
+	CounterDefs []PerfCounterDef
 
 	Frequency int64
-
-	rawData *perfObjectType
 }
 
 // PerfInstance Each object can have multiple instances. For example,
 // In case the object has no instances, we return one single PerfInstance with an empty name.
 type PerfInstance struct {
 	// *not* resolved using a name table
-	Name     string
-	Counters []*PerfCounter
-
-	rawData         *perfInstanceDefinition
-	rawCounterBlock *perfCounterBlock
+	Name string
+	// Counters holds one counter per counter definition of the object, in the
+	// same order. The counters of all instances share one backing array.
+	Counters []PerfCounter
 }
 
 type PerfCounterDef struct {
@@ -187,7 +182,7 @@ type PerfCounterDef struct {
 	IsNanosecondCounter bool
 	HasSecondValue      bool
 
-	rawData *perfCounterDefinition
+	rawData perfCounterDefinition
 }
 
 type PerfCounter struct {
@@ -345,15 +340,17 @@ var errMalformedPerformanceData = errors.New("malformed performance data")
 
 // parsePerformanceData parses a PERF_DATA_BLOCK as returned by queryRawData.
 // If counterName is set, only the first object with that name is returned.
+//
+// The structures are decoded straight from buffer. The parsed objects copy
+// everything they keep, so they do not reference buffer.
 func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, error) {
 	bufLen := int64(len(buffer))
-	r := bytes.NewReader(buffer)
 
 	// Read global header
 
-	header := new(perfDataBlock)
+	var header perfDataBlock
 
-	if err := header.BinaryReadFrom(r); err != nil {
+	if err := header.decode(buffer); err != nil {
 		return nil, fmt.Errorf("failed to read performance data block: %w", err)
 	}
 
@@ -375,16 +372,13 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 
 	objects := make([]*PerfObject, numObjects)
 
+	// nameScratch holds the UTF16 code units of an instance name while it is decoded.
+	var nameScratch []uint16
+
 	for i := range numObjects {
-		_, err := r.Seek(objOffset, io.SeekStart)
-		if err != nil {
-			return nil, err
-		}
+		var obj perfObjectType
 
-		obj := new(perfObjectType)
-
-		err = obj.BinaryReadFrom(r)
-		if err != nil {
+		if err := obj.decode(bytesAt(buffer, objOffset)); err != nil {
 			return nil, err
 		}
 
@@ -415,8 +409,17 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 			numInstances = 1
 		}
 
-		instances := make([]*PerfInstance, numInstances)
-		counterDefs := make([]*PerfCounterDef, numCounterDefs)
+		// The counters of all instances share one backing array below, which
+		// grows with instances × counters. Counter values do not overlap, so
+		// each takes at least one byte of the object. Checking that keeps the
+		// allocation linear in the buffer size for malformed data.
+		if numInstances*numCounterDefs > bufLen-objOffset {
+			return nil, fmt.Errorf("%w: %d instances with %d counters each do not fit into the buffer", errMalformedPerformanceData, numInstances, numCounterDefs)
+		}
+
+		instances := make([]PerfInstance, numInstances)
+		counterDefs := make([]PerfCounterDef, numCounterDefs)
+		counters := make([]PerfCounter, numInstances*numCounterDefs)
 
 		objects[i] = &PerfObject{
 			Name:        perfCounterName,
@@ -424,18 +427,19 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 			Instances:   instances,
 			CounterDefs: counterDefs,
 			Frequency:   obj.PerfFreq,
-			rawData:     obj,
 		}
 
-		for i := range numCounterDefs {
-			def := new(perfCounterDefinition)
+		// The counter definitions follow the object header.
+		defOffset := objOffset + perfObjectTypeSize
 
-			err := def.BinaryReadFrom(r)
-			if err != nil {
+		for j := range numCounterDefs {
+			var def perfCounterDefinition
+
+			if err := def.decode(bytesAt(buffer, defOffset+j*perfCounterDefinitionSize)); err != nil {
 				return nil, err
 			}
 
-			counterDefs[i] = &PerfCounterDef{
+			counterDefs[j] = PerfCounterDef{
 				Name:      def.LookupName(),
 				NameIndex: uint(def.CounterNameTitleIndex),
 				rawData:   def,
@@ -452,32 +456,18 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 		if obj.NumInstances <= 0 { //nolint:nestif
 			blockOffset := objOffset + int64(obj.DefinitionLength)
 
-			if _, err := r.Seek(blockOffset, io.SeekStart); err != nil {
+			if _, err := parseCounterBlock(buffer, blockOffset, counterDefs, counters); err != nil {
 				return nil, err
 			}
 
-			_, counters, err := parseCounterBlock(buffer, r, blockOffset, counterDefs)
-			if err != nil {
-				return nil, err
-			}
-
-			instances[0] = &PerfInstance{
-				Name:            "",
-				Counters:        counters,
-				rawData:         nil,
-				rawCounterBlock: nil,
-			}
+			instances[0].Counters = counters
 		} else {
 			instOffset := objOffset + int64(obj.DefinitionLength)
 
-			for i := range numInstances {
-				if _, err := r.Seek(instOffset, io.SeekStart); err != nil {
-					return nil, err
-				}
+			for j := range numInstances {
+				var inst perfInstanceDefinition
 
-				inst := new(perfInstanceDefinition)
-
-				if err = inst.BinaryReadFrom(r); err != nil {
+				if err := inst.decode(bytesAt(buffer, instOffset)); err != nil {
 					return nil, err
 				}
 
@@ -486,21 +476,18 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 					return nil, fmt.Errorf("%w: instance name at offset %d with length %d exceeds the buffer", errMalformedPerformanceData, namePos, inst.NameLength)
 				}
 
-				name, _ := readUTF16StringAtPos(r, namePos, inst.NameLength)
-				pos := instOffset + int64(inst.ByteLength)
+				instances[j].Name, nameScratch = decodeUTF16String(buffer[namePos:namePos+int64(inst.NameLength)], nameScratch)
 
-				offset, counters, err := parseCounterBlock(buffer, r, pos, counterDefs)
+				blockOffset := instOffset + int64(inst.ByteLength)
+				instCounters := counters[j*numCounterDefs : (j+1)*numCounterDefs : (j+1)*numCounterDefs]
+
+				blockLength, err := parseCounterBlock(buffer, blockOffset, counterDefs, instCounters)
 				if err != nil {
 					return nil, err
 				}
 
-				instances[i] = &PerfInstance{
-					Name:     name,
-					Counters: counters,
-					rawData:  inst,
-				}
-
-				instOffset = pos + offset
+				instances[j].Counters = instCounters
+				instOffset = blockOffset + blockLength
 			}
 		}
 
@@ -516,30 +503,35 @@ func parsePerformanceData(buffer []byte, counterName string) ([]*PerfObject, err
 	return objects[:numFilteredObjects], nil
 }
 
-func parseCounterBlock(b []byte, r io.ReadSeeker, pos int64, defs []*PerfCounterDef) (int64, []*PerfCounter, error) {
-	_, err := r.Seek(pos, io.SeekStart)
-	if err != nil {
-		return 0, nil, err
+// bytesAt returns b from offset pos on, or nil if pos lies outside of b.
+func bytesAt(b []byte, pos int64) []byte {
+	if pos < 0 || pos > int64(len(b)) {
+		return nil
 	}
 
-	block := new(perfCounterBlock)
+	return b[pos:]
+}
 
-	err = block.BinaryReadFrom(r)
-	if err != nil {
-		return 0, nil, err
+// parseCounterBlock reads the PERF_COUNTER_BLOCK at pos into counters, which
+// has one element per definition in defs. It returns the length of the block.
+func parseCounterBlock(b []byte, pos int64, defs []PerfCounterDef, counters []PerfCounter) (int64, error) {
+	var block perfCounterBlock
+
+	if err := block.decode(bytesAt(b, pos)); err != nil {
+		return 0, err
 	}
 
 	// The block length includes the length field itself. A shorter block would
 	// make the next instance overlap this one.
 	if int64(block.ByteLength) < perfCounterBlockSize {
-		return 0, nil, fmt.Errorf("%w: counter block at offset %d has length %d", errMalformedPerformanceData, pos, block.ByteLength)
+		return 0, fmt.Errorf("%w: counter block at offset %d has length %d", errMalformedPerformanceData, pos, block.ByteLength)
 	}
 
-	counters := make([]*PerfCounter, len(defs))
+	for i := range defs {
+		def := &defs[i]
 
-	for i, def := range defs {
 		valueOffset := pos + int64(def.rawData.CounterOffset)
-		size := counterValueSize(def.rawData)
+		size := counterValueSize(&def.rawData)
 
 		end := valueOffset + size
 		if def.HasSecondValue {
@@ -547,24 +539,24 @@ func parseCounterBlock(b []byte, r io.ReadSeeker, pos int64, defs []*PerfCounter
 		}
 
 		if end > int64(len(b)) {
-			return 0, nil, fmt.Errorf("%w: value of counter %s at offset %d exceeds the buffer", errMalformedPerformanceData, def.Name, valueOffset)
+			return 0, fmt.Errorf("%w: value of counter %s at offset %d exceeds the buffer", errMalformedPerformanceData, def.Name, valueOffset)
 		}
 
-		value := convertCounterValue(def.rawData, b, valueOffset)
+		value := convertCounterValue(&def.rawData, b, valueOffset)
 		secondValue := int64(0)
 
 		if def.HasSecondValue {
-			secondValue = convertCounterValue(def.rawData, b, valueOffset+8)
+			secondValue = convertCounterValue(&def.rawData, b, valueOffset+8)
 		}
 
-		counters[i] = &PerfCounter{
+		counters[i] = PerfCounter{
 			Value:       value,
 			Def:         def,
 			SecondValue: secondValue,
 		}
 	}
 
-	return int64(block.ByteLength), counters, nil
+	return int64(block.ByteLength), nil
 }
 
 // counterValueSize returns the number of bytes convertCounterValue reads.

internal/pdh/registry/perflib_parse_test.go

@@ -20,6 +20,7 @@ package registry
 import (
 	"bytes"
 	"encoding/binary"
+	"io"
 	"testing"
 	"unsafe"
 
@@ -190,6 +191,73 @@ func TestParsePerformanceDataTruncated(t *testing.T) {
 	}
 }
 
+// sharedValuePerfData returns a consistent PERF_DATA_BLOCK with one object
+// whose counters all read the same 4 bytes of each minimal counter block.
+func sharedValuePerfData(tb testing.TB, numCounters, numInstances uint32) []byte {
+	tb.Helper()
+
+	var body bytes.Buffer
+
+	write := func(v any) {
+		require.NoError(tb, binary.Write(&body, bo, v))
+	}
+
+	for range numCounters {
+		write(perfCounterDefinition{
+			ByteLength:  uint32(perfCounterDefinitionSize),
+			CounterType: pdh.PERF_COUNTER_RAWCOUNT,
+			CounterSize: 4,
+		})
+	}
+
+	for range numInstances {
+		write(perfInstanceDefinition{ByteLength: uint32(perfInstanceDefinitionSize)})
+		write(perfCounterBlock{ByteLength: uint32(perfCounterBlockSize)})
+	}
+
+	object := perfObjectType{
+		TotalByteLength:  uint32(perfObjectTypeSize) + uint32(body.Len()),
+		DefinitionLength: uint32(perfObjectTypeSize) + numCounters*uint32(perfCounterDefinitionSize),
+		HeaderLength:     uint32(perfObjectTypeSize),
+		NumCounters:      numCounters,
+		NumInstances:     int32(numInstances),
+	}
+
+	header := perfDataBlock{
+		Signature:       [4]uint16{'P', 'E', 'R', 'F'},
+		LittleEndian:    1,
+		TotalByteLength: uint32(perfDataBlockSize) + object.TotalByteLength,
+		HeaderLength:    uint32(perfDataBlockSize),
+		NumObjectTypes:  1,
+	}
+
+	var buf bytes.Buffer
+
+	require.NoError(tb, binary.Write(&buf, bo, header))
+	require.NoError(tb, binary.Write(&buf, bo, object))
+	buf.Write(body.Bytes())
+
+	return buf.Bytes()
+}
+
+// TestParsePerformanceDataCounterProduct checks that the counters of all
+// instances, which share one allocation, are bounded by the buffer size. Data
+// whose counter values overlap could otherwise request an allocation that
+// grows with the square of the buffer size.
+func TestParsePerformanceDataCounterProduct(t *testing.T) {
+	t.Parallel()
+
+	objects, err := parsePerformanceData(sharedValuePerfData(t, 200, 10), "")
+	require.NoError(t, err)
+	require.Len(t, objects[0].Instances, 10)
+	require.Len(t, objects[0].Instances[9].Counters, 200)
+	require.Equal(t, perfCounterBlockSize, objects[0].Instances[9].Counters[199].Value)
+
+	_, err = parsePerformanceData(sharedValuePerfData(t, 200, 200), "")
+	require.ErrorIs(t, err, errMalformedPerformanceData)
+	require.ErrorContains(t, err, "200 instances with 200 counters each")
+}
+
 // TestParsePerformanceDataLive parses the data of a real object.
 func TestParsePerformanceDataLive(t *testing.T) {
 	t.Parallel()
@@ -309,3 +377,49 @@ func TestNewCollectorUnknownObject(t *testing.T) {
 	_, err := NewCollector[counterValues]("Object That Does Not Exist", nil)
 	require.ErrorContains(t, err, "not found in the counter name table")
 }
+
+// BenchmarkParsePerformanceData measures the parser alone, without the
+// RegQueryValueEx call that dominates BenchmarkCollectorCollect.
+func BenchmarkParsePerformanceData(b *testing.B) {
+	buffer := livePerformanceData(b, "Process")
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		if _, err := parsePerformanceData(buffer, "Process"); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// TestDecodeMatchesBinaryRead checks the hand-written decoders of the raw
+// structures against binary.Read.
+func TestDecodeMatchesBinaryRead(t *testing.T) {
+	t.Parallel()
+
+	type decoder interface {
+		decode(b []byte) error
+	}
+
+	for _, tc := range []struct {
+		decoded, expected decoder
+	}{
+		{new(perfDataBlock), new(perfDataBlock)},
+		{new(perfObjectType), new(perfObjectType)},
+		{new(perfCounterDefinition), new(perfCounterDefinition)},
+		{new(perfCounterBlock), new(perfCounterBlock)},
+		{new(perfInstanceDefinition), new(perfInstanceDefinition)},
+	} {
+		// Distinct bytes, so a field read from the wrong offset gets a different value.
+		buffer := make([]byte, binary.Size(tc.expected))
+		for i := range buffer {
+			buffer[i] = byte(i + 1)
+		}
+
+		require.NoError(t, binary.Read(bytes.NewReader(buffer), bo, tc.expected))
+		require.NoError(t, tc.decoded.decode(buffer))
+		require.Equal(t, tc.expected, tc.decoded)
+
+		require.ErrorIs(t, tc.decoded.decode(buffer[:len(buffer)-1]), io.ErrUnexpectedEOF)
+	}
+}

internal/pdh/registry/raw_types.go

@@ -18,7 +18,6 @@
 package registry
 
 import (
-	"encoding/binary"
 	"io"
 
 	"golang.org/x/sys/windows"
@@ -63,8 +62,40 @@ type perfDataBlock struct {
 	SystemNameOffset uint32
 }
 
-func (p *perfDataBlock) BinaryReadFrom(r io.Reader) error {
-	return binary.Read(r, bo, p)
+// decode reads p from the start of b. The fields are decoded directly, since
+// binary.Read allocates and walks the struct by reflection on every call.
+func (p *perfDataBlock) decode(b []byte) error {
+	if int64(len(b)) < perfDataBlockSize {
+		return io.ErrUnexpectedEOF
+	}
+
+	*p = perfDataBlock{
+		Signature:       [4]uint16{bo.Uint16(b[0:]), bo.Uint16(b[2:]), bo.Uint16(b[4:]), bo.Uint16(b[6:])},
+		LittleEndian:    bo.Uint32(b[8:]),
+		Version:         bo.Uint32(b[12:]),
+		Revision:        bo.Uint32(b[16:]),
+		TotalByteLength: bo.Uint32(b[20:]),
+		HeaderLength:    bo.Uint32(b[24:]),
+		NumObjectTypes:  bo.Uint32(b[28:]),
+		DefaultObject:   int32(bo.Uint32(b[32:])),
+		SystemTime: windows.Systemtime{
+			Year:         bo.Uint16(b[36:]),
+			Month:        bo.Uint16(b[38:]),
+			DayOfWeek:    bo.Uint16(b[40:]),
+			Day:          bo.Uint16(b[42:]),
+			Hour:         bo.Uint16(b[44:]),
+			Minute:       bo.Uint16(b[46:]),
+			Second:       bo.Uint16(b[48:]),
+			Milliseconds: bo.Uint16(b[50:]),
+		},
+		PerfTime:         int64(bo.Uint64(b[56:])),
+		PerfFreq:         int64(bo.Uint64(b[64:])),
+		PerfTime100nSec:  int64(bo.Uint64(b[72:])),
+		SystemNameLength: bo.Uint32(b[80:]),
+		SystemNameOffset: bo.Uint32(b[84:]),
+	}
+
+	return nil
 }
 
 /*
@@ -105,8 +136,30 @@ type perfObjectType struct {
 	PerfFreq             int64
 }
 
-func (p *perfObjectType) BinaryReadFrom(r io.Reader) error {
-	return binary.Read(r, bo, p)
+// decode reads p from the start of b.
+func (p *perfObjectType) decode(b []byte) error {
+	if int64(len(b)) < perfObjectTypeSize {
+		return io.ErrUnexpectedEOF
+	}
+
+	*p = perfObjectType{
+		TotalByteLength:      bo.Uint32(b[0:]),
+		DefinitionLength:     bo.Uint32(b[4:]),
+		HeaderLength:         bo.Uint32(b[8:]),
+		ObjectNameTitleIndex: bo.Uint32(b[12:]),
+		ObjectNameTitle:      bo.Uint32(b[16:]),
+		ObjectHelpTitleIndex: bo.Uint32(b[20:]),
+		ObjectHelpTitle:      bo.Uint32(b[24:]),
+		DetailLevel:          bo.Uint32(b[28:]),
+		NumCounters:          bo.Uint32(b[32:]),
+		DefaultCounter:       int32(bo.Uint32(b[36:])),
+		NumInstances:         int32(bo.Uint32(b[40:])),
+		CodePage:             bo.Uint32(b[44:]),
+		PerfTime:             int64(bo.Uint64(b[48:])),
+		PerfFreq:             int64(bo.Uint64(b[56:])),
+	}
+
+	return nil
 }
 
 /*
@@ -139,8 +192,26 @@ type perfCounterDefinition struct {
 	CounterOffset         uint32
 }
 
-func (p *perfCounterDefinition) BinaryReadFrom(r io.Reader) error {
-	return binary.Read(r, bo, p)
+// decode reads p from the start of b.
+func (p *perfCounterDefinition) decode(b []byte) error {
+	if int64(len(b)) < perfCounterDefinitionSize {
+		return io.ErrUnexpectedEOF
+	}
+
+	*p = perfCounterDefinition{
+		ByteLength:            bo.Uint32(b[0:]),
+		CounterNameTitleIndex: bo.Uint32(b[4:]),
+		CounterNameTitle:      bo.Uint32(b[8:]),
+		CounterHelpTitleIndex: bo.Uint32(b[12:]),
+		CounterHelpTitle:      bo.Uint32(b[16:]),
+		DefaultScale:          int32(bo.Uint32(b[20:])),
+		DetailLevel:           bo.Uint32(b[24:]),
+		CounterType:           bo.Uint32(b[28:]),
+		CounterSize:           bo.Uint32(b[32:]),
+		CounterOffset:         bo.Uint32(b[36:]),
+	}
+
+	return nil
 }
 
 func (p *perfCounterDefinition) LookupName() string {
@@ -159,8 +230,15 @@ type perfCounterBlock struct {
 	ByteLength uint32
 }
 
-func (p *perfCounterBlock) BinaryReadFrom(r io.Reader) error {
-	return binary.Read(r, bo, p)
+// decode reads p from the start of b.
+func (p *perfCounterBlock) decode(b []byte) error {
+	if int64(len(b)) < perfCounterBlockSize {
+		return io.ErrUnexpectedEOF
+	}
+
+	p.ByteLength = bo.Uint32(b)
+
+	return nil
 }
 
 /*
@@ -185,6 +263,20 @@ type perfInstanceDefinition struct {
 	NameLength             uint32
 }
 
-func (p *perfInstanceDefinition) BinaryReadFrom(r io.Reader) error {
-	return binary.Read(r, bo, p)
+// decode reads p from the start of b.
+func (p *perfInstanceDefinition) decode(b []byte) error {
+	if int64(len(b)) < perfInstanceDefinitionSize {
+		return io.ErrUnexpectedEOF
+	}
+
+	*p = perfInstanceDefinition{
+		ByteLength:             bo.Uint32(b[0:]),
+		ParentObjectTitleIndex: bo.Uint32(b[4:]),
+		ParentObjectInstance:   bo.Uint32(b[8:]),
+		UniqueID:               bo.Uint32(b[12:]),
+		NameOffset:             bo.Uint32(b[16:]),
+		NameLength:             bo.Uint32(b[20:]),
+	}
+
+	return nil
 }

internal/pdh/registry/utf16.go

@@ -18,27 +18,23 @@
 package registry
 
 import (
-	"encoding/binary"
 	"io"
 
 	"golang.org/x/sys/windows"
 )
 
-// readUTF16StringAtPos Read an unterminated UTF16 string at a given position, specifying its length.
-func readUTF16StringAtPos(r io.ReadSeeker, absPos int64, length uint32) (string, error) {
-	value := make([]uint16, length/2)
+// decodeUTF16String decodes the unterminated UTF16 string b up to the first NUL.
+// A trailing odd byte is ignored. scratch holds the code units; it is returned
+// with any growth so the caller can reuse it for the next string. The returned
+// string does not reference b.
+func decodeUTF16String(b []byte, scratch []uint16) (string, []uint16) {
+	scratch = scratch[:0]
 
-	_, err := r.Seek(absPos, io.SeekStart)
-	if err != nil {
-		return "", err
+	for i := 0; i+1 < len(b); i += 2 {
+		scratch = append(scratch, bo.Uint16(b[i:]))
 	}
 
-	err = binary.Read(r, bo, value)
-	if err != nil {
-		return "", err
-	}
-
-	return windows.UTF16ToString(value), nil
+	return windows.UTF16ToString(scratch), scratch
 }
 
 // readUTF16String Reads a null-terminated UTF16 string at the current offset.