feat: return all spans for flamegraph under a limit

* fix: add missing filtering for ip address for scalar data

In domain listing api for external api monitoring,
we have option to filter out the IP address but
it only handles timeseries and raw type data while
domain list handler returns scalar data.

* fix: switch to new derived attributes for ip filtering

---------

Co-authored-by: Nityananda Gohain <nityanandagohain@gmail.com>

pkg/modules/thirdpartyapi/translator.go

@@ -120,6 +120,8 @@ func FilterResponse(results []*qbtypes.QueryRangeResponse) []*qbtypes.QueryRange
 					}
 				}
 				resultData.Rows = filteredRows
+			case *qbtypes.ScalarData:
+				resultData.Data = filterScalarDataIPs(resultData.Columns, resultData.Data)
 			}
 
 			filteredData = append(filteredData, result)
@@ -145,6 +147,39 @@ func shouldIncludeSeries(series *qbtypes.TimeSeries) bool {
 	return true
 }
 
+func filterScalarDataIPs(columns []*qbtypes.ColumnDescriptor, data [][]any) [][]any {
+	// Find column indices for server address fields
+	serverColIndices := make([]int, 0)
+	for i, col := range columns {
+		if col.Name == derivedKeyHTTPHost {
+			serverColIndices = append(serverColIndices, i)
+		}
+	}
+
+	if len(serverColIndices) == 0 {
+		return data
+	}
+
+	filtered := make([][]any, 0, len(data))
+	for _, row := range data {
+		includeRow := true
+		for _, colIdx := range serverColIndices {
+			if colIdx < len(row) {
+				if strVal, ok := row[colIdx].(string); ok {
+					if net.ParseIP(strVal) != nil {
+						includeRow = false
+						break
+					}
+				}
+			}
+		}
+		if includeRow {
+			filtered = append(filtered, row)
+		}
+	}
+	return filtered
+}
+
 func shouldIncludeRow(row *qbtypes.RawRow) bool {
 	if row.Data != nil {
 		if domainVal, ok := row.Data[derivedKeyHTTPHost]; ok {

pkg/modules/thirdpartyapi/translator_test.go

@@ -117,6 +117,59 @@ func TestFilterResponse(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "should filter out IP addresses from scalar data",
+			input: []*qbtypes.QueryRangeResponse{
+				{
+					Data: qbtypes.QueryData{
+						Results: []any{
+							&qbtypes.ScalarData{
+								QueryName: "endpoints",
+								Columns: []*qbtypes.ColumnDescriptor{
+									{
+										TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: derivedKeyHTTPHost},
+										Type:              qbtypes.ColumnTypeGroup,
+									},
+									{
+										TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "endpoints"},
+										Type:              qbtypes.ColumnTypeAggregation,
+									},
+								},
+								Data: [][]any{
+									{"192.168.1.1", 10},
+									{"example.com", 20},
+									{"10.0.0.1", 5},
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: []*qbtypes.QueryRangeResponse{
+				{
+					Data: qbtypes.QueryData{
+						Results: []any{
+							&qbtypes.ScalarData{
+								QueryName: "endpoints",
+								Columns: []*qbtypes.ColumnDescriptor{
+									{
+										TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: derivedKeyHTTPHost},
+										Type:              qbtypes.ColumnTypeGroup,
+									},
+									{
+										TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "endpoints"},
+										Type:              qbtypes.ColumnTypeAggregation,
+									},
+								},
+								Data: [][]any{
+									{"example.com", 20},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {

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

@@ -1046,7 +1046,19 @@ func (r *ClickHouseReader) GetWaterfallSpansForTraceWithMetadata(ctx context.Con
 	}
 
 	processingPostCache := time.Now()
-	selectedSpans, uncollapsedSpans, rootServiceName, rootServiceEntryPoint := tracedetail.GetSelectedSpans(req.UncollapsedSpans, req.SelectedSpanID, traceRoots, spanIdToSpanNodeMap, req.IsSelectedSpanIDUnCollapsed)
+	limit := min(req.Limit, tracedetail.MaxLimitToSelectAllSpans)
+	selectAllSpans := totalSpans <= uint64(limit)
+
+	var (
+		selectedSpans                          []*model.Span
+		uncollapsedSpans                       []string
+		rootServiceName, rootServiceEntryPoint string
+	)
+	if selectAllSpans {
+		selectedSpans, rootServiceName, rootServiceEntryPoint = tracedetail.GetAllSpans(traceRoots)
+	} else {
+		selectedSpans, uncollapsedSpans, rootServiceName, rootServiceEntryPoint = tracedetail.GetSelectedSpans(req.UncollapsedSpans, req.SelectedSpanID, traceRoots, spanIdToSpanNodeMap, req.IsSelectedSpanIDUnCollapsed)
+	}
 	zap.L().Info("getWaterfallSpansForTraceWithMetadata: processing post cache", zap.Duration("duration", time.Since(processingPostCache)), zap.String("traceID", traceID))
 
 	// convert start timestamp to millis because right now frontend is expecting it in millis
@@ -1059,7 +1071,7 @@ func (r *ClickHouseReader) GetWaterfallSpansForTraceWithMetadata(ctx context.Con
 	}
 
 	response.Spans = selectedSpans
-	response.UncollapsedSpans = uncollapsedSpans
+	response.UncollapsedSpans = uncollapsedSpans // ignoring if all spans are returning
 	response.StartTimestampMillis = startTime / 1000000
 	response.EndTimestampMillis = endTime / 1000000
 	response.TotalSpansCount = totalSpans
@@ -1068,6 +1080,7 @@ func (r *ClickHouseReader) GetWaterfallSpansForTraceWithMetadata(ctx context.Con
 	response.RootServiceEntryPoint = rootServiceEntryPoint
 	response.ServiceNameToTotalDurationMap = serviceNameToTotalDurationMap
 	response.HasMissingSpans = hasMissingSpans
+	response.HasMore = !selectAllSpans
 	return response, nil
 }
 
@@ -1216,12 +1229,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/app/traces/tracedetail/waterfall.go

@@ -9,6 +9,9 @@ import (
 
 var (
 	SPAN_LIMIT_PER_REQUEST_FOR_WATERFALL float64 = 500
+
+	maxDepthForSelectedSpanChildren int  = 5
+	MaxLimitToSelectAllSpans        uint = 10_000
 )
 
 type Interval struct {
@@ -88,8 +91,11 @@ func getPathFromRootToSelectedSpanId(node *model.Span, selectedSpanId string, un
 	return isPresentInSubtreeForTheNode, spansFromRootToNode
 }
 
-func traverseTrace(span *model.Span, uncollapsedSpans []string, level uint64, isPartOfPreOrder bool, hasSibling bool, selectedSpanId string) []*model.Span {
+func traverseTrace(span *model.Span, uncollapsedSpans []string, level uint64, isPartOfPreOrder bool, hasSibling bool, selectedSpanId string,
+	depthFromSelectedSpan int, isSelectedSpanIDUnCollapsed bool, selectAllSpan bool) ([]*model.Span, []string) {
+
 	preOrderTraversal := []*model.Span{}
+	autoExpandedSpans := []string{}
 
 	// sort the children to maintain the order across requests
 	sort.Slice(span.Children, func(i, j int) bool {
@@ -126,15 +132,40 @@ func traverseTrace(span *model.Span, uncollapsedSpans []string, level uint64, is
 		preOrderTraversal = append(preOrderTraversal, &nodeWithoutChildren)
 	}
 
+	nextDepthFromSelectedSpan := -1
+	if span.SpanID == selectedSpanId && isSelectedSpanIDUnCollapsed {
+		nextDepthFromSelectedSpan = 1
+	} else if depthFromSelectedSpan >= 1 && depthFromSelectedSpan < maxDepthForSelectedSpanChildren {
+		nextDepthFromSelectedSpan = depthFromSelectedSpan + 1
+	}
+
 	for index, child := range span.Children {
-		_childTraversal := traverseTrace(child, uncollapsedSpans, level+1, isPartOfPreOrder && slices.Contains(uncollapsedSpans, span.SpanID), index != (len(span.Children)-1), selectedSpanId)
+		// A child is included in the pre-order output if its parent is uncollapsed
+		// OR if the child falls within MAX_DEPTH_FOR_SELECTED_SPAN_CHILDREN levels
+		// below the selected span.
+		isChildWithinMaxDepth := nextDepthFromSelectedSpan >= 1
+		isAlreadyUncollapsed := slices.Contains(uncollapsedSpans, span.SpanID)
+		childIsPartOfPreOrder := isPartOfPreOrder && (isAlreadyUncollapsed || isChildWithinMaxDepth)
+		if selectAllSpan {
+			childIsPartOfPreOrder = true
+		}
+
+		if isPartOfPreOrder && isChildWithinMaxDepth && !isAlreadyUncollapsed {
+			if !slices.Contains(autoExpandedSpans, span.SpanID) {
+				autoExpandedSpans = append(autoExpandedSpans, span.SpanID)
+			}
+		}
+
+		_childTraversal, _autoExpanded := traverseTrace(child, uncollapsedSpans, level+1, childIsPartOfPreOrder, index != (len(span.Children)-1), selectedSpanId,
+			nextDepthFromSelectedSpan, isSelectedSpanIDUnCollapsed, selectAllSpan)
 		preOrderTraversal = append(preOrderTraversal, _childTraversal...)
+		autoExpandedSpans = append(autoExpandedSpans, _autoExpanded...)
 		nodeWithoutChildren.SubTreeNodeCount += child.SubTreeNodeCount + 1
 		span.SubTreeNodeCount += child.SubTreeNodeCount + 1
 	}
 
 	nodeWithoutChildren.SubTreeNodeCount += 1
-	return preOrderTraversal
+	return preOrderTraversal, autoExpandedSpans
 
 }
 
@@ -168,7 +199,13 @@ func GetSelectedSpans(uncollapsedSpans []string, selectedSpanID string, traceRoo
 			_, spansFromRootToNode := getPathFromRootToSelectedSpanId(rootNode, selectedSpanID, updatedUncollapsedSpans, isSelectedSpanIDUnCollapsed)
 			updatedUncollapsedSpans = append(updatedUncollapsedSpans, spansFromRootToNode...)
 
-			_preOrderTraversal := traverseTrace(rootNode, updatedUncollapsedSpans, 0, true, false, selectedSpanID)
+			_preOrderTraversal, _autoExpanded := traverseTrace(rootNode, updatedUncollapsedSpans, 0, true, false, selectedSpanID, -1, isSelectedSpanIDUnCollapsed, false)
+			// Merge auto-expanded spans into updatedUncollapsedSpans for returning in response
+			for _, spanID := range _autoExpanded {
+				if !slices.Contains(updatedUncollapsedSpans, spanID) {
+					updatedUncollapsedSpans = append(updatedUncollapsedSpans, spanID)
+				}
+			}
 			_selectedSpanIndex := findIndexForSelectedSpanFromPreOrder(_preOrderTraversal, selectedSpanID)
 
 			if _selectedSpanIndex != -1 {
@@ -212,3 +249,17 @@ func GetSelectedSpans(uncollapsedSpans []string, selectedSpanID string, traceRoo
 
 	return preOrderTraversal[startIndex:endIndex], updatedUncollapsedSpans, rootServiceName, rootServiceEntryPoint
 }
+
+func GetAllSpans(traceRoots []*model.Span) (spans []*model.Span, rootServiceName, rootEntryPoint string) {
+	for _, root := range traceRoots {
+		childSpans, _ := traverseTrace(root, nil, 0, true, false, "", -1, false, true)
+		spans = append(spans, childSpans...)
+		if rootServiceName == "" {
+			rootServiceName = root.ServiceName
+		}
+		if rootEntryPoint == "" {
+			rootEntryPoint = root.Name
+		}
+	}
+	return
+}

pkg/query-service/model/queryParams.go

@@ -333,10 +333,12 @@ type GetWaterfallSpansForTraceWithMetadataParams struct {
 	SelectedSpanID              string   `json:"selectedSpanId"`
 	IsSelectedSpanIDUnCollapsed bool     `json:"isSelectedSpanIDUnCollapsed"`
 	UncollapsedSpans            []string `json:"uncollapsedSpans"`
+	Limit                       uint     `json:"limit"`
 }
 
 type GetFlamegraphSpansForTraceParams struct {
 	SelectedSpanID string `json:"selectedSpanId"`
+	Limit          uint   `json:"limit"`
 }
 
 type SpanFilterParams struct {

pkg/query-service/model/response.go

@@ -329,6 +329,7 @@ type GetWaterfallSpansForTraceWithMetadataResponse struct {
 	HasMissingSpans               bool              `json:"hasMissingSpans"`
 	// this is needed for frontend and query service sync
 	UncollapsedSpans []string `json:"uncollapsedSpans"`
+	HasMore          bool     `json:"hasMore"`
 }
 
 type GetFlamegraphSpansForTraceResponse struct {
@@ -336,6 +337,7 @@ type GetFlamegraphSpansForTraceResponse struct {
 	EndTimestampMillis   uint64              `json:"endTimestampMillis"`
 	DurationNano         uint64              `json:"durationNano"`
 	Spans                [][]*FlamegraphSpan `json:"spans"`
+	HasMore              bool                `json:"hasMore"`
 }
 
 type OtelSpanRef struct {