scheduled_task: run COM work on long-lived workers (#2741)

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

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

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

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)

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()
+}

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)
 				}
 			}

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
+	}
+}

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()
 

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)

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))
+	}
+}

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")
+}