Fix real live flamegraph key handling and startup viewport sync

2 files changed, 104 insertions, 9 deletions

diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 0d42b6b..13d7de9 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -12,7 +12,11 @@ import (
 	"ior/internal/event"
 )
 
-const liveTrieMinFraction = 0.001
+const (
+	liveTrieMinFraction                     = 0.001
+	liveTrieMinVisibleChildrenWhenPruned    = 8
+	liveTrieVisibleChildrenFallbackMaxDepth = 1
+)
 
 type trieSnapshot struct {
 	Name     string          `json:"n"`
@@ -244,29 +248,45 @@ func subtreeTotal(node *trieNode) uint64 {
 }
 
 func buildSnapshot(node *trieNode, depth int, minFraction float64, rootTotal uint64) *trieSnapshot {
-	snapshot, _ := buildSnapshotWithTotal(node, depth, minFraction, rootTotal)
+	snapshot, _ := buildSnapshotWithTotal(node, depth, minFraction, rootTotal, false)
 	return snapshot
 }
 
-func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, rootTotal uint64) (*trieSnapshot, uint64) {
+type childSnapshotState struct {
+	node     *trieNode
+	snapshot *trieSnapshot
+	total    uint64
+}
+
+func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, rootTotal uint64, forceKeep bool) (*trieSnapshot, uint64) {
 	total := node.value
 	children := slices.Clone(node.children)
 	sort.Slice(children, func(i, j int) bool {
 		return children[i].name < children[j].name
 	})
 
-	childSnapshots := make([]*trieSnapshot, 0, len(children))
+	childStates := make([]childSnapshotState, 0, len(children))
 	for _, child := range children {
-		childSnapshot, childTotal := buildSnapshotWithTotal(child, depth+1, minFraction, rootTotal)
+		childSnapshot, childTotal := buildSnapshotWithTotal(child, depth+1, minFraction, rootTotal, false)
 		total += childTotal
-		if childSnapshot != nil {
-			childSnapshots = append(childSnapshots, childSnapshot)
-		}
+		childStates = append(childStates, childSnapshotState{
+			node:     child,
+			snapshot: childSnapshot,
+			total:    childTotal,
+		})
 	}
 
-	if depth > 0 && rootTotal > 0 && float64(total)/float64(rootTotal) < minFraction {
+	if !forceKeep && depth > 0 && rootTotal > 0 && float64(total)/float64(rootTotal) < minFraction {
 		return nil, total
 	}
+	ensureFallbackVisibleChildren(childStates, depth, minFraction, rootTotal)
+
+	childSnapshots := make([]*trieSnapshot, 0, len(childStates))
+	for _, child := range childStates {
+		if child.snapshot != nil {
+			childSnapshots = append(childSnapshots, child.snapshot)
+		}
+	}
 
 	snapshot := &trieSnapshot{
 		Name:  node.name,
@@ -278,3 +298,43 @@ func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, root
 	}
 	return snapshot, total
 }
+
+func ensureFallbackVisibleChildren(children []childSnapshotState, depth int, minFraction float64, rootTotal uint64) {
+	if depth > liveTrieVisibleChildrenFallbackMaxDepth {
+		return
+	}
+	visible := 0
+	for _, child := range children {
+		if child.snapshot != nil {
+			visible++
+		}
+	}
+	if visible > 0 {
+		return
+	}
+
+	candidates := make([]int, 0, len(children))
+	for idx, child := range children {
+		if child.total > 0 {
+			candidates = append(candidates, idx)
+		}
+	}
+	sort.Slice(candidates, func(i, j int) bool {
+		left := children[candidates[i]]
+		right := children[candidates[j]]
+		if left.total == right.total {
+			return left.node.name < right.node.name
+		}
+		return left.total > right.total
+	})
+
+	limit := liveTrieMinVisibleChildrenWhenPruned
+	if len(candidates) < limit {
+		limit = len(candidates)
+	}
+	for i := 0; i < limit; i++ {
+		idx := candidates[i]
+		forced, _ := buildSnapshotWithTotal(children[idx].node, depth+1, minFraction, rootTotal, true)
+		children[idx].snapshot = forced
+	}
+}
diff --git a/internal/flamegraph/livetrie_test.go b/internal/flamegraph/livetrie_test.go
index 632f668..c5ed32c 100644
--- a/internal/flamegraph/livetrie_test.go
+++ b/internal/flamegraph/livetrie_test.go
@@ -221,6 +221,41 @@ func TestLiveTrieSnapshotJSONPrunesTinyNodes(t *testing.T) {
 	}
 }
 
+func TestLiveTrieSnapshotJSONKeepsFallbackChildrenWhenAllAreTinyAtRoot(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm"}, "count")
+	const total = 6000
+	for i := 0; i < total; i++ {
+		comm := fmt.Sprintf("svc-%04d", i)
+		lt.Ingest(newTestPair(comm, 42, uint32(100000+i), "/tmp/a", 1, 1, 1))
+	}
+
+	snap := decodeLiveSnapshot(t, lt)
+	if len(snap.Children) == 0 {
+		t.Fatalf("expected fallback root children when pruning would hide every branch")
+	}
+	if got, want := len(snap.Children), liveTrieMinVisibleChildrenWhenPruned; got != want {
+		t.Fatalf("expected fallback to keep %d root children, got %d", want, got)
+	}
+}
+
+func TestLiveTrieSnapshotJSONKeepsFallbackChildrenAtDepthOne(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm", "pid"}, "count")
+	const total = 6000
+	for i := 0; i < total; i++ {
+		pid := uint32(100000 + i)
+		lt.Ingest(newTestPair("svc", pid, pid, "/tmp/a", 1, 1, 1))
+	}
+
+	snap := decodeLiveSnapshot(t, lt)
+	commNode := findSnapshotPath(t, &snap, "svc")
+	if len(commNode.Children) == 0 {
+		t.Fatalf("expected fallback depth-one children for pid branches")
+	}
+	if got, want := len(commNode.Children), liveTrieMinVisibleChildrenWhenPruned; got != want {
+		t.Fatalf("expected fallback to keep %d depth-one children, got %d", want, got)
+	}
+}
+
 func TestLiveTrieConcurrentIngestAndSnapshot(t *testing.T) {
 	lt := NewLiveTrie([]string{"comm", "pid"}, "count")