diff --git a/AGENTS.md b/AGENTS.md index b10afaf0..fd46fae6 100644 --- a/AGENTS.md +++ b/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 diff --git a/UPGRADE.md b/UPGRADE.md index bc5568d7..492b5741 100644 --- a/UPGRADE.md +++ b/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 diff --git a/internal/collector/wmi/wmi.go b/internal/collector/wmi/wmi.go index 4020a14e..0425db74 100644 --- a/internal/collector/wmi/wmi.go +++ b/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() { diff --git a/internal/mi/AGENTS.md b/internal/mi/AGENTS.md index a5367fe6..fd2f5314 100644 --- a/internal/mi/AGENTS.md +++ b/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. diff --git a/internal/mi/application.go b/internal/mi/application.go index 0b9ef16d..4503397a 100644 --- a/internal/mi/application.go +++ b/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() diff --git a/internal/mi/callbacks.go b/internal/mi/callbacks.go index 526da631..703fd737 100644 --- a/internal/mi/callbacks.go +++ b/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) } diff --git a/internal/mi/callbacks_test.go b/internal/mi/callbacks_test.go index eabcaff9..3943ec0c 100644 --- a/internal/mi/callbacks_test.go +++ b/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) +} diff --git a/internal/mi/session.go b/internal/mi/session.go index 0846dacc..364a0b3e 100644 --- a/internal/mi/session.go +++ b/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 diff --git a/internal/mi/session_internal_test.go b/internal/mi/session_internal_test.go new file mode 100644 index 00000000..99c0bbbd --- /dev/null +++ b/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() + } + }) + } +}