flamegraph: guard SnapshotJSON cache writes (5p)

2 files changed, 33 insertions, 2 deletions

diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 47299e6..4554780 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -234,8 +234,11 @@ func (lt *LiveTrie) SnapshotJSON() ([]byte, uint64) {
 	}
 
 	lt.cacheMu.Lock()
-	lt.cacheVersion = version
-	lt.cacheJSON = slices.Clone(payload)
+	// Only commit if no concurrent caller stored a newer version.
+	if version >= lt.cacheVersion {
+		lt.cacheVersion = version
+		lt.cacheJSON = slices.Clone(payload)
+	}
 	lt.cacheMu.Unlock()
 
 	return payload, version
diff --git a/internal/flamegraph/livetrie_test.go b/internal/flamegraph/livetrie_test.go
index fc9e6a6..a52b72e 100644
--- a/internal/flamegraph/livetrie_test.go
+++ b/internal/flamegraph/livetrie_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"runtime"
+	"slices"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -439,6 +440,33 @@ func TestLiveTrieSnapshotJSONCaching(t *testing.T) {
 	}
 }
 
+func TestLiveTrieSnapshotJSONSkipsStaleCacheWrite(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm"}, "count", "count")
+	lt.Ingest(newTestPair("svc", 42, 1001, "/tmp/a", 1, 1, 1))
+
+	_, version := lt.SnapshotJSON()
+	newerPayload := []byte(`{"n":"newer","v":7}`)
+
+	lt.cacheMu.Lock()
+	lt.cacheVersion = version + 1
+	lt.cacheJSON = slices.Clone(newerPayload)
+	lt.cacheMu.Unlock()
+
+	_, _ = lt.SnapshotJSON()
+
+	lt.cacheMu.Lock()
+	gotVersion := lt.cacheVersion
+	gotPayload := slices.Clone(lt.cacheJSON)
+	lt.cacheMu.Unlock()
+
+	if gotVersion != version+1 {
+		t.Fatalf("cache version overwritten by stale snapshot: got %d want %d", gotVersion, version+1)
+	}
+	if !bytes.Equal(gotPayload, newerPayload) {
+		t.Fatalf("cache payload overwritten by stale snapshot: got %q want %q", gotPayload, newerPayload)
+	}
+}
+
 func TestLiveTrieSnapshotJSONPrunesTinyNodes(t *testing.T) {
 	lt := NewLiveTrie([]string{"comm"}, "count", "count")
 	for i := 0; i < 2000; i++ {