From 1cbae72a9e1506a416cffe53ba4c61fe7ca95644 Mon Sep 17 00:00:00 2001
From: Andrej Ocenas <mr.ocenas@gmail.com>
Date: Fri, 6 Oct 2023 15:03:20 +0200
Subject: [PATCH] Pyroscope: Fix backend panic when querying out of bounds
 (#76053)

---
 .../pyroscopeClient.go                        |  5 +++
 .../pyroscopeClient_test.go                   | 16 ++++++-
 .../grafana-pyroscope-datasource/query.go     | 44 ++++++++++++++-----
 3 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go b/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go
index 46196f156a9..b63879f2c58 100644
--- a/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go
+++ b/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go
@@ -149,6 +149,11 @@ func (c *PyroscopeClient) GetProfile(ctx context.Context, profileTypeID, labelSe
 		return nil, err
 	}
 
+	if resp.Msg.Flamegraph == nil {
+		// Not an error, can happen when querying data oout of range.
+		return nil, nil
+	}
+
 	levels := make([]*Level, len(resp.Msg.Flamegraph.Levels))
 	for i, level := range resp.Msg.Flamegraph.Levels {
 		levels[i] = &Level{
diff --git a/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient_test.go b/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient_test.go
index 7f1534b3d37..4b07ff9cfd2 100644
--- a/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient_test.go
+++ b/pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient_test.go
@@ -51,10 +51,21 @@ func Test_PyroscopeClient(t *testing.T) {
 		}
 		require.Equal(t, series, resp)
 	})
+
+	t.Run("GetProfile with empty response", func(t *testing.T) {
+		connectClient.SendEmptyProfileResponse = true
+		maxNodes := int64(-1)
+		resp, err := client.GetProfile(context.Background(), "memory:alloc_objects:count:space:bytes", "{}", 0, 100, &maxNodes)
+		require.Nil(t, err)
+		// Mainly ensuring this does not panic like before
+		require.Nil(t, resp)
+		connectClient.SendEmptyProfileResponse = false
+	})
 }
 
 type FakePyroscopeConnectClient struct {
-	Req any
+	Req                      any
+	SendEmptyProfileResponse bool
 }
 
 func (f *FakePyroscopeConnectClient) ProfileTypes(ctx context.Context, c *connect.Request[querierv1.ProfileTypesRequest]) (*connect.Response[querierv1.ProfileTypesResponse], error) {
@@ -75,6 +86,9 @@ func (f *FakePyroscopeConnectClient) Series(ctx context.Context, c *connect.Requ
 
 func (f *FakePyroscopeConnectClient) SelectMergeStacktraces(ctx context.Context, c *connect.Request[querierv1.SelectMergeStacktracesRequest]) (*connect.Response[querierv1.SelectMergeStacktracesResponse], error) {
 	f.Req = c
+	if f.SendEmptyProfileResponse {
+		return &connect.Response[querierv1.SelectMergeStacktracesResponse]{Msg: &querierv1.SelectMergeStacktracesResponse{}}, nil
+	}
 	return &connect.Response[querierv1.SelectMergeStacktracesResponse]{
 		Msg: &querierv1.SelectMergeStacktracesResponse{
 			Flamegraph: &querierv1.FlameGraph{
diff --git a/pkg/tsdb/grafana-pyroscope-datasource/query.go b/pkg/tsdb/grafana-pyroscope-datasource/query.go
index 931f23fdbb7..0bbd30a6e4c 100644
--- a/pkg/tsdb/grafana-pyroscope-datasource/query.go
+++ b/pkg/tsdb/grafana-pyroscope-datasource/query.go
@@ -91,22 +91,30 @@ func (d *PyroscopeDatasource) query(ctx context.Context, pCtx backend.PluginCont
 				logger.Error("Error GetProfile()", "err", err)
 				return err
 			}
-			frame := responseToDataFrames(prof)
+
+			var frame *data.Frame
+			if prof != nil {
+				frame = responseToDataFrames(prof)
+
+				// If query called with streaming on then return a channel
+				// to subscribe on a client-side and consume updates from a plugin.
+				// Feel free to remove this if you don't need streaming for your datasource.
+				if qm.WithStreaming {
+					channel := live.Channel{
+						Scope:     live.ScopeDatasource,
+						Namespace: pCtx.DataSourceInstanceSettings.UID,
+						Path:      "stream",
+					}
+					frame.SetMeta(&data.FrameMeta{Channel: channel.String()})
+				}
+			} else {
+				// We still send empty data frame to give feedback that query really run, just didn't return any data.
+				frame = getEmptyDataFrame()
+			}
 			responseMutex.Lock()
 			response.Frames = append(response.Frames, frame)
 			responseMutex.Unlock()
 
-			// If query called with streaming on then return a channel
-			// to subscribe on a client-side and consume updates from a plugin.
-			// Feel free to remove this if you don't need streaming for your datasource.
-			if qm.WithStreaming {
-				channel := live.Channel{
-					Scope:     live.ScopeDatasource,
-					Namespace: pCtx.DataSourceInstanceSettings.UID,
-					Path:      "stream",
-				}
-				frame.SetMeta(&data.FrameMeta{Channel: channel.String()})
-			}
 			return nil
 		})
 	}
@@ -273,6 +281,18 @@ func (pt *ProfileTree) String() string {
 	return tree.String()
 }
 
+func getEmptyDataFrame() *data.Frame {
+	var emptyProfileDataFrame = data.NewFrame("response")
+	emptyProfileDataFrame.Meta = &data.FrameMeta{PreferredVisualization: "flamegraph"}
+	emptyProfileDataFrame.Fields = data.Fields{
+		data.NewField("level", nil, []int64{}),
+		data.NewField("value", nil, []int64{}),
+		data.NewField("self", nil, []int64{}),
+		data.NewField("label", nil, []string{}),
+	}
+	return emptyProfileDataFrame
+}
+
 type CustomMeta struct {
 	ProfileTypeID string
 }