mi: read elements without per-call allocations (#2652)

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

internal/collector/wmi/types.go

@@ -42,6 +42,9 @@ type Query struct {
 type LabelProperty struct {
 	Name  string `json:"name"  yaml:"name"`
 	Label string `json:"label" yaml:"label"`
+
+	// Resolved at Build time.
+	elementName mi.ElementName
 }
 
 type Property struct {
@@ -52,8 +55,9 @@ type Property struct {
 	Labels map[string]string `json:"labels" yaml:"labels"`
 
 	// Resolved at Build time.
-	desc       *prometheus.Desc
-	metricType prometheus.ValueType
+	elementName mi.ElementName
+	desc        *prometheus.Desc
+	metricType  prometheus.ValueType
 }
 
 // UnmarshalYAML decodes the wmi block of the configuration file, where the

internal/collector/wmi/wmi.go

@@ -264,6 +264,12 @@ func (c *Collector) Build(logger *slog.Logger, miSession *mi.Session) error {
 				continue
 			}
 
+			if labelProperty.elementName, err = mi.NewElementName(labelProperty.Name); err != nil {
+				errs = append(errs, fmt.Errorf("query %s: label property %q: %w", query.Name, labelProperty.Name, err))
+
+				continue
+			}
+
 			if labelProperty.Label == "" {
 				labelProperty.Label = sanitizeMetricName(labelProperty.Name)
 			}
@@ -290,6 +296,12 @@ func (c *Collector) Build(logger *slog.Logger, miSession *mi.Session) error {
 				continue
 			}
 
+			if property.elementName, err = mi.NewElementName(property.Name); err != nil {
+				errs = append(errs, fmt.Errorf("query %s: property %q: %w", query.Name, property.Name, err))
+
+				continue
+			}
+
 			if slices.Contains(propertyNames, strings.ToLower(property.Name)) {
 				errs = append(errs, fmt.Errorf("query %s: property %q is duplicated", query.Name, property.Name))
 
@@ -515,7 +527,7 @@ func (c *Collector) collectQuery(ch chan<- prometheus.Metric, query Query, deadl
 		}
 
 		for i, labelProperty := range query.LabelProperties {
-			element, err := instance.GetElement(labelProperty.Name)
+			element, err := instance.GetElementByName(labelProperty.elementName)
 			if err != nil {
 				return fmt.Errorf("failed to read label property %s: %w", labelProperty.Name, err)
 			}
@@ -537,7 +549,7 @@ func (c *Collector) collectQuery(ch chan<- prometheus.Metric, query Query, deadl
 				continue
 			}
 
-			element, err := instance.GetElement(property.Name)
+			element, err := instance.GetElementByName(property.elementName)
 			if err != nil {
 				propertyErrs[property.Name] = fmt.Errorf("failed to read property %s: %w", property.Name, err)
 

internal/mi/application.go

@@ -183,6 +183,7 @@ func (application *Application) NewSession(options *DestinationOptions) (*Sessio
 		return nil, fmt.Errorf("failed to set timeout: %w", err)
 	}
 
+	session.application = application
 	session.defaultOperationOptions = defaultOperationOptions
 
 	return session, nil

internal/mi/errors.go

@@ -20,6 +20,7 @@ package mi
 import "errors"
 
 var (
-	ErrNotInitialized    = errors.New("not initialized")
-	ErrInvalidEntityType = errors.New("invalid entity type")
+	ErrNotInitialized     = errors.New("not initialized")
+	ErrInvalidEntityType  = errors.New("invalid entity type")
+	ErrInvalidElementName = errors.New("invalid element name")
 )

internal/mi/instance.go

@@ -88,14 +88,41 @@ func (instance *Instance) Delete() error {
 	return nil
 }
 
-func (instance *Instance) GetElement(elementName string) (*Element, error) {
+// ElementName is the UTF-16 name of an instance element.
+type ElementName *uint16
+
+func NewElementName(elementName string) (ElementName, error) {
+	return windows.UTF16PtrFromString(elementName)
+}
+
+// GetElement gets the element with the given name. To read the same element
+// from many instances, convert the name once with [NewElementName] and use
+// [Instance.GetElementByName].
+func (instance *Instance) GetElement(elementName string) (Element, error) {
 	if instance == nil || instance.ft == nil {
-		return nil, ErrNotInitialized
+		return Element{}, ErrNotInitialized
 	}
 
-	elementNameUTF16, err := windows.UTF16PtrFromString(elementName)
+	elementNameUTF16, err := NewElementName(elementName)
 	if err != nil {
-		return nil, fmt.Errorf("failed to convert element name %s to UTF-16: %w", elementName, err)
+		return Element{}, fmt.Errorf("failed to convert element name %s to UTF-16: %w", elementName, err)
+	}
+
+	return instance.GetElementByName(elementNameUTF16)
+}
+
+// GetElementByName gets the element with the given name. Strings and arrays
+// in the returned element point into the instance and are only valid while
+// the instance is.
+//
+// https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_instance_getelement
+func (instance *Instance) GetElementByName(elementName ElementName) (Element, error) {
+	if instance == nil || instance.ft == nil {
+		return Element{}, ErrNotInitialized
+	}
+
+	if elementName == nil {
+		return Element{}, ErrInvalidElementName
 	}
 
 	// MI_Value is a union sized to its largest member. On 64-bit MI_Datetime is
@@ -113,7 +140,7 @@ func (instance *Instance) GetElement(elementName string) (*Element, error) {
 	r0, _, _ := syscall.SyscallN(
 		instance.ft.GetElement,
 		uintptr(unsafe.Pointer(instance)),
-		uintptr(unsafe.Pointer(elementNameUTF16)),
+		uintptr(unsafe.Pointer(elementName)),
 		uintptr(unsafe.Pointer(&valueBuf)),
 		uintptr(unsafe.Pointer(&valueType)),
 		uintptr(unsafe.Pointer(&flags)),
@@ -121,10 +148,10 @@ func (instance *Instance) GetElement(elementName string) (*Element, error) {
 	)
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
-		return nil, result
+		return Element{}, result
 	}
 
-	return &Element{
+	return Element{
 		value:     uintptr(valueBuf[0]),
 		arrayLen:  uint32(valueBuf[1]),
 		valueType: valueType,

internal/mi/mi_bench_test.go

@@ -56,3 +56,108 @@ func Benchmark_MI_Query_Unmarshal(b *testing.B) {
 
 	b.ReportAllocs()
 }
+
+func Benchmark_MI_QueryFunc_GetElement(b *testing.B) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(b, err)
+
+	session, err := application.NewSession(nil)
+	require.NoError(b, err)
+
+	b.Cleanup(func() {
+		require.NoError(b, session.Close())
+		require.NoError(b, application.Close())
+	})
+
+	query, err := mi.NewQuery("SELECT Name, PercentIdleTime, PercentProcessorTime, InterruptsPersec FROM Win32_PerfRawData_PerfOS_Processor")
+	require.NoError(b, err)
+
+	properties := []string{"Name", "PercentIdleTime", "PercentProcessorTime", "InterruptsPersec"}
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		err := session.QueryFunc(mi.NamespaceRootCIMv2, query, -1, func(instance *mi.Instance) error {
+			for _, property := range properties {
+				if _, err := instance.GetElement(property); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		})
+		require.NoError(b, err)
+	}
+}
+
+func Benchmark_MI_Query_Unmarshal_Processor(b *testing.B) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(b, err)
+
+	session, err := application.NewSession(nil)
+	require.NoError(b, err)
+
+	b.Cleanup(func() {
+		require.NoError(b, session.Close())
+		require.NoError(b, application.Close())
+	})
+
+	query, err := mi.NewQuery("SELECT Name, PercentIdleTime, PercentProcessorTime, InterruptsPersec FROM Win32_PerfRawData_PerfOS_Processor")
+	require.NoError(b, err)
+
+	var processors []struct {
+		Name                 string `mi:"Name"`
+		PercentIdleTime      uint64 `mi:"PercentIdleTime"`
+		PercentProcessorTime uint64 `mi:"PercentProcessorTime"`
+		InterruptsPersec     uint32 `mi:"InterruptsPersec"`
+	}
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		require.NoError(b, session.Query(&processors, mi.NamespaceRootCIMv2, query, -1))
+	}
+}
+
+// Benchmark_MI_QueryFunc_GetElementByName reads elements by names converted
+// once, as the wmi collector does. The zero timeout creates per-query
+// operation options.
+func Benchmark_MI_QueryFunc_GetElementByName(b *testing.B) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(b, err)
+
+	session, err := application.NewSession(nil)
+	require.NoError(b, err)
+
+	b.Cleanup(func() {
+		require.NoError(b, session.Close())
+		require.NoError(b, application.Close())
+	})
+
+	query, err := mi.NewQuery("SELECT Name, PercentIdleTime, PercentProcessorTime, InterruptsPersec FROM Win32_PerfRawData_PerfOS_Processor")
+	require.NoError(b, err)
+
+	properties := make([]mi.ElementName, 0, 4)
+
+	for _, property := range []string{"Name", "PercentIdleTime", "PercentProcessorTime", "InterruptsPersec"} {
+		name, err := mi.NewElementName(property)
+		require.NoError(b, err)
+
+		properties = append(properties, name)
+	}
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		err := session.QueryFunc(mi.NamespaceRootCIMv2, query, 0, func(instance *mi.Instance) error {
+			for _, property := range properties {
+				if _, err := instance.GetElementByName(property); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		})
+		require.NoError(b, err)
+	}
+}

internal/mi/mi_test.go

@@ -1079,3 +1079,59 @@ func Test_MI_QueryUnmarshal_InvalidClass(t *testing.T) {
 		require.Equal(t, []win32Process{{Name: "System Idle Process"}}, processes)
 	}
 }
+
+func Test_MI_Query_NilSession(t *testing.T) {
+	t.Parallel()
+
+	var (
+		session   *mi.Session
+		processes []win32Process
+	)
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process WHERE Handle = 0")
+	require.NoError(t, err)
+
+	require.ErrorIs(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, 0), mi.ErrNotInitialized)
+	require.ErrorIs(t, session.QueryFunc(mi.NamespaceRootCIMv2, query, 0, func(*mi.Instance) error { return nil }), mi.ErrNotInitialized)
+}
+
+func Test_MI_GetElementByName(t *testing.T) {
+	application, err := mi.ApplicationInitialize()
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		require.NoError(t, application.Close())
+	})
+
+	session, err := application.NewSession(nil)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		require.NoError(t, session.Close())
+	})
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process WHERE Handle = 4")
+	require.NoError(t, err)
+
+	name, err := mi.NewElementName("Name")
+	require.NoError(t, err)
+
+	var names []string
+
+	err = session.QueryFunc(mi.NamespaceRootCIMv2, query, time.Second, func(instance *mi.Instance) error {
+		_, err := instance.GetElementByName(nil)
+		require.ErrorIs(t, err, mi.ErrInvalidElementName)
+
+		element, err := instance.GetElementByName(name)
+		if err != nil {
+			return err
+		}
+
+		value, err := element.String()
+		names = append(names, value)
+
+		return err
+	})
+	require.NoError(t, err)
+	require.Equal(t, []string{"System"}, names)
+}

internal/mi/operation.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 	"unsafe"
@@ -232,8 +233,15 @@ func (o *Operation) Unmarshal[T any](dst *[]T) error {
 type miField struct {
 	index int
 	tag   string
+	name  ElementName
 }
 
+// miFieldsCache maps a struct type to its []miField. Collectors query the same
+// types on every scrape, so the tags are parsed and converted to UTF-16 once.
+//
+//nolint:gochecknoglobals
+var miFieldsCache sync.Map
+
 // prepareUnmarshal checks that dst is non-nil and T is a struct, resets *dst to empty
 // and returns the `mi`-tagged fields of T. Callers run it before starting a query,
 // so invalid input is rejected without a WMI round trip.
@@ -254,8 +262,14 @@ func prepareUnmarshal[T any](dst *[]T) ([]miField, error) {
 
 // miFieldsOf returns the fields of T that carry an `mi` tag.
 // It returns ErrInvalidEntityType if T is not a struct.
+// The returned slice is shared and must not be modified.
 func miFieldsOf[T any]() ([]miField, error) {
 	structType := reflect.TypeFor[T]()
+
+	if fields, ok := miFieldsCache.Load(structType); ok {
+		return fields.([]miField), nil //nolint:forcetypeassert
+	}
+
 	if structType.Kind() != reflect.Struct {
 		return nil, ErrInvalidEntityType
 	}
@@ -263,12 +277,23 @@ func miFieldsOf[T any]() ([]miField, error) {
 	fields := make([]miField, 0, structType.NumField())
 
 	for i := range structType.NumField() {
-		if miTag := structType.Field(i).Tag.Get("mi"); miTag != "" {
-			fields = append(fields, miField{index: i, tag: miTag})
+		miTag := structType.Field(i).Tag.Get("mi")
+		if miTag == "" {
+			continue
 		}
+
+		name, err := NewElementName(miTag)
+		if err != nil {
+			return nil, fmt.Errorf("invalid mi tag %q: %w", miTag, err)
+		}
+
+		fields = append(fields, miField{index: i, tag: miTag, name: name})
 	}
 
-	return fields, nil
+	// Concurrent first queries of a type agree on one shared slice.
+	cached, _ := miFieldsCache.LoadOrStore(structType, fields)
+
+	return cached.([]miField), nil //nolint:forcetypeassert
 }
 
 // unmarshal iterates over the operation's instances and appends them to dst.

internal/mi/session.go

@@ -37,6 +37,11 @@ type Session struct {
 	reserved2 uintptr
 	ft        *SessionFT
 
+	// Go-only state after the MI_Session fields, which MI never writes past.
+	//
+	// application created the session and outlives it, so per-query
+	// operation options are created without MI_Session_GetApplication.
+	application             *Application
 	defaultOperationOptions *OperationOptions
 }
 
@@ -202,7 +207,7 @@ func (s *Session) queryInstances(flags OperationFlags, operationOptions *Operati
 // instead of being silently left at its zero value.
 func unmarshalInstance(instance *Instance, fields []miField, structValue reflect.Value, skipMissing bool) error {
 	for _, f := range fields {
-		element, err := instance.GetElement(f.tag)
+		element, err := instance.GetElementByName(f.name)
 		if err != nil {
 			if skipMissing && errors.Is(err, MI_RESULT_NO_SUCH_PROPERTY) {
 				continue
@@ -211,7 +216,7 @@ func unmarshalInstance(instance *Instance, fields []miField, structValue reflect
 			return fmt.Errorf("failed to get element %s: %w", f.tag, err)
 		}
 
-		if err := setField(f.tag, structValue.Field(f.index), element); err != nil {
+		if err := setField(f.tag, structValue.Field(f.index), &element); err != nil {
 			return err
 		}
 	}
@@ -457,16 +462,15 @@ func (s *Session) QueryFunc(namespaceName Namespace, queryExpression Query, quer
 // options, so the session defaults are used. The caller must delete non-nil
 // options once the operation is closed.
 func (s *Session) newOperationOptions(queryTimeout time.Duration) (*OperationOptions, error) {
+	if s == nil || s.ft == nil {
+		return nil, ErrNotInitialized
+	}
+
 	if queryTimeout < 0 {
 		return nil, nil //nolint:nilnil
 	}
 
-	app, err := s.GetApplication()
-	if err != nil {
-		return nil, fmt.Errorf("failed to get application: %w", err)
-	}
-
-	operationOptions, err := app.NewOperationOptions()
+	operationOptions, err := s.application.NewOperationOptions()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create operation options: %w", err)
 	}

internal/mi/unmarshal_test.go

@@ -23,6 +23,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows"
 )
 
 // newField returns a settable reflect.Value holding the zero value of T.
@@ -112,7 +113,34 @@ func TestPrepareUnmarshal(t *testing.T) {
 		fields, err := prepareUnmarshal(&dst)
 		require.NoError(t, err)
 		require.Empty(t, dst)
-		require.Equal(t, []miField{{index: 0, tag: "Name"}, {index: 2, tag: "ProcessId"}}, fields)
+		require.Len(t, fields, 2)
+		require.Equal(t, 0, fields[0].index)
+		require.Equal(t, "Name", fields[0].tag)
+		require.Equal(t, "Name", windows.UTF16PtrToString(fields[0].name))
+		require.Equal(t, 2, fields[1].index)
+		require.Equal(t, "ProcessId", fields[1].tag)
+		require.Equal(t, "ProcessId", windows.UTF16PtrToString(fields[1].name))
+	})
+
+	t.Run("fields_are_cached", func(t *testing.T) {
+		t.Parallel()
+
+		first, err := prepareUnmarshal(new([]row))
+		require.NoError(t, err)
+
+		second, err := prepareUnmarshal(new([]row))
+		require.NoError(t, err)
+
+		require.Same(t, &first[0], &second[0])
+	})
+
+	t.Run("tag_with_nul", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := prepareUnmarshal(new([]struct {
+			Name string `mi:"Na\x00me"`
+		}))
+		require.Error(t, err)
 	})
 
 	t.Run("nil_pointer", func(t *testing.T) {