mi: deliver query results through MI callbacks (#2666)

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

internal/mi/AGENTS.md

@@ -0,0 +1,41 @@
+# MI package guidance
+
+`Session.Query`, `QueryUnmarshal` and `QueryFunc` receive results through
+`MI_OperationCallbacks` ([`callbacks.go`](callbacks.go)). Do not move them back
+to the synchronous `MI_Operation_GetInstance` loop: when its result hand-over
+waits, e.g. for parallel queries or slow providers such as `MSFT_StoragePool`,
+the WMI client leaks Event handles (`miutils!RtlInterlockedCompareWait`).
+`Test_MI_ParallelQuery_HandleGrowth` guards this. `QueryInstances` and
+`Operation.GetInstance` remain as the synchronous low-level API for tests.
+
+Rules for the callback path:
+
+- Create native callbacks once (`sync.OnceValue`) and identify the query by
+  the callback context. `windows.NewCallback` never frees a callback and a
+  process can create only about 2000.
+- Keep the callback context and the `MI_OperationCallbacks` pinned until
+  `MI_Operation_Close` has returned.
+- Never call `Cancel` or `Close` from a callback. The goroutine that started
+  the query cancels on request and closes after the final callback
+  (`moreResults == MI_FALSE`).
+- Recover panics in callbacks; they must not unwind into MI.
+- Never block a callback on a goroutine. `QueryFunc` copies each instance
+  (`MI_Instance_Clone`) in the callback and runs `fn` for the copies on the
+  calling goroutine. Handing the original over and waiting for `fn` costs two
+  thread wake-ups per instance, about 25% more process CPU than the copy
+  (`Benchmark_MI_ProcessCPU`).
+- MI reports parameter errors from within `MI_Session_QueryInstances`; the
+  runtime runs that callback on the calling goroutine. Never run caller code
+  such as the `QueryFunc` handler inside a callback: a panic cannot reach the
+  caller there and `runtime.Goexit` is fatal.
+- While waiting for callbacks, keep a timer pending. Otherwise the runtime can
+  declare a false deadlock (golang/go#55015, still open in Go 1.27);
+  `Test_MI_Query_DeadlockDetector` reproduces it. Do not use a goroutine that
+  sleeps forever.
+- Guard shared callback state with a mutex: results of one operation may
+  arrive on different MI threads. `-race` cannot reveal a missing lock there,
+  because the runtime treats every callback as synchronized with every native
+  call.
+
+Run `CGO_ENABLED=1 go test -race` (gcc from w64devkit) and a
+`-gcflags=all=-d=checkptr` run for changes in this package.

internal/mi/callbacks.go

@@ -0,0 +1,438 @@
+// 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 mi
+
+import (
+	"errors"
+	"fmt"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"time"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+// Queries deliver their results through MI_OperationCallbacks instead of the
+// synchronous MI_Operation_GetInstance. With GetInstance, the WMI client hands
+// every result from its delivery thread to the reading thread, and creates
+// Event handles it never closes whenever that hand-over has to wait, e.g. for
+// parallel queries or slow providers. A callback receives the result on the
+// delivery thread, so there is no such hand-over.
+//
+// The callback runs on MI threads, which are foreign to the Go runtime:
+//
+//   - It is created once. windows.NewCallback never frees a callback and
+//     supports only a limited number of them per process.
+//   - The asyncQuery, which holds the MI_OperationCallbacks and is the
+//     callback context, is pinned until the operation is closed.
+//   - It never calls into the operation. Cancel and Close are called by the
+//     goroutine that started the query, after the final callback for Close.
+//   - It never runs caller code and never blocks on a goroutine. A panic in
+//     it is recovered, as it must not unwind into native code.
+//
+// MI reports some failures, e.g. invalid parameters, from within
+// MI_Session_QueryInstances. The runtime runs such a callback on the
+// goroutine that started the query.
+
+// operationCallbacks is MI_OperationCallbacks.
+//
+// https://learn.microsoft.com/en-us/windows/win32/api/mi/ns-mi-mi_operationcallbacks
+type operationCallbacks struct {
+	callbackContext         unsafe.Pointer
+	promptUser              uintptr
+	writeError              uintptr
+	writeMessage            uintptr
+	writeProgress           uintptr
+	instanceResult          uintptr
+	indicationResult        uintptr
+	classResult             uintptr
+	streamedParameterResult uintptr
+}
+
+// instanceResultCallback is the MI_OperationCallback_Instance of all queries.
+// The callback context identifies the query.
+//
+// https://learn.microsoft.com/en-us/windows/win32/api/mi/nc-mi-mi_operationcallback_instance
+//
+//nolint:gochecknoglobals
+var instanceResultCallback = sync.OnceValue(func() uintptr {
+	return windows.NewCallback(func(
+		_ *Operation,
+		query *asyncQuery,
+		instance *Instance,
+		moreResults Boolean,
+		resultCode ResultError,
+		errorMessage *uint16,
+		errorDetails *Instance,
+		_ uintptr, // resultAcknowledgement, only set with MI_OPERATIONFLAGS_MANUAL_ACK_RESULTS.
+	) uintptr {
+		query.instanceResult(instance, moreResults == True, resultCode, errorMessage, errorDetails)
+
+		return 0
+	})
+})
+
+// deadlockGuardInterval is the period of the timer a query keeps pending
+// while it waits for its callbacks.
+//
+// The runtime does not know that an MI thread will call back, so while every
+// goroutine waits for MI, it declares a deadlock and exits
+// (https://github.com/golang/go/issues/55015). A pending timer suppresses
+// that check. The interval does not bound the query; MI timeouts do.
+const deadlockGuardInterval = time.Minute
+
+// guardTimers holds stopped deadlock guard timers for reuse.
+//
+//nolint:gochecknoglobals
+var guardTimers = sync.Pool{
+	New: func() any {
+		timer := time.NewTimer(deadlockGuardInterval)
+		timer.Stop()
+
+		return timer
+	},
+}
+
+// newGuardTimer returns a running deadlock guard timer. Return it with
+// releaseGuardTimer.
+func newGuardTimer() *time.Timer {
+	timer := guardTimers.Get().(*time.Timer) //nolint:forcetypeassert
+	timer.Reset(deadlockGuardInterval)
+
+	return timer
+}
+
+func releaseGuardTimer(timer *time.Timer) {
+	// Since Go 1.23, a stopped timer delivers no stale value after Reset.
+	timer.Stop()
+	guardTimers.Put(timer)
+}
+
+// errStopQuery cancels a query without an error.
+var errStopQuery = errors.New("stop query")
+
+// callbackPanicError reports a panic recovered in an MI callback. The value
+// is formatted by Error, outside the callback, because formatting can panic
+// as well.
+type callbackPanicError struct {
+	value any
+}
+
+func (e *callbackPanicError) Error() string {
+	return fmt.Sprintf("panic in MI instance callback: %v", e.value)
+}
+
+// asyncQuery is the state of one MI_Session_QueryInstances operation that
+// delivers its results to instanceResultCallback. It is pinned while MI may
+// use it.
+type asyncQuery struct {
+	callbacks operationCallbacks
+	operation Operation
+	pinner    runtime.Pinner
+
+	// onInstance is called for every instance, inside the callback. An error
+	// cancels the query and is returned by wait; errStopQuery cancels it
+	// without an error.
+	onInstance func(*Instance) error
+
+	// handler is called on the waiting goroutine for copies of the instances,
+	// which the callbacks queue in pending and announce on queued. It is nil
+	// unless created by newCopyQuery. batch is reused for pending; only the
+	// waiting goroutine uses it.
+	handler func(*Instance) error
+	pending []*Instance
+	batch   []*Instance
+	queued  chan struct{}
+	// handlerErr is the error handler returned, if any. Only the waiting
+	// goroutine uses it. It takes precedence over err: the handler only sees
+	// instances that arrived before any failed result.
+	handlerErr error
+
+	// cancelRequested asks the waiting goroutine to cancel the operation.
+	cancelRequested chan struct{}
+	// done is closed by the final callback. MI does not use the query any
+	// more afterward, and err is final. finished guards against closing it
+	// twice, which would panic into MI.
+	done     chan struct{}
+	finished atomic.Bool
+
+	// mu serializes the callbacks. MI delivers the results of an operation
+	// one at a time, but possibly on different threads, and its own ordering
+	// is invisible to Go. The race detector cannot point out a missing lock
+	// here: it treats every callback as synchronized with every native call.
+	// mu also guards pending against the waiting goroutine.
+	mu  sync.Mutex
+	err error
+}
+
+// newAsyncQuery returns a query that calls onInstance inside the callback.
+func newAsyncQuery(onInstance func(*Instance) error) *asyncQuery {
+	return &asyncQuery{
+		onInstance:      onInstance,
+		cancelRequested: make(chan struct{}, 1),
+		done:            make(chan struct{}),
+	}
+}
+
+// newCopyQuery returns a query that calls handler on the goroutine running
+// [asyncQuery.wait], for a copy of every instance. The callback queues the
+// copy and returns right away; the copy is deleted when handler returns.
+func newCopyQuery(handler func(*Instance) error) *asyncQuery {
+	query := newAsyncQuery(nil)
+	query.onInstance = query.queueCopy
+	query.handler = handler
+	query.queued = make(chan struct{}, 1)
+
+	return query
+}
+
+// start starts the query as an MI_Session_QueryInstances operation. The
+// caller must call close afterward. Nil operationOptions fall back to the
+// session defaults.
+//
+// MI_Session_QueryInstances returns void: a failed query still yields an
+// operation, and its status, e.g. a timeout or an invalid class, is reported
+// by the final callback.
+func (q *asyncQuery) start(session *Session, flags OperationFlags, operationOptions *OperationOptions,
+	namespaceName Namespace, queryDialect QueryDialect, queryExpression Query,
+) {
+	q.callbacks.callbackContext = unsafe.Pointer(q)
+	q.callbacks.instanceResult = instanceResultCallback()
+
+	q.pinner.Pin(q)
+
+	if operationOptions == nil {
+		operationOptions = session.defaultOperationOptions
+	}
+
+	_, _, _ = syscall.SyscallN(
+		session.ft.QueryInstances,
+		uintptr(unsafe.Pointer(session)),
+		uintptr(flags),
+		uintptr(unsafe.Pointer(operationOptions)),
+		uintptr(unsafe.Pointer(namespaceName)),
+		uintptr(unsafe.Pointer(queryDialect)),
+		uintptr(unsafe.Pointer(queryExpression)),
+		uintptr(unsafe.Pointer(&q.callbacks)),
+		uintptr(unsafe.Pointer(&q.operation)),
+	)
+}
+
+// instanceResult handles one MI_OperationCallback_Instance call.
+func (q *asyncQuery) instanceResult(instance *Instance, moreResults bool, resultCode ResultError,
+	errorMessage *uint16, errorDetails *Instance,
+) {
+	q.mu.Lock()
+
+	defer func() {
+		// A panic must not unwind into MI. It fails the query instead.
+		if r := recover(); r != nil {
+			q.fail(&callbackPanicError{value: r})
+		}
+
+		q.mu.Unlock()
+
+		// The final result is the last call MI makes with this context.
+		if !moreResults && q.finished.CompareAndSwap(false, true) {
+			close(q.done)
+		}
+	}()
+
+	if !errors.Is(resultCode, MI_RESULT_OK) {
+		// The first error wins. A query cancelled because of an error ends
+		// with a failed result, which must not replace that error.
+		if q.err == nil {
+			q.err = fmt.Errorf("failed to get instance: %w", instanceResultError(resultCode, errorMessage, errorDetails))
+		}
+
+		return
+	}
+
+	// After an error or an early stop, the remaining results are skipped
+	// until the cancellation takes effect.
+	if q.err != nil || instance == nil {
+		return
+	}
+
+	if err := q.onInstance(instance); err != nil {
+		q.fail(err)
+	}
+}
+
+// fail records err and asks the waiting goroutine to cancel the operation.
+// q.mu must be held.
+func (q *asyncQuery) fail(err error) {
+	if q.err == nil {
+		q.err = err
+	}
+
+	select {
+	case q.cancelRequested <- struct{}{}:
+	default:
+	}
+}
+
+// queueCopy queues a copy of instance for the handler. q.mu must be held.
+func (q *asyncQuery) queueCopy(instance *Instance) error {
+	clone, err := instance.clone()
+	if err != nil {
+		return fmt.Errorf("failed to copy instance: %w", err)
+	}
+
+	q.pending = append(q.pending, clone)
+
+	select {
+	case q.queued <- struct{}{}:
+	default:
+	}
+
+	return nil
+}
+
+// wait blocks until the final callback and returns the query result. It
+// cancels the operation when a callback asks for it, and runs the handler of
+// a copy query.
+func (q *asyncQuery) wait() error {
+	guard := newGuardTimer()
+	defer releaseGuardTimer(guard)
+
+	for {
+		select {
+		case <-q.done:
+			// Copies queued by the last callbacks.
+			q.handleQueued()
+
+			if q.handlerErr != nil {
+				return q.handlerErr
+			}
+
+			if errors.Is(q.err, errStopQuery) {
+				return nil
+			}
+
+			return q.err
+		case <-q.cancelRequested:
+			_ = q.operation.Cancel()
+		case <-q.queued:
+			q.handleQueued()
+		case <-guard.C:
+			guard.Reset(deadlockGuardInterval)
+		}
+	}
+}
+
+// handleQueued calls handler for the queued copies and deletes them. After a
+// handler error, it only deletes them. If handler panics or exits the
+// goroutine, the copies it has not handled are deleted by close.
+func (q *asyncQuery) handleQueued() {
+	for {
+		q.mu.Lock()
+		batch := q.pending
+		q.pending, q.batch = q.batch, nil
+		q.mu.Unlock()
+
+		if len(batch) == 0 {
+			q.batch = batch
+
+			return
+		}
+
+		// Until handled, the copies stay in q.batch for close.
+		q.batch = batch
+
+		for i, instance := range batch {
+			if q.handlerErr == nil {
+				if err := q.handler(instance); err != nil {
+					q.handlerErr = err
+
+					// Stops the callbacks from copying further instances.
+					q.mu.Lock()
+					q.fail(err)
+					q.mu.Unlock()
+				}
+			}
+
+			_ = instance.Delete()
+			batch[i] = nil
+		}
+
+		q.batch = batch[:0]
+	}
+}
+
+// deleteQueued deletes the copies that were not handled. The callbacks must
+// have finished.
+func (q *asyncQuery) deleteQueued() {
+	for _, instances := range [][]*Instance{q.batch, q.pending} {
+		for _, instance := range instances {
+			if instance != nil {
+				_ = instance.Delete()
+			}
+		}
+	}
+
+	q.batch, q.pending = nil, nil
+}
+
+// close cancels the operation unless it has finished, waits for the final
+// callback, deletes unhandled copies, closes the operation and unpins the
+// query. The goroutine that started the query must call it, also when wait
+// did not return because a handler panicked.
+func (q *asyncQuery) close() {
+	select {
+	case <-q.done:
+	default:
+		_ = q.operation.Cancel()
+
+		guard := newGuardTimer()
+		defer releaseGuardTimer(guard)
+
+		for finished := false; !finished; {
+			select {
+			case <-q.done:
+				finished = true
+			case <-guard.C:
+				guard.Reset(deadlockGuardInterval)
+			}
+		}
+	}
+
+	q.deleteQueued()
+
+	// MI_Operation_Close blocks until the final result has been delivered,
+	// which has happened.
+	_ = q.operation.close()
+
+	q.pinner.Unpin()
+}
+
+// query runs a query whose instances are passed to onInstance inside the MI
+// callback.
+func (s *Session) query(flags OperationFlags, operationOptions *OperationOptions, namespaceName Namespace,
+	queryDialect QueryDialect, queryExpression Query, onInstance func(*Instance) error,
+) error {
+	query := newAsyncQuery(onInstance)
+
+	query.start(s, flags, operationOptions, namespaceName, queryDialect, queryExpression)
+	defer query.close()
+
+	return query.wait()
+}

internal/mi/callbacks_test.go

@@ -0,0 +1,276 @@
+// 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 mi_test
+
+import (
+	"errors"
+	"os"
+	"os/exec"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/prometheus-community/windows_exporter/internal/mi"
+	"github.com/prometheus-community/windows_exporter/internal/utils/testutils"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows"
+)
+
+type msftStoragePool struct {
+	FriendlyName string `mi:"FriendlyName"`
+}
+
+// newTestSession returns a session that is closed with the test.
+func newTestSession(tb testing.TB) *mi.Session {
+	tb.Helper()
+
+	application, err := mi.ApplicationInitialize()
+	require.NoError(tb, err)
+
+	tb.Cleanup(func() { require.NoError(tb, application.Close()) })
+
+	session, err := application.NewSession(nil)
+	require.NoError(tb, err)
+
+	tb.Cleanup(func() { require.NoError(tb, session.Close()) })
+
+	return session
+}
+
+// storagePoolQueryTimeout is generous: the storage provider is slow, and CI
+// runs several test binaries at once. The tests check handles, not latency.
+const storagePoolQueryTimeout = 30 * time.Second
+
+// storagePoolQuery returns the MSFT_StoragePool query, or skips if the storage
+// provider is not installed or not accessible, e.g. without elevation. Other
+// errors fail the test.
+func storagePoolQuery(tb testing.TB, session *mi.Session) mi.Query {
+	tb.Helper()
+
+	query, err := mi.NewQuery("SELECT FriendlyName FROM MSFT_StoragePool")
+	require.NoError(tb, err)
+
+	var pools []msftStoragePool
+
+	err = session.Query(&pools, mi.NamespaceRootStorage, query, storagePoolQueryTimeout)
+	if errors.Is(err, mi.MI_RESULT_INVALID_NAMESPACE) || errors.Is(err, mi.MI_RESULT_INVALID_CLASS) ||
+		errors.Is(err, mi.MI_RESULT_NOT_SUPPORTED) || errors.Is(err, mi.MI_RESULT_ACCESS_DENIED) {
+		tb.Skipf("MSFT_StoragePool is not available: %v", err)
+	}
+
+	require.NoError(tb, err)
+
+	return query
+}
+
+// Test_MI_ParallelQuery_HandleGrowth runs parallel queries against the slow
+// MSFT_StoragePool provider. With MI_Operation_GetInstance, the WMI client
+// leaked Event handles whenever its result hand-over had to wait, which
+// parallel queries against a slow provider do: about one per three queries.
+func Test_MI_ParallelQuery_HandleGrowth(t *testing.T) {
+	const (
+		workers           = 8
+		queriesPerWorker  = 50
+		rounds            = 4
+		maxHandleIncrease = 100
+	)
+
+	session := newTestSession(t)
+	query := storagePoolQuery(t, session)
+
+	round := func() {
+		var wg sync.WaitGroup
+
+		for range workers {
+			wg.Go(func() {
+				for i := range queriesPerWorker {
+					// Both query paths deliver results through the callback.
+					if i%2 == 0 {
+						var pools []msftStoragePool
+
+						if err := session.Query(&pools, mi.NamespaceRootStorage, query, storagePoolQueryTimeout); err != nil {
+							t.Error(err)
+						}
+
+						continue
+					}
+
+					err := session.QueryFunc(mi.NamespaceRootStorage, query, storagePoolQueryTimeout, func(*mi.Instance) error {
+						return nil
+					})
+					if err != nil {
+						t.Error(err)
+					}
+				}
+			})
+		}
+
+		wg.Wait()
+	}
+
+	handleCount := func() int64 {
+		count, err := testutils.GetProcessHandleCount(windows.CurrentProcess())
+		require.NoError(t, err)
+
+		return int64(count)
+	}
+
+	// The first parallel round grows MI and Go thread pools.
+	round()
+
+	baseline := handleCount()
+
+	for range rounds {
+		round()
+	}
+
+	current := handleCount()
+
+	t.Logf("handle count after warm-up %d, after %d queries %d", baseline, workers*queriesPerWorker*rounds, current)
+
+	require.LessOrEqual(t, current, baseline+maxHandleIncrease,
+		"handle count grew from %d to %d after %d parallel queries", baseline, current, workers*queriesPerWorker*rounds)
+}
+
+// Test_MI_QueryFunc_Panic checks that a panic in fn reaches the caller, after
+// the query has been cancelled and closed, and leaves the session usable.
+func Test_MI_QueryFunc_Panic(t *testing.T) {
+	session := newTestSession(t)
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	for range 3 {
+		var calls int
+
+		require.PanicsWithValue(t, "fn panicked", func() {
+			_ = session.QueryFunc(mi.NamespaceRootCIMv2, query, 5*time.Second, func(*mi.Instance) error {
+				calls++
+
+				panic("fn panicked")
+			})
+		})
+		require.Equal(t, 1, calls)
+	}
+
+	var processes []win32Process
+
+	require.NoError(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, 5*time.Second))
+	require.NotEmpty(t, processes)
+}
+
+// Test_MI_QueryFunc_Timeout checks that a timeout reaches QueryFunc as the
+// native MI result.
+func Test_MI_QueryFunc_Timeout(t *testing.T) {
+	session := newTestSession(t)
+
+	query, err := mi.NewQuery("SELECT * FROM Win32_Process")
+	require.NoError(t, err)
+
+	err = session.QueryFunc(mi.NamespaceRootCIMv2, query, time.Millisecond, func(*mi.Instance) error {
+		return nil
+	})
+	require.ErrorIs(t, err, mi.MI_RESULT_INVALID_OPERATION_TIMEOUT)
+}
+
+// Test_MI_QueryFunc_Nested runs a query from within fn, which runs on the
+// calling goroutine. The instance passed to fn stays valid across it.
+func Test_MI_QueryFunc_Nested(t *testing.T) {
+	session := newTestSession(t)
+
+	outer, err := mi.NewQuery("SELECT Name FROM Win32_Process WHERE Handle = 0 OR Handle = 4")
+	require.NoError(t, err)
+
+	inner, err := mi.NewQuery("SELECT Name FROM Win32_Process WHERE Handle = 4")
+	require.NoError(t, err)
+
+	var calls int
+
+	err = session.QueryFunc(mi.NamespaceRootCIMv2, outer, 5*time.Second, func(instance *mi.Instance) error {
+		calls++
+
+		var processes []win32Process
+
+		if err := session.Query(&processes, mi.NamespaceRootCIMv2, inner, 5*time.Second); err != nil {
+			return err
+		}
+
+		if len(processes) != 1 || processes[0].Name != "System" {
+			return errors.New("unexpected inner result")
+		}
+
+		// The outer instance is still valid.
+		element, err := instance.GetElement("Name")
+		if err != nil {
+			return err
+		}
+
+		_, err = element.String()
+
+		return err
+	})
+	require.NoError(t, err)
+	require.Equal(t, 2, calls)
+}
+
+// Test_MI_Query_InvalidParameter checks a failure MI reports from within
+// MI_Session_QueryInstances, on the calling goroutine.
+func Test_MI_Query_InvalidParameter(t *testing.T) {
+	session := newTestSession(t)
+
+	var processes []win32Process
+
+	require.ErrorIs(t, session.Query(&processes, mi.NamespaceRootCIMv2, nil, 5*time.Second), mi.MI_RESULT_INVALID_PARAMETER)
+
+	err := session.QueryFunc(nil, nil, 5*time.Second, func(*mi.Instance) error {
+		return errors.New("unexpected instance")
+	})
+	require.ErrorIs(t, err, mi.MI_RESULT_INVALID_PARAMETER)
+}
+
+// Test_MI_Query_DeadlockDetector runs a query in a test binary without the
+// test timeout, so no timer is pending while the only goroutine waits for MI
+// callbacks. The runtime must not declare a deadlock
+// (https://github.com/golang/go/issues/55015).
+func Test_MI_Query_DeadlockDetector(t *testing.T) {
+	if os.Getenv("MI_TEST_DEADLOCK_DETECTOR") == "1" {
+		session := newTestSession(t)
+
+		query, err := mi.NewQuery("SELECT * FROM Win32_Process")
+		require.NoError(t, err)
+
+		for range 5 {
+			var processes []win32Process
+
+			require.NoError(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, -1))
+			require.NoError(t, session.QueryFunc(mi.NamespaceRootCIMv2, query, -1, func(*mi.Instance) error { return nil }))
+		}
+
+		return
+	}
+
+	executable, err := os.Executable()
+	require.NoError(t, err)
+
+	cmd := exec.CommandContext(t.Context(), executable, "-test.run=^Test_MI_Query_DeadlockDetector$", "-test.timeout=0", "-test.count=1")
+
+	cmd.Env = append(os.Environ(), "MI_TEST_DEADLOCK_DETECTOR=1")
+
+	output, err := cmd.CombinedOutput()
+	require.NoError(t, err, "%s", output)
+}

internal/mi/instance.go

@@ -74,6 +74,25 @@ type ClassDecl struct {
 	OwningClass uintptr
 }
 
+// clone returns a copy of the instance that stays valid until it is deleted.
+//
+// https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_instance_clone
+func (instance *Instance) clone() (*Instance, error) {
+	if instance == nil || instance.ft == nil {
+		return nil, ErrNotInitialized
+	}
+
+	var clone *Instance
+
+	r0, _, _ := syscall.SyscallN(instance.ft.Clone, uintptr(unsafe.Pointer(instance)), uintptr(unsafe.Pointer(&clone)))
+
+	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
+		return nil, result
+	}
+
+	return clone, nil
+}
+
 func (instance *Instance) Delete() error {
 	if instance == nil || instance.ft == nil {
 		return ErrNotInitialized

internal/mi/mi_bench_test.go

@@ -18,10 +18,16 @@
 package mi_test
 
 import (
+	"errors"
+	"runtime"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/prometheus-community/windows_exporter/internal/mi"
+	"github.com/prometheus-community/windows_exporter/internal/utils/testutils"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows"
 )
 
 func Benchmark_MI_Query_Unmarshal(b *testing.B) {
@@ -161,3 +167,227 @@ func Benchmark_MI_QueryFunc_GetElementByName(b *testing.B) {
 		require.NoError(b, err)
 	}
 }
+
+// The Sync benchmarks read the same results with MI_Operation_GetInstance,
+// the path the high-level queries used before they switched to callbacks.
+
+func Benchmark_MI_Query_Unmarshal_Sync(b *testing.B) {
+	session := newTestSession(b)
+
+	b.ReportAllocs()
+
+	var processes []win32Process
+
+	for b.Loop() {
+		operation, err := session.QueryInstances(mi.OperationFlagsStandardRTTI, nil, mi.NamespaceRootCIMv2, mi.QueryDialectWQL,
+			"SELECT Name FROM Win32_Process WHERE Handle = 0 OR Handle = 4")
+		require.NoError(b, err)
+		require.NoError(b, operation.Unmarshal(&processes))
+		require.NoError(b, operation.Close())
+		require.Equal(b, []win32Process{{Name: "System Idle Process"}, {Name: "System"}}, processes)
+	}
+}
+
+func Benchmark_MI_QueryFunc_GetElement_Sync(b *testing.B) {
+	session := newTestSession(b)
+
+	properties := []string{"Name", "PercentIdleTime", "PercentProcessorTime", "InterruptsPersec"}
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		operation, err := session.QueryInstances(mi.OperationFlagsStandardRTTI, nil, mi.NamespaceRootCIMv2, mi.QueryDialectWQL,
+			"SELECT Name, PercentIdleTime, PercentProcessorTime, InterruptsPersec FROM Win32_PerfRawData_PerfOS_Processor")
+		require.NoError(b, err)
+
+		for {
+			instance, moreResults, err := operation.GetInstance()
+			require.NoError(b, err)
+
+			if instance == nil {
+				break
+			}
+
+			for _, property := range properties {
+				_, err := instance.GetElement(property)
+				require.NoError(b, err)
+			}
+
+			if !moreResults {
+				break
+			}
+		}
+
+		require.NoError(b, operation.Close())
+	}
+}
+
+// Benchmark_MI_Parallel_StoragePool runs parallel queries against a slow
+// provider and reports the handle count growth per 1000 queries. Both modes
+// use the session's default operation options.
+func Benchmark_MI_Parallel_StoragePool(b *testing.B) {
+	for _, mode := range []string{"Async", "Sync"} {
+		// b.Run calls the function once per b.N step, so the session is set
+		// up and warmed up once per mode outside of it.
+		session := newTestSession(b)
+		query := storagePoolQuery(b, session)
+
+		queryOnce := func() error {
+			var pools []msftStoragePool
+
+			if mode == "Async" {
+				return session.Query(&pools, mi.NamespaceRootStorage, query, -1)
+			}
+
+			operation, err := session.QueryInstances(mi.OperationFlagsStandardRTTI, nil, mi.NamespaceRootStorage, mi.QueryDialectWQL,
+				"SELECT FriendlyName FROM MSFT_StoragePool")
+			if err != nil {
+				return err
+			}
+
+			err = operation.Unmarshal(&pools)
+
+			return errors.Join(err, operation.Close())
+		}
+
+		// Grow the MI and Go thread pools before the handle baseline.
+		var wg sync.WaitGroup
+
+		for range 2 * runtime.GOMAXPROCS(0) {
+			wg.Go(func() {
+				for range 4 {
+					if err := queryOnce(); err != nil {
+						b.Error(err)
+					}
+				}
+			})
+		}
+
+		wg.Wait()
+
+		b.Run(mode, func(b *testing.B) {
+			start, err := testutils.GetProcessHandleCount(windows.CurrentProcess())
+			require.NoError(b, err)
+
+			b.ReportAllocs()
+			b.SetParallelism(2)
+			b.ResetTimer()
+
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					if err := queryOnce(); err != nil {
+						b.Error(err)
+
+						return
+					}
+				}
+			})
+
+			b.StopTimer()
+
+			end, err := testutils.GetProcessHandleCount(windows.CurrentProcess())
+			require.NoError(b, err)
+
+			b.ReportMetric(float64(int64(end)-int64(start))*1000/float64(b.N), "handles/1k-queries")
+		})
+	}
+}
+
+// processCPU returns the CPU time of the process, including MI threads.
+func processCPU(tb testing.TB) time.Duration {
+	tb.Helper()
+
+	var creation, exit, kernel, user windows.Filetime
+
+	require.NoError(tb, windows.GetProcessTimes(windows.CurrentProcess(), &creation, &exit, &kernel, &user))
+
+	return time.Duration(kernel.Nanoseconds() + user.Nanoseconds())
+}
+
+// Benchmark_MI_ProcessCPU compares the process CPU time per query of the
+// callback paths with MI_Operation_GetInstance. GetProcessTimes counts in
+// 15.6ms steps, so run it with a fixed -benchtime of a few thousand queries.
+func Benchmark_MI_ProcessCPU(b *testing.B) {
+	const queryString = "SELECT Name, PercentProcessorTime FROM Win32_PerfRawData_PerfOS_Processor"
+
+	type processor struct {
+		Name                 string `mi:"Name"`
+		PercentProcessorTime uint64 `mi:"PercentProcessorTime"`
+	}
+
+	query, err := mi.NewQuery(queryString)
+	require.NoError(b, err)
+
+	name, err := mi.NewElementName("PercentProcessorTime")
+	require.NoError(b, err)
+
+	readSync := func(session *mi.Session, read func(*mi.Operation) error) error {
+		operation, err := session.QueryInstances(mi.OperationFlagsStandardRTTI, nil, mi.NamespaceRootCIMv2, mi.QueryDialectWQL, queryString)
+		if err != nil {
+			return err
+		}
+
+		return errors.Join(read(operation), operation.Close())
+	}
+
+	cases := []struct {
+		name  string
+		query func(*mi.Session) error
+	}{
+		{"Query/Async", func(session *mi.Session) error {
+			var processors []processor
+
+			return session.Query(&processors, mi.NamespaceRootCIMv2, query, -1)
+		}},
+		{"Query/Sync", func(session *mi.Session) error {
+			var processors []processor
+
+			return readSync(session, func(operation *mi.Operation) error { return operation.Unmarshal(&processors) })
+		}},
+		{"QueryFunc/Async", func(session *mi.Session) error {
+			return session.QueryFunc(mi.NamespaceRootCIMv2, query, -1, func(instance *mi.Instance) error {
+				_, err := instance.GetElementByName(name)
+
+				return err
+			})
+		}},
+		{"QueryFunc/Sync", func(session *mi.Session) error {
+			return readSync(session, func(operation *mi.Operation) error {
+				for {
+					instance, moreResults, err := operation.GetInstance()
+					if err != nil || instance == nil {
+						return err
+					}
+
+					if _, err := instance.GetElementByName(name); err != nil {
+						return err
+					}
+
+					if !moreResults {
+						return nil
+					}
+				}
+			})
+		}},
+	}
+
+	for _, tc := range cases {
+		b.Run(tc.name, func(b *testing.B) {
+			session := newTestSession(b)
+
+			for range 20 {
+				require.NoError(b, tc.query(session))
+			}
+
+			b.ReportAllocs()
+
+			start := processCPU(b)
+
+			for b.Loop() {
+				require.NoError(b, tc.query(session))
+			}
+
+			b.ReportMetric(float64((processCPU(b)-start).Microseconds())/float64(b.N), "cpu-us/op")
+		})
+	}
+}

internal/mi/operation.go

@@ -114,6 +114,12 @@ func (o *Operation) Close() error {
 		}
 	}
 
+	return o.close()
+}
+
+// close calls MI_Operation_Close, which blocks until the final result has
+// been delivered.
+func (o *Operation) close() error {
 	r0, _, _ := syscall.SyscallN(o.ft.Close, uintptr(unsafe.Pointer(o)))
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
@@ -160,37 +166,11 @@ func (o *Operation) GetInstance() (*Instance, bool, error) {
 		uintptr(unsafe.Pointer(&errorDetails)),
 	)
 
-	//nolint:nestif
 	if !errors.Is(instanceResult, MI_RESULT_OK) {
-		errorMessage := strings.TrimSpace(windows.UTF16PtrToString(errorMessageUTF16))
-
-		// We need a language neutral way to detect an operation timeout, because MI_RESULT_OPERATION_TIMED_OUT
-		// is not returned by the API, but instead we get MI_RESULT_INVALID_OPERATION_TIMEOUT with a specific error code
-		// in the error details.
-		if errorDetails != nil {
-			count, _ := errorDetails.GetElementCount()
-			if count != 0 {
-				errorCodeRaw, err := errorDetails.GetElement("error_Code")
-				if err == nil {
-					errorCodeValue, _ := errorCodeRaw.GetValue()
-
-					errorCode, ok := errorCodeValue.(uint32)
-					if ok && errorCode == 262148 {
-						instanceResult = MI_RESULT_INVALID_OPERATION_TIMEOUT
-						errorMessage = ""
-					}
-				}
-			}
-		}
-
-		if errorMessage != "" {
-			errorMessage = fmt.Sprintf(" (%s)", errorMessage)
-		}
-
 		// A failed result is always the final one.
 		o.completed = true
 
-		return nil, false, fmt.Errorf("instance result: %w%s", instanceResult, errorMessage)
+		return nil, false, instanceResultError(instanceResult, errorMessageUTF16, errorDetails)
 	}
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
@@ -204,6 +184,37 @@ func (o *Operation) GetInstance() (*Instance, bool, error) {
 	return instance, moreResults == True, nil
 }
 
+// instanceResultError returns the error for a failed instance result.
+func instanceResultError(instanceResult ResultError, errorMessageUTF16 *uint16, errorDetails *Instance) error {
+	errorMessage := strings.TrimSpace(windows.UTF16PtrToString(errorMessageUTF16))
+
+	// We need a language neutral way to detect an operation timeout, because MI_RESULT_OPERATION_TIMED_OUT
+	// is not returned by the API, but instead we get MI_RESULT_INVALID_OPERATION_TIMEOUT with a specific error code
+	// in the error details.
+	//nolint:nestif
+	if errorDetails != nil {
+		count, _ := errorDetails.GetElementCount()
+		if count != 0 {
+			errorCodeRaw, err := errorDetails.GetElement("error_Code")
+			if err == nil {
+				errorCodeValue, _ := errorCodeRaw.GetValue()
+
+				errorCode, ok := errorCodeValue.(uint32)
+				if ok && errorCode == 262148 {
+					instanceResult = MI_RESULT_INVALID_OPERATION_TIMEOUT
+					errorMessage = ""
+				}
+			}
+		}
+	}
+
+	if errorMessage != "" {
+		errorMessage = fmt.Sprintf(" (%s)", errorMessage)
+	}
+
+	return fmt.Errorf("instance result: %w%s", instanceResult, errorMessage)
+}
+
 func (o *Operation) Unmarshal[T any](dst *[]T) error {
 	if o == nil || o.ft == nil {
 		return ErrNotInitialized
@@ -299,23 +310,14 @@ func (o *Operation) unmarshal[T any](dst *[]T, fields []miField, skipMissing boo
 			break
 		}
 
-		counter, err := instance.GetElementCount()
-		if err != nil {
-			return fmt.Errorf("failed to get element count: %w", err)
-		}
+		if err := appendInstance(dst, instance, fields, skipMissing); err != nil {
+			if errors.Is(err, errStopQuery) {
+				break
+			}
 
-		if counter == 0 {
-			break
-		}
-
-		var elem T
-
-		if err := unmarshalInstance(instance, fields, reflect.ValueOf(&elem).Elem(), skipMissing); err != nil {
 			return err
 		}
 
-		*dst = append(*dst, elem)
-
 		if !moreResults {
 			break
 		}
@@ -324,6 +326,30 @@ func (o *Operation) unmarshal[T any](dst *[]T, fields []miField, skipMissing boo
 	return nil
 }
 
+// appendInstance unmarshals instance and appends it to dst. An instance
+// without elements ends the result set; appendInstance then returns
+// errStopQuery.
+func appendInstance[T any](dst *[]T, instance *Instance, fields []miField, skipMissing bool) error {
+	counter, err := instance.GetElementCount()
+	if err != nil {
+		return fmt.Errorf("failed to get element count: %w", err)
+	}
+
+	if counter == 0 {
+		return errStopQuery
+	}
+
+	var elem T
+
+	if err := unmarshalInstance(instance, fields, reflect.ValueOf(&elem).Elem(), skipMissing); err != nil {
+		return err
+	}
+
+	*dst = append(*dst, elem)
+
+	return nil
+}
+
 func (o *OperationOptions) SetTimeout(timeout time.Duration) error {
 	if o == nil || o.ft == nil {
 		return ErrNotInitialized

internal/mi/session.go

@@ -162,6 +162,9 @@ func (s *Session) GetApplication() (*Application, error) {
 }
 
 // QueryInstances queries for a set of instances based on a query expression.
+// The results are read synchronously with [Operation.GetInstance]. Prefer
+// [Session.Query] or [Session.QueryFunc], which receive them through callbacks;
+// the synchronous path leaks WMI client handles under parallel load.
 //
 // https://learn.microsoft.com/en-us/windows/win32/api/mi/nf-mi-mi_session_queryinstances
 func (s *Session) QueryInstances(flags OperationFlags, operationOptions *OperationOptions, namespaceName Namespace,
@@ -391,13 +394,9 @@ func (s *Session) QueryUnmarshal[T any](dst *[]T,
 		return err
 	}
 
-	operation := s.queryInstances(flags, operationOptions, namespaceName, queryDialect, queryExpression)
-
-	defer func() {
-		_ = operation.Close()
-	}()
-
-	return operation.unmarshal(dst, fields, true)
+	return s.query(flags, operationOptions, namespaceName, queryDialect, queryExpression, func(instance *Instance) error {
+		return appendInstance(dst, instance, fields, true)
+	})
 }
 
 // Query queries for a set of instances based on a query expression.
@@ -420,9 +419,18 @@ func (s *Session) Query[T any](dst *[]T, namespaceName Namespace, queryExpressio
 // order. It is meant for queries whose result shape is only known at runtime,
 // where [Session.Query] cannot unmarshal into a struct.
 //
-// The instance, and every element read from it, is owned by the operation and
-// only valid until fn returns. If fn returns an error, the query is cancelled
-// and that error is returned.
+// The instance passed to fn, and every element read from it, is only valid
+// until fn returns. If fn returns an error, the query is cancelled and that
+// error is returned.
+//
+// fn runs on the calling goroutine, for a copy of each instance: MI delivers
+// the instances to a callback on one of its threads, which copies them and
+// returns right away. A panic in fn cancels the query and propagates to the
+// caller once MI has finished the operation.
+//
+// MI does not wait for fn. With a slow fn, copies of the remaining result set
+// queue up in native memory, and the query timeout does not stop fn; check a
+// deadline in fn where that matters.
 func (s *Session) QueryFunc(namespaceName Namespace, queryExpression Query, queryTimeout time.Duration, fn func(*Instance) error) error {
 	if s == nil || s.ft == nil {
 		return ErrNotInitialized
@@ -440,31 +448,14 @@ func (s *Session) QueryFunc(namespaceName Namespace, queryExpression Query, quer
 		}()
 	}
 
-	operation := s.queryInstances(OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
+	query := newCopyQuery(fn)
 
-	// Close cancels the operation if fn aborted it early.
-	defer func() {
-		_ = operation.Close()
-	}()
+	query.start(s, OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
 
-	for {
-		instance, moreResults, err := operation.GetInstance()
-		if err != nil {
-			return fmt.Errorf("failed to get instance: %w", err)
-		}
+	// close cancels the operation if fn aborted it early or panicked.
+	defer query.close()
 
-		if instance == nil {
-			return nil
-		}
-
-		if err := fn(instance); err != nil {
-			return err
-		}
-
-		if !moreResults {
-			return nil
-		}
-	}
+	return query.wait()
 }
 
 // newOperationOptions creates operation options carrying queryTimeout. A zero