From 8aa5d470ca34831757e7aa4b47822210505a6db5 Mon Sep 17 00:00:00 2001
From: Andre Pereira <adrapereira@gmail.com>
Date: Mon, 6 Nov 2023 16:29:59 +0000
Subject: [PATCH] Tempo: Fix streaming query restart after Grafana server
 reboot (#77614)

* Fix streaming query restart after Grafana server reboot

* TraceQL Search filter name improvements

* Add flag to enable streaming in tempo docker block

---------

Co-authored-by: Andrej Ocenas <mr.ocenas@gmail.com>
---
 devenv/docker/blocks/tempo/tempo.yaml         |   2 +
 .../SpanFilters/SpanFilters.tsx               |   4 +-
 .../SearchTraceQLEditor/TraceQLSearch.tsx     |   2 +-
 .../tempo/SearchTraceQLEditor/utils.ts        |   4 +
 .../app/plugins/datasource/tempo/streaming.ts | 113 ++++++++++--------
 5 files changed, 69 insertions(+), 56 deletions(-)

diff --git a/devenv/docker/blocks/tempo/tempo.yaml b/devenv/docker/blocks/tempo/tempo.yaml
index 6e94ff11a8d..e8378e9b9a6 100644
--- a/devenv/docker/blocks/tempo/tempo.yaml
+++ b/devenv/docker/blocks/tempo/tempo.yaml
@@ -52,3 +52,5 @@ storage:
 
 overrides:
   metrics_generator_processors: [local-blocks, service-graphs, span-metrics]
+
+stream_over_http_enabled: true
diff --git a/public/app/features/explore/TraceView/components/TracePageHeader/SpanFilters/SpanFilters.tsx b/public/app/features/explore/TraceView/components/TracePageHeader/SpanFilters/SpanFilters.tsx
index f29caaca8b6..7a27e19a20c 100644
--- a/public/app/features/explore/TraceView/components/TracePageHeader/SpanFilters/SpanFilters.tsx
+++ b/public/app/features/explore/TraceView/components/TracePageHeader/SpanFilters/SpanFilters.tsx
@@ -342,9 +342,9 @@ export const SpanFilters = memo((props: SpanFilterProps) => {
         </InlineFieldRow>
         <InlineFieldRow>
           <InlineField
-            label="Span Duration"
+            label="Duration"
             labelWidth={16}
-            tooltip="Filter by span duration. Accepted units are ns, us, ms, s, m, h"
+            tooltip="Filter by duration. Accepted units are ns, us, ms, s, m, h"
           >
             <HorizontalGroup spacing="xs" align="flex-start">
               <Select
diff --git a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/TraceQLSearch.tsx b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/TraceQLSearch.tsx
index 409b9688d4f..d65f7b28f15 100644
--- a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/TraceQLSearch.tsx
+++ b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/TraceQLSearch.tsx
@@ -150,7 +150,7 @@ const TraceQLSearch = ({ datasource, query, onChange }: Props) => {
             />
           </InlineSearchField>
           <InlineSearchField
-            label={'Duration'}
+            label={'Span Duration'}
             tooltip="The span duration, i.e.	end - start time of the span. Accepted units are ns, ms, s, m, h"
           >
             <HorizontalGroup spacing={'sm'}>
diff --git a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/utils.ts b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/utils.ts
index c7106b4a1af..9406a1c92e2 100644
--- a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/utils.ts
+++ b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/utils.ts
@@ -41,6 +41,10 @@ export const filterTitle = (f: TraceqlFilter) => {
   if (f.tag === 'name') {
     return 'Span Name';
   }
+  // Special case for the resource service name
+  if (f.tag === 'service.name' && f.scope === TraceqlSearchScope.Resource) {
+    return 'Service Name';
+  }
   return startCase(filterScopedTag(f));
 };
 
diff --git a/public/app/plugins/datasource/tempo/streaming.ts b/public/app/plugins/datasource/tempo/streaming.ts
index baf88073d41..23eef3182ce 100644
--- a/public/app/plugins/datasource/tempo/streaming.ts
+++ b/public/app/plugins/datasource/tempo/streaming.ts
@@ -1,5 +1,5 @@
 import { capitalize } from 'lodash';
-import { map, Observable, defer, mergeMap } from 'rxjs';
+import { map, Observable, takeWhile } from 'rxjs';
 import { v4 as uuidv4 } from 'uuid';
 
 import {
@@ -20,7 +20,7 @@ import { SearchStreamingState } from './dataquery.gen';
 import { DEFAULT_SPSS, TempoDatasource } from './datasource';
 import { formatTraceQLResponse } from './resultTransformer';
 import { SearchMetrics, TempoJsonData, TempoQuery } from './types';
-export async function getLiveStreamKey(): Promise<string> {
+function getLiveStreamKey(): string {
   return uuidv4();
 }
 
@@ -34,59 +34,66 @@ export function doTempoChannelStream(
 
   let frames: DataFrame[] | undefined = undefined;
   let state: LoadingState = LoadingState.NotStarted;
+  const requestTime = performance.now();
 
-  return defer(() => getLiveStreamKey()).pipe(
-    mergeMap((key) => {
-      const requestTime = performance.now();
-      return getGrafanaLiveSrv()
-        .getStream<MutableDataFrame>({
-          scope: LiveChannelScope.DataSource,
-          namespace: ds.uid,
-          path: `search/${key}`,
-          data: {
-            ...query,
-            SpansPerSpanSet: query.spss ?? DEFAULT_SPSS,
-            timeRange: {
-              from: range.from.toISOString(),
-              to: range.to.toISOString(),
-            },
-          },
-        })
-        .pipe(
-          map((evt) => {
-            if ('message' in evt && evt?.message) {
-              const currentTime = performance.now();
-              const elapsedTime = currentTime - requestTime;
-              // Schema should be [traces, metrics, state, error]
-              const traces = evt.message.data.values[0][0];
-              const metrics = evt.message.data.values[1][0];
-              const frameState: SearchStreamingState = evt.message.data.values[2][0];
-              const error = evt.message.data.values[3][0];
-
-              switch (frameState) {
-                case SearchStreamingState.Done:
-                  state = LoadingState.Done;
-                  break;
-                case SearchStreamingState.Streaming:
-                  state = LoadingState.Streaming;
-                  break;
-                case SearchStreamingState.Error:
-                  throw new Error(error);
-              }
-
-              frames = [
-                metricsDataFrame(metrics, frameState, elapsedTime),
-                ...formatTraceQLResponse(traces, instanceSettings, query.tableType),
-              ];
-            }
-            return {
-              data: frames || [],
-              state,
-            };
-          })
-        );
+  return getGrafanaLiveSrv()
+    .getStream<MutableDataFrame>({
+      scope: LiveChannelScope.DataSource,
+      namespace: ds.uid,
+      path: `search/${getLiveStreamKey()}`,
+      data: {
+        ...query,
+        SpansPerSpanSet: query.spss ?? DEFAULT_SPSS,
+        timeRange: {
+          from: range.from.toISOString(),
+          to: range.to.toISOString(),
+        },
+      },
     })
-  );
+    .pipe(
+      takeWhile((evt) => {
+        if ('message' in evt && evt?.message) {
+          const frameState: SearchStreamingState = evt.message.data.values[2][0];
+          if (frameState === SearchStreamingState.Done || frameState === SearchStreamingState.Error) {
+            return false;
+          }
+        }
+        return true;
+      }, true)
+    )
+    .pipe(
+      map((evt) => {
+        if ('message' in evt && evt?.message) {
+          const currentTime = performance.now();
+          const elapsedTime = currentTime - requestTime;
+          // Schema should be [traces, metrics, state, error]
+          const traces = evt.message.data.values[0][0];
+          const metrics = evt.message.data.values[1][0];
+          const frameState: SearchStreamingState = evt.message.data.values[2][0];
+          const error = evt.message.data.values[3][0];
+
+          switch (frameState) {
+            case SearchStreamingState.Done:
+              state = LoadingState.Done;
+              break;
+            case SearchStreamingState.Streaming:
+              state = LoadingState.Streaming;
+              break;
+            case SearchStreamingState.Error:
+              throw new Error(error);
+          }
+
+          frames = [
+            metricsDataFrame(metrics, frameState, elapsedTime),
+            ...formatTraceQLResponse(traces, instanceSettings, query.tableType),
+          ];
+        }
+        return {
+          data: frames || [],
+          state,
+        };
+      })
+    );
 }
 
 function metricsDataFrame(metrics: SearchMetrics, state: SearchStreamingState, elapsedTime: number) {