Phlare: ability to set maxNodes parameter (#67017)

pkg/tsdb/phlare/client.go

@@ -10,7 +10,7 @@ type ProfilingClient interface {
 	LabelNames(ctx context.Context, query string, start int64, end int64) ([]string, error)
 	LabelValues(ctx context.Context, query string, label string, start int64, end int64) ([]string, error)
 	GetSeries(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, groupBy []string, step float64) (*SeriesResponse, error)
-	GetProfile(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64) (*ProfileResponse, error)
+	GetProfile(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, maxNodes *int64) (*ProfileResponse, error)
 }
 
 type ProfileType struct {

pkg/tsdb/phlare/kinds/dataquery/types_dataquery_gen.go

@@ -35,6 +35,9 @@ type GrafanaPyroscopeDataQuery struct {
 	// Specifies the query label selectors.
 	LabelSelector string `json:"labelSelector"`
 
+	// Sets the maximum number of nodes in the flamegraph.
+	MaxNodes *int64 `json:"maxNodes,omitempty"`
+
 	// Specifies the type of profile to query.
 	ProfileTypeId string `json:"profileTypeId"`
 

pkg/tsdb/phlare/phlareClient.go

@@ -90,13 +90,14 @@ func (c *PhlareClient) GetSeries(ctx context.Context, profileTypeID string, labe
 	}, nil
 }
 
-func (c *PhlareClient) GetProfile(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64) (*ProfileResponse, error) {
+func (c *PhlareClient) GetProfile(ctx context.Context, profileTypeID, labelSelector string, start, end int64, maxNodes *int64) (*ProfileResponse, error) {
 	req := &connect.Request[querierv1.SelectMergeStacktracesRequest]{
 		Msg: &querierv1.SelectMergeStacktracesRequest{
 			ProfileTypeID: profileTypeID,
 			LabelSelector: labelSelector,
 			Start:         start,
 			End:           end,
+			MaxNodes:      maxNodes,
 		},
 	}
 

pkg/tsdb/phlare/phlareClient_test.go

@@ -32,7 +32,8 @@ func Test_PhlareClient(t *testing.T) {
 	})
 
 	t.Run("GetProfile", func(t *testing.T) {
-		resp, err := client.GetProfile(context.Background(), "memory:alloc_objects:count:space:bytes", "{}", 0, 100)
+		maxNodes := int64(-1)
+		resp, err := client.GetProfile(context.Background(), "memory:alloc_objects:count:space:bytes", "{}", 0, 100, &maxNodes)
 		require.Nil(t, err)
 
 		series := &ProfileResponse{

pkg/tsdb/phlare/pyroscopeClient.go

@@ -83,11 +83,14 @@ type PyroFlamebearer struct {
 	Names    []string  `json:"names"`
 }
 
-func (c *PyroscopeClient) getProfileData(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, groupBy []string) (*PyroscopeProfileResponse, error) {
+func (c *PyroscopeClient) getProfileData(ctx context.Context, profileTypeID, labelSelector string, start, end int64, maxNodes *int64, groupBy []string) (*PyroscopeProfileResponse, error) {
 	params := url.Values{}
 	params.Add("from", strconv.FormatInt(start, 10))
 	params.Add("until", strconv.FormatInt(end, 10))
 	params.Add("query", profileTypeID+labelSelector)
+	if maxNodes != nil {
+		params.Add("maxNodes", strconv.FormatInt(*maxNodes, 10))
+	}
 	params.Add("format", "json")
 	if len(groupBy) > 0 {
 		params.Add("groupBy", groupBy[0])
@@ -122,8 +125,8 @@ func (c *PyroscopeClient) getProfileData(ctx context.Context, profileTypeID stri
 	return respData, nil
 }
 
-func (c *PyroscopeClient) GetProfile(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64) (*ProfileResponse, error) {
-	respData, err := c.getProfileData(ctx, profileTypeID, labelSelector, start, end, nil)
+func (c *PyroscopeClient) GetProfile(ctx context.Context, profileTypeID, labelSelector string, start, end int64, maxNodes *int64) (*ProfileResponse, error) {
+	respData, err := c.getProfileData(ctx, profileTypeID, labelSelector, start, end, maxNodes, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -154,11 +157,11 @@ func (c *PyroscopeClient) GetProfile(ctx context.Context, profileTypeID string,
 	}, nil
 }
 
-func (c *PyroscopeClient) GetSeries(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, groupBy []string, step float64) (*SeriesResponse, error) {
+func (c *PyroscopeClient) GetSeries(ctx context.Context, profileTypeID string, labelSelector string, start, end int64, groupBy []string, step float64) (*SeriesResponse, error) {
 	// This is super ineffective at the moment. We need 2 different APIs one for profile one for series (timeline) data
 	// but Pyro returns all in single response. This currently does the simplest thing and calls the same API 2 times
 	// and gets the part of the response it needs.
-	respData, err := c.getProfileData(ctx, profileTypeID, labelSelector, start, end, groupBy)
+	respData, err := c.getProfileData(ctx, profileTypeID, labelSelector, start, end, nil, groupBy)
 	if err != nil {
 		return nil, err
 	}

pkg/tsdb/phlare/query.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"math"
+	"sync"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -13,6 +14,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/live"
 	"github.com/grafana/grafana/pkg/tsdb/phlare/kinds/dataquery"
 	"github.com/xlab/treeprint"
+	"golang.org/x/sync/errgroup"
 )
 
 type queryModel struct {
@@ -41,63 +43,76 @@ func (d *PhlareDatasource) query(ctx context.Context, pCtx backend.PluginContext
 		return response
 	}
 
+	responseMutex := sync.Mutex{}
+	g, gCtx := errgroup.WithContext(ctx)
 	if query.QueryType == queryTypeMetrics || query.QueryType == queryTypeBoth {
-		var dsJson dsJsonModel
-		err = json.Unmarshal(pCtx.DataSourceInstanceSettings.JSONData, &dsJson)
-		if err != nil {
-			response.Error = fmt.Errorf("error unmarshaling datasource json model: %v", err)
-			return response
-		}
-
-		parsedInterval := time.Second * 15
-		if dsJson.MinStep != "" {
-			parsedInterval, err = gtime.ParseDuration(dsJson.MinStep)
+		g.Go(func() error {
+			var dsJson dsJsonModel
+			err = json.Unmarshal(pCtx.DataSourceInstanceSettings.JSONData, &dsJson)
 			if err != nil {
-				parsedInterval = time.Second * 15
-				logger.Debug("Failed to parse the MinStep using default", "MinStep", dsJson.MinStep)
+				return fmt.Errorf("error unmarshaling datasource json model: %v", err)
 			}
-		}
-		logger.Debug("Sending SelectSeriesRequest", "queryModel", qm)
-		seriesResp, err := d.client.GetSeries(
-			ctx,
-			qm.ProfileTypeId,
-			qm.LabelSelector,
-			query.TimeRange.From.UnixMilli(),
-			query.TimeRange.To.UnixMilli(),
-			qm.GroupBy,
-			math.Max(query.Interval.Seconds(), parsedInterval.Seconds()),
-		)
-		if err != nil {
-			logger.Error("Querying SelectSeries()", "err", err)
-			response.Error = err
-			return response
-		}
-		// add the frames to the response.
-		response.Frames = append(response.Frames, seriesToDataFrames(seriesResp)...)
+
+			parsedInterval := time.Second * 15
+			if dsJson.MinStep != "" {
+				parsedInterval, err = gtime.ParseDuration(dsJson.MinStep)
+				if err != nil {
+					parsedInterval = time.Second * 15
+					logger.Debug("Failed to parse the MinStep using default", "MinStep", dsJson.MinStep)
+				}
+			}
+			logger.Debug("Sending SelectSeriesRequest", "queryModel", qm)
+			seriesResp, err := d.client.GetSeries(
+				gCtx,
+				qm.ProfileTypeId,
+				qm.LabelSelector,
+				query.TimeRange.From.UnixMilli(),
+				query.TimeRange.To.UnixMilli(),
+				qm.GroupBy,
+				math.Max(query.Interval.Seconds(), parsedInterval.Seconds()),
+			)
+			if err != nil {
+				logger.Error("Querying SelectSeries()", "err", err)
+				return err
+			}
+			// add the frames to the response.
+			responseMutex.Lock()
+			response.Frames = append(response.Frames, seriesToDataFrames(seriesResp)...)
+			responseMutex.Unlock()
+			return nil
+		})
 	}
 
 	if query.QueryType == queryTypeProfile || query.QueryType == queryTypeBoth {
-		logger.Debug("Calling GetProfile", "queryModel", qm)
-		prof, err := d.client.GetProfile(ctx, qm.ProfileTypeId, qm.LabelSelector, query.TimeRange.From.UnixMilli(), query.TimeRange.To.UnixMilli())
-		if err != nil {
-			logger.Error("Error GetProfile()", "err", err)
-			response.Error = err
-			return response
-		}
-		frame := responseToDataFrames(prof)
-		response.Frames = append(response.Frames, frame)
-
-		// 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",
+		g.Go(func() error {
+			logger.Debug("Calling GetProfile", "queryModel", qm)
+			prof, err := d.client.GetProfile(gCtx, qm.ProfileTypeId, qm.LabelSelector, query.TimeRange.From.UnixMilli(), query.TimeRange.To.UnixMilli(), qm.MaxNodes)
+			if err != nil {
+				logger.Error("Error GetProfile()", "err", err)
+				return err
 			}
-			frame.SetMeta(&data.FrameMeta{Channel: channel.String()})
-		}
+			frame := responseToDataFrames(prof)
+			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
+		})
+	}
+
+	if err := g.Wait(); err != nil {
+		response.Error = g.Wait()
 	}
 
 	return response

pkg/tsdb/phlare/query_test.go

@@ -28,6 +28,12 @@ func Test_query(t *testing.T) {
 		resp := ds.query(context.Background(), pCtx, *dataQuery)
 		require.Nil(t, resp.Error)
 		require.Equal(t, 2, len(resp.Frames))
+
+		// The order of the frames is not guaranteed, so we normalize it
+		if resp.Frames[0].Fields[0].Name == "level" {
+			resp.Frames[1], resp.Frames[0] = resp.Frames[0], resp.Frames[1]
+		}
+
 		require.Equal(t, "time", resp.Frames[0].Fields[0].Name)
 		require.Equal(t, data.NewField("level", nil, []int64{0, 1, 2}), resp.Frames[1].Fields[0])
 	})
@@ -289,7 +295,7 @@ func (f *FakeClient) LabelNames(ctx context.Context, query string, start int64,
 	panic("implement me")
 }
 
-func (f *FakeClient) GetProfile(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64) (*ProfileResponse, error) {
+func (f *FakeClient) GetProfile(ctx context.Context, profileTypeID, labelSelector string, start, end int64, maxNodes *int64) (*ProfileResponse, error) {
 	return &ProfileResponse{
 		Flamebearer: &Flamebearer{
 			Names: []string{"foo", "bar", "baz"},
@@ -305,7 +311,7 @@ func (f *FakeClient) GetProfile(ctx context.Context, profileTypeID string, label
 	}, nil
 }
 
-func (f *FakeClient) GetSeries(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, groupBy []string, step float64) (*SeriesResponse, error) {
+func (f *FakeClient) GetSeries(ctx context.Context, profileTypeID, labelSelector string, start, end int64, groupBy []string, step float64) (*SeriesResponse, error) {
 	f.Args = []interface{}{profileTypeID, labelSelector, start, end, groupBy, step}
 	return &SeriesResponse{
 		Series: []*Series{