fix(frontend): apply validated max-nodes default to v1 flame graph output

pkg/frontend/frontend_diff.go

@@ -53,12 +53,13 @@ func (f *Frontend) Diff(
 	var left, right *phlaremodel.FunctionNameTree
 	g.Go(func() error {
 		var leftErr error
-		left, leftErr = f.selectMergeStacktracesTree(ctx, connect.NewRequest(c.Msg.Left))
+		// maxNodes is validated above and applied via NewFlamegraphDiff below.
+		left, _, leftErr = f.selectMergeStacktracesTree(ctx, connect.NewRequest(c.Msg.Left))
 		return leftErr
 	})
 	g.Go(func() error {
 		var rightErr error
-		right, rightErr = f.selectMergeStacktracesTree(ctx, connect.NewRequest(c.Msg.Right))
+		right, _, rightErr = f.selectMergeStacktracesTree(ctx, connect.NewRequest(c.Msg.Right))
 		return rightErr
 	})
 	if err = g.Wait(); err != nil {

pkg/frontend/frontend_select_merge_span_profile.go

@@ -84,9 +84,12 @@ func (f *Frontend) SelectMergeSpanProfile(
 	var resp querierv1.SelectMergeSpanProfileResponse
 	switch c.Msg.Format {
 	default:
-		resp.Flamegraph = phlaremodel.NewFlameGraph(t, c.Msg.GetMaxNodes())
+		// Use the validated maxNodes (default applied, clamped to the max), not
+		// c.Msg.GetMaxNodes(): the request value is 0 when omitted, which would
+		// disable truncation of the final flame graph and bypass both limits.
+		resp.Flamegraph = phlaremodel.NewFlameGraph(t, maxNodes)
 	case querierv1.ProfileFormat_PROFILE_FORMAT_TREE:
-		resp.Tree = t.Bytes(c.Msg.GetMaxNodes(), nil)
+		resp.Tree = t.Bytes(maxNodes, nil)
 	}
 	return connect.NewResponse(&resp), nil
 }

pkg/frontend/frontend_select_merge_stacktraces.go

@@ -29,40 +29,47 @@ func (f *Frontend) SelectMergeStacktraces(
 	if len(c.Msg.TraceIdSelector) > 0 {
 		return nil, connect.NewError(connect.CodeUnimplemented, errors.New("trace_id_selector is only supported with the v2 query backend"))
 	}
-	t, err := f.selectMergeStacktracesTree(ctx, c)
+	t, maxNodes, err := f.selectMergeStacktracesTree(ctx, c)
 	if err != nil {
 		return nil, err
 	}
 	var resp querierv1.SelectMergeStacktracesResponse
 	switch c.Msg.Format {
 	default:
-		resp.Flamegraph = phlaremodel.NewFlameGraph(t, c.Msg.GetMaxNodes())
+		resp.Flamegraph = phlaremodel.NewFlameGraph(t, maxNodes)
 	case querierv1.ProfileFormat_PROFILE_FORMAT_TREE:
-		resp.Tree = t.Bytes(c.Msg.GetMaxNodes(), nil)
+		resp.Tree = t.Bytes(maxNodes, nil)
 	}
 	return connect.NewResponse(&resp), nil
 }
 
+// selectMergeStacktracesTree merges the matching profiles into a single tree.
+// It returns the validated maxNodes (with the per-tenant default applied when
+// the request omits it, and clamped to the configured maximum) so that callers
+// truncate the final flame graph / tree with the same limit used for the
+// per-query fan-out. Callers MUST use the returned maxNodes rather than
+// c.Msg.GetMaxNodes(): the request value is 0 when the client omits it, which
+// disables truncation entirely and bypasses both the default and the max limit.
 func (f *Frontend) selectMergeStacktracesTree(
 	ctx context.Context,
 	c *connect.Request[querierv1.SelectMergeStacktracesRequest],
-) (*phlaremodel.FunctionNameTree, error) {
+) (*phlaremodel.FunctionNameTree, int64, error) {
 	ctx = connectgrpc.WithProcedure(ctx, querierv1connect.QuerierServiceSelectMergeStacktracesProcedure)
 	tenantIDs, err := tenant.TenantIDs(ctx)
 	if err != nil {
-		return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		return nil, 0, connect.NewError(connect.CodeInvalidArgument, err)
 	}
 
 	validated, err := validation.ValidateRangeRequest(f.limits, tenantIDs, model.Interval{Start: model.Time(c.Msg.Start), End: model.Time(c.Msg.End)}, model.Now())
 	if err != nil {
-		return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		return nil, 0, connect.NewError(connect.CodeInvalidArgument, err)
 	}
 	if validated.IsEmpty {
-		return new(phlaremodel.FunctionNameTree), nil
+		return new(phlaremodel.FunctionNameTree), 0, nil
 	}
 	maxNodes, err := validation.ValidateMaxNodes(f.limits, tenantIDs, c.Msg.GetMaxNodes())
 	if err != nil {
-		return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		return nil, 0, connect.NewError(connect.CodeInvalidArgument, err)
 	}
 
 	g, ctx := errgroup.WithContext(ctx)
@@ -102,8 +109,8 @@ func (f *Frontend) selectMergeStacktracesTree(
 	}
 
 	if err = g.Wait(); err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	return m.Tree(), nil
+	return m.Tree(), maxNodes, nil
 }

pkg/frontend/frontend_select_merge_stacktraces_test.go

@@ -0,0 +1,114 @@
+package frontend
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"connectrpc.com/connect"
+	"github.com/grafana/dskit/user"
+	"github.com/stretchr/testify/require"
+
+	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
+	"github.com/grafana/pyroscope/v2/pkg/model"
+	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend"
+	"github.com/grafana/pyroscope/v2/pkg/util/connectgrpc"
+	"github.com/grafana/pyroscope/v2/pkg/util/httpgrpc"
+)
+
+// wideTreeRoundTripper answers SelectMergeStacktraces sub-requests with a tree
+// of `leaves` sibling leaves under the root, returned untruncated (maxNodes=-1)
+// regardless of the per-query maxNodes. This mimics the real fan-out/merge: each
+// shard truncates to the limit, but their union is much wider, so the final
+// merged tree the frontend assembles is larger than the node limit.
+func wideTreeRoundTripper(leaves int) *mockRoundTripper {
+	return &mockRoundTripper{callback: func(ctx context.Context, req *httpgrpc.HTTPRequest) (*httpgrpc.HTTPResponse, error) {
+		return connectgrpc.HandleUnary[querierv1.SelectMergeStacktracesRequest, querierv1.SelectMergeStacktracesResponse](ctx, req, func(ctx context.Context, req *connect.Request[querierv1.SelectMergeStacktracesRequest]) (*connect.Response[querierv1.SelectMergeStacktracesResponse], error) {
+			s := new(model.FunctionNameTree)
+			// Distinct, descending weights so truncation to the node limit pushes
+			// the long tail of small leaves into an "other" bucket.
+			for i := 0; i < leaves; i++ {
+				s.InsertStack(int64(leaves-i), "root", model.FunctionName("leaf"+string(rune('a'+i))))
+			}
+			return connect.NewResponse(&querierv1.SelectMergeStacktracesResponse{
+				Flamegraph: model.NewFlameGraph(s, -1),
+			}), nil
+		})
+	}}
+}
+
+func flameGraphNodeCount(fg *querierv1.FlameGraph) int {
+	var n int
+	for _, l := range fg.GetLevels() {
+		n += len(l.GetValues()) / 4
+	}
+	return n
+}
+
+func flameGraphHasName(fg *querierv1.FlameGraph, name string) bool {
+	for _, n := range fg.GetNames() {
+		if n == name {
+			return true
+		}
+	}
+	return false
+}
+
+// Test_Frontend_SelectMergeStacktraces_MaxNodesDefault is a regression test for
+// the bug where an omitted maxNodes disabled truncation of the final merged
+// flame graph: SelectMergeStacktraces built the response with
+// c.Msg.GetMaxNodes() (0 when omitted) instead of the validated maxNodes, so
+// the configured default (and the max ceiling) were bypassed and the full tree
+// was returned. The default must be applied when the client omits maxNodes.
+func Test_Frontend_SelectMergeStacktraces_MaxNodesDefault(t *testing.T) {
+	const (
+		defaultNodes = 4
+		maxNodes     = 100_000
+		leaves       = 32
+	)
+
+	limits := mockfrontend.NewMockLimits(t)
+	limits.On("MaxFlameGraphNodesDefault", "test").Return(defaultNodes).Maybe()
+	limits.On("MaxFlameGraphNodesMax", "test").Return(maxNodes).Maybe()
+	limits.On("MaxQueryLookback", "test").Return(time.Hour * 24).Maybe()
+	limits.On("MaxQueryLength", "test").Return(time.Hour).Maybe()
+	limits.On("MaxQueryParallelism", "test").Return(100).Maybe()
+	limits.On("QuerySplitDuration", "test").Return(time.Hour).Maybe()
+
+	frontend := Frontend{limits: limits, GRPCRoundTripper: wideTreeRoundTripper(leaves)}
+	ctx := user.InjectOrgID(context.Background(), "test")
+	now := time.Now().UnixMilli()
+	profileType := "memory:inuse_space:bytes:space:byte"
+
+	newReq := func(maxNodes *int64) *connect.Request[querierv1.SelectMergeStacktracesRequest] {
+		return connect.NewRequest(&querierv1.SelectMergeStacktracesRequest{
+			ProfileTypeID: profileType,
+			LabelSelector: "{}",
+			Start:         now,
+			End:           now + 1000,
+			MaxNodes:      maxNodes,
+		})
+	}
+
+	t.Run("omitted maxNodes applies the configured default", func(t *testing.T) {
+		resp, err := frontend.SelectMergeStacktraces(ctx, newReq(nil))
+		require.NoError(t, err)
+		fg := resp.Msg.GetFlamegraph()
+		require.NotNil(t, fg)
+		// Truncation to the default must collapse the surplus leaves into "other"
+		// and bound the node count, instead of returning the full wide tree.
+		require.True(t, flameGraphHasName(fg, "other"),
+			"expected an 'other' node from default truncation; got names=%v", fg.GetNames())
+		require.LessOrEqual(t, flameGraphNodeCount(fg), defaultNodes+2,
+			"node count must be bounded by the default max-nodes")
+		require.Less(t, flameGraphNodeCount(fg), leaves,
+			"omitted maxNodes must not return the full untruncated tree")
+	})
+
+	t.Run("explicit maxNodes above the max is rejected", func(t *testing.T) {
+		n := int64(maxNodes + 1)
+		_, err := frontend.SelectMergeStacktraces(ctx, newReq(&n))
+		require.Error(t, err)
+		require.ErrorContains(t, err, "max flamegraph nodes")
+	})
+}