feat(agent): runtime memory gauges on heartbeat + on-demand pprof capture (#2389)

agent/internal/collectors/runtime_stats.go

@@ -0,0 +1,31 @@
+package collectors
+
+import "runtime"
+
+// RuntimeStats is a snapshot of the agent's own Go runtime memory state,
+// reported on every heartbeat so fleet-wide agent memory leaks are visible
+// from the server without shell access to the device (issue #2389).
+type RuntimeStats struct {
+	HeapAllocBytes    uint64 `json:"heapAllocBytes"`
+	HeapInuseBytes    uint64 `json:"heapInuseBytes"`
+	HeapReleasedBytes uint64 `json:"heapReleasedBytes"`
+	SysBytes          uint64 `json:"sysBytes"`
+	NumGC             uint32 `json:"numGc"`
+	Goroutines        int    `json:"goroutines"`
+}
+
+// CollectRuntimeStats reads the Go runtime's memory statistics for this
+// process. runtime.ReadMemStats is a brief stop-the-world read (microseconds
+// for a typical agent heap) — cheap enough to call on every heartbeat.
+func CollectRuntimeStats() *RuntimeStats {
+	var ms runtime.MemStats
+	runtime.ReadMemStats(&ms)
+	return &RuntimeStats{
+		HeapAllocBytes:    ms.HeapAlloc,
+		HeapInuseBytes:    ms.HeapInuse,
+		HeapReleasedBytes: ms.HeapReleased,
+		SysBytes:          ms.Sys,
+		NumGC:             ms.NumGC,
+		Goroutines:        runtime.NumGoroutine(),
+	}
+}

agent/internal/collectors/runtime_stats_test.go

@@ -0,0 +1,54 @@
+package collectors
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestCollectRuntimeStats(t *testing.T) {
+	stats := CollectRuntimeStats()
+	if stats == nil {
+		t.Fatal("CollectRuntimeStats returned nil")
+	}
+
+	// A running Go process always has a non-empty heap and at least the
+	// current goroutine.
+	if stats.HeapAllocBytes == 0 {
+		t.Error("HeapAllocBytes should be non-zero for a live process")
+	}
+	if stats.HeapInuseBytes == 0 {
+		t.Error("HeapInuseBytes should be non-zero for a live process")
+	}
+	if stats.SysBytes == 0 {
+		t.Error("SysBytes should be non-zero for a live process")
+	}
+	if stats.Goroutines < 1 {
+		t.Errorf("Goroutines = %d, want >= 1", stats.Goroutines)
+	}
+	// HeapInuse counts whole spans, so it can never be below HeapAlloc.
+	if stats.HeapInuseBytes < stats.HeapAllocBytes {
+		t.Errorf("HeapInuseBytes (%d) < HeapAllocBytes (%d)", stats.HeapInuseBytes, stats.HeapAllocBytes)
+	}
+}
+
+func TestRuntimeStatsJSONShape(t *testing.T) {
+	// The wire field names are a contract with the API heartbeatSchema
+	// (apps/api/src/routes/agents/schemas.ts agentRuntime) — renaming a JSON
+	// key silently drops the gauge server-side.
+	data, err := json.Marshal(CollectRuntimeStats())
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	var m map[string]any
+	if err := json.Unmarshal(data, &m); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	for _, key := range []string{
+		"heapAllocBytes", "heapInuseBytes", "heapReleasedBytes",
+		"sysBytes", "numGc", "goroutines",
+	} {
+		if _, ok := m[key]; !ok {
+			t.Errorf("expected JSON key %q missing", key)
+		}
+	}
+}

agent/internal/heartbeat/handlers.go

@@ -96,6 +96,9 @@ var handlerRegistry = map[string]CommandHandler{
 	// Log shipping
 	tools.CmdSetLogLevel: handleSetLogLevel,
 
+	// Runtime diagnostics — on-demand pprof capture (#2389)
+	tools.CmdCapturePprof: handleCapturePprof,
+
 	// Auto-update management
 	tools.CmdSetAutoUpdate: handleSetAutoUpdate,
 }

agent/internal/heartbeat/handlers_diag.go

@@ -0,0 +1,108 @@
+package heartbeat
+
+// On-demand runtime diagnostics (#2389). capture_pprof captures heap and/or
+// goroutine profiles in-process and returns them base64-encoded in the command
+// result. There is deliberately NO pprof HTTP listener: the agent is a root
+// daemon and must expose nothing reachable off-box, so the signed/queued
+// command path is the only trigger.
+
+import (
+	"bytes"
+	"encoding/base64"
+	"fmt"
+	"runtime"
+	"runtime/pprof"
+	"time"
+
+	"github.com/breeze-rmm/agent/internal/collectors"
+	"github.com/breeze-rmm/agent/internal/remote/tools"
+)
+
+// maxProfileBytes caps a single captured profile's raw size. Heap and
+// goroutine profiles are sampled/compact — typically well under 1 MB even for
+// large processes — and the server-side command result caps stdout at 5 MB,
+// so 1 MiB raw per profile (~1.37 MiB base64) keeps the combined result
+// comfortably inside that. Var (not const) so tests can exercise the cap
+// without allocating a gigantic heap.
+var maxProfileBytes = 1 << 20
+
+func handleCapturePprof(_ *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).
+	profile := "all"
+	if raw, ok := cmd.Payload["profile"]; ok {
+		s, isString := raw.(string)
+		if !isString {
+			return tools.CommandResult{
+				Status: "failed",
+				Error:  fmt.Sprintf("profile must be a string, got %T", raw),
+			}
+		}
+		profile = s
+	}
+
+	var wantHeap, wantGoroutine bool
+	switch profile {
+	case "all":
+		wantHeap, wantGoroutine = true, true
+	case "heap":
+		wantHeap = true
+	case "goroutine":
+		wantGoroutine = true
+	default:
+		return tools.CommandResult{
+			Status: "failed",
+			Error:  fmt.Sprintf("invalid profile %q: must be heap, goroutine, or all", profile),
+		}
+	}
+
+	result := map[string]any{
+		"capturedAt": time.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(),
+	}
+
+	if wantHeap {
+		// Force a GC first so the heap profile reflects live (in-use) objects
+		// rather than garbage awaiting collection — same effect as
+		// net/http/pprof's ?gc=1.
+		runtime.GC()
+		b64, size, err := capturePprofProfile("heap")
+		if err != nil {
+			return tools.CommandResult{Status: "failed", Error: err.Error()}
+		}
+		result["heapProfileBase64"] = b64
+		result["heapProfileBytes"] = size
+	}
+
+	if wantGoroutine {
+		b64, size, err := capturePprofProfile("goroutine")
+		if err != nil {
+			return tools.CommandResult{Status: "failed", Error: err.Error()}
+		}
+		result["goroutineProfileBase64"] = b64
+		result["goroutineProfileBytes"] = size
+	}
+
+	return tools.NewSuccessResult(result, 0)
+}
+
+// capturePprofProfile writes the named runtime/pprof profile (debug=0 →
+// gzip-compressed protobuf, the format `go tool pprof` consumes) and returns
+// it base64-encoded along with its raw byte size.
+func capturePprofProfile(name string) (string, int, error) {
+	p := pprof.Lookup(name)
+	if p == nil {
+		return "", 0, fmt.Errorf("profile %q not found", name)
+	}
+	var buf bytes.Buffer
+	if err := p.WriteTo(&buf, 0); err != nil {
+		return "", 0, fmt.Errorf("failed to write %s profile: %v", name, err)
+	}
+	if buf.Len() > maxProfileBytes {
+		return "", 0, fmt.Errorf("%s profile is %d bytes, exceeds the %d byte result cap", name, buf.Len(), maxProfileBytes)
+	}
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), buf.Len(), nil
+}

agent/internal/heartbeat/handlers_diag_test.go

@@ -0,0 +1,164 @@
+package heartbeat
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"testing"
+
+	"github.com/breeze-rmm/agent/internal/collectors"
+)
+
+// decodeDiagResult parses the JSON payload NewSuccessResult marshals into
+// Stdout.
+func decodeDiagResult(t *testing.T, stdout string) map[string]any {
+	t.Helper()
+	var result map[string]any
+	if err := json.Unmarshal([]byte(stdout), &result); err != nil {
+		t.Fatalf("result stdout is not valid JSON: %v", err)
+	}
+	return result
+}
+
+// assertValidPprofBlob base64-decodes the named field and checks it looks
+// like a debug=0 runtime/pprof profile (gzip-compressed protobuf, magic
+// bytes 0x1f 0x8b).
+func assertValidPprofBlob(t *testing.T, result map[string]any, field string) {
+	t.Helper()
+	b64, ok := result[field].(string)
+	if !ok || b64 == "" {
+		t.Fatalf("%s missing or not a string", field)
+	}
+	raw, err := base64.StdEncoding.DecodeString(b64)
+	if err != nil {
+		t.Fatalf("%s is not valid base64: %v", field, err)
+	}
+	if len(raw) < 2 || raw[0] != 0x1f || raw[1] != 0x8b {
+		t.Fatalf("%s does not start with gzip magic bytes (got % x)", field, raw[:min(2, len(raw))])
+	}
+}
+
+func TestHandleCapturePprofDefaultCapturesBoth(t *testing.T) {
+	res := handleCapturePprof(nil, Command{ID: "c1", Type: "capture_pprof"})
+	if res.Status != "completed" {
+		t.Fatalf("status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+	result := decodeDiagResult(t, res.Stdout)
+	assertValidPprofBlob(t, result, "heapProfileBase64")
+	assertValidPprofBlob(t, result, "goroutineProfileBase64")
+
+	if _, ok := result["capturedAt"].(string); !ok {
+		t.Error("capturedAt missing")
+	}
+	rt, ok := result["runtime"].(map[string]any)
+	if !ok {
+		t.Fatal("runtime stats snapshot missing")
+	}
+	if goroutines, _ := rt["goroutines"].(float64); goroutines < 1 {
+		t.Errorf("runtime.goroutines = %v, want >= 1", rt["goroutines"])
+	}
+}
+
+func TestHandleCapturePprofHeapOnly(t *testing.T) {
+	res := handleCapturePprof(nil, Command{
+		ID: "c2", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "heap"},
+	})
+	if res.Status != "completed" {
+		t.Fatalf("status = %q (error: %q), want completed", res.Status, res.Error)
+	}
+	result := decodeDiagResult(t, res.Stdout)
+	assertValidPprofBlob(t, result, "heapProfileBase64")
+	if _, present := result["goroutineProfileBase64"]; present {
+		t.Error("goroutineProfileBase64 should not be present for profile=heap")
+	}
+}
+
+func TestHandleCapturePprofGoroutineOnly(t *testing.T) {
+	res := handleCapturePprof(nil, Command{
+		ID: "c3", 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)
+	assertValidPprofBlob(t, result, "goroutineProfileBase64")
+	if _, present := result["heapProfileBase64"]; present {
+		t.Error("heapProfileBase64 should not be present for profile=goroutine")
+	}
+}
+
+func TestHandleCapturePprofRejectsUnknownProfile(t *testing.T) {
+	res := handleCapturePprof(nil, Command{
+		ID: "c4", Type: "capture_pprof",
+		Payload: map[string]any{"profile": "cpu"},
+	})
+	if res.Status != "failed" {
+		t.Fatalf("status = %q, want failed for unknown profile", res.Status)
+	}
+	if res.Error == "" {
+		t.Error("expected a descriptive error for an unknown profile")
+	}
+}
+
+func TestHandleCapturePprofSizeCap(t *testing.T) {
+	orig := maxProfileBytes
+	maxProfileBytes = 1 // every real profile exceeds one byte
+	defer func() { maxProfileBytes = orig }()
+
+	res := handleCapturePprof(nil, Command{ID: "c5", Type: "capture_pprof"})
+	if res.Status != "failed" {
+		t.Fatalf("status = %q, want failed when profile exceeds size cap", res.Status)
+	}
+	if res.Error == "" {
+		t.Error("expected size-cap error message")
+	}
+}
+
+func TestHandleCapturePprofRejectsNonStringProfile(t *testing.T) {
+	res := handleCapturePprof(nil, Command{
+		ID: "c6", Type: "capture_pprof",
+		Payload: map[string]any{"profile": 123},
+	})
+	if res.Status != "failed" {
+		t.Fatalf("status = %q, want failed for non-string profile", res.Status)
+	}
+	if res.Error == "" {
+		t.Error("expected a descriptive error for a non-string profile")
+	}
+}
+
+// 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
+// renaming the struct tag — or dropping the assignment in sendHeartbeat —
+// would silently darken the gauges fleet-wide while every other test on both
+// sides stayed green.
+func TestHeartbeatPayloadAgentRuntimeWireKey(t *testing.T) {
+	withStats, err := json.Marshal(HeartbeatPayload{
+		Status:       "ok",
+		AgentRuntime: collectors.CollectRuntimeStats(),
+	})
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	var m map[string]any
+	if err := json.Unmarshal(withStats, &m); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	if _, ok := m["agentRuntime"]; !ok {
+		t.Error(`payload with AgentRuntime set must serialize an "agentRuntime" key`)
+	}
+
+	withoutStats, err := json.Marshal(HeartbeatPayload{Status: "ok"})
+	if err != nil {
+		t.Fatalf("marshal: %v", err)
+	}
+	m = map[string]any{}
+	if err := json.Unmarshal(withoutStats, &m); err != nil {
+		t.Fatalf("unmarshal: %v", err)
+	}
+	if _, ok := m["agentRuntime"]; ok {
+		t.Error("nil AgentRuntime must be omitted (omitempty) so old-server compat is preserved")
+	}
+}

agent/internal/heartbeat/handlers_test.go

@@ -91,6 +91,9 @@ var allCommandTypes = []string{
 	// handlers.go — log shipping
 	tools.CmdSetLogLevel,
 
+	// handlers.go — runtime diagnostics (handlers_diag.go)
+	tools.CmdCapturePprof,
+
 	// handlers_autoupdate.go
 	tools.CmdSetAutoUpdate,