feat: Allow to optionally enforce maxNodes through limits (#4436)

Author: simonswine

cmd/pyroscope/help-all.txt.tmpl

@@ -579,6 +579,8 @@ Usage of ./pyroscope:
     	Maximum number of flame graph nodes by default. 0 to disable. (default 8192)
   -querier.max-flamegraph-nodes-max int
     	Maximum number of flame graph nodes allowed. 0 to disable. (default 1048576)
+  -querier.max-flamegraph-nodes-on-select-merge-profile
+    	Enforce the max nodes limits and defaults on SelectMergeProfile API. Historically this limit was not enforced to enable to gather full pprof profiles without truncation.
   -querier.max-query-length duration
     	The limit to length of queries. 0 to disable. (default 1d)
   -querier.max-query-lookback duration

pkg/frontend/frontend_diff_test.go

@@ -51,6 +51,10 @@ func (m *mockLimits) MaxFlameGraphNodesMax(_ string) int {
 	return 100_000
 }
 
+func (m *mockLimits) MaxFlameGraphNodesOnSelectMergeProfile(_ string) bool {
+	return true
+}
+
 func (m *mockLimits) SymbolizerEnabled(s string) bool { return true }
 
 type mockRoundTripper struct {

pkg/frontend/frontend_select_merge_profile.go

@@ -48,6 +48,21 @@ func (f *Frontend) SelectMergeProfile(
 	// NOTE: Max nodes limit is not set by default:
 	//   the method is used for pprof export and
 	//   truncation is not applicable for that.
+	maxNodesEnabled := false
+	for _, tenantID := range tenantIDs {
+		if f.limits.MaxFlameGraphNodesOnSelectMergeProfile(tenantID) {
+			maxNodesEnabled = true
+		}
+	}
+
+	maxNodes := c.Msg.MaxNodes
+	if maxNodesEnabled {
+		maxNodesV, err := validation.ValidateMaxNodes(f.limits, tenantIDs, c.Msg.GetMaxNodes())
+		if err != nil {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+		maxNodes = &maxNodesV
+	}
 
 	var m pprof.ProfileMerge
 	for intervals.Next() {
@@ -58,7 +73,7 @@ func (f *Frontend) SelectMergeProfile(
 				LabelSelector:      c.Msg.LabelSelector,
 				Start:              r.Start.UnixMilli(),
 				End:                r.End.UnixMilli(),
-				MaxNodes:           c.Msg.MaxNodes,
+				MaxNodes:           maxNodes,
 				StackTraceSelector: c.Msg.StackTraceSelector,
 			})
 			resp, err := connectgrpc.RoundTripUnary[

pkg/test/mocks/mockfrontend/mock_limits.go

@@ -87,9 +87,9 @@ type Limits struct {
 	QueryAnalysisSeriesEnabled bool           `yaml:"query_analysis_series_enabled" json:"query_analysis_series_enabled"`
 
 	// Flame graph enforced limits.
-	MaxFlameGraphNodesDefault int `yaml:"max_flamegraph_nodes_default" json:"max_flamegraph_nodes_default"`
-	MaxFlameGraphNodesMax     int `yaml:"max_flamegraph_nodes_max" json:"max_flamegraph_nodes_max"`
-
+	MaxFlameGraphNodesDefault              int  `yaml:"max_flamegraph_nodes_default" json:"max_flamegraph_nodes_default"`
+	MaxFlameGraphNodesMax                  int  `yaml:"max_flamegraph_nodes_max" json:"max_flamegraph_nodes_max"`
+	MaxFlameGraphNodesOnSelectMergeProfile bool `yaml:"max_flamegraph_nodes_on_select_merge_profile" json:"max_flamegraph_nodes_on_select_merge_profile" category:"advanced" doc:"hidden"`
 	// Store-gateway.
 	StoreGatewayTenantShardSize int `yaml:"store_gateway_tenant_shard_size" json:"store_gateway_tenant_shard_size"`
 
@@ -186,6 +186,7 @@ func (l *Limits) RegisterFlags(f *flag.FlagSet) {
 
 	f.IntVar(&l.MaxFlameGraphNodesDefault, "querier.max-flamegraph-nodes-default", 8<<10, "Maximum number of flame graph nodes by default. 0 to disable.")
 	f.IntVar(&l.MaxFlameGraphNodesMax, "querier.max-flamegraph-nodes-max", 1<<20, "Maximum number of flame graph nodes allowed. 0 to disable.")
+	f.BoolVar(&l.MaxFlameGraphNodesOnSelectMergeProfile, "querier.max-flamegraph-nodes-on-select-merge-profile", false, "Enforce the max nodes limits and defaults on SelectMergeProfile API. Historically this limit was not enforced to enable to gather full pprof profiles without truncation.")
 
 	f.Var(&l.DistributorAggregationWindow, "distributor.aggregation-window", "Duration of the distributor aggregation window. Requires aggregation period to be specified. 0 to disable.")
 	f.Var(&l.DistributorAggregationPeriod, "distributor.aggregation-period", "Duration of the distributor aggregation period. Requires aggregation window to be specified. 0 to disable.")
@@ -428,6 +429,11 @@ func (o *Overrides) MaxFlameGraphNodesMax(tenantID string) int {
 	return o.getOverridesForTenant(tenantID).MaxFlameGraphNodesMax
 }
 
+// MaxFlameGraphNodesOnSelectMergeProfiles returns if the max flame graph nodes should be enforced for the SelectMergeProfile API.
+func (o *Overrides) MaxFlameGraphNodesOnSelectMergeProfile(tenantID string) bool {
+	return o.getOverridesForTenant(tenantID).MaxFlameGraphNodesOnSelectMergeProfile
+}
+
 // StoreGatewayTenantShardSize returns the store-gateway shard size for a given user.
 func (o *Overrides) StoreGatewayTenantShardSize(userID string) int {
 	return o.getOverridesForTenant(userID).StoreGatewayTenantShardSize

pkg/test/mocks/mockobjstore/mock_bucket.go

@@ -16,8 +16,9 @@ type MockLimits struct {
 	MaxLabelValueLengthValue        int
 	MaxLabelNamesPerSeriesValue     int
 
-	MaxFlameGraphNodesDefaultValue int
-	MaxFlameGraphNodesMaxValue     int
+	MaxFlameGraphNodesDefaultValue              int
+	MaxFlameGraphNodesMaxValue                  int
+	MaxFlameGraphNodesOnSelectMergeProfileValue bool
 
 	DistributorAggregationWindowValue time.Duration
 	DistributorAggregationPeriodValue time.Duration
@@ -51,6 +52,9 @@ func (m MockLimits) QuerySanitizeOnMerge(tenantID string) bool {
 }
 func (m MockLimits) MaxFlameGraphNodesDefault(string) int { return m.MaxFlameGraphNodesDefaultValue }
 func (m MockLimits) MaxFlameGraphNodesMax(string) int     { return m.MaxFlameGraphNodesMaxValue }
+func (m MockLimits) MaxFlameGraphNodesOnSelectMergeProfile(string) bool {
+	return m.MaxFlameGraphNodesOnSelectMergeProfileValue
+}
 
 func (m MockLimits) MaxLabelNameLength(userID string) int     { return m.MaxLabelNameLengthValue }
 func (m MockLimits) MaxLabelValueLength(userID string) int    { return m.MaxLabelValueLengthValue }

pkg/validation/limits.go

@@ -517,6 +517,7 @@ func SanitizeTimeRange(limits RangeRequestLimits, tenant []string, start, end *i
 type FlameGraphLimits interface {
 	MaxFlameGraphNodesDefault(string) int
 	MaxFlameGraphNodesMax(string) int
+	MaxFlameGraphNodesOnSelectMergeProfile(string) bool
 }
 
 func ValidateMaxNodes(l FlameGraphLimits, tenantIDs []string, n int64) (int64, error) {