Revert "chore: extract out flamegraph building logic for easier review"

This reverts commit f9222c9930.

docs/api/openapi.yml

@@ -5663,58 +5663,6 @@ components:
       - attribute
       - resource
       type: string
-    SpantypesFlamegraphSpan:
-      properties:
-        attributes:
-          additionalProperties: {}
-          type: object
-        durationNano:
-          minimum: 0
-          type: integer
-        event:
-          items:
-            $ref: '#/components/schemas/SpantypesEvent'
-          nullable: true
-          type: array
-        hasError:
-          type: boolean
-        level:
-          format: int64
-          type: integer
-        name:
-          type: string
-        parentSpanId:
-          type: string
-        resource:
-          additionalProperties:
-            type: string
-          type: object
-        serviceName:
-          type: string
-        spanId:
-          type: string
-        timestamp:
-          minimum: 0
-          type: integer
-      type: object
-    SpantypesGettableFlamegraphTrace:
-      properties:
-        endTimestampMillis:
-          format: int64
-          type: integer
-        hasMore:
-          type: boolean
-        spans:
-          items:
-            items:
-              $ref: '#/components/schemas/SpantypesFlamegraphSpan'
-            type: array
-          nullable: true
-          type: array
-        startTimestampMillis:
-          format: int64
-          type: integer
-      type: object
     SpantypesGettableSpanMapperGroups:
       properties:
         items:
@@ -5762,15 +5710,6 @@ components:
           nullable: true
           type: array
       type: object
-    SpantypesPostableFlamegraph:
-      properties:
-        selectFields:
-          items:
-            $ref: '#/components/schemas/TelemetrytypesTelemetryFieldKey'
-          type: array
-        selectedSpanId:
-          type: string
-      type: object
     SpantypesPostableSpanMapper:
       properties:
         config:
@@ -18939,75 +18878,6 @@ paths:
       summary: Put profile in Zeus for a deployment.
       tags:
       - zeus
-  /api/v3/traces/{traceID}/flamegraph:
-    post:
-      deprecated: false
-      description: Returns the flamegraph view of spans for a given trace ID.
-      operationId: GetFlamegraph
-      parameters:
-      - in: path
-        name: traceID
-        required: true
-        schema:
-          type: string
-      requestBody:
-        content:
-          application/json:
-            schema:
-              $ref: '#/components/schemas/SpantypesPostableFlamegraph'
-      responses:
-        "200":
-          content:
-            application/json:
-              schema:
-                properties:
-                  data:
-                    $ref: '#/components/schemas/SpantypesGettableFlamegraphTrace'
-                  status:
-                    type: string
-                required:
-                - status
-                - data
-                type: object
-          description: OK
-        "400":
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/RenderErrorResponse'
-          description: Bad Request
-        "401":
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/RenderErrorResponse'
-          description: Unauthorized
-        "403":
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/RenderErrorResponse'
-          description: Forbidden
-        "404":
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/RenderErrorResponse'
-          description: Not Found
-        "500":
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/RenderErrorResponse'
-          description: Internal Server Error
-      security:
-      - api_key:
-        - VIEWER
-      - tokenizer:
-        - VIEWER
-      summary: Get flamegraph view for a trace
-      tags:
-      - tracedetail
   /api/v3/traces/{traceID}/waterfall:
     post:
       deprecated: false

frontend/src/api/generated/services/sigNoz.schemas.ts

@@ -6689,81 +6689,6 @@ export enum SpantypesFieldContextDTO {
 	attribute = 'attribute',
 	resource = 'resource',
 }
-export type SpantypesFlamegraphSpanDTOAttributes = { [key: string]: unknown };
-
-export type SpantypesFlamegraphSpanDTOResource = { [key: string]: string };
-
-export interface SpantypesFlamegraphSpanDTO {
-	/**
-	 * @type object
-	 */
-	attributes?: SpantypesFlamegraphSpanDTOAttributes;
-	/**
-	 * @type integer
-	 * @minimum 0
-	 */
-	durationNano?: number;
-	/**
-	 * @type array,null
-	 */
-	event?: SpantypesEventDTO[] | null;
-	/**
-	 * @type boolean
-	 */
-	hasError?: boolean;
-	/**
-	 * @type integer
-	 * @format int64
-	 */
-	level?: number;
-	/**
-	 * @type string
-	 */
-	name?: string;
-	/**
-	 * @type string
-	 */
-	parentSpanId?: string;
-	/**
-	 * @type object
-	 */
-	resource?: SpantypesFlamegraphSpanDTOResource;
-	/**
-	 * @type string
-	 */
-	serviceName?: string;
-	/**
-	 * @type string
-	 */
-	spanId?: string;
-	/**
-	 * @type integer
-	 * @minimum 0
-	 */
-	timestamp?: number;
-}
-
-export interface SpantypesGettableFlamegraphTraceDTO {
-	/**
-	 * @type integer
-	 * @format int64
-	 */
-	endTimestampMillis?: number;
-	/**
-	 * @type boolean
-	 */
-	hasMore?: boolean;
-	/**
-	 * @type array,null
-	 */
-	spans?: SpantypesFlamegraphSpanDTO[][] | null;
-	/**
-	 * @type integer
-	 * @format int64
-	 */
-	startTimestampMillis?: number;
-}
-
 export type SpantypesSpanMapperGroupConditionDTOAnyOf = {
 	/**
 	 * @type array,null
@@ -7039,17 +6964,6 @@ export interface SpantypesGettableWaterfallTraceDTO {
 	uncollapsedSpans?: string[] | null;
 }
 
-export interface SpantypesPostableFlamegraphDTO {
-	/**
-	 * @type array
-	 */
-	selectFields?: TelemetrytypesTelemetryFieldKeyDTO[];
-	/**
-	 * @type string
-	 */
-	selectedSpanId?: string;
-}
-
 export enum SpantypesSpanMapperOperationDTO {
 	move = 'move',
 	copy = 'copy',
@@ -9307,17 +9221,6 @@ export type GetHosts200 = {
 	status: string;
 };
 
-export type GetFlamegraphPathParameters = {
-	traceID: string;
-};
-export type GetFlamegraph200 = {
-	data: SpantypesGettableFlamegraphTraceDTO;
-	/**
-	 * @type string
-	 */
-	status: string;
-};
-
 export type GetWaterfallPathParameters = {
 	traceID: string;
 };

frontend/src/api/generated/services/tracedetail/index.ts

@@ -12,119 +12,17 @@ import type {
 } from 'react-query';
 
 import type {
-	GetFlamegraph200,
-	GetFlamegraphPathParameters,
 	GetWaterfall200,
 	GetWaterfallPathParameters,
 	GetWaterfallV4200,
 	GetWaterfallV4PathParameters,
 	RenderErrorResponseDTO,
-	SpantypesPostableFlamegraphDTO,
 	SpantypesPostableWaterfallDTO,
 } from '../sigNoz.schemas';
 
 import { GeneratedAPIInstance } from '../../../generatedAPIInstance';
 import type { ErrorType, BodyType } from '../../../generatedAPIInstance';
 
-/**
- * Returns the flamegraph view of spans for a given trace ID.
- * @summary Get flamegraph view for a trace
- */
-export const getFlamegraph = (
-	{ traceID }: GetFlamegraphPathParameters,
-	spantypesPostableFlamegraphDTO?: BodyType<SpantypesPostableFlamegraphDTO>,
-	signal?: AbortSignal,
-) => {
-	return GeneratedAPIInstance<GetFlamegraph200>({
-		url: `/api/v3/traces/${traceID}/flamegraph`,
-		method: 'POST',
-		headers: { 'Content-Type': 'application/json' },
-		data: spantypesPostableFlamegraphDTO,
-		signal,
-	});
-};
-
-export const getGetFlamegraphMutationOptions = <
-	TError = ErrorType<RenderErrorResponseDTO>,
-	TContext = unknown,
->(options?: {
-	mutation?: UseMutationOptions<
-		Awaited<ReturnType<typeof getFlamegraph>>,
-		TError,
-		{
-			pathParams: GetFlamegraphPathParameters;
-			data?: BodyType<SpantypesPostableFlamegraphDTO>;
-		},
-		TContext
-	>;
-}): UseMutationOptions<
-	Awaited<ReturnType<typeof getFlamegraph>>,
-	TError,
-	{
-		pathParams: GetFlamegraphPathParameters;
-		data?: BodyType<SpantypesPostableFlamegraphDTO>;
-	},
-	TContext
-> => {
-	const mutationKey = ['getFlamegraph'];
-	const { mutation: mutationOptions } = options
-		? options.mutation &&
-			'mutationKey' in options.mutation &&
-			options.mutation.mutationKey
-			? options
-			: { ...options, mutation: { ...options.mutation, mutationKey } }
-		: { mutation: { mutationKey } };
-
-	const mutationFn: MutationFunction<
-		Awaited<ReturnType<typeof getFlamegraph>>,
-		{
-			pathParams: GetFlamegraphPathParameters;
-			data?: BodyType<SpantypesPostableFlamegraphDTO>;
-		}
-	> = (props) => {
-		const { pathParams, data } = props ?? {};
-
-		return getFlamegraph(pathParams, data);
-	};
-
-	return { mutationFn, ...mutationOptions };
-};
-
-export type GetFlamegraphMutationResult = NonNullable<
-	Awaited<ReturnType<typeof getFlamegraph>>
->;
-export type GetFlamegraphMutationBody =
-	| BodyType<SpantypesPostableFlamegraphDTO>
-	| undefined;
-export type GetFlamegraphMutationError = ErrorType<RenderErrorResponseDTO>;
-
-/**
- * @summary Get flamegraph view for a trace
- */
-export const useGetFlamegraph = <
-	TError = ErrorType<RenderErrorResponseDTO>,
-	TContext = unknown,
->(options?: {
-	mutation?: UseMutationOptions<
-		Awaited<ReturnType<typeof getFlamegraph>>,
-		TError,
-		{
-			pathParams: GetFlamegraphPathParameters;
-			data?: BodyType<SpantypesPostableFlamegraphDTO>;
-		},
-		TContext
-	>;
-}): UseMutationResult<
-	Awaited<ReturnType<typeof getFlamegraph>>,
-	TError,
-	{
-		pathParams: GetFlamegraphPathParameters;
-		data?: BodyType<SpantypesPostableFlamegraphDTO>;
-	},
-	TContext
-> => {
-	return useMutation(getGetFlamegraphMutationOptions(options));
-};
 /**
  * Returns the waterfall view of spans for a given trace ID with tree structure, metadata, and windowed pagination
  * @summary Get waterfall view for a trace

pkg/types/spantypes/flamegraph_trace.go

@@ -39,15 +39,81 @@ func NewFlamegraphTraceFromStorable(spans []StorableSpan) *FlamegraphTrace {
 }
 
 func (t *FlamegraphTrace) GetAllLevels() [][]*FlamegraphSpan {
-	return nil
+	allLevels := t.buildAllLevels()
+	for _, node := range t.nodeByID {
+		node.Children = nil // children not required after building tree
+	}
+	return allLevels
 }
 
-// GetSelectedLevels returns the window of levels around selectedSpanID with sampling applied to dense levels.
+// GetSelectedLevels returns the level window for selectedSpanID with sampling applied to
+// dense levels. It always applies windowing — callers should only invoke this when the
+// trace is known to exceed the select-all limit.
+// Children are cleared after traversal so the tree can be GC'd.
 func (t *FlamegraphTrace) GetSelectedLevels(
 	selectedSpanID string,
 	levelLimit, spansPerLevel, topLatencyCount, bucketCount int,
 ) []FlamegraphLevel {
-	return nil
+	allLevels := t.buildAllLevels()
+	for _, node := range t.nodeByID {
+		node.Children = nil
+	}
+
+	selectedIndex := 0
+	if selectedSpanID != "" {
+	outer:
+		for i, lvl := range allLevels {
+			for _, span := range lvl {
+				if span.SpanID == selectedSpanID {
+					selectedIndex = i
+					break outer
+				}
+			}
+		}
+	}
+
+	lowerLimit := selectedIndex - int(float64(levelLimit)*0.4)
+	upperLimit := selectedIndex + int(float64(levelLimit)*0.6)
+
+	if lowerLimit < 0 {
+		upperLimit -= lowerLimit
+		lowerLimit = 0
+	}
+	if upperLimit > len(allLevels) {
+		lowerLimit -= upperLimit - len(allLevels)
+		upperLimit = len(allLevels)
+	}
+	if lowerLimit < 0 {
+		lowerLimit = 0
+	}
+
+	result := make([]FlamegraphLevel, 0, upperLimit-lowerLimit)
+	for i := lowerLimit; i < upperLimit; i++ {
+		lvl := allLevels[i]
+		if len(lvl) == 0 {
+			continue
+		}
+		var sampled []*FlamegraphSpan
+		if len(lvl) > spansPerLevel {
+			sampled = sampleFlamegraphLevel(lvl, selectedSpanID, i == selectedIndex,
+				t.startTime, t.endTime, topLatencyCount, bucketCount)
+		} else {
+			sampled = lvl
+		}
+		if len(sampled) == 0 {
+			continue
+		}
+		spanIDs := make([]string, len(sampled))
+		for j, s := range sampled {
+			spanIDs[j] = s.SpanID
+		}
+		result = append(result, FlamegraphLevel{
+			Level:   sampled[0].Level,
+			SpanIDs: spanIDs,
+		})
+	}
+
+	return result
 }
 
 func (t *FlamegraphTrace) EnrichSelectedSpans(selectedSpans []FlamegraphLevel, fullSpans []StorableSpan) [][]*FlamegraphSpan {
@@ -108,6 +174,98 @@ func (t *FlamegraphTrace) buildSpanTree() {
 	})
 }
 
+func (t *FlamegraphTrace) buildAllLevels() [][]*FlamegraphSpan {
+	var result [][]*FlamegraphSpan
+
+	type entry struct {
+		node  *FlamegraphSpan
+		depth int64
+	}
+
+	for _, root := range t.roots {
+		levelMap := make(map[int64][]*FlamegraphSpan)
+		maxDepth := int64(-1)
+
+		queue := []entry{{root, 0}}
+		for len(queue) > 0 {
+			curr := queue[0]
+			queue = queue[1:]
+			curr.node.Level = curr.depth
+			levelMap[curr.depth] = append(levelMap[curr.depth], curr.node)
+			if curr.depth > maxDepth {
+				maxDepth = curr.depth
+			}
+			for _, child := range curr.node.Children {
+				queue = append(queue, entry{child, curr.depth + 1})
+			}
+		}
+
+		for depth := int64(0); depth <= maxDepth; depth++ {
+			if spans, ok := levelMap[depth]; ok {
+				result = append(result, spans)
+			}
+		}
+	}
+
+	return result
+}
+
+func sampleFlamegraphLevel(
+	spans []*FlamegraphSpan,
+	selectedSpanID string,
+	isSelectedLevel bool,
+	startTime, endTime uint64,
+	topLatencyCount, bucketCount int,
+) []*FlamegraphSpan {
+	sorted := make([]*FlamegraphSpan, len(spans))
+	copy(sorted, spans)
+	sort.Slice(sorted, func(i, j int) bool {
+		return sorted[i].DurationNano > sorted[j].DurationNano
+	})
+
+	var sampled []*FlamegraphSpan
+
+	topK := min(topLatencyCount, len(sorted))
+	sampled = append(sampled, sorted[:topK]...)
+
+	if isSelectedLevel {
+		for _, span := range sorted {
+			if span.SpanID == selectedSpanID {
+				sampled = append(sampled, span)
+				break
+			}
+		}
+	}
+
+	bucketSize := (endTime - startTime) / uint64(bucketCount)
+	if bucketSize == 0 {
+		bucketSize = 1
+	}
+	buckets := make([][]*FlamegraphSpan, bucketCount)
+	for _, span := range sorted {
+		if span.Timestamp < startTime || span.Timestamp > endTime {
+			continue
+		}
+		idx := int((span.Timestamp - startTime) / bucketSize)
+		if idx < 0 {
+			idx = 0
+		} else if idx >= bucketCount {
+			idx = bucketCount - 1
+		}
+		buckets[idx] = append(buckets[idx], span)
+	}
+	for i := range buckets {
+		if len(buckets[i]) > 2 {
+			buckets[i] = buckets[i][:2]
+		}
+	}
+	for _, bucket := range buckets {
+		sampled = append(sampled, bucket...)
+	}
+
+	return sampled
+}
+
 func flamegraphSpanIndex(spans []*FlamegraphSpan, spanID string) int {
 	for i, s := range spans {
 		if s != nil && s.SpanID == spanID {