feat: return all spans for flamegraph under a limit

pkg/query-service/app/clickhouseReader/reader.go

@@ -1216,12 +1216,19 @@ func (r *ClickHouseReader) GetFlamegraphSpansForTrace(ctx context.Context, orgID
 	}
 
 	processingPostCache := time.Now()
-	selectedSpansForRequest := tracedetail.GetSelectedSpansForFlamegraphForRequest(req.SelectedSpanID, selectedSpans, startTime, endTime)
-	zap.L().Info("getFlamegraphSpansForTrace: processing post cache", zap.Duration("duration", time.Since(processingPostCache)), zap.String("traceID", traceID))
+	selectedSpansForRequest := selectedSpans
+	limit := min(req.Limit, tracedetail.MaxLimitWithoutSampling)
+	totalSpanCount := tracedetail.GetTotalSpanCount(selectedSpans)
+	if totalSpanCount > uint64(limit) {
+		selectedSpansForRequest = tracedetail.GetSelectedSpansForFlamegraphForRequest(req.SelectedSpanID, selectedSpans, startTime, endTime)
+	}
+	zap.L().Info("getFlamegraphSpansForTrace: processing post cache", zap.Duration("duration", time.Since(processingPostCache)), zap.String("traceID", traceID),
+		zap.Uint64("totalSpanCount", totalSpanCount))
 
 	trace.Spans = selectedSpansForRequest
 	trace.StartTimestampMillis = startTime / 1000000
 	trace.EndTimestampMillis = endTime / 1000000
+	trace.HasMore = totalSpanCount > uint64(limit)
 	return trace, nil
 }
 

pkg/query-service/app/traces/tracedetail/flamegraph.go

@@ -10,6 +10,8 @@ var (
 	SPAN_LIMIT_PER_REQUEST_FOR_FLAMEGRAPH float64 = 50
 	SPAN_LIMIT_PER_LEVEL                  int     = 100
 	TIMESTAMP_SAMPLING_BUCKET_COUNT       int     = 50
+
+	MaxLimitWithoutSampling uint = 120_000
 )
 
 func ContainsFlamegraphSpan(slice []*model.FlamegraphSpan, item *model.FlamegraphSpan) bool {
@@ -184,3 +186,12 @@ func GetSelectedSpansForFlamegraphForRequest(selectedSpanID string, selectedSpan
 
 	return selectedSpansForRequest
 }
+
+func GetTotalSpanCount(spans [][]*model.FlamegraphSpan) uint64 {
+	levelCount := len(spans)
+	spanCount := uint64(0)
+	for i := range levelCount {
+		spanCount += uint64(len(spans[i]))
+	}
+	return spanCount
+}

pkg/query-service/model/queryParams.go

@@ -337,6 +337,7 @@ type GetWaterfallSpansForTraceWithMetadataParams struct {
 
 type GetFlamegraphSpansForTraceParams struct {
 	SelectedSpanID string `json:"selectedSpanId"`
+	Limit          uint   `json:"limit"`
 }
 
 type SpanFilterParams struct {

pkg/query-service/model/response.go

@@ -336,6 +336,7 @@ type GetFlamegraphSpansForTraceResponse struct {
 	EndTimestampMillis   uint64              `json:"endTimestampMillis"`
 	DurationNano         uint64              `json:"durationNano"`
 	Spans                [][]*FlamegraphSpan `json:"spans"`
+	HasMore              bool                `json:"hasMore"`
 }
 
 type OtelSpanRef struct {