diff --git a/agent/internal/heartbeat/handlers_diag.go b/agent/internal/heartbeat/handlers_diag.go index 6d8c295acf..119e2a35ac 100644 --- a/agent/internal/heartbeat/handlers_diag.go +++ b/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 { diff --git a/agent/internal/heartbeat/handlers_diag_test.go b/agent/internal/heartbeat/handlers_diag_test.go index 2f6b9ca48d..f27332b202 100644 --- a/agent/internal/heartbeat/handlers_diag_test.go +++ b/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