Use pyroscope tree representation in MergeProfilesStacktraces API (#702)

* Add pyroscope tree to the MergeProfilesStacktracesResult type

* Add support for multiple merge result formats

* Refactor tree and flamegraph types to the model package

* Update reference-configuration-parameters

* Add stack truncation implementation notes

* Implement tree merge

* Fix tree merge

* Fix stacktrace rewrite

* Refine tree merge

* Fix tests

* Refactor unused functions

* Remove dead code

* Fix tree min value counting

Author: kolesnikovae

api/gen/proto/go/ingester/v1/ingester.pb.go

@@ -42,18 +42,29 @@ message SelectProfilesRequest {
   int64 start = 3;
   int64 end = 4;
 }
+
 message MergeProfilesStacktracesRequest {
   // The client starts the stream with a request containing the profile type and the labels.
   SelectProfilesRequest request = 1;
-
+  // Max nodes in the resulting tree.
+  optional int64 max_nodes = 3;
   // On a batch of profiles, the client sends the profiles to keep for merging.
   repeated bool profiles = 2;
 }
 
 message MergeProfilesStacktracesResult {
+  StacktracesMergeFormat format = 3;
   // The list of stracktraces with their respective value
   repeated StacktraceSample stacktraces = 1;
   repeated string function_names = 2;
+  // Merge result marshaled to pyroscope tree bytes.
+  bytes tree_bytes = 4;
+}
+
+enum StacktracesMergeFormat {
+  MERGE_FORMAT_UNSPECIFIED = 0;
+  MERGE_FORMAT_STACKTRACES = 1;
+  MERGE_FORMAT_TREE = 2;
 }
 
 message MergeProfilesStacktracesResponse {

api/gen/proto/go/ingester/v1/ingester_vtproto.pb.go

@@ -722,6 +722,9 @@
     "v1MergeProfilesStacktracesResult": {
       "type": "object",
       "properties": {
+        "format": {
+          "$ref": "#/definitions/v1StacktracesMergeFormat"
+        },
         "stacktraces": {
           "type": "array",
           "items": {
@@ -734,6 +737,11 @@
           "items": {
             "type": "string"
           }
+        },
+        "treeBytes": {
+          "type": "string",
+          "format": "byte",
+          "description": "Merge result marshaled to pyroscope tree bytes."
         }
       }
     },
@@ -967,6 +975,15 @@
         }
       }
     },
+    "v1StacktracesMergeFormat": {
+      "type": "string",
+      "enum": [
+        "MERGE_FORMAT_UNSPECIFIED",
+        "MERGE_FORMAT_STACKTRACES",
+        "MERGE_FORMAT_TREE"
+      ],
+      "default": "MERGE_FORMAT_UNSPECIFIED"
+    },
     "v1State": {
       "type": "string",
       "enum": [

api/ingester/v1/ingester.proto

@@ -153,7 +153,7 @@ client:
 
 api:
   # base URL for when the server is behind a reverse proxy with a different path
-  # CLI flag: -base-url
+  # CLI flag: -api.base-url
   [base-url: <string> | default = ""]
 
 # The server block configures the HTTP and gRPC server of the launched

api/openapiv2/gen/phlare.swagger.json

@@ -6,6 +6,7 @@ require (
 	github.com/bufbuild/connect-go v1.4.1
 	github.com/bufbuild/connect-grpchealth-go v1.0.0
 	github.com/cespare/xxhash/v2 v2.2.0
+	github.com/dennwc/varint v1.0.0
 	github.com/drone/envsubst v1.0.3
 	github.com/dustin/go-humanize v1.0.0
 	github.com/felixge/fgprof v0.9.4-0.20221116204635-ececf7638e93
@@ -115,7 +116,6 @@ require (
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.5.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/dennwc/varint v1.0.0 // indirect
 	github.com/dgraph-io/badger/v2 v2.2007.2 // indirect
 	github.com/dgraph-io/ristretto v0.1.1 // indirect
 	github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 // indirect

docs/sources/configure-server/reference-configuration-parameters/index.md

@@ -1,8 +1,6 @@
-package querier
+package model
 
 import (
-	"container/heap"
-
 	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
 	"github.com/pyroscope-io/pyroscope/pkg/structs/flamebearer"
 	"github.com/samber/lo"
@@ -11,13 +9,7 @@ import (
 	typesv1 "github.com/grafana/phlare/api/gen/proto/go/types/v1"
 )
 
-type stackNode struct {
-	xOffset int
-	level   int
-	node    *node
-}
-
-func NewFlameGraph(t *tree, maxNodes int64) *querierv1.FlameGraph {
+func NewFlameGraph(t *Tree, maxNodes int64) *querierv1.FlameGraph {
 	var total, max int64
 	for _, node := range t.root {
 		total += node.total
@@ -175,37 +167,3 @@ func ExportDiffToFlamebearer(fg *querierv1.FlameGraphDiff, profileType *typesv1.
 
 	return fb
 }
-
-// minValue returns the minimum "total" value a node in a tree has to have to show up in
-// the resulting flamegraph
-func (t *tree) minValue(maxNodes int64) int64 {
-	if maxNodes == -1 { // -1 means show all nodes
-		return 0
-	}
-	nodes := t.root
-
-	mh := &minHeap{}
-	heap.Init(mh)
-
-	for len(nodes) > 0 {
-		node := nodes[0]
-		nodes = nodes[1:]
-		number := node.total
-
-		if mh.Len() < int(maxNodes) {
-			heap.Push(mh, number)
-		} else {
-			if number > (*mh)[0] {
-				heap.Pop(mh)
-				heap.Push(mh, number)
-				nodes = append(node.children, nodes...)
-			}
-		}
-	}
-
-	if mh.Len() < int(maxNodes) {
-		return 0
-	}
-
-	return (*mh)[0]
-}

go.mod

@@ -1,4 +1,4 @@
-package querier
+package model
 
 import (
 	"bytes"
@@ -25,7 +25,7 @@ const MaxNodes = 8192
 //	i+4 = total   , right tree
 //	i+5 = self    , right tree
 //	i+6 = index in the names array
-func NewFlamegraphDiff(left, right *tree, maxNodes int) (*querierv1.FlameGraphDiff, error) {
+func NewFlamegraphDiff(left, right *Tree, maxNodes int) (*querierv1.FlameGraphDiff, error) {
 	// The algorithm doesn't work properly with negative nodes
 	// Although it's possible to silently drop these nodes
 	// Let's fail early and analyze properly with real data when the issue happens
@@ -148,7 +148,7 @@ func NewFlamegraphDiff(left, right *tree, maxNodes int) (*querierv1.FlameGraphDi
 // combineTree aligns 2 trees by making them having the same structure with the
 // same number of nodes
 // It also makes the tree have a single root
-func combineTree(leftTree, rightTree *tree) (*tree, *tree) {
+func combineTree(leftTree, rightTree *Tree) (*Tree, *Tree) {
 	leftTotal := int64(0)
 	for _, l := range leftTree.root {
 		leftTotal = leftTotal + l.total
@@ -161,15 +161,15 @@ func combineTree(leftTree, rightTree *tree) (*tree, *tree) {
 
 	// differently from pyroscope, there could be multiple roots
 	// so we add a fake root as expected
-	leftTree = &tree{
+	leftTree = &Tree{
 		root: []*node{{
 			children: leftTree.root,
 			total:    leftTotal,
 			self:     0,
 		}},
 	}
 
-	rightTree = &tree{
+	rightTree = &Tree{
 		root: []*node{{
 			children: rightTree.root,
 			total:    rightTotal,
@@ -263,7 +263,7 @@ func nextPow2(a int) int {
 	return a
 }
 
-func combineMinValues(leftTree, rightTree *tree, maxNodes int) uint64 {
+func combineMinValues(leftTree, rightTree *Tree, maxNodes int) uint64 {
 	c := cappedarr.New(maxNodes)
 	combineIterateWithTotal(leftTree, rightTree, func(left uint64, right uint64) bool {
 		return c.Push(maxUint64(left, right))
@@ -272,7 +272,7 @@ func combineMinValues(leftTree, rightTree *tree, maxNodes int) uint64 {
 }
 
 // iterate both trees, both trees must be returned from combineTree
-func combineIterateWithTotal(leftTree, rightTree *tree, cb func(uint64, uint64) bool) {
+func combineIterateWithTotal(leftTree, rightTree *Tree, cb func(uint64, uint64) bool) {
 	leftNodes, rghtNodes := leftTree.root, rightTree.root
 	i := 0
 	for len(leftNodes) > 0 {
@@ -289,7 +289,7 @@ func combineIterateWithTotal(leftTree, rightTree *tree, cb func(uint64, uint64)
 }
 
 // isPositiveTree returns whether a tree only contain positive values
-func isPositiveTree(t *tree) bool {
+func isPositiveTree(t *Tree) bool {
 	stack := Stack[*node]{}
 	for _, node := range t.root {
 		stack.Push(node)
@@ -313,7 +313,7 @@ func isPositiveTree(t *tree) bool {
 	return true
 }
 
-func assertPositiveTrees(left *tree, right *tree) error {
+func assertPositiveTrees(left *Tree, right *Tree) error {
 	leftRes := isPositiveTree(left)
 	rightRes := isPositiveTree(right)
 

pkg/querier/flamegraph.go renamed to pkg/model/flamegraph.go

@@ -1,4 +1,4 @@
-package querier
+package model
 
 import (
 	"testing"

pkg/querier/flamegraph_diff.go renamed to pkg/model/flamegraph_diff.go

@@ -1,4 +1,4 @@
-package querier
+package model
 
 import (
 	"fmt"