From 3f1cbdd54b65f66ebf188f3d269af0271a39be3b Mon Sep 17 00:00:00 2001 From: Dimitry Linkov Date: Tue, 18 Aug 2026 12:30:04 -0700 Subject: [PATCH] fix(frontend): apply validated max-nodes default to v1 flame graph output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The v1 frontend SelectMergeStacktraces and SelectMergeSpanProfile handlers validated max-nodes (applying the per-tenant `max_flamegraph_nodes_default` and clamping to `max_flamegraph_nodes_max`) and used the validated value for the per-query fan-out — but then built the final merged flame graph / tree with the raw request value `c.Msg.GetMaxNodes()`. When a client omits max-nodes (e.g. Grafana Explore Profiles, which only sends the parameter when set), the raw value is 0. `Tree.minValue` treats maxNodes < 1 as "no limit", so the final flame graph was emitted untruncated: the configured default was ignored and the max ceiling was bypassed (the merged tree could exceed `max_flamegraph_nodes_max`, since that limit lives in ValidateMaxNodes which the final assembly skipped). The user-visible effect is that an omitted max-nodes renders a far larger tree than an explicit one, so "other" appears to grow when you *raise* max-nodes toward the default. Fix: return the validated maxNodes from selectMergeStacktracesTree and use it for both NewFlameGraph and Tree.Bytes; use the in-scope validated maxNodes in the span-profile handler. The v2 read path already threads the validated value through TreeQuery, so this aligns v1 with v2. Diff was already correct. Adds a regression test asserting an omitted max-nodes truncates to the configured default (produces an "other" node, bounded node count) and that an explicit value above the max is rejected. Co-Authored-By: Claude Opus 5 --- pkg/frontend/frontend_diff.go | 5 +- .../frontend_select_merge_span_profile.go | 7 +- .../frontend_select_merge_stacktraces.go | 27 +++-- .../frontend_select_merge_stacktraces_test.go | 114 ++++++++++++++++++ 4 files changed, 139 insertions(+), 14 deletions(-) create mode 100644 pkg/frontend/frontend_select_merge_stacktraces_test.go diff --git a/pkg/frontend/frontend_diff.go b/pkg/frontend/frontend_diff.go index 23493cd897..8d71b5f2d8 100644 --- a/pkg/frontend/frontend_diff.go +++ b/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 { diff --git a/pkg/frontend/frontend_select_merge_span_profile.go b/pkg/frontend/frontend_select_merge_span_profile.go index 11721e5cd9..ffa8804d52 100644 --- a/pkg/frontend/frontend_select_merge_span_profile.go +++ b/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 } diff --git a/pkg/frontend/frontend_select_merge_stacktraces.go b/pkg/frontend/frontend_select_merge_stacktraces.go index ea6d65bd1c..ec90059135 100644 --- a/pkg/frontend/frontend_select_merge_stacktraces.go +++ b/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 } diff --git a/pkg/frontend/frontend_select_merge_stacktraces_test.go b/pkg/frontend/frontend_select_merge_stacktraces_test.go new file mode 100644 index 0000000000..f487c990dd --- /dev/null +++ b/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") + }) +}