pdh: collect synchronously under mutex (#2559)

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

internal/pdh/collector.go

@@ -59,8 +59,7 @@ type Collector[T any] struct {
 	// left out of the sample, so a missing value is never reported as zero.
 	partialRows bool
 
-	collectCh chan *[]T
+	state *collectState
-	errorCh   chan error
 }
 
 // Row holds the values of one instance collected by a collector from [NewDynamicCollector].
@@ -326,10 +325,7 @@ func newCollector[T any](logger *slog.Logger, resultType CounterType, object str
 		return collector, fmt.Errorf("failed to initialize collector: %w", err)
 	}
 
-	collector.collectCh = make(chan *[]T)
+	collector.state = &collectState{}
-	collector.errorCh = make(chan error)
-
-	go collector.collectWorker()
 
 	// Collect initial data because some counters need to be read twice to get the correct value.
 	var collectValues []T
@@ -400,27 +396,17 @@ func (c *Collector[T]) Collect(dst *[]T) error {
 		return errors.New("dst must not be nil")
 	}
 
-	c.mu.RLock()
+	c.mu.Lock()
-	defer c.mu.RUnlock()
+	defer c.mu.Unlock()
 
-	if len(c.counters) == 0 || c.handle == 0 || c.collectCh == nil || c.errorCh == nil {
+	if len(c.counters) == 0 || c.handle == 0 || c.state == nil {
 		return ErrPerformanceCounterNotInitialized
 	}
 
-	c.collectCh <- dst
+	return c.collect(dst, c.state)
-
-	return <-c.errorCh
 }
 
-func (c *Collector[T]) collectWorker() {
+// collectState holds the buffers that the collector reuses between samples.
-	var state collectState
-
-	for dst := range c.collectCh {
-		c.errorCh <- c.collect(dst, &state)
-	}
-}
-
-// collectState holds the buffers that the collect worker reuses between samples.
 type collectState struct {
 	// buf starts empty so that the first call only queries the required size with a nil buffer.
 	// PdhGetRawCounterArrayW writes 8 bytes into a non-nil buffer, even if lpdwBufferSize is smaller,
@@ -690,8 +676,7 @@ func (r *rowSet[T]) finish() {
 	}
 }
 
-// panicError logs a panic recovered in the collect worker and returns it as an
+// panicError logs a panic recovered while collecting and returns it as an error.
-// error. The worker goroutine has no other recovery, so a panic would end the process.
 func (c *Collector[T]) panicError(r any) error {
 	c.logger.Error("recovered from panic while collecting performance counters",
 		slog.String("object", c.object),
@@ -715,17 +700,7 @@ func (c *Collector[T]) Close() {
 	}
 
 	c.handle = 0
-
+	c.state = nil
-	if c.collectCh != nil {
-		close(c.collectCh)
-	}
-
-	if c.errorCh != nil {
-		close(c.errorCh)
-	}
-
-	c.collectCh = nil
-	c.errorCh = nil
 }
 
 func formatCounterPath(object, instance, counterName string) string {

internal/pdh/collector_internal_test.go

@@ -155,8 +155,7 @@ type processThreads struct {
 	ThreadCount float64 `perfdata:"Thread Count"`
 }
 
-// TestCollectRecoversPanic checks that a panic in the worker goroutine is
+// TestCollectRecoversPanic checks that a panic while collecting is returned as an error.
-// returned as an error instead of terminating the process.
 func TestCollectRecoversPanic(t *testing.T) {
 	t.Parallel()
 
@@ -177,7 +176,7 @@ func TestCollectRecoversPanic(t *testing.T) {
 	err = c.Collect(&dst)
 	require.ErrorContains(t, err, "panic while collecting performance counters of Process")
 
-	// The worker survives and serves the next request.
+	// Collection still works after the panic.
 	c.mu.Lock()
 	c.counters[0].FieldIndexValue = 1
 	c.mu.Unlock()

internal/pdh/collector_test.go

@@ -19,6 +19,7 @@ package pdh_test
 
 import (
 	"log/slog"
+	"sync"
 	"testing"
 	"time"
 
@@ -73,6 +74,113 @@ func TestCollector(t *testing.T) {
 	}
 }
 
+func TestCollectorConcurrentCollect(t *testing.T) {
+	t.Parallel()
+
+	for _, counterType := range []pdh.CounterType{pdh.CounterTypeRaw, pdh.CounterTypeFormatted} {
+		t.Run(string(counterType), func(t *testing.T) {
+			t.Parallel()
+
+			collector, err := pdh.NewCollector[process](slog.New(slog.DiscardHandler), counterType, "Process", pdh.InstancesAll)
+			require.NoError(t, err)
+
+			t.Cleanup(collector.Close)
+
+			const callers = 8
+
+			data := make([][]process, callers)
+			errs := make([]error, callers)
+			start := make(chan struct{})
+
+			var wg sync.WaitGroup
+
+			for i := range callers {
+				wg.Go(func() {
+					<-start
+
+					for range 4 {
+						if errs[i] = collector.Collect(&data[i]); errs[i] != nil {
+							return
+						}
+					}
+				})
+			}
+
+			close(start)
+			wg.Wait()
+
+			for i := range callers {
+				require.NoError(t, errs[i])
+				require.NotEmpty(t, data[i])
+			}
+		})
+	}
+}
+
+func TestCollectorConcurrentClose(t *testing.T) {
+	t.Parallel()
+
+	collector, err := pdh.NewCollector[process](slog.New(slog.DiscardHandler), pdh.CounterTypeRaw, "Process", pdh.InstancesAll)
+	require.NoError(t, err)
+
+	t.Cleanup(collector.Close)
+
+	const callers = 8
+
+	errs := make([]error, callers)
+	start := make(chan struct{})
+
+	var wg sync.WaitGroup
+
+	for i := range callers {
+		wg.Go(func() {
+			<-start
+
+			var data []process
+
+			errs[i] = collector.Collect(&data)
+		})
+	}
+
+	wg.Go(func() {
+		<-start
+		collector.Close()
+	})
+
+	close(start)
+	wg.Wait()
+
+	for _, err := range errs {
+		if err != nil {
+			require.ErrorIs(t, err, pdh.ErrPerformanceCounterNotInitialized)
+		}
+	}
+
+	var data []process
+
+	require.ErrorIs(t, collector.Collect(&data), pdh.ErrPerformanceCounterNotInitialized)
+	collector.Close()
+}
+
+func TestCollectorInitializationError(t *testing.T) {
+	t.Parallel()
+
+	type invalidProcess struct {
+		ThreadCount float64 `perfdata:"Thread Count"`
+		HandleCount string  `perfdata:"Handle Count"`
+	}
+
+	collector, err := pdh.NewCollector[invalidProcess](slog.New(slog.DiscardHandler), pdh.CounterTypeRaw, "Process", pdh.InstancesAll)
+	require.ErrorContains(t, err, "field HandleCount must be a float64")
+	require.NotNil(t, collector)
+
+	t.Cleanup(collector.Close)
+
+	var data []invalidProcess
+
+	require.ErrorIs(t, collector.Collect(&data), pdh.ErrPerformanceCounterNotInitialized)
+}
+
 type processorInformation struct {
 	Name          string
 	ProcessorTime float64 `perfdata:"% Processor Time"`