diff --git a/AGENTS.md b/AGENTS.md index fd46fae6..f7554672 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -159,6 +159,15 @@ exercised a live Windows role. patterns, and MI NULL values from populated unions. Respect borrowed/owned pointers, allocator-specific cleanup and COM thread initialization; use the native COM helpers rather than restoring the removed `go-ole` dependency. + Run COM calls on an `ole.Worker` or another long-lived COM thread that + `Build` starts and `Close` stops; initializing and uninitializing COM per + scrape unloads in-process COM servers each time. A COM call cannot be + interrupted: if `Collect` waits for one, wait without a deadline of its own + (`ole.Do` with `context.Background()`), so that the collection's busy-skip + handles a hung call instead of every later scrape queueing behind it. +- A collector's `Close` waits at most 3 seconds (`closeTimeout`) for its own + workers, below the 5 seconds the collection gives all collectors to close. + Document what a worker still running after that does with its resources. - Keep CLI and YAML validation consistent. Reject unknown fields, unsupported sub-collector names and invalid values; preserve CLI precedence and repeated flag/list behavior. Changes to config shapes also need schema, flag parsing diff --git a/internal/collector/scheduled_task/AGENTS.md b/internal/collector/scheduled_task/AGENTS.md index ea47aa13..7264fcdc 100644 --- a/internal/collector/scheduled_task/AGENTS.md +++ b/internal/collector/scheduled_task/AGENTS.md @@ -65,7 +65,10 @@ of them empty), about 150 ms per scrape without filters. disabled. Skipping the read for disabled tasks would change published values. - A persistent COM connection (dedicated thread, service kept connected) saves ~1.7 ms per scrape. `CoCreateInstance` and `Connect` are cheap; the service's - per-call time dominates. + per-call time dominates. The collector reads tasks on an `ole.Worker` that + keeps COM initialized from `Build` to `Close`, but still creates and connects + the `TaskService` per scrape, so a restarted Schedule service needs no + reconnect logic. - Direct MS-TSCH RPC over `ncalrpc` (hand-marshalled NDR, `I_RpcSendReceive`) is only 5–8 µs faster per call. It would only save the duplicate `SchRpcGetLastRunInfo` (~18 ms per scrape, ~12%) plus ~9 ms of client CPU, at diff --git a/internal/collector/scheduled_task/scheduled_task.go b/internal/collector/scheduled_task/scheduled_task.go index 72216bca..f3dd9e32 100644 --- a/internal/collector/scheduled_task/scheduled_task.go +++ b/internal/collector/scheduled_task/scheduled_task.go @@ -18,11 +18,11 @@ package scheduled_task import ( + "context" "errors" "fmt" "log/slog" "regexp" - "runtime" "strconv" "strings" "time" @@ -73,8 +73,16 @@ type Collector struct { // the last scrape. Each task publishes 19 series, and building them is most // of the allocations of a scrape. metricCache metriccache.Cache[string, taskValues] + + // worker reads the tasks on a thread that keeps COM initialized between + // scrapes, so that the Task Scheduler client is not unloaded after each one. + worker *ole.Worker } +// closeTimeout bounds how long Close waits for a running task read. It stays +// below the 5 seconds that the collection waits for its collectors to close. +const closeTimeout = 3 * time.Second + // TaskState ... // https://docs.microsoft.com/en-us/windows/desktop/api/taskschd/ne-taskschd-task_state type TaskState uint @@ -178,7 +186,22 @@ func (c *Collector) ScrapeCaches() []scrapecache.Cache { return []scrapecache.Cache{&c.metricCache} } +// Close stops the COM worker. It waits up to closeTimeout for a task read that +// outlived its scrape. If the read is still running after that, Close returns +// an error and the worker stops and uninitializes COM on its own once the read +// returns. func (c *Collector) Close() error { + if c.worker == nil { + return nil + } + + worker := c.worker + c.worker = nil + + if err := worker.Close(closeTimeout); err != nil { + return fmt.Errorf("close Task Scheduler worker: %w", err) + } + return nil } @@ -216,11 +239,31 @@ func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error { c.folderFilter = newFolderFilter(c.config.TaskInclude, c.config.TaskExclude) + if c.worker != nil { + if err := c.Close(); err != nil { + return err + } + } + + worker, err := ole.NewWorker() + if err != nil { + return fmt.Errorf("start Task Scheduler worker: %w", err) + } + + c.worker = worker + return nil } +// Collect reads the tasks on the COM worker and waits for the result without +// a deadline of its own. A Task Scheduler call cannot be interrupted, so a +// read that outlives the scrape budget keeps Collect running: the collection +// reports the scrape as timed out, keeps the collector's slot until Collect +// returns and skips the collector in later scrapes meanwhile. Waiting for the +// worker with the scrape budget instead would let every later scrape queue +// behind a hung call for its whole budget. func (c *Collector) Collect(ch chan<- prometheus.Metric, _ time.Duration) error { - return c.collect(ch) + return c.collect(context.Background(), ch) } //nolint:gochecknoglobals @@ -232,11 +275,15 @@ var TASK_RESULT_STATUSES = []string{ "not_scheduled", "terminated", "no_valid_triggers", "event_trigger", "queued", "unknown", } -func (c *Collector) collect(ch chan<- prometheus.Metric) error { - scheduledTasks, err := c.taskReader().getScheduledTasks() +func (c *Collector) collect(ctx context.Context, ch chan<- prometheus.Metric) error { + if c.worker == nil { + return ole.ErrWorkerStopped + } + + scheduledTasks, err := ole.Do(ctx, c.worker, c.taskReader().getScheduledTasks) if errors.Is(err, errTasksSkipped) { // Tasks and folders that can't be read are skipped, the other tasks are still collected. - c.logger.Warn("failed to read some scheduled tasks", + c.logger.WarnContext(ctx, "failed to read some scheduled tasks", slog.Any("err", err), ) } else if err != nil { @@ -359,19 +406,10 @@ type taskReader struct { } // getScheduledTasks returns the tasks that pass includeTask. It does not walk -// the subfolders that fail includeFolder. +// the subfolders that fail includeFolder. It must run on a thread that joined +// the multithreaded apartment, such as an [ole.Worker], and releases every COM +// object it creates before it returns. func (r taskReader) getScheduledTasks() (ScheduledTasks, error) { - // COM initialization and every interface call stay on the same OS thread. - runtime.LockOSThread() - - defer runtime.UnlockOSThread() - - if err := ole.Initialize(); err != nil { - return nil, err - } - - defer ole.Uninitialize() - service, err := taskschd.NewTaskService() if err != nil { return nil, fmt.Errorf("create Task Scheduler service: %w", err) diff --git a/internal/collector/scheduled_task/scheduled_task_hang_test.go b/internal/collector/scheduled_task/scheduled_task_hang_test.go new file mode 100644 index 00000000..4227ee15 --- /dev/null +++ b/internal/collector/scheduled_task/scheduled_task_hang_test.go @@ -0,0 +1,91 @@ +// 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 scheduled_task_test + +import ( + "context" + "log/slog" + "strings" + "testing" + "time" + + "github.com/prometheus-community/windows_exporter/internal/collector/scheduled_task" + "github.com/prometheus-community/windows_exporter/pkg/collector" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/common/expfmt" + "github.com/stretchr/testify/require" +) + +// TestHungWorkerSkipsLaterScrapes checks that a Task Scheduler call that does +// not return times out one scrape, and that later scrapes skip the collector +// at once instead of waiting for the hung worker for their whole budget. +func TestHungWorkerSkipsLaterScrapes(t *testing.T) { + c := scheduled_task.New(nil) + collection := collector.New(collector.Map{scheduled_task.Name: c}) + + require.NoError(t, collection.Build(t.Context(), slog.New(slog.DiscardHandler))) + + release := scheduled_task.HangWorker(t, c) + + handler, err := collection.NewHandler(300*time.Millisecond, slog.New(slog.DiscardHandler), nil) + require.NoError(t, err) + + // The first scrape starts Collect, which waits for the hung worker until the budget ends. + first := gather(t, handler) + require.Contains(t, first, `windows_exporter_collector_timeout{collector="scheduled_task"} 1`) + + // The second scrape finds the first Collect still running and skips the collector. + start := time.Now() + second := gather(t, handler) + require.Less(t, time.Since(start), 150*time.Millisecond, "a later scrape does not wait for the hung worker") + require.Contains(t, second, `windows_exporter_collector_timeout{collector="scheduled_task"} 1`) + + release() + + // Once the worker answers, the pending Collect returns and the collector succeeds again. + require.Eventually(t, func() bool { + ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) + defer cancel() + + handler, err := collection.NewHandlerWithContext(ctx, 30*time.Second, slog.New(slog.DiscardHandler), nil) + require.NoError(t, err) + + return strings.Contains(gather(t, handler), `windows_exporter_collector_success{collector="scheduled_task"} 1`) + }, time.Minute, 100*time.Millisecond) + + require.NoError(t, collection.Close()) +} + +func gather(t *testing.T, handler prometheus.Collector) string { + t.Helper() + + registry := prometheus.NewRegistry() + require.NoError(t, registry.Register(handler)) + + families, err := registry.Gather() + require.NoError(t, err) + + var out strings.Builder + + for _, family := range families { + _, err := expfmt.MetricFamilyToText(&out, family) + require.NoError(t, err) + } + + return out.String() +} diff --git a/internal/collector/scheduled_task/scheduled_task_native_test.go b/internal/collector/scheduled_task/scheduled_task_native_test.go index 641e3c56..bb321903 100644 --- a/internal/collector/scheduled_task/scheduled_task_native_test.go +++ b/internal/collector/scheduled_task/scheduled_task_native_test.go @@ -501,8 +501,12 @@ func BenchmarkGetScheduledTasks(b *testing.B) { c := New(&Config{TaskExclude: regexp.MustCompile(fmt.Sprintf("^(?:%s)$", bc.exclude))}) require.NoError(b, c.Build(slog.New(slog.DiscardHandler), nil)) + b.Cleanup(func() { require.NoError(b, c.Close()) }) + + reader := c.taskReader() + for b.Loop() { - if _, err := c.taskReader().getScheduledTasks(); err != nil { + if _, err := ole.Do(b.Context(), c.worker, reader.getScheduledTasks); err != nil { b.Fatal(err) } } diff --git a/internal/collector/scheduled_task/worker_test.go b/internal/collector/scheduled_task/worker_test.go new file mode 100644 index 00000000..91200497 --- /dev/null +++ b/internal/collector/scheduled_task/worker_test.go @@ -0,0 +1,54 @@ +// 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 scheduled_task + +import ( + "context" + "testing" + + "github.com/prometheus-community/windows_exporter/internal/ole" +) + +// HangWorker blocks the COM worker of c, like a Task Scheduler call that does +// not return, until the returned function is called. It returns once the +// worker is blocked. +func HangWorker(t *testing.T, c *Collector) func() { + t.Helper() + + running := make(chan struct{}) + unblock := make(chan struct{}) + done := make(chan struct{}) + + go func() { + defer close(done) + + _, _ = ole.Do(context.Background(), c.worker, func() (struct{}, error) { + close(running) + <-unblock + + return struct{}{}, nil + }) + }() + + <-running + + return func() { + close(unblock) + <-done + } +} diff --git a/internal/collector/update/update.go b/internal/collector/update/update.go index 5098ba2d..a098fc2d 100644 --- a/internal/collector/update/update.go +++ b/internal/collector/update/update.go @@ -67,6 +67,8 @@ type Collector struct { mu sync.RWMutex ctxCancelFn context.CancelFunc + // workerDone is closed when the Windows Update worker returned and released its COM objects. + workerDone chan struct{} logger *slog.Logger @@ -108,16 +110,51 @@ func NewWithFlags(app *kingpin.Application) *Collector { return c } +// closeTimeout bounds how long Close waits for the Windows Update worker. It +// stays below the 5 seconds that the collection waits for its collectors to close. +const closeTimeout = 3 * time.Second + +// Close stops the Windows Update worker and waits up to closeTimeout for it. +// +// The worker observes the cancellation between searches only: a running +// search is a synchronous WUA call that cannot be interrupted, and an online +// search can take minutes. If the worker does not return in time, Close logs +// a warning and returns. The worker keeps running until the search returns. +// It owns its COM objects and only shares metricsBuf, under mu, with Collect. +// After the cancellation it does not store the result of the running search: +// it releases its COM objects, uninitializes COM and exits on its own without +// starting another search. If the process exits first, the search is abandoned. func (c *Collector) Close() error { - if c.ctxCancelFn != nil { - c.ctxCancelFn() - c.ctxCancelFn = nil + if c.ctxCancelFn == nil { + return nil + } + + c.ctxCancelFn() + c.ctxCancelFn = nil + + timer := time.NewTimer(closeTimeout) + defer timer.Stop() + + select { + case <-c.workerDone: + case <-timer.C: + c.logger.Warn("Windows Update search still running at close, it finishes in the background", + slog.Duration("waited", closeTimeout), + ) } return nil } func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error { + // A Build without a Close in between stops the worker of the earlier + // Build, which would otherwise keep searching until the process exits. + if c.ctxCancelFn != nil { + if err := c.Close(); err != nil { + return err + } + } + c.logger = logger.With(slog.String("collector", Name)) c.logger.Info("update collector is in an experimental state! The configuration and metrics may change in future. Please report any issues.") @@ -154,8 +191,11 @@ func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error { // scheduleUpdateStatus sends exactly one value to initErrCh, nil after a successful initialization. initErrCh := make(chan error, 1) + workerDone := make(chan struct{}) go func() { + defer close(workerDone) + err := recovery.Run(func() error { c.scheduleUpdateStatus(ctx, logger, initErrCh, c.config.Online) @@ -175,6 +215,7 @@ func (c *Collector) Build(logger *slog.Logger, _ *mi.Session) error { }() c.ctxCancelFn = cancel + c.workerDone = workerDone if err := <-initErrCh; err != nil { return fmt.Errorf("failed to initialize Windows Update collector: %w", err) @@ -264,10 +305,16 @@ func (c *Collector) scheduleUpdateStatus(ctx context.Context, logger *slog.Logge err = recovery.Run(func() error { var err error - metricsBuf, err = c.fetchUpdates(logger, searcher) + metricsBuf, err = c.fetchUpdates(logger, searcher, len(metricsBuf)) return err }) + + // A search that outlived Close publishes nothing. + if ctx.Err() != nil { + return + } + if err != nil { logger.ErrorContext(ctx, "failed to fetch updates", slog.Any("err", err), @@ -291,8 +338,11 @@ func (c *Collector) scheduleUpdateStatus(ctx context.Context, logger *slog.Logge } } -func (c *Collector) fetchUpdates(logger *slog.Logger, searcher *wuapi.UpdateSearcher) ([]prometheus.Metric, error) { - metricsBuf := make([]prometheus.Metric, 0, len(c.metricsBuf)*2+1) +// fetchUpdates searches for pending updates. previousLen, the number of +// metrics of the previous search, sizes the buffer. fetchUpdates reads no +// collector state that Collect or Close change. +func (c *Collector) fetchUpdates(logger *slog.Logger, searcher *wuapi.UpdateSearcher, previousLen int) ([]prometheus.Metric, error) { + metricsBuf := make([]prometheus.Metric, 0, previousLen+1) timeStart := time.Now() diff --git a/internal/collector/update/update_internal_test.go b/internal/collector/update/update_internal_test.go index 2bfafe14..74987e83 100644 --- a/internal/collector/update/update_internal_test.go +++ b/internal/collector/update/update_internal_test.go @@ -19,9 +19,11 @@ package update import ( "bytes" + "context" "fmt" "log/slog" "testing" + "time" "unicode/utf16" "unsafe" @@ -33,6 +35,86 @@ import ( "golang.org/x/sys/windows" ) +func TestCloseWaitsForWorker(t *testing.T) { + for _, tc := range []struct { + name string + stopsIn time.Duration + timedOut bool + }{ + {name: "worker stops in time", stopsIn: 200 * time.Millisecond}, + {name: "search outlives the close bound", stopsIn: closeTimeout + time.Second, timedOut: true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + + // The fake worker models a search that returns stopsIn after the cancellation. + go func() { + defer close(done) + + <-ctx.Done() + time.Sleep(tc.stopsIn) + }() + + var logs bytes.Buffer + + c := &Collector{ + logger: slog.New(slog.NewTextHandler(&logs, nil)), + ctxCancelFn: cancel, + workerDone: done, + } + start := time.Now() + + require.NoError(t, c.Close()) + + if tc.timedOut { + require.Less(t, time.Since(start), closeTimeout+time.Second/2, "Close is bounded") + require.Contains(t, logs.String(), "Windows Update search still running at close") + <-done + } else { + require.GreaterOrEqual(t, time.Since(start), tc.stopsIn, "Close waits for the worker") + require.Empty(t, logs.String()) + } + + require.NoError(t, c.Close(), "Close is idempotent") + }) + } +} + +// TestBuildStopsPreviousWorker checks that a second Build stops the worker of +// the first one instead of leaking it. +func TestBuildStopsPreviousWorker(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + + go func() { + defer close(done) + + <-ctx.Done() + }() + + c := New(nil) + c.logger = slog.New(slog.DiscardHandler) + c.ctxCancelFn = cancel + c.workerDone = done + + err := c.Build(slog.New(slog.DiscardHandler), nil) + + t.Cleanup(func() { _ = c.Close() }) + + select { + case <-done: + default: + t.Fatal("Build did not stop the worker of the previous Build") + } + + if err != nil { + // The Windows Update API may be unavailable on this host; the old + // worker must be stopped either way. + t.Logf("Build: %v", err) + } +} + func TestFetchUpdatesPartialCollection(t *testing.T) { for _, tc := range []struct { name string @@ -55,7 +137,7 @@ func TestFetchUpdatesPartialCollection(t *testing.T) { searcher, attempted, criteria := fakeUpdateSearch(tc.countError) - metrics, err := c.fetchUpdates(logger, searcher) + metrics, err := c.fetchUpdates(logger, searcher, 0) if tc.countError { require.ErrorIs(t, err, ole.HRESULT(0x80070005)) require.Nil(t, metrics) diff --git a/internal/ole/worker_windows.go b/internal/ole/worker_windows.go new file mode 100644 index 00000000..a59370f1 --- /dev/null +++ b/internal/ole/worker_windows.go @@ -0,0 +1,194 @@ +// 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 ole + +import ( + "context" + "errors" + "fmt" + "runtime" + "sync" + "time" + + "github.com/prometheus-community/windows_exporter/internal/utils/recovery" +) + +var ( + // ErrWorkerStopped is returned for work submitted after Worker.Close. + ErrWorkerStopped = errors.New("COM worker stopped") + // ErrWorkerTimeout is returned when a worker does not answer within the caller's budget. + ErrWorkerTimeout = errors.New("COM worker did not answer in time") +) + +// Worker runs functions on one goroutine that is locked to its OS thread and +// has joined the multithreaded apartment for its whole lifetime. +// +// Initializing and uninitializing COM for each call unloads in-process COM +// servers whenever the apartment's last reference goes away, so every call +// would load them again. A Worker initializes COM once in NewWorker and +// uninitializes it after it stops. +// +// Functions run one at a time in submission order. COM objects created by a +// function belong to the worker's thread: release them before the function +// returns and never hand them to the caller. +type Worker struct { + jobs chan func() + stop chan struct{} + // done is closed after the worker's goroutine returned and uninitialized COM. + done chan struct{} + closeOnce sync.Once +} + +// NewWorker starts a worker and returns the error of COM initialization. +// A worker returned without an error must be closed. +func NewWorker() (*Worker, error) { + w := &Worker{ + jobs: make(chan func()), + stop: make(chan struct{}), + done: make(chan struct{}), + } + + initErrCh := make(chan error, 1) + + go w.run(initErrCh) + + if err := <-initErrCh; err != nil { + <-w.done + + return nil, err + } + + return w, nil +} + +func (w *Worker) run(initErrCh chan<- error) { + defer close(w.done) + + // COM initialization and every interface call stay on the same OS thread. + // The thread is not unlocked: if Uninitialize fails to balance an + // initialization, the thread exits with the goroutine instead of carrying + // apartment state to other goroutines. + runtime.LockOSThread() + + if err := Initialize(); err != nil { + initErrCh <- err + + return + } + + defer Uninitialize() + + initErrCh <- nil + + for { + select { + case job := <-w.jobs: + job() + case <-w.stop: + return + } + } +} + +// Close stops the worker. A function that is still running finishes first, +// then the worker uninitializes COM. Close waits up to timeout for that and +// returns ErrWorkerTimeout if the worker is still busy. The worker then stops +// and uninitializes COM on its own once the function returns. Close is safe to +// call more than once. +func (w *Worker) Close(timeout time.Duration) error { + w.closeOnce.Do(func() { close(w.stop) }) + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case <-w.done: + return nil + case <-timer.C: + return fmt.Errorf("stop: %w", ErrWorkerTimeout) + } +} + +// Do runs fn on the worker and returns its result. A panic in fn is returned +// as an error wrapping [recovery.ErrPanic], and the worker keeps running. +// COM objects that fn created before the panic leak unless fn releases them +// in deferred calls. If fn calls runtime.Goexit, the worker uninitializes COM +// and exits, and this and every later call return ErrWorkerStopped. +// +// If ctx is done before the worker starts or finishes fn, Do returns an error +// wrapping ErrWorkerTimeout and the cause of ctx. A function that already +// started keeps running on the worker; its result is discarded, so fn must not +// write to memory that the caller reads after Do returns. Functions submitted +// later wait until it returns. +// +//nolint:ireturn // T is the concrete result type of fn. +func Do[T any](ctx context.Context, w *Worker, fn func() (T, error)) (T, error) { + type result struct { + value T + err error + } + + // Buffered, so that the worker does not block on an abandoned call. + resultCh := make(chan result, 1) + + job := func() { + var value T + + err := recovery.Run(func() error { + var err error + + value, err = fn() + + return err + }) + + resultCh <- result{value: value, err: err} + } + + var zero T + + if ctx.Err() != nil { + return zero, fmt.Errorf("%w: %w", ErrWorkerTimeout, context.Cause(ctx)) + } + + select { + case w.jobs <- job: + case <-w.stop: + return zero, ErrWorkerStopped + case <-w.done: + return zero, ErrWorkerStopped + case <-ctx.Done(): + return zero, fmt.Errorf("%w: worker busy: %w", ErrWorkerTimeout, context.Cause(ctx)) + } + + select { + case r := <-resultCh: + return r.value, r.err + case <-w.done: + // The worker exited, normally after fn sent its result, but also if fn + // called runtime.Goexit, which recovery.Run does not stop. + select { + case r := <-resultCh: + return r.value, r.err + default: + return zero, ErrWorkerStopped + } + case <-ctx.Done(): + return zero, fmt.Errorf("%w: %w", ErrWorkerTimeout, context.Cause(ctx)) + } +} diff --git a/internal/ole/worker_windows_test.go b/internal/ole/worker_windows_test.go new file mode 100644 index 00000000..d9d8dd5c --- /dev/null +++ b/internal/ole/worker_windows_test.go @@ -0,0 +1,166 @@ +// 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 ole + +import ( + "context" + "errors" + "runtime" + "testing" + "time" + "unsafe" + + "github.com/prometheus-community/windows_exporter/internal/utils/recovery" + "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" +) + +type apartment struct { + threadID uint32 + kind int32 + qualifier int32 +} + +func currentApartment() (apartment, error) { + a := apartment{threadID: windows.GetCurrentThreadId()} + + hr, _, _ := ole32.NewProc("CoGetApartmentType").Call( + uintptr(unsafe.Pointer(&a.kind)), + uintptr(unsafe.Pointer(&a.qualifier)), + ) + + return a, ResultError(hr) +} + +func TestWorkerKeepsApartment(t *testing.T) { + w, err := NewWorker() + require.NoError(t, err) + + first, err := Do(t.Context(), w, currentApartment) + require.NoError(t, err) + require.Equal(t, int32(1), first.kind, "APTTYPE_MTA") + require.Equal(t, int32(0), first.qualifier, "explicitly initialized, not the implicit MTA") + + for range 10 { + next, err := Do(t.Context(), w, currentApartment) + require.NoError(t, err) + require.Equal(t, first, next, "every call runs on the same initialized thread") + } + + require.NoError(t, w.Close(time.Second)) + require.NoError(t, w.Close(time.Second), "Close is idempotent") + + _, err = Do(t.Context(), w, currentApartment) + require.ErrorIs(t, err, ErrWorkerStopped) +} + +func TestWorkerRecoversPanics(t *testing.T) { + w, err := NewWorker() + require.NoError(t, err) + + defer func() { require.NoError(t, w.Close(time.Second)) }() + + _, err = Do(t.Context(), w, func() (int, error) { panic("boom") }) + require.ErrorIs(t, err, recovery.ErrPanic) + + value, err := Do(t.Context(), w, func() (int, error) { return 42, nil }) + require.NoError(t, err) + require.Equal(t, 42, value, "the worker keeps running after a panic") + + errTest := errors.New("test") + value, err = Do(t.Context(), w, func() (int, error) { return 7, errTest }) + require.ErrorIs(t, err, errTest) + require.Equal(t, 7, value, "partial results are returned with the error") +} + +// TestWorkerGoexit checks that a function calling runtime.Goexit, which +// recovery.Run does not stop, ends the worker instead of hanging Do. +func TestWorkerGoexit(t *testing.T) { + w, err := NewWorker() + require.NoError(t, err) + + done := make(chan error, 1) + + go func() { + _, err := Do(context.Background(), w, func() (int, error) { + runtime.Goexit() + + return 0, nil + }) + done <- err + }() + + select { + case err := <-done: + require.ErrorIs(t, err, ErrWorkerStopped) + case <-time.After(10 * time.Second): + t.Fatal("Do did not return after the worker exited") + } + + _, err = Do(context.Background(), w, func() (int, error) { return 1, nil }) + require.ErrorIs(t, err, ErrWorkerStopped, "later calls do not block") + + require.NoError(t, w.Close(time.Second)) +} + +func TestWorkerBudget(t *testing.T) { + w, err := NewWorker() + require.NoError(t, err) + + release := make(chan struct{}) + running := make(chan struct{}) + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + _, err = Do(ctx, w, func() (struct{}, error) { + close(running) + <-release + + return struct{}{}, nil + }) + require.ErrorIs(t, err, ErrWorkerTimeout) + require.ErrorIs(t, err, context.DeadlineExceeded) + + <-running + + // The abandoned call still occupies the worker. + ctx, cancel = context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + _, err = Do(ctx, w, func() (struct{}, error) { return struct{}{}, nil }) + require.ErrorIs(t, err, ErrWorkerTimeout) + + // A done context fails before the call is submitted. + canceled, cancelNow := context.WithCancel(t.Context()) + cancelNow() + + _, err = Do(canceled, w, func() (struct{}, error) { + t.Error("a call with a done context must not run") + + return struct{}{}, nil + }) + require.ErrorIs(t, err, ErrWorkerTimeout) + + // Close waits for the running call within its bound. + require.ErrorIs(t, w.Close(50*time.Millisecond), ErrWorkerTimeout) + + close(release) + + require.NoError(t, w.Close(time.Second), "the worker stops once the call returns") +}