mi: bound every query with a timeout and add context cancellation (#2738)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

AGENTS.md

@@ -130,8 +130,9 @@ exercised a live Windows role.
   the original `Collection` owns collectors and MI resources. A timed-out call
   may still be running: prevent overlap, drain late metrics and do not free its
   resources while it runs. Honor the supplied scrape budget in blocking APIs.
-  `Build` has no budget: bound its MI probe queries with `mi.BuildQueryTimeout`,
-  never `0` (no MI timeout), so one stalled provider cannot block startup.
+  `Build` has no budget: bound its MI probe queries with `mi.BuildQueryTimeout`.
+  A non-positive MI query timeout selects `mi.DefaultQueryTimeout` (5s), which
+  is too short for slow providers.
 - Collectors that reuse metrics or label pairs between scrapes keep an
   `internal/metriccache` or `internal/labelcache` field and list it in
   `ScrapeCaches()` (`internal/scrapecache`). The collection resets these caches

UPGRADE.md

@@ -24,7 +24,7 @@ Most installations need no changes. Go through the checklist below and read the
 | You set scrape timeouts above 5 minutes or below the timeout margin                             | Review the new limits                                                   | [Scrape timeouts](#scrape-timeouts)                               |
 | You use the sample Grafana dashboard                                                            | Use Grafana 13 or later and import the new dashboard                    | [Dashboard and alerts](#dashboard-and-alerts)                     |
 | You install with the MSI                                                                        | Note the one-time process stop during the upgrade                       | [Installer](#installer)                                           |
-| You embed `pkg/collector` as a Go library                                                       | Rename `Config.Filetime` to `Config.File`                               | [Library users](#library-users)                                   |
+| You embed `pkg/collector` as a Go library                                                       | Rename `Config.Filetime`; pass a non-zero `Collect` timeout             | [Library users](#library-users)                                   |
 | You build from source                                                                           | Use Go 1.27.2                                                           | [Building from source](#building-from-source)                     |
 
 ## filetime collector removed
@@ -438,6 +438,7 @@ These changes affect Go programs that embed `github.com/prometheus-community/win
 - Views created by `Collection.WithCollectors` no longer own the collectors. Their `Close` does nothing; close the original `Collection`.
 - `Collector.Collect` is never called concurrently on one collector instance, and not at all after its `Build` failed.
 - A collector can return an error wrapping `errors.ErrUnsupported` from `Build` to mark itself unavailable on the host. `Collection.Build` logs a warning and continues.
+- WMI (MI) queries always have a timeout. Calling `Collector.Collect` with a `maxScrapeDuration` of 0 or less used to run the queries of MI-based collectors, such as `wmi`, `cpu_info` or `diskdrive`, without a timeout; they now time out after 5 seconds. Pass the scrape budget you want instead of 0.
 
 ## Building from source
 

internal/collector/wmi/wmi.go

@@ -516,7 +516,7 @@ func (c *Collector) Collect(ch chan<- prometheus.Metric, maxScrapeDuration time.
 }
 
 func (c *Collector) collectQuery(ch chan<- prometheus.Metric, scrape *labelcache.Scrape, query Query, deadline time.Time) error {
-	// Zero means no timeout, which leaves the MI default in place.
+	// Zero selects mi.DefaultQueryTimeout.
 	var timeout time.Duration
 
 	if !deadline.IsZero() {

internal/mi/AGENTS.md

@@ -37,5 +37,31 @@ Rules for the callback path:
   because the runtime treats every callback as synchronized with every native
   call.
 
+Query timeouts and options:
+
+- Every query runs with an MI timeout. MI treats a zero timeout as unset, so
+  a non-positive timeout selects `DefaultQueryTimeout` (the session default
+  options) and a positive one below 1ms is raised to 1ms.
+- Operation options must outlive every operation that uses them. Options for
+  whole-second timeouts are cached per session and shared by concurrent
+  operations; `Session.Close` deletes them after `MI_Session_Close` has
+  closed every operation. Other timeouts, e.g. the remaining scrape budget,
+  get per-query options that are deleted after the operation is closed.
+- `startQuery` holds the read lock of `Session.lifecycle` from the `closed`
+  check until `MI_Session_QueryInstances` has returned; `Close` holds the
+  write lock across `MI_Session_Close` and the option deletion. Start every
+  callback query through `runQuery`/`startQuery`, never by calling
+  `asyncQuery.start` with session options directly.
+
+Cancellation: `QueryContext` and `QueryFuncContext` do not start a query
+for a done context. When it is done during the query, `wait` fails the query
+through `fail`/`cancelRequested`, so `Cancel` runs on the waiting goroutine,
+and still waits for the final callback before `close` frees anything. The
+handler is not called after the context is done; a copy it skips fails the
+query with the context's cause, even after the final result, so partial
+results never look complete. A context that is done after the final result
+(`complete`, set under `mu`) does not fail a query that has nothing left to
+hand over.
+
 Run `CGO_ENABLED=1 go test -race` (gcc from w64devkit) and a
 `-gcflags=all=-d=checkptr` run for changes in this package.

internal/mi/application.go

@@ -176,7 +176,7 @@ func (application *Application) NewSession(options *DestinationOptions) (*Sessio
 		return nil, fmt.Errorf("failed to create default operation options: %w", err)
 	}
 
-	if err = defaultOperationOptions.SetTimeout(5 * time.Second); err != nil {
+	if err = defaultOperationOptions.SetTimeout(DefaultQueryTimeout); err != nil {
 		_ = defaultOperationOptions.Delete()
 		_ = session.Close()
 

internal/mi/callbacks.go

@@ -18,6 +18,7 @@
 package mi
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"runtime"
@@ -166,6 +167,9 @@ type asyncQuery struct {
 	// instances that arrived before any failed result.
 	handlerErr error
 
+	// ctx is the context passed to wait. Only the waiting goroutine uses it.
+	ctx context.Context //nolint:containedctx // lives for the duration of one wait.
+
 	// 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
@@ -181,6 +185,8 @@ type asyncQuery struct {
 	// mu also guards pending against the waiting goroutine.
 	mu  sync.Mutex
 	err error
+	// complete is set under mu by the final callback, before done is closed.
+	complete bool
 }
 
 // newAsyncQuery returns a query that calls onInstance inside the callback.
@@ -248,6 +254,10 @@ func (q *asyncQuery) instanceResult(instance *Instance, moreResults bool, result
 			q.fail(&callbackPanicError{value: r})
 		}
 
+		if !moreResults {
+			q.complete = true
+		}
+
 		q.mu.Unlock()
 
 		// The final result is the last call MI makes with this context.
@@ -308,14 +318,35 @@ func (q *asyncQuery) queueCopy(instance *Instance) error {
 }
 
 // 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 {
+// cancels the operation when a callback asks for it or ctx is done, and runs
+// the handler of a copy query.
+//
+// A done ctx fails the query with the context's cause, unless it has failed
+// before, and stops calling the handler. wait still returns only after the
+// final callback, so close never releases memory MI may use.
+func (q *asyncQuery) wait(ctx context.Context) error {
 	guard := newGuardTimer()
 	defer releaseGuardTimer(guard)
 
+	q.ctx = ctx
+	ctxDone := ctx.Done()
+
 	for {
 		select {
+		case <-ctxDone:
+			// Fires once. The cancellation goes through cancelRequested,
+			// like one a callback asks for.
+			ctxDone = nil
+
+			// A query whose final result has arrived is complete; a late
+			// cancellation must not turn it into an error.
+			q.mu.Lock()
+
+			if !q.complete {
+				q.fail(context.Cause(ctx))
+			}
+
+			q.mu.Unlock()
 		case <-q.done:
 			// Copies queued by the last callbacks.
 			q.handleQueued()
@@ -360,7 +391,16 @@ func (q *asyncQuery) handleQueued() {
 
 		for i, instance := range batch {
 			if q.handlerErr == nil {
-				if err := q.handler(instance); err != nil {
+				err := context.Cause(q.ctx)
+				if err == nil {
+					err = q.handler(instance)
+				}
+
+				// After ctx is done, the copies are only deleted. That fails
+				// the query, also when MI has already delivered the final
+				// result: skipped instances must not look like a complete
+				// result.
+				if err != nil {
 					q.handlerErr = err
 
 					// Stops the callbacks from copying further instances.
@@ -429,10 +469,25 @@ func (q *asyncQuery) close() {
 func (s *Session) query(flags OperationFlags, operationOptions *OperationOptions, namespaceName Namespace,
 	queryDialect QueryDialect, queryExpression Query, onInstance func(*Instance) error,
 ) error {
-	query := newAsyncQuery(onInstance)
+	return s.runQuery(context.Background(), newAsyncQuery(onInstance), flags, operationOptions, 0,
+		namespaceName, queryDialect, queryExpression)
+}
 
-	query.start(s, flags, operationOptions, namespaceName, queryDialect, queryExpression)
+// runQuery starts query, see [Session.startQuery], waits for it with
+// cancellation, see [asyncQuery.wait], and closes it.
+func (s *Session) runQuery(ctx context.Context, query *asyncQuery, flags OperationFlags, operationOptions *OperationOptions,
+	queryTimeout time.Duration, namespaceName Namespace, queryDialect QueryDialect, queryExpression Query,
+) error {
+	release, err := s.startQuery(query, flags, operationOptions, queryTimeout, namespaceName, queryDialect, queryExpression)
+	if err != nil {
+		return err
+	}
+
+	// Deferred first so the options outlive the operation that uses them.
+	defer release()
+
+	// close cancels the operation if a handler aborted it early or panicked.
 	defer query.close()
 
-	return query.wait()
+	return query.wait(ctx)
 }

internal/mi/callbacks_test.go

@@ -18,6 +18,7 @@
 package mi_test
 
 import (
+	"context"
 	"errors"
 	"os"
 	"os/exec"
@@ -274,3 +275,97 @@ func Test_MI_Query_DeadlockDetector(t *testing.T) {
 	output, err := cmd.CombinedOutput()
 	require.NoError(t, err, "%s", output)
 }
+
+// Test_MI_QueryContext_CancelBeforeStart checks that a query with a done
+// context is not started.
+func Test_MI_QueryContext_CancelBeforeStart(t *testing.T) {
+	t.Parallel()
+
+	session := newTestSession(t)
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	processes := []win32Process{{Name: "stale"}}
+
+	require.ErrorIs(t, session.QueryContext(ctx, &processes, mi.NamespaceRootCIMv2, query, 5*time.Second), context.Canceled)
+	require.Empty(t, processes)
+
+	err = session.QueryFuncContext(ctx, mi.NamespaceRootCIMv2, query, 5*time.Second, func(*mi.Instance) error {
+		return errors.New("unexpected instance")
+	})
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+// Test_MI_QueryContext_CancelDuringQuery checks that a context that is done
+// while the query runs cancels the operation and is reported, and that the
+// session stays usable.
+func Test_MI_QueryContext_CancelDuringQuery(t *testing.T) {
+	session := newTestSession(t)
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	t.Run("QueryFuncContext", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+
+		var handled int
+
+		err := session.QueryFuncContext(ctx, mi.NamespaceRootCIMv2, query, 5*time.Second, func(*mi.Instance) error {
+			handled++
+
+			cancel()
+
+			return nil
+		})
+		require.ErrorIs(t, err, context.Canceled)
+		require.Equal(t, 1, handled)
+	})
+
+	t.Run("QueryContext", func(t *testing.T) {
+		// CIM_DataFile enumerates every file of the drive, which takes far
+		// longer than the context.
+		slowQuery, err := mi.NewQuery("SELECT Name FROM CIM_DataFile WHERE Drive = 'C:'")
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
+		defer cancel()
+
+		type cimDataFile struct {
+			Name string `mi:"Name"`
+		}
+
+		var files []cimDataFile
+
+		start := time.Now()
+
+		err = session.QueryContext(ctx, &files, mi.NamespaceRootCIMv2, slowQuery, time.Minute)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Less(t, time.Since(start), 30*time.Second)
+	})
+
+	var processes []win32Process
+
+	require.NoError(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, 5*time.Second))
+	require.NotEmpty(t, processes)
+}
+
+// Test_MI_Query_ZeroTimeout checks that a zero timeout runs the query with
+// mi.DefaultQueryTimeout instead of failing.
+func Test_MI_Query_ZeroTimeout(t *testing.T) {
+	t.Parallel()
+
+	session := newTestSession(t)
+
+	query, err := mi.NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	var processes []win32Process
+
+	require.NoError(t, session.Query(&processes, mi.NamespaceRootCIMv2, query, 0))
+	require.NotEmpty(t, processes)
+}

internal/mi/session.go

@@ -18,10 +18,12 @@
 package mi
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math"
 	"reflect"
+	"sync"
 	"syscall"
 	"time"
 	"unsafe"
@@ -30,11 +32,27 @@ import (
 )
 
 // BuildQueryTimeout bounds the probe queries collectors run in Build. Build
-// has no scrape budget, and a zero timeout sets no MI timeout, so a stalled
-// provider would otherwise block startup. A timed-out query fails with
-// MI_RESULT_INVALID_OPERATION_TIMEOUT. MI ignores timeouts below 1ms.
+// has no scrape budget, and the session default, [DefaultQueryTimeout], is
+// too short for slow providers. A timed-out query fails with
+// MI_RESULT_INVALID_OPERATION_TIMEOUT.
 const BuildQueryTimeout = 30 * time.Second
 
+// DefaultQueryTimeout is the MI timeout of a query that is started with a
+// non-positive timeout or with nil operation options. Every query started
+// through [Session.Query] or [Session.QueryFunc] has an MI timeout: MI treats
+// a zero timeout as unset, which would let a stalled provider block the query
+// forever.
+const DefaultQueryTimeout = 5 * time.Second
+
+// minQueryTimeout is the shortest MI timeout. MI ignores timeouts below 1ms,
+// so shorter positive timeouts are raised to it instead of disabling the
+// timeout.
+const minQueryTimeout = time.Millisecond
+
+// maxCachedOperationOptions bounds the number of fixed query timeouts whose
+// operation options a session keeps.
+const maxCachedOperationOptions = 16
+
 // Session represents a session.
 //
 // https://learn.microsoft.com/en-us/windows/win32/api/mi/ns-mi-mi_session
@@ -47,8 +65,25 @@ type Session struct {
 	//
 	// application created the session and outlives it, so per-query
 	// operation options are created without MI_Session_GetApplication.
-	application             *Application
+	application *Application
+	// defaultOperationOptions carry DefaultQueryTimeout. They are used for
+	// nil operation options and non-positive query timeouts.
 	defaultOperationOptions *OperationOptions
+
+	// lifecycle orders starting queries against Close. startQuery holds a
+	// read lock from checking closed until MI_Session_QueryInstances has
+	// returned, so Close, which holds the write lock, never deletes options
+	// that a query is about to start with. closed is guarded by it.
+	lifecycle sync.RWMutex
+	closed    bool
+
+	// optionsMu guards cachedOperationOptions, which concurrent queries fill
+	// under the read lock of lifecycle. The cached options carry fixed
+	// timeouts, see operationOptions. Like the default options, they are
+	// shared by concurrent operations and deleted by Close after
+	// MI_Session_Close has closed every operation.
+	optionsMu              sync.Mutex
+	cachedOperationOptions map[time.Duration]*OperationOptions
 }
 
 // SessionFT represents the function table for Session.
@@ -80,20 +115,43 @@ func (s *Session) Close() error {
 		return ErrNotInitialized
 	}
 
+	// The write lock waits for queries that are being started, see
+	// startQuery, and makes queries started afterwards fail instead of using
+	// options that are about to be deleted.
+	s.lifecycle.Lock()
+	defer s.lifecycle.Unlock()
+
+	if s.closed {
+		return ErrNotInitialized
+	}
+
+	s.closed = true
+
 	// MI_Session_Close cancels running operations and blocks until their
-	// handles are closed. The default options are deleted afterwards, so an
-	// operation that is still open never uses freed options.
+	// handles are closed. The default and cached options are deleted
+	// afterwards, so an operation that is still open never uses freed options.
+	// The goroutines that close the operations do not take lifecycle.
 	r0, _, _ := syscall.SyscallN(s.ft.Close,
 		uintptr(unsafe.Pointer(s)),
 		0,
 		0,
 	)
 
+	s.optionsMu.Lock()
+
+	for _, operationOptions := range s.cachedOperationOptions {
+		_ = operationOptions.Delete()
+	}
+
+	s.cachedOperationOptions = nil
+
 	if s.defaultOperationOptions != nil {
 		_ = s.defaultOperationOptions.Delete()
 		s.defaultOperationOptions = nil
 	}
 
+	s.optionsMu.Unlock()
+
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
 		return result
 	}
@@ -405,25 +463,59 @@ func (s *Session) QueryUnmarshal[T any](dst *[]T,
 	})
 }
 
-// Query queries for a set of instances based on a query expression.
+// Query queries for a set of instances based on a query expression and
+// unmarshals them into dst. It is [Session.QueryContext] without
+// cancellation.
 func (s *Session) Query[T any](dst *[]T, namespaceName Namespace, queryExpression Query, queryTimeout time.Duration) error {
-	operationOptions, err := s.newOperationOptions(queryTimeout)
+	return s.QueryContext(context.Background(), dst, namespaceName, queryExpression, queryTimeout)
+}
+
+// QueryContext queries for a set of instances based on a query expression and
+// unmarshals them into dst.
+//
+// queryTimeout is the MI timeout of the query. A non-positive timeout selects
+// [DefaultQueryTimeout], so MI never runs the query without a timeout. A
+// positive timeout below 1ms is raised to 1ms, the shortest one MI honors.
+//
+// If ctx is done before the query starts, the query is not started. If ctx is
+// done while the query runs, the operation is cancelled. Either way the
+// context's cause is returned, unless the query failed before. QueryContext
+// returns only after MI has delivered the final result of a started query.
+func (s *Session) QueryContext[T any](ctx context.Context, dst *[]T, namespaceName Namespace, queryExpression Query,
+	queryTimeout time.Duration,
+) error {
+	if s == nil || s.ft == nil {
+		return ErrNotInitialized
+	}
+
+	fields, err := prepareUnmarshal(dst)
 	if err != nil {
 		return err
 	}
 
-	if operationOptions != nil {
-		defer func() {
-			_ = operationOptions.Delete()
-		}()
+	if ctx.Err() != nil {
+		return context.Cause(ctx)
 	}
 
-	return s.QueryUnmarshal(dst, OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
+	query := newAsyncQuery(func(instance *Instance) error {
+		return appendInstance(dst, instance, fields, true)
+	})
+
+	return s.runQuery(ctx, query, OperationFlagsStandardRTTI, nil, queryTimeout, namespaceName, QueryDialectWQL, queryExpression)
 }
 
-// QueryFunc runs a WQL query and calls fn for every returned instance, in
-// order. It is meant for queries whose result shape is only known at runtime,
-// where [Session.Query] cannot unmarshal into a struct.
+// QueryFunc runs a WQL query and calls fn for every returned instance. It is
+// [Session.QueryFuncContext] without cancellation.
+func (s *Session) QueryFunc(namespaceName Namespace, queryExpression Query, queryTimeout time.Duration, fn func(*Instance) error) error {
+	return s.QueryFuncContext(context.Background(), namespaceName, queryExpression, queryTimeout, fn)
+}
+
+// QueryFuncContext runs a WQL query and calls fn for every returned instance,
+// in order. It is meant for queries whose result shape is only known at
+// runtime, where [Session.QueryContext] cannot unmarshal into a struct.
+//
+// queryTimeout and ctx behave as for [Session.QueryContext]. Once ctx is
+// done, fn is not called any more.
 //
 // 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
@@ -437,57 +529,133 @@ func (s *Session) Query[T any](dst *[]T, namespaceName Namespace, queryExpressio
 // 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 {
+func (s *Session) QueryFuncContext(ctx context.Context, namespaceName Namespace, queryExpression Query,
+	queryTimeout time.Duration, fn func(*Instance) error,
+) error {
 	if s == nil || s.ft == nil {
 		return ErrNotInitialized
 	}
 
-	operationOptions, err := s.newOperationOptions(queryTimeout)
-	if err != nil {
-		return err
+	if ctx.Err() != nil {
+		return context.Cause(ctx)
 	}
 
-	if operationOptions != nil {
-		// Deferred first so the options outlive the operation that uses them.
-		defer func() {
-			_ = operationOptions.Delete()
-		}()
-	}
-
-	query := newCopyQuery(fn)
-
-	query.start(s, OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
-
-	// close cancels the operation if fn aborted it early or panicked.
-	defer query.close()
-
-	return query.wait()
+	return s.runQuery(ctx, newCopyQuery(fn), OperationFlagsStandardRTTI, nil, queryTimeout, namespaceName, QueryDialectWQL, queryExpression)
 }
 
-// newOperationOptions creates operation options carrying queryTimeout. A zero
-// timeout leaves the MI default in place; a negative timeout returns nil
-// 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 {
+// startQuery starts query with operationOptions or, if they are nil, with
+// the options for queryTimeout, see operationOptions. It returns a function
+// that releases the options once the operation is closed.
+//
+// The read lock of lifecycle is held from checking closed until
+// MI_Session_QueryInstances has returned and the operation exists. Close
+// takes the write lock, so it either fails this query first or waits until
+// the operation exists, which MI_Session_Close then cancels and waits for,
+// before it deletes the options.
+func (s *Session) startQuery(query *asyncQuery, flags OperationFlags, operationOptions *OperationOptions, queryTimeout time.Duration,
+	namespaceName Namespace, queryDialect QueryDialect, queryExpression Query,
+) (func(), error) {
+	s.lifecycle.RLock()
+	defer s.lifecycle.RUnlock()
+
+	if s.closed {
 		return nil, ErrNotInitialized
 	}
 
-	if queryTimeout < 0 {
-		return nil, nil //nolint:nilnil
+	release := func() {}
+
+	if operationOptions == nil {
+		var err error
+
+		operationOptions, release, err = s.operationOptions(queryTimeout)
+		if err != nil {
+			return nil, err
+		}
 	}
 
+	query.start(s, flags, operationOptions, namespaceName, queryDialect, queryExpression)
+
+	return release, nil
+}
+
+// operationOptions returns the operation options for queryTimeout and a
+// function that releases them once the operation that uses them is closed.
+//
+//   - A non-positive timeout, or DefaultQueryTimeout itself, returns the
+//     session defaults.
+//   - A timeout of whole seconds is taken as a fixed, configured timeout, e.g.
+//     BuildQueryTimeout. Its options are created once and kept until Close,
+//     for up to maxCachedOperationOptions distinct timeouts. Every query with
+//     such a timeout saves three native calls.
+//   - Any other timeout, typically the remaining scrape budget, gets options
+//     of its own, which release deletes.
+//
+// The caller holds the read lock of lifecycle and has checked closed, see
+// startQuery. After Close it fails with ErrNotInitialized.
+func (s *Session) operationOptions(queryTimeout time.Duration) (*OperationOptions, func(), error) {
+	s.optionsMu.Lock()
+	defer s.optionsMu.Unlock()
+
+	if s.defaultOperationOptions == nil {
+		return nil, nil, ErrNotInitialized
+	}
+
+	if queryTimeout <= 0 || queryTimeout == DefaultQueryTimeout {
+		return s.defaultOperationOptions, func() {}, nil
+	}
+
+	queryTimeout = max(queryTimeout, minQueryTimeout)
+
+	if queryTimeout%time.Second != 0 {
+		return s.ownOperationOptions(queryTimeout)
+	}
+
+	if operationOptions, ok := s.cachedOperationOptions[queryTimeout]; ok {
+		return operationOptions, func() {}, nil
+	}
+
+	if len(s.cachedOperationOptions) >= maxCachedOperationOptions {
+		return s.ownOperationOptions(queryTimeout)
+	}
+
+	operationOptions, err := s.newOperationOptions(queryTimeout)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if s.cachedOperationOptions == nil {
+		s.cachedOperationOptions = make(map[time.Duration]*OperationOptions, 1)
+	}
+
+	s.cachedOperationOptions[queryTimeout] = operationOptions
+
+	return operationOptions, func() {}, nil
+}
+
+// ownOperationOptions returns new operation options for queryTimeout and a
+// function that deletes them.
+func (s *Session) ownOperationOptions(queryTimeout time.Duration) (*OperationOptions, func(), error) {
+	operationOptions, err := s.newOperationOptions(queryTimeout)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return operationOptions, func() { _ = operationOptions.Delete() }, nil
+}
+
+// newOperationOptions creates operation options carrying queryTimeout, which
+// must be at least minQueryTimeout. The caller must delete them once the
+// operation that uses them is closed.
+func (s *Session) newOperationOptions(queryTimeout time.Duration) (*OperationOptions, error) {
 	operationOptions, err := s.application.NewOperationOptions()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create operation options: %w", err)
 	}
 
-	if queryTimeout > 0 {
-		if err = operationOptions.SetTimeout(queryTimeout); err != nil {
-			_ = operationOptions.Delete()
+	if err = operationOptions.SetTimeout(queryTimeout); err != nil {
+		_ = operationOptions.Delete()
 
-			return nil, fmt.Errorf("failed to set timeout: %w", err)
-		}
+		return nil, fmt.Errorf("failed to set timeout: %w", err)
 	}
 
 	return operationOptions, nil

internal/mi/session_internal_test.go

@@ -0,0 +1,316 @@
+// 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 (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// newInternalTestSession returns a session and its application. The caller
+// closes the session; the application is closed with the test.
+func newInternalTestSession(tb testing.TB) *Session {
+	tb.Helper()
+
+	application, err := ApplicationInitialize()
+	require.NoError(tb, err)
+
+	tb.Cleanup(func() { require.NoError(tb, application.Close()) })
+
+	session, err := application.NewSession(nil)
+	require.NoError(tb, err)
+
+	return session
+}
+
+func TestSession_OperationOptions(t *testing.T) {
+	session := newInternalTestSession(t)
+
+	t.Run("non-positive and default timeouts use the session defaults", func(t *testing.T) {
+		for _, timeout := range []time.Duration{0, -1, -time.Hour, DefaultQueryTimeout} {
+			operationOptions, release, err := session.operationOptions(timeout)
+			require.NoError(t, err)
+			require.Same(t, session.defaultOperationOptions, operationOptions)
+
+			release()
+		}
+	})
+
+	t.Run("fixed timeouts are cached", func(t *testing.T) {
+		first, release, err := session.operationOptions(BuildQueryTimeout)
+		require.NoError(t, err)
+		release()
+
+		second, release, err := session.operationOptions(BuildQueryTimeout)
+		require.NoError(t, err)
+		release()
+
+		require.Same(t, first, second)
+		require.NotSame(t, session.defaultOperationOptions, first)
+		require.Same(t, first, session.cachedOperationOptions[BuildQueryTimeout])
+	})
+
+	t.Run("other timeouts get options of their own", func(t *testing.T) {
+		for _, timeout := range []time.Duration{time.Microsecond, 1500 * time.Millisecond} {
+			operationOptions, release, err := session.operationOptions(timeout)
+			require.NoError(t, err)
+			require.NotSame(t, session.defaultOperationOptions, operationOptions)
+			require.NotContains(t, session.cachedOperationOptions, timeout)
+
+			release()
+		}
+	})
+
+	t.Run("the cache is bounded", func(t *testing.T) {
+		for i := range maxCachedOperationOptions + 2 {
+			_, release, err := session.operationOptions(time.Duration(100+i) * time.Second)
+			require.NoError(t, err)
+
+			release()
+		}
+
+		require.Len(t, session.cachedOperationOptions, maxCachedOperationOptions)
+	})
+
+	require.NoError(t, session.Close())
+	require.Nil(t, session.cachedOperationOptions)
+	require.Nil(t, session.defaultOperationOptions)
+
+	for _, timeout := range []time.Duration{0, BuildQueryTimeout, 1500 * time.Millisecond} {
+		_, _, err := session.operationOptions(timeout)
+		require.ErrorIs(t, err, ErrNotInitialized, timeout)
+	}
+}
+
+// TestSession_CloseWaitsForStartingQuery checks that Close does not delete
+// operation options while a query is being started, and that queries fail
+// after Close instead of starting with deleted options.
+func TestSession_CloseWaitsForStartingQuery(t *testing.T) {
+	session := newInternalTestSession(t)
+
+	queryExpression, err := NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	// A query between checking closed and MI_Session_QueryInstances holds
+	// this read lock.
+	session.lifecycle.RLock()
+
+	closed := make(chan error, 1)
+
+	go func() { closed <- session.Close() }()
+
+	select {
+	case err := <-closed:
+		session.lifecycle.RUnlock()
+		require.Failf(t, "Close returned while a query was being started", "err: %v", err)
+	case <-time.After(200 * time.Millisecond):
+	}
+
+	require.NotNil(t, session.defaultOperationOptions, "options deleted while a query was being started")
+
+	session.lifecycle.RUnlock()
+
+	select {
+	case err := <-closed:
+		require.NoError(t, err)
+	case <-time.After(time.Minute):
+		require.FailNow(t, "Close did not return")
+	}
+
+	require.ErrorIs(t, session.QueryFunc(NamespaceRootCIMv2, queryExpression, BuildQueryTimeout, func(*Instance) error { return nil }),
+		ErrNotInitialized)
+	require.ErrorIs(t, session.QueryUnmarshal(&[]struct {
+		Name string `mi:"Name"`
+	}{}, OperationFlagsStandardRTTI, nil, NamespaceRootCIMv2, QueryDialectWQL, queryExpression), ErrNotInitialized)
+	require.ErrorIs(t, session.Close(), ErrNotInitialized)
+}
+
+// TestSession_CloseDuringQueries closes a session while queries with
+// default, cached and per-query options start and run on other goroutines.
+// Every query either runs, is cancelled by MI_Session_Close, or fails with
+// ErrNotInitialized; none may start with deleted options.
+func TestSession_CloseDuringQueries(t *testing.T) {
+	session := newInternalTestSession(t)
+
+	queryExpression, err := NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	var (
+		wg      sync.WaitGroup
+		started = make(chan struct{})
+		once    sync.Once
+	)
+
+	for _, timeout := range []time.Duration{0, BuildQueryTimeout, 1500 * time.Millisecond} {
+		for range 2 {
+			wg.Go(func() {
+				for {
+					err := session.QueryFunc(NamespaceRootCIMv2, queryExpression, timeout, func(*Instance) error {
+						once.Do(func() { close(started) })
+
+						return nil
+					})
+					if errors.Is(err, ErrNotInitialized) {
+						return
+					}
+				}
+			})
+		}
+	}
+
+	<-started
+	time.Sleep(50 * time.Millisecond)
+
+	require.NoError(t, session.Close())
+
+	wg.Wait()
+}
+
+// Test_asyncQuery_ContextDoneAfterFinalResult checks a context that is done
+// after MI's final callback. Instances the handler did not see fail the
+// query, while a result that was complete stays successful.
+func Test_asyncQuery_ContextDoneAfterFinalResult(t *testing.T) {
+	session := newInternalTestSession(t)
+
+	t.Cleanup(func() { require.NoError(t, session.Close()) })
+
+	queryExpression, err := NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	t.Run("unhandled copies", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+
+		var handled int
+
+		query := newCopyQuery(func(*Instance) error {
+			handled++
+
+			cancel()
+
+			return nil
+		})
+
+		query.start(session, OperationFlagsStandardRTTI, nil, NamespaceRootCIMv2, QueryDialectWQL, queryExpression)
+		defer query.close()
+
+		// Every copy is queued and the final callback has run before wait.
+		<-query.done
+
+		require.ErrorIs(t, query.wait(ctx), context.Canceled)
+		require.Equal(t, 1, handled)
+	})
+
+	t.Run("complete result", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+
+		var instances int
+
+		query := newAsyncQuery(func(*Instance) error {
+			instances++
+
+			return nil
+		})
+
+		query.start(session, OperationFlagsStandardRTTI, nil, NamespaceRootCIMv2, QueryDialectWQL, queryExpression)
+		defer query.close()
+
+		<-query.done
+		cancel()
+
+		require.NoError(t, query.wait(ctx))
+		require.Positive(t, instances)
+	})
+}
+
+// Test_asyncQuery_ContextCancel checks that a cancelled query returns only
+// after MI's final callback, stops calling the handler, and leaves no queued
+// copies behind once closed.
+func Test_asyncQuery_ContextCancel(t *testing.T) {
+	session := newInternalTestSession(t)
+
+	t.Cleanup(func() { require.NoError(t, session.Close()) })
+
+	queryExpression, err := NewQuery("SELECT Name FROM Win32_Process")
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+
+	var handled int
+
+	query := newCopyQuery(func(*Instance) error {
+		handled++
+
+		cancel()
+
+		return nil
+	})
+
+	query.start(session, OperationFlagsStandardRTTI, nil, NamespaceRootCIMv2, QueryDialectWQL, queryExpression)
+
+	err = query.wait(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+	require.Equal(t, 1, handled, "the handler must not run after the context is done")
+	require.True(t, query.finished.Load(), "wait returned before the final callback")
+
+	select {
+	case <-query.done:
+	default:
+		require.Fail(t, "wait returned before the final callback")
+	}
+
+	query.close()
+
+	require.Nil(t, query.pending)
+	require.Nil(t, query.batch)
+}
+
+// Benchmark_Session_OperationOptions compares cached operation options for a
+// fixed timeout with options created for every query.
+func Benchmark_Session_OperationOptions(b *testing.B) {
+	session := newInternalTestSession(b)
+
+	b.Cleanup(func() { require.NoError(b, session.Close()) })
+
+	for _, bc := range []struct {
+		name    string
+		timeout time.Duration
+	}{
+		{name: "cached", timeout: 5 * time.Second},
+		{name: "per-query", timeout: 5*time.Second + time.Microsecond},
+	} {
+		b.Run(bc.name, func(b *testing.B) {
+			b.ReportAllocs()
+
+			for b.Loop() {
+				_, release, err := session.operationOptions(bc.timeout)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				release()
+			}
+		})
+	}
+}