feat: Support stacktrace selector for SelectMergeStacktraces requests (#4380)

* Add stack_trace_selector to SelectMergeStacktracesRequest

* Implement stacktrace selectors for SelectMergeStacktraces

Only implements for the small case

* Implement the tree single pass

* Use truncation mark instead of sentinel values

* Make lint happy

* Add tests and fix implementation for tree resolver

* Test with max nodes active

* Refactor tree matching for readability

---------

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

Author: bryanhuhta

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

@@ -2108,6 +2108,10 @@
         "format": {
           "$ref": "#/definitions/v1ProfileFormat",
           "description": "Profile format specifies the format of profile to be returned.\nIf not specified, the profile will be returned in flame graph format."
+        },
+        "stackTraceSelector": {
+          "$ref": "#/definitions/v1StackTraceSelector",
+          "description": "Select stack traces that match the provided selector."
         }
       }
     },
@@ -2432,6 +2436,9 @@
           "items": {
             "type": "string"
           }
+        },
+        "stackTraceSelector": {
+          "$ref": "#/definitions/v1StackTraceSelector"
         }
       }
     },

api/gen/proto/go/querier/v1/querier_vtproto.pb.go

@@ -71,6 +71,8 @@ message SelectMergeStacktracesRequest {
   // Profile format specifies the format of profile to be returned.
   // If not specified, the profile will be returned in flame graph format.
   ProfileFormat format = 6;
+  // Select stack traces that match the provided selector.
+  optional types.v1.StackTraceSelector stack_trace_selector = 7;
 }
 
 enum ProfileFormat {

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

@@ -171,6 +171,7 @@ message TimeSeriesReport {
 message TreeQuery {
   int64 max_nodes = 1;
   repeated string span_selector = 2;
+  optional types.v1.StackTraceSelector stack_trace_selector = 3;
 }
 
 message TreeReport {

api/gen/proto/go/query/v1/query_vtproto.pb.go

@@ -64,7 +64,10 @@ func (q *QueryFrontend) selectMergeStacktracesTree(
 		LabelSelector: labelSelector,
 		Query: []*queryv1.Query{{
 			QueryType: queryv1.QueryType_QUERY_TREE,
-			Tree:      &queryv1.TreeQuery{MaxNodes: maxNodes},
+			Tree: &queryv1.TreeQuery{
+				MaxNodes:           maxNodes,
+				StackTraceSelector: c.Msg.StackTraceSelector,
+			},
 		}},
 	})
 	if err != nil {

api/openapiv2/gen/phlare.swagger.json

@@ -242,7 +242,7 @@ func (r *Resolver) Tree() (*model.Tree, error) {
 	var lock sync.Mutex
 	tree := new(model.Tree)
 	err := r.withSymbols(ctx, func(symbols *Symbols, appender *SampleAppender) error {
-		resolved, err := symbols.Tree(ctx, appender, r.maxNodes)
+		resolved, err := symbols.Tree(ctx, appender, r.maxNodes, SelectStackTraces(symbols, r.sts))
 		if err != nil {
 			return err
 		}
@@ -336,6 +336,7 @@ func (r *Symbols) Tree(
 	ctx context.Context,
 	appender *SampleAppender,
 	maxNodes int64,
+	selection *SelectedStackTraces,
 ) (*model.Tree, error) {
-	return buildTree(ctx, r, appender, maxNodes)
+	return buildTree(ctx, r, appender, maxNodes, selection)
 }