From 96fdb6f7792ac9fe390eddc33c60a979f18c9650 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan-Otto=20Kr=C3=B6pke?= Date: Sat, 10 Oct 2026 23:15:42 +0200 Subject: [PATCH] mi: deliver query results through MI callbacks (#2666) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Claude Opus 5.5 Signed-off-by: Jan-Otto Kröpke --- internal/mi/AGENTS.md | 41 ++++ internal/mi/callbacks.go | 438 ++++++++++++++++++++++++++++++++++ internal/mi/callbacks_test.go | 276 +++++++++++++++++++++ internal/mi/instance.go | 19 ++ internal/mi/mi_bench_test.go | 230 ++++++++++++++++++ internal/mi/operation.go | 106 ++++---- internal/mi/session.go | 55 ++--- 7 files changed, 1093 insertions(+), 72 deletions(-) create mode 100644 internal/mi/AGENTS.md create mode 100644 internal/mi/callbacks.go create mode 100644 internal/mi/callbacks_test.go diff --git a/internal/mi/AGENTS.md b/internal/mi/AGENTS.md new file mode 100644 index 00000000..a5367fe6 --- /dev/null +++ b/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. diff --git a/internal/mi/callbacks.go b/internal/mi/callbacks.go new file mode 100644 index 00000000..526da631 --- /dev/null +++ b/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() +} diff --git a/internal/mi/callbacks_test.go b/internal/mi/callbacks_test.go new file mode 100644 index 00000000..eabcaff9 --- /dev/null +++ b/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) +} diff --git a/internal/mi/instance.go b/internal/mi/instance.go index d2b6cedb..a7881d3e 100644 --- a/internal/mi/instance.go +++ b/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 diff --git a/internal/mi/mi_bench_test.go b/internal/mi/mi_bench_test.go index 7d9088ac..115fc8d3 100644 --- a/internal/mi/mi_bench_test.go +++ b/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") + }) + } +} diff --git a/internal/mi/operation.go b/internal/mi/operation.go index e2f348a3..0679c657 100644 --- a/internal/mi/operation.go +++ b/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 diff --git a/internal/mi/session.go b/internal/mi/session.go index f8624bd4..5b4e11d7 100644 --- a/internal/mi/session.go +++ b/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