feat: show spinner on expanding span and fix autoscroll

Two improvemnts are included in this commit:
- Show a loading spiner in UI when a span is expanded
  to signal api call.
- Don't auto scroll the waterfall to bring the selected
  span to centre in case it's expanded or scrolled up
  because it looks very flickering behaviour to user

frontend/src/container/TraceWaterfall/TraceWaterfall.tsx

@@ -15,6 +15,7 @@ import './TraceWaterfall.styles.scss';
 export interface IInterestedSpan {
 	spanId: string;
 	isUncollapsed: boolean;
+	shouldScrollToSpan?: boolean;
 }
 
 interface ITraceWaterfallProps {
@@ -109,6 +110,7 @@ function TraceWaterfall(props: ITraceWaterfallProps): JSX.Element {
 						setTraceFlamegraphStatsWidth={setTraceFlamegraphStatsWidth}
 						selectedSpan={selectedSpan}
 						setSelectedSpan={setSelectedSpan}
+						isFetchingTraceData={isFetchingTraceData}
 					/>
 				);
 			default:
@@ -117,6 +119,7 @@ function TraceWaterfall(props: ITraceWaterfallProps): JSX.Element {
 	}, [
 		errorFetchingTraceData,
 		interestedSpanId,
+		isFetchingTraceData,
 		selectedSpan,
 		setInterestedSpanId,
 		setSelectedSpan,

frontend/src/container/TraceWaterfall/TraceWaterfallStates/Success/Success.styles.scss

@@ -1,3 +1,16 @@
+@keyframes spin {
+	from {
+		transform: rotate(0deg);
+	}
+	to {
+		transform: rotate(360deg);
+	}
+}
+
+.spin {
+	animation: spin 1s linear infinite;
+}
+
 .success-content {
 	overflow-y: hidden;
 	overflow-x: hidden;

frontend/src/container/TraceWaterfall/TraceWaterfallStates/Success/Success.tsx

@@ -30,6 +30,7 @@ import {
 	ChevronDown,
 	ChevronRight,
 	Leaf,
+	Loader2,
 } from 'lucide-react';
 import { useAppContext } from 'providers/App/App';
 import { Span } from 'types/api/trace/getTraceV2';
@@ -43,6 +44,20 @@ import './Success.styles.scss';
 const CONNECTOR_WIDTH = 28;
 const VERTICAL_CONNECTOR_WIDTH = 1;
 
+// Component to render chevron icon with spinner
+function ChevronIcon({
+	isFetching,
+	isCollapsed,
+}: {
+	isFetching: boolean;
+	isCollapsed: boolean;
+}): JSX.Element {
+	if (isFetching) {
+		return <Loader2 size={14} className="spin" />;
+	}
+	return isCollapsed ? <ChevronRight size={14} /> : <ChevronDown size={14} />;
+}
+
 interface ITraceMetadata {
 	traceId: string;
 	startTime: number;
@@ -58,6 +73,7 @@ interface ISuccessProps {
 	setTraceFlamegraphStatsWidth: Dispatch<SetStateAction<number>>;
 	selectedSpan: Span | undefined;
 	setSelectedSpan: Dispatch<SetStateAction<Span | undefined>>;
+	isFetchingTraceData: boolean;
 }
 
 function SpanOverview({
@@ -70,6 +86,8 @@ function SpanOverview({
 	filteredSpanIds,
 	isFilterActive,
 	traceMetadata,
+	isFetchingTraceData,
+	interestedSpanId,
 }: {
 	span: Span;
 	isSpanCollapsed: boolean;
@@ -80,6 +98,8 @@ function SpanOverview({
 	filteredSpanIds: string[];
 	isFilterActive: boolean;
 	traceMetadata: ITraceMetadata;
+	isFetchingTraceData: boolean;
+	interestedSpanId: IInterestedSpan;
 }): JSX.Element {
 	const isRootSpan = span.level === 0;
 	const { hasEditPermission } = useAppContext();
@@ -145,11 +165,12 @@ function SpanOverview({
 									}}
 									className="collapse-uncollapse-button"
 								>
-									{isSpanCollapsed ? (
-										<ChevronRight size={14} />
-									) : (
-										<ChevronDown size={14} />
-									)}
+									<ChevronIcon
+										isFetching={
+											isFetchingTraceData && interestedSpanId.spanId === span.spanId
+										}
+										isCollapsed={isSpanCollapsed}
+									/>
 									<Typography.Text className="children-count">
 										{span.subTreeNodeCount}
 									</Typography.Text>
@@ -342,6 +363,8 @@ function getWaterfallColumns({
 	handleAddSpanToFunnel,
 	filteredSpanIds,
 	isFilterActive,
+	isFetchingTraceData,
+	interestedSpanId,
 }: {
 	handleCollapseUncollapse: (id: string, collapse: boolean) => void;
 	uncollapsedNodes: string[];
@@ -351,8 +374,10 @@ function getWaterfallColumns({
 	handleAddSpanToFunnel: (span: Span) => void;
 	filteredSpanIds: string[];
 	isFilterActive: boolean;
-}): ColumnDef<Span, any>[] {
-	const waterfallColumns: ColumnDef<Span, any>[] = [
+	isFetchingTraceData: boolean;
+	interestedSpanId: IInterestedSpan;
+}): ColumnDef<Span, string>[] {
+	const waterfallColumns: ColumnDef<Span, string>[] = [
 		columnDefHelper.display({
 			id: 'span-name',
 			header: '',
@@ -367,6 +392,8 @@ function getWaterfallColumns({
 					traceMetadata={traceMetadata}
 					filteredSpanIds={filteredSpanIds}
 					isFilterActive={isFilterActive}
+					isFetchingTraceData={isFetchingTraceData}
+					interestedSpanId={interestedSpanId}
 				/>
 			),
 			size: 450,
@@ -411,6 +438,7 @@ function Success(props: ISuccessProps): JSX.Element {
 		setTraceFlamegraphStatsWidth,
 		setSelectedSpan,
 		selectedSpan,
+		isFetchingTraceData,
 	} = props;
 
 	const [filteredSpanIds, setFilteredSpanIds] = useState<string[]>([]);
@@ -427,7 +455,11 @@ function Success(props: ISuccessProps): JSX.Element {
 
 	const handleCollapseUncollapse = useCallback(
 		(spanId: string, collapse: boolean) => {
-			setInterestedSpanId({ spanId, isUncollapsed: !collapse });
+			setInterestedSpanId({
+				spanId,
+				isUncollapsed: !collapse,
+				shouldScrollToSpan: false,
+			});
 		},
 		[setInterestedSpanId],
 	);
@@ -445,7 +477,11 @@ function Success(props: ISuccessProps): JSX.Element {
 		if (range?.startIndex === 0 && instance.isScrolling) {
 			// do not trigger for trace root as nothing to fetch above
 			if (spans[0].level !== 0) {
-				setInterestedSpanId({ spanId: spans[0].spanId, isUncollapsed: false });
+				setInterestedSpanId({
+					spanId: spans[0].spanId,
+					isUncollapsed: false,
+					shouldScrollToSpan: false,
+				});
 			}
 			return;
 		}
@@ -454,6 +490,7 @@ function Success(props: ISuccessProps): JSX.Element {
 			setInterestedSpanId({
 				spanId: spans[spans.length - 1].spanId,
 				isUncollapsed: false,
+				shouldScrollToSpan: false,
 			});
 		}
 	};
@@ -495,6 +532,8 @@ function Success(props: ISuccessProps): JSX.Element {
 				handleAddSpanToFunnel,
 				filteredSpanIds,
 				isFilterActive,
+				isFetchingTraceData,
+				interestedSpanId,
 			}),
 		[
 			handleCollapseUncollapse,
@@ -505,6 +544,8 @@ function Success(props: ISuccessProps): JSX.Element {
 			handleAddSpanToFunnel,
 			filteredSpanIds,
 			isFilterActive,
+			isFetchingTraceData,
+			interestedSpanId,
 		],
 	);
 
@@ -514,12 +555,16 @@ function Success(props: ISuccessProps): JSX.Element {
 				(span) => span.spanId === interestedSpanId.spanId,
 			);
 			if (idx !== -1) {
-				setTimeout(() => {
-					virtualizerRef.current?.scrollToIndex(idx, {
-						align: 'center',
-						behavior: 'auto',
-					});
-				}, 400);
+				// Only scroll to center when navigating (URL/flamegraph), not on
+				// expand/collapse or virtualizer boundary fetches.
+				if (interestedSpanId.shouldScrollToSpan) {
+					setTimeout(() => {
+						virtualizerRef.current?.scrollToIndex(idx, {
+							align: 'center',
+							behavior: 'auto',
+						});
+					}, 400);
+				}
 
 				setSelectedSpan(spans[idx]);
 			}

frontend/src/pages/TraceDetailV2/TraceDetailV2.tsx

@@ -30,6 +30,7 @@ function TraceDetailsV2(): JSX.Element {
 		() => ({
 			spanId: urlQuery.get('spanId') || '',
 			isUncollapsed: urlQuery.get('spanId') !== '',
+			shouldScrollToSpan: urlQuery.get('spanId') !== '',
 		}),
 	);
 	const [
@@ -43,6 +44,7 @@ function TraceDetailsV2(): JSX.Element {
 		setInterestedSpanId({
 			spanId: urlQuery.get('spanId') || '',
 			isUncollapsed: urlQuery.get('spanId') !== '',
+			shouldScrollToSpan: true,
 		});
 	}, [urlQuery]);
 

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

@@ -9,6 +9,7 @@ import (
 
 var (
 	SPAN_LIMIT_PER_REQUEST_FOR_WATERFALL float64 = 500
+	MAX_DEPTH_FOR_SELECTED_SPAN_CHILDREN int     = 5
 )
 
 type Interval struct {
@@ -88,8 +89,9 @@ 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) ([]*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 +128,36 @@ 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 < MAX_DEPTH_FOR_SELECTED_SPAN_CHILDREN {
+		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 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)
 		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 +191,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)
+			// 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 {