Reduce memory allocations when rendering diff flamegraphs (#4430)

* Reduce memory allocations when rendering diff flamegraphs

* Add benchmark

---------

Co-authored-by: Christian Simon <[email protected]>

Author: aleks-p

pkg/model/flamegraph_diff.go

@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"fmt"
 
-	"github.com/grafana/pyroscope/pkg/og/structs/cappedarr"
+	"github.com/grafana/pyroscope/pkg/util/minheap"
 
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
 )
@@ -95,7 +95,8 @@ func NewFlamegraphDiff(left, right *Tree, maxNodes int64) (*querierv1.FlameGraph
 				int64(i),
 			}
 
-			res.Levels[level].Values = append(values, res.Levels[level].Values...)
+			// We need to prepend values here, but this is expensive. We'll reverse the order later.
+			res.Levels[level].Values = append(res.Levels[level].Values, values...)
 			xLeftOffset += left.self
 			xRghtOffset += rght.self
 			otherLeftTotal, otherRghtTotal := int64(0), int64(0)
@@ -140,6 +141,11 @@ func NewFlamegraphDiff(left, right *Tree, maxNodes int64) (*querierv1.FlameGraph
 		}
 	}
 
+	// reverse each level's values since we appended values instead of prepending them
+	for _, level := range res.Levels {
+		reverseSliceInChunks(level.Values, 7)
+	}
+
 	deltaEncoding(res.Levels, 0, 7)
 	deltaEncoding(res.Levels, 3, 7)
 
@@ -259,18 +265,29 @@ func combineMinValues(leftTree, rightTree *Tree, maxNodes int) uint64 {
 	if maxNodes < 1 {
 		return 0
 	}
-	// Trees are combined, meaning that their structures are
-	// identical, therefore the resulting tree can not have
-	// more nodes than any of them.
 	treeSize := leftTree.size(make([]*node, 0, defaultDFSSize))
 	if treeSize <= int64(maxNodes) {
 		return 0
 	}
-	c := cappedarr.New(maxNodes)
+
+	h := make([]int64, 0, maxNodes)
 	combineIterateWithTotal(leftTree, rightTree, func(left uint64, right uint64) bool {
-		return c.Push(max(left, right))
+		maxVal := int64(max(left, right))
+		if len(h) >= maxNodes {
+			if maxVal > h[0] {
+				h = minheap.Pop(h)
+				h = minheap.Push(h, maxVal)
+			}
+		} else {
+			h = minheap.Push(h, maxVal)
+		}
+		return true
 	})
-	return c.MinValue()
+
+	if len(h) < maxNodes {
+		return 0
+	}
+	return uint64(h[0])
 }
 
 // iterate both trees, both trees must be returned from combineTree
@@ -364,3 +381,16 @@ func prependTreeNode(s []*node, x *node) []*node {
 	s[0] = x
 	return s
 }
+
+// reverseSliceInChunks reverses a slice in chunks of the given size
+func reverseSliceInChunks(slice []int64, chunkSize int) {
+	if len(slice) < chunkSize {
+		return
+	}
+
+	for i, j := 0, len(slice)-chunkSize; i < j; i, j = i+chunkSize, j-chunkSize {
+		for k := 0; k < chunkSize; k++ {
+			slice[i+k], slice[j+k] = slice[j+k], slice[i+k]
+		}
+	}
+}

pkg/model/flamegraph_diff_bench_test.go

@@ -0,0 +1,30 @@
+package model
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Benchmark_NewFlamegraphDiff(b *testing.B) {
+	leftTreeBytes, err := os.ReadFile("testdata/diff_left_tree.bin")
+	require.NoError(b, err)
+	rightTreeBytes, err := os.ReadFile("testdata/diff_right_tree.bin")
+	require.NoError(b, err)
+
+	leftTree, err := UnmarshalTree(leftTreeBytes)
+	require.NoError(b, err)
+
+	rightTree, err := UnmarshalTree(rightTreeBytes)
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		diff, err := NewFlamegraphDiff(leftTree, rightTree, 163840)
+		require.NoError(b, err)
+		require.NotNil(b, diff)
+	}
+}