Tempo: Add template vars to embedded flamegraph (#80363)

* Add template vars to embedded flamegraph

* Update location of templateSrv.replace in spanFlameGraph

public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/SpanFlameGraph.tsx

@@ -13,7 +13,7 @@ import {
   TimeZone,
 } from '@grafana/data';
 import { FlameGraph } from '@grafana/flamegraph';
-import { config } from '@grafana/runtime';
+import { config, getTemplateSrv } from '@grafana/runtime';
 import { useStyles2 } from '@grafana/ui';
 import { TraceToProfilesOptions } from 'app/core/components/TraceToProfiles/TraceToProfilesSettings';
 import { getDatasourceSrv } from 'app/features/plugins/datasource_srv';
@@ -21,7 +21,14 @@ import { PyroscopeQueryType } from 'app/plugins/datasource/grafana-pyroscope-dat
 import { PyroscopeDataSource } from 'app/plugins/datasource/grafana-pyroscope-datasource/datasource';
 import { Query } from 'app/plugins/datasource/grafana-pyroscope-datasource/types';
 
-import { defaultProfilingKeys, getFormattedTags, pyroscopeProfileIdTagKey } from '../../../createSpanLink';
+import {
+  defaultProfilingKeys,
+  getFormattedTags,
+  pyroscopeProfileIdTagKey,
+  scopedVarsFromSpan,
+  scopedVarsFromTags,
+  scopedVarsFromTrace,
+} from '../../../createSpanLink';
 import { TraceSpan } from '../../types/trace';
 
 import { TraceFlameGraphs } from '.';
@@ -33,10 +40,21 @@ export type SpanFlameGraphProps = {
   traceFlameGraphs: TraceFlameGraphs;
   setTraceFlameGraphs: (flameGraphs: TraceFlameGraphs) => void;
   setRedrawListView: (redraw: {}) => void;
+  traceDuration: number;
+  traceName: string;
 };
 
 export default function SpanFlameGraph(props: SpanFlameGraphProps) {
-  const { span, traceToProfilesOptions, timeZone, traceFlameGraphs, setTraceFlameGraphs, setRedrawListView } = props;
+  const {
+    span,
+    traceToProfilesOptions,
+    timeZone,
+    traceFlameGraphs,
+    setTraceFlameGraphs,
+    setRedrawListView,
+    traceDuration,
+    traceName,
+  } = props;
   const [sizeRef, { height: containerHeight }] = useMeasure<HTMLDivElement>();
   const styles = useStyles2(getStyles);
 
@@ -80,7 +98,12 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
     ) => {
       let labelSelector = '{}';
       if (traceToProfilesOptions.customQuery && traceToProfilesOptions.query) {
-        labelSelector = traceToProfilesOptions.query;
+        const scopedVars = {
+          ...scopedVarsFromTrace(traceDuration, traceName, span.traceID),
+          ...scopedVarsFromSpan(span),
+          ...scopedVarsFromTags(span, traceToProfilesOptions),
+        };
+        labelSelector = getTemplateSrv().replace(traceToProfilesOptions.query, scopedVars);
       } else {
         const tags =
           traceToProfilesOptions.tags && traceToProfilesOptions.tags.length > 0
@@ -119,7 +142,7 @@ export default function SpanFlameGraph(props: SpanFlameGraphProps) {
         setTraceFlameGraphs({ ...traceFlameGraphs, [profileTagValue]: flameGraph });
       }
     },
-    [getTimeRangeForProfile, profileTagValue, setTraceFlameGraphs, timeZone, traceFlameGraphs]
+    [getTimeRangeForProfile, profileTagValue, setTraceFlameGraphs, timeZone, traceDuration, traceFlameGraphs, traceName]
   );
 
   useEffect(() => {

public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetail/index.tsx

@@ -125,6 +125,8 @@ export type SpanDetailProps = {
   timeZone: TimeZone;
   tagsToggle: (spanID: string) => void;
   traceStartTime: number;
+  traceDuration: number;
+  traceName: string;
   warningsToggle: (spanID: string) => void;
   stackTracesToggle: (spanID: string) => void;
   referenceItemToggle: (spanID: string, reference: TraceSpanReference) => void;
@@ -148,6 +150,8 @@ export default function SpanDetail(props: SpanDetailProps) {
     span,
     tagsToggle,
     traceStartTime,
+    traceDuration,
+    traceName,
     warningsToggle,
     stackTracesToggle,
     referencesToggle,
@@ -402,6 +406,8 @@ export default function SpanDetail(props: SpanDetailProps) {
               setTraceFlameGraphs={setTraceFlameGraphs}
               traceToProfilesOptions={traceToProfilesOptions}
               setRedrawListView={setRedrawListView}
+              traceDuration={traceDuration}
+              traceName={traceName}
             />
           )}
         <small className={styles.debugInfo}>

public/app/features/explore/TraceView/components/TraceTimelineViewer/SpanDetailRow.tsx

@@ -90,6 +90,8 @@ export type SpanDetailRowProps = {
   timeZone: TimeZone;
   tagsToggle: (spanID: string) => void;
   traceStartTime: number;
+  traceDuration: number;
+  traceName: string;
   hoverIndentGuideIds: Set<string>;
   addHoverIndentGuideId: (spanID: string) => void;
   removeHoverIndentGuideId: (spanID: string) => void;
@@ -131,6 +133,8 @@ export class UnthemedSpanDetailRow extends React.PureComponent<SpanDetailRowProp
       timeZone,
       tagsToggle,
       traceStartTime,
+      traceDuration,
+      traceName,
       hoverIndentGuideIds,
       addHoverIndentGuideId,
       removeHoverIndentGuideId,
@@ -181,6 +185,8 @@ export class UnthemedSpanDetailRow extends React.PureComponent<SpanDetailRowProp
               timeZone={timeZone}
               tagsToggle={tagsToggle}
               traceStartTime={traceStartTime}
+              traceDuration={traceDuration}
+              traceName={traceName}
               createSpanLink={createSpanLink}
               focusedSpanId={focusedSpanId}
               createFocusSpanLink={createFocusSpanLink}

public/app/features/explore/TraceView/components/TraceTimelineViewer/VirtualizedTraceView.tsx

@@ -587,6 +587,8 @@ export class UnthemedVirtualizedTraceView extends React.Component<VirtualizedTra
           timeZone={timeZone}
           tagsToggle={detailTagsToggle}
           traceStartTime={trace.startTime}
+          traceDuration={trace.duration}
+          traceName={trace.traceName}
           hoverIndentGuideIds={hoverIndentGuideIds}
           addHoverIndentGuideId={addHoverIndentGuideId}
           removeHoverIndentGuideId={removeHoverIndentGuideId}

public/app/features/explore/TraceView/createSpanLink.tsx

@@ -55,7 +55,7 @@ export function createSpanLinkFactory({
     return undefined;
   }
 
-  let scopedVars = scopedVarsFromTrace(trace);
+  let scopedVars = scopedVarsFromTrace(trace.duration, trace.traceName, trace.traceID);
   const hasLinks = dataFrame.fields.some((f) => Boolean(f.config.links?.length));
 
   const createSpanLinks = legacyCreateSpanLinkFactory(
@@ -596,14 +596,14 @@ function buildMetricsQuery(
  * Variables from trace that can be used in the query
  * @param trace
  */
-function scopedVarsFromTrace(trace: Trace): ScopedVars {
+export function scopedVarsFromTrace(duration: number, name: string, traceId: string): ScopedVars {
   return {
     __trace: {
       text: 'Trace',
       value: {
-        duration: trace.duration,
-        name: trace.traceName,
-        traceId: trace.traceID,
+        duration,
+        name,
+        traceId,
       },
     },
   };
@@ -613,7 +613,7 @@ function scopedVarsFromTrace(trace: Trace): ScopedVars {
  * Variables from span that can be used in the query
  * @param span
  */
-function scopedVarsFromSpan(span: TraceSpan): ScopedVars {
+export function scopedVarsFromSpan(span: TraceSpan): ScopedVars {
   const tags: ScopedVars = {};
 
   // We put all these tags together similar way we do for the __tags variable. This means there can be some overriding
@@ -643,7 +643,10 @@ function scopedVarsFromSpan(span: TraceSpan): ScopedVars {
  * Variables from tags that can be used in the query
  * @param span
  */
-function scopedVarsFromTags(span: TraceSpan, traceToProfilesOptions: TraceToProfilesOptions | undefined): ScopedVars {
+export function scopedVarsFromTags(
+  span: TraceSpan,
+  traceToProfilesOptions: TraceToProfilesOptions | undefined
+): ScopedVars {
   let tags: ScopedVars = {};
 
   if (traceToProfilesOptions) {