fix(agent): capture_pprof reports worker-wedge gauges + 30s capture throttle (#2422)

agent/internal/heartbeat/handlers_diag.go

@@ -12,9 +12,9 @@ import (
 	"fmt"
 	"runtime"
 	"runtime/pprof"
+	"sync/atomic"
 	"time"
 
-	"github.com/breeze-rmm/agent/internal/collectors"
 	"github.com/breeze-rmm/agent/internal/remote/tools"
 )
 
@@ -26,7 +26,58 @@ import (
 // without allocating a gigantic heap.
 var maxProfileBytes = 1 << 20
 
-func handleCapturePprof(_ *Heartbeat, cmd Command) tools.CommandResult {
+// capturePprofMinIntervalNs rate-limits captures. capture_pprof is a
+// server-queued command (up to 10 concurrent, 100 queued) and every heap/all
+// capture forces a stop-the-world runtime.GC(), so without a floor a burst of
+// queued captures degenerates into back-to-back GC pauses (#2422). Same
+// pattern as the heartbeat watchdog dump throttle (#2392). Atomic so tests
+// can shrink it.
+var capturePprofMinIntervalNs atomic.Int64
+
+// capturePprofLastNs is the unix-nano timestamp of the last admitted capture
+// (0 = never).
+var capturePprofLastNs atomic.Int64
+
+func init() {
+	capturePprofMinIntervalNs.Store(int64(30 * time.Second))
+}
+
+// capturePprofMinInterval returns the current minimum interval between
+// admitted captures.
+func capturePprofMinInterval() time.Duration {
+	return time.Duration(capturePprofMinIntervalNs.Load())
+}
+
+// setCapturePprofMinInterval overrides the capture rate-limit interval and
+// returns the previous value. Intended for tests.
+func setCapturePprofMinInterval(d time.Duration) time.Duration {
+	return time.Duration(capturePprofMinIntervalNs.Swap(int64(d)))
+}
+
+// resetCapturePprofThrottle clears the cross-invocation rate-limit state.
+// Intended for tests.
+func resetCapturePprofThrottle() {
+	capturePprofLastNs.Store(0)
+}
+
+// capturePprofTryAcquire reports whether a capture may run now, atomically
+// claiming the slot if so. Safe for concurrent handler invocations
+// (overlapping pool workers race for one slot). The slot is consumed even if
+// the capture itself later fails — acceptable, because the expensive part
+// (runtime.GC + profile serialization) may already have run by then.
+func capturePprofTryAcquire(now time.Time, interval time.Duration) bool {
+	for {
+		last := capturePprofLastNs.Load()
+		if last != 0 && now.UnixNano()-last < int64(interval) {
+			return false
+		}
+		if capturePprofLastNs.CompareAndSwap(last, now.UnixNano()) {
+			return true
+		}
+	}
+}
+
+func handleCapturePprof(h *Heartbeat, cmd Command) tools.CommandResult {
 	// Strict payload validation: key absent → default "all"; key present but
 	// not a string → error (don't let a malformed payload silently force a
 	// GC + double capture the caller never asked for).
@@ -57,11 +108,27 @@ func handleCapturePprof(_ *Heartbeat, cmd Command) tools.CommandResult {
 		}
 	}
 
+	// Rate-limit only after payload validation, so a malformed request is
+	// rejected on its own merits without burning the capture slot.
+	now := time.Now()
+	if interval := capturePprofMinInterval(); !capturePprofTryAcquire(now, interval) {
+		return tools.CommandResult{
+			Status: "failed",
+			Error: fmt.Sprintf(
+				"capture_pprof rate-limited: at most one capture per %s (heap captures force a stop-the-world GC); retry later",
+				interval),
+		}
+	}
+
 	result := map[string]any{
-		"capturedAt": time.Now().UTC().Format(time.RFC3339),
+		"capturedAt": now.UTC().Format(time.RFC3339),
 		// Snapshot of the runtime gauges at capture time, so the profile can
 		// be correlated with the heartbeat trend without a second command.
-		"runtime": collectors.CollectRuntimeStats(),
+		// Must go through h.collectAgentRuntime — the raw
+		// collectors.CollectRuntimeStats() never populates the worker-pool
+		// wedge gauges (commandsInFlight/commandsOverdue), which are exactly
+		// what an operator chasing an overdue-commands trend needs (#2422).
+		"runtime": h.collectAgentRuntime(now),
 	}
 
 	if wantHeap {

agent/internal/heartbeat/handlers_diag_test.go

@@ -3,11 +3,25 @@ package heartbeat
 import (
 	"encoding/base64"
 	"encoding/json"
+	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/breeze-rmm/agent/internal/collectors"
 )
 
+// newTestHeartbeatForDiag returns a zero-value Heartbeat for exercising
+// handleCapturePprof and clears the package-level capture throttle so each
+// test starts with a fresh rate-limit slot.
+func newTestHeartbeatForDiag(t *testing.T) *Heartbeat {
+	t.Helper()
+	resetCapturePprofThrottle()
+	t.Cleanup(resetCapturePprofThrottle)
+	return &Heartbeat{}
+}
+
 // decodeDiagResult parses the JSON payload NewSuccessResult marshals into
 // Stdout.
 func decodeDiagResult(t *testing.T, stdout string) map[string]any {
@@ -38,7 +52,7 @@ func assertValidPprofBlob(t *testing.T, result map[string]any, field string) {
 }
 
 func TestHandleCapturePprofDefaultCapturesBoth(t *testing.T) {
-	res := handleCapturePprof(nil, Command{ID: "c1", Type: "capture_pprof"})
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{ID: "c1", Type: "capture_pprof"})
 	if res.Status != "completed" {
 		t.Fatalf("status = %q (error: %q), want completed", res.Status, res.Error)
 	}
@@ -59,7 +73,7 @@ func TestHandleCapturePprofDefaultCapturesBoth(t *testing.T) {
 }
 
 func TestHandleCapturePprofHeapOnly(t *testing.T) {
-	res := handleCapturePprof(nil, Command{
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{
 		ID: "c2", Type: "capture_pprof",
 		Payload: map[string]any{"profile": "heap"},
 	})
@@ -74,7 +88,7 @@ func TestHandleCapturePprofHeapOnly(t *testing.T) {
 }
 
 func TestHandleCapturePprofGoroutineOnly(t *testing.T) {
-	res := handleCapturePprof(nil, Command{
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{
 		ID: "c3", Type: "capture_pprof",
 		Payload: map[string]any{"profile": "goroutine"},
 	})
@@ -89,7 +103,7 @@ func TestHandleCapturePprofGoroutineOnly(t *testing.T) {
 }
 
 func TestHandleCapturePprofRejectsUnknownProfile(t *testing.T) {
-	res := handleCapturePprof(nil, Command{
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{
 		ID: "c4", Type: "capture_pprof",
 		Payload: map[string]any{"profile": "cpu"},
 	})
@@ -106,7 +120,7 @@ func TestHandleCapturePprofSizeCap(t *testing.T) {
 	maxProfileBytes = 1 // every real profile exceeds one byte
 	defer func() { maxProfileBytes = orig }()
 
-	res := handleCapturePprof(nil, Command{ID: "c5", Type: "capture_pprof"})
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{ID: "c5", Type: "capture_pprof"})
 	if res.Status != "failed" {
 		t.Fatalf("status = %q, want failed when profile exceeds size cap", res.Status)
 	}
@@ -116,7 +130,7 @@ func TestHandleCapturePprofSizeCap(t *testing.T) {
 }
 
 func TestHandleCapturePprofRejectsNonStringProfile(t *testing.T) {
-	res := handleCapturePprof(nil, Command{
+	res := handleCapturePprof(newTestHeartbeatForDiag(t), Command{
 		ID: "c6", Type: "capture_pprof",
 		Payload: map[string]any{"profile": 123},
 	})
@@ -128,6 +142,150 @@ func TestHandleCapturePprofRejectsNonStringProfile(t *testing.T) {
 	}
 }
 
+// TestHandleCapturePprofIncludesWedgeGauges verifies the runtime snapshot in
+// the capture result carries the worker-pool wedge gauges — i.e. that the
+// handler goes through h.collectAgentRuntime, not the raw collector, which
+// would report a permanently-plausible 0/0 (#2422).
+func TestHandleCapturePprofIncludesWedgeGauges(t *testing.T) {
+	h := newTestHeartbeatForDiag(t)
+	// One command that started an hour ago with a 1s watchdog tier:
+	// in flight AND overdue at capture time.
+	key := h.trackInFlight(time.Now().Add(-time.Hour), time.Second)
+	defer h.untrackInFlight(key)
+
+	res := handleCapturePprof(h, Command{
+		ID: "g1", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "goroutine"},
+	})
+	if res.Status != "completed" {
+		t.Fatalf("status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+	result := decodeDiagResult(t, res.Stdout)
+	rt, ok := result["runtime"].(map[string]any)
+	if !ok {
+		t.Fatal("runtime stats snapshot missing")
+	}
+	if got, _ := rt["commandsInFlight"].(float64); got != 1 {
+		t.Errorf("runtime.commandsInFlight = %v, want 1", rt["commandsInFlight"])
+	}
+	if got, _ := rt["commandsOverdue"].(float64); got != 1 {
+		t.Errorf("runtime.commandsOverdue = %v, want 1", rt["commandsOverdue"])
+	}
+}
+
+// TestHandleCapturePprofThrottled verifies back-to-back captures are
+// rate-limited (each heap/all capture forces a stop-the-world GC) and that
+// the slot frees up once the interval elapses.
+func TestHandleCapturePprofThrottled(t *testing.T) {
+	h := newTestHeartbeatForDiag(t)
+
+	res := handleCapturePprof(h, Command{
+		ID: "t1", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "goroutine"},
+	})
+	if res.Status != "completed" {
+		t.Fatalf("first capture: status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+
+	res = handleCapturePprof(h, Command{
+		ID: "t2", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "goroutine"},
+	})
+	if res.Status != "failed" {
+		t.Fatalf("second capture inside the interval: status = %q, want failed", res.Status)
+	}
+	if !strings.Contains(res.Error, "rate-limited") {
+		t.Errorf("rate-limit rejection error = %q, want it to mention rate-limited", res.Error)
+	}
+
+	// Shrink the interval below the elapsed time — the slot must free up.
+	prev := setCapturePprofMinInterval(time.Nanosecond)
+	t.Cleanup(func() { setCapturePprofMinInterval(prev) })
+	time.Sleep(time.Millisecond)
+
+	res = handleCapturePprof(h, Command{
+		ID: "t3", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "goroutine"},
+	})
+	if res.Status != "completed" {
+		t.Fatalf("capture after interval elapsed: status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+}
+
+// TestCapturePprofTryAcquireConcurrent verifies that overlapping pool
+// workers racing for one capture slot yield exactly one winner — the burst
+// scenario (#2422: up to 10 concurrent queued captures) the throttle exists
+// to prevent. Mirrors TestHeartbeatWatchdogTryAcquireDumpConcurrent; the CAS
+// loop is a separate copy, so the watchdog's test does not cover it.
+func TestCapturePprofTryAcquireConcurrent(t *testing.T) {
+	resetCapturePprofThrottle()
+	t.Cleanup(resetCapturePprofThrottle)
+
+	now := time.Now()
+	const n = 32
+	var winners atomic.Int64
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			if capturePprofTryAcquire(now, time.Hour) {
+				winners.Add(1)
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	if got := winners.Load(); got != 1 {
+		t.Fatalf("expected exactly 1 winner among %d concurrent acquisitions, got %d", n, got)
+	}
+}
+
+// TestCapturePprofTryAcquireInterval pins the interval boundary semantics
+// with injected times (no sleeps).
+func TestCapturePprofTryAcquireInterval(t *testing.T) {
+	resetCapturePprofThrottle()
+	t.Cleanup(resetCapturePprofThrottle)
+
+	base := time.Now()
+	interval := 30 * time.Second
+
+	if !capturePprofTryAcquire(base, interval) {
+		t.Fatal("first acquisition must succeed")
+	}
+	if capturePprofTryAcquire(base.Add(interval-time.Nanosecond), interval) {
+		t.Fatal("acquisition 1ns before the interval elapses must be rejected")
+	}
+	if !capturePprofTryAcquire(base.Add(interval), interval) {
+		t.Fatal("acquisition exactly at the interval must succeed")
+	}
+}
+
+// TestHandleCapturePprofValidationDoesNotConsumeSlot verifies a malformed
+// payload is rejected without burning the capture rate-limit slot.
+func TestHandleCapturePprofValidationDoesNotConsumeSlot(t *testing.T) {
+	h := newTestHeartbeatForDiag(t)
+
+	res := handleCapturePprof(h, Command{
+		ID: "v1", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "cpu"},
+	})
+	if res.Status != "failed" {
+		t.Fatalf("invalid profile: status = %q, want failed", res.Status)
+	}
+
+	res = handleCapturePprof(h, Command{
+		ID: "v2", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "goroutine"},
+	})
+	if res.Status != "completed" {
+		t.Fatalf("valid capture after validation failure: status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+}
+
 // TestHeartbeatPayloadAgentRuntimeWireKey pins the OUTER wire key of the
 // runtime gauges. The API heartbeatSchema matches on the literal
 // "agentRuntime" (apps/api/src/routes/agents/schemas.ts) with .optional(), so