Tempo: TraceQL metrics streaming (#99037)

* TraceQL metrics streaming POC

* Reduce duplicate frames by using scan() and combineResponses()

* Trying to remove samples outside of time range

* Remove code to clean out of range

* Metrics streaming config toggle

* Sync opening the search and metrics options

* Fix tests

* Fix issues after conflicts

* Fix tests

* Use absolute value when computing minXDelta

* Revert last commit

* Fix frame sorting

* Remove all duplicates

* Use fields from schema to get the frames

* Use FieldCache

* Address PR comments
This commit is contained in:
Andre Pereira
2025-02-11 11:11:01 +00:00
committed by GitHub
parent cbe5741096
commit d48802cdfb
16 changed files with 453 additions and 77 deletions
+103
View File
@@ -0,0 +1,103 @@
package tempo
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"github.com/grafana/grafana/pkg/tsdb/tempo/traceql"
"google.golang.org/grpc/metadata"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana-plugin-sdk-go/backend/tracing"
"github.com/grafana/grafana/pkg/tsdb/tempo/kinds/dataquery"
"github.com/grafana/tempo/pkg/tempopb"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
)
const MetricsPathPrefix = "metrics/"
func (s *Service) runMetricsStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender, datasource *Datasource) error {
ctx, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.runMetricsStream")
defer span.End()
response := &backend.DataResponse{}
var backendQuery *backend.DataQuery
err := json.Unmarshal(req.Data, &backendQuery)
if err != nil {
response.Error = fmt.Errorf("error unmarshaling backend query model: %v", err)
span.RecordError(response.Error)
span.SetStatus(codes.Error, response.Error.Error())
return err
}
var qrr *tempopb.QueryRangeRequest
err = json.Unmarshal(req.Data, &qrr)
if err != nil {
response.Error = fmt.Errorf("error unmarshaling Tempo query model: %v", err)
span.RecordError(response.Error)
span.SetStatus(codes.Error, response.Error.Error())
return err
}
if qrr.GetQuery() == "" {
return fmt.Errorf("query is empty")
}
qrr.Start = uint64(backendQuery.TimeRange.From.UnixNano())
qrr.End = uint64(backendQuery.TimeRange.To.UnixNano())
// Setting the user agent for the gRPC call. When DS is decoupled we don't recreate instance when grafana config
// changes or updates, so we have to get it from context.
// Ideally this would be pushed higher, so it's set once for all rpc calls, but we have only one now.
ctx = metadata.AppendToOutgoingContext(ctx, "User-Agent", backend.UserAgentFromContext(ctx).String())
stream, err := datasource.StreamingClient.MetricsQueryRange(ctx, qrr)
if err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
s.logger.Error("Error Search()", "err", err)
return err
}
return s.processMetricsStream(ctx, qrr.Query, stream, sender)
}
func (s *Service) processMetricsStream(ctx context.Context, query string, stream tempopb.StreamingQuerier_MetricsQueryRangeClient, sender StreamSender) error {
ctx, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.processStream")
defer span.End()
messageCount := 0
for {
msg, err := stream.Recv()
messageCount++
span.SetAttributes(attribute.Int("message_count", messageCount))
if errors.Is(err, io.EOF) {
if err := s.sendResponse(ctx, nil, nil, dataquery.SearchStreamingStateDone, sender); err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
return err
}
break
}
if err != nil {
s.logger.Error("Error receiving message", "err", err)
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
return err
}
transformed := traceql.TransformMetricsResponse(query, *msg)
if err := s.sendResponse(ctx, transformed, msg.Metrics, dataquery.SearchStreamingStateStreaming, sender); err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, err.Error())
return err
}
}
return nil
}
+30 -21
View File
@@ -89,7 +89,7 @@ func (s *Service) processStream(ctx context.Context, stream tempopb.StreamingQue
messageCount++
span.SetAttributes(attribute.Int("message_count", messageCount))
if errors.Is(err, io.EOF) {
if err := s.sendResponse(ctx, &ExtendedResponse{
if err := s.sendSearchResponse(ctx, &ExtendedResponse{
State: dataquery.SearchStreamingStateDone,
SearchResponse: &tempopb.SearchResponse{
Metrics: metrics,
@@ -114,7 +114,7 @@ func (s *Service) processStream(ctx context.Context, stream tempopb.StreamingQue
traceList = removeDuplicates(traceList)
span.SetAttributes(attribute.Int("traces_count", len(traceList)))
if err := s.sendResponse(ctx, &ExtendedResponse{
if err := s.sendSearchResponse(ctx, &ExtendedResponse{
State: dataquery.SearchStreamingStateStreaming,
SearchResponse: &tempopb.SearchResponse{
Metrics: metrics,
@@ -130,34 +130,43 @@ func (s *Service) processStream(ctx context.Context, stream tempopb.StreamingQue
return nil
}
func (s *Service) sendResponse(ctx context.Context, response *ExtendedResponse, sender StreamSender) error {
_, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.sendResponse")
func (s *Service) sendSearchResponse(ctx context.Context, response *ExtendedResponse, sender StreamSender) error {
_, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.sendSearchResponse")
defer span.End()
frame := createResponseDataFrame()
if response != nil {
span.SetAttributes(attribute.Int("trace_count", len(response.Traces)), attribute.String("state", string(response.State)))
tracesAsJson, err := json.Marshal(response.Traces)
if err != nil {
return err
}
tracesRawMessage := json.RawMessage(tracesAsJson)
frame.Fields[0].Append(tracesRawMessage)
metricsAsJson, err := json.Marshal(response.Metrics)
if err != nil {
return err
}
metricsRawMessage := json.RawMessage(metricsAsJson)
frame.Fields[1].Append(metricsRawMessage)
frame.Fields[2].Append(string(response.State))
frame.Fields[3].Append("")
return s.sendResponse(ctx, response.Traces, response.Metrics, response.State, sender)
}
return sender.SendFrame(frame, data.IncludeAll)
}
func (s *Service) sendResponse(ctx context.Context, result interface{}, metrics *tempopb.SearchMetrics, state dataquery.SearchStreamingState, sender StreamSender) error {
_, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.sendResponse")
defer span.End()
frame := createResponseDataFrame()
tracesAsJson, err := json.Marshal(result)
if err != nil {
return err
}
tracesRawMessage := json.RawMessage(tracesAsJson)
frame.Fields[0].Append(tracesRawMessage)
metricsAsJson, err := json.Marshal(metrics)
if err != nil {
return err
}
metricsRawMessage := json.RawMessage(metricsAsJson)
frame.Fields[1].Append(metricsRawMessage)
frame.Fields[2].Append(string(state))
frame.Fields[3].Append("")
return sender.SendFrame(frame, data.IncludeAll)
}
func sendError(searchErr error, sender StreamSender) error {
frame := createResponseDataFrame()
@@ -173,7 +182,7 @@ func sendError(searchErr error, sender StreamSender) error {
func createResponseDataFrame() *data.Frame {
frame := data.NewFrame("response")
frame.Fields = append(frame.Fields, data.NewField("traces", nil, []json.RawMessage{}))
frame.Fields = append(frame.Fields, data.NewField("result", nil, []json.RawMessage{}))
frame.Fields = append(frame.Fields, data.NewField("metrics", nil, []json.RawMessage{}))
frame.Fields = append(frame.Fields, data.NewField("state", nil, []string{}))
frame.Fields = append(frame.Fields, data.NewField("error", nil, []string{}))
+16 -3
View File
@@ -8,19 +8,22 @@ import (
"github.com/grafana/grafana-plugin-sdk-go/backend"
)
func (s *Service) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
func (s *Service) SubscribeStream(_ context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
s.logger.Debug("Allowing access to stream", "path", req.Path, "user", req.PluginContext.User)
status := backend.SubscribeStreamStatusPermissionDenied
if strings.HasPrefix(req.Path, SearchPathPrefix) {
status = backend.SubscribeStreamStatusOK
}
if strings.HasPrefix(req.Path, MetricsPathPrefix) {
status = backend.SubscribeStreamStatusOK
}
return &backend.SubscribeStreamResponse{
Status: status,
}, nil
}
func (s *Service) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
func (s *Service) PublishStream(_ context.Context, _ *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
s.logger.Debug("PublishStream called")
// Do not allow publishing at all.
@@ -31,9 +34,9 @@ func (s *Service) PublishStream(ctx context.Context, req *backend.PublishStreamR
func (s *Service) RunStream(ctx context.Context, request *backend.RunStreamRequest, sender *backend.StreamSender) error {
s.logger.Debug("New stream call", "path", request.Path)
tempoDatasource, err := s.getDSInfo(ctx, request.PluginContext)
if strings.HasPrefix(request.Path, SearchPathPrefix) {
tempoDatasource, err := s.getDSInfo(ctx, request.PluginContext)
if err != nil {
return err
}
@@ -43,6 +46,16 @@ func (s *Service) RunStream(ctx context.Context, request *backend.RunStreamReque
return nil
}
}
if strings.HasPrefix(request.Path, MetricsPathPrefix) {
if err != nil {
return err
}
if err = s.runMetricsStream(ctx, request, sender, tempoDatasource); err != nil {
return sendError(err, sender)
} else {
return nil
}
}
return fmt.Errorf("unknown path %s", request.Path)
}
+4 -2
View File
@@ -155,8 +155,10 @@ func (s *Service) performTraceRequest(ctx context.Context, dsInfo *Datasource, a
}
defer func() {
if err := resp.Body.Close(); err != nil {
ctxLogger.Error("Failed to close response body", "error", err, "function", logEntrypoint())
if resp != nil && resp.Body != nil {
if err := resp.Body.Close(); err != nil {
ctxLogger.Error("Failed to close response body", "error", err, "function", logEntrypoint())
}
}
}()
+15 -5
View File
@@ -3,6 +3,7 @@ package traceql
import (
"fmt"
"regexp"
"sort"
"strconv"
"strings"
"time"
@@ -13,7 +14,7 @@ import (
v1 "github.com/grafana/tempo/pkg/tempopb/common/v1"
)
func TransformMetricsResponse(query *dataquery.TempoQuery, resp tempopb.QueryRangeResponse) []*data.Frame {
func TransformMetricsResponse(query string, resp tempopb.QueryRangeResponse) []*data.Frame {
// prealloc frames
frames := make([]*data.Frame, len(resp.Series))
var exemplarFrames []*data.Frame
@@ -37,10 +38,11 @@ func TransformMetricsResponse(query *dataquery.TempoQuery, resp tempopb.QueryRan
},
Meta: &data.FrameMeta{
PreferredVisualization: data.VisTypeGraph,
Type: data.FrameTypeTimeSeriesMulti,
},
}
isHistogram := isHistogramQuery(*query.Query)
isHistogram := isHistogramQuery(query)
if isHistogram {
frame.Meta.PreferredVisualizationPluginID = "heatmap"
}
@@ -128,10 +130,18 @@ func transformLabelsAndGetName(seriesLabels []v1.KeyValue) (string, data.Labels)
if len(seriesLabels) == 1 {
_, name = metricsValueToString(seriesLabels[0].GetValue())
} else {
var labelStrings []string
for key, val := range labels {
labelStrings = append(labelStrings, fmt.Sprintf("%s=%s", key, val))
keys := make([]string, 0, len(labels))
for k := range labels {
keys = append(keys, k)
}
sort.Strings(keys)
var labelStrings []string
for _, key := range keys {
labelStrings = append(labelStrings, fmt.Sprintf("%s=%s", key, labels[key]))
}
name = fmt.Sprintf("{%s}", strings.Join(labelStrings, ", "))
}
}
+4 -15
View File
@@ -13,9 +13,7 @@ import (
func TestTransformMetricsResponse_EmptyResponse(t *testing.T) {
resp := tempopb.QueryRangeResponse{}
queryStr := ""
query := &dataquery.TempoQuery{Query: &queryStr}
frames := TransformMetricsResponse(query, resp)
frames := TransformMetricsResponse("", resp)
assert.Empty(t, frames)
}
@@ -32,9 +30,7 @@ func TestTransformMetricsResponse_SingleSeriesSingleLabel(t *testing.T) {
},
},
}
queryStr := ""
query := &dataquery.TempoQuery{Query: &queryStr}
frames := TransformMetricsResponse(query, resp)
frames := TransformMetricsResponse("", resp)
assert.Len(t, frames, 1)
assert.Equal(t, "value1", frames[0].RefID)
assert.Equal(t, "value1", frames[0].Name)
@@ -47,9 +43,6 @@ func TestTransformMetricsResponse_SingleSeriesSingleLabel(t *testing.T) {
}
func TestTransformMetricsResponse_SingleSeriesMultipleLabels(t *testing.T) {
// Skipping for now because this test is broken.
t.Skip()
resp := tempopb.QueryRangeResponse{
Series: []*tempopb.TimeSeries{
{
@@ -65,9 +58,7 @@ func TestTransformMetricsResponse_SingleSeriesMultipleLabels(t *testing.T) {
},
},
}
queryStr := ""
query := &dataquery.TempoQuery{Query: &queryStr}
frames := TransformMetricsResponse(query, resp)
frames := TransformMetricsResponse("", resp)
assert.Len(t, frames, 1)
assert.Equal(t, "{label1=\"value1\", label2=123, label3=123.456, label4=true}", frames[0].RefID)
assert.Equal(t, "{label1=\"value1\", label2=123, label3=123.456, label4=true}", frames[0].Name)
@@ -100,9 +91,7 @@ func TestTransformMetricsResponse_MultipleSeries(t *testing.T) {
},
},
}
queryStr := ""
query := &dataquery.TempoQuery{Query: &queryStr}
frames := TransformMetricsResponse(query, resp)
frames := TransformMetricsResponse("", resp)
assert.Len(t, frames, 2)
assert.Equal(t, "value1", frames[0].RefID)
assert.Equal(t, "value1", frames[0].Name)
+5 -3
View File
@@ -71,8 +71,10 @@ func (s *Service) runTraceQlQueryMetrics(ctx context.Context, pCtx backend.Plugi
resp, responseBody, err := s.performMetricsQuery(ctx, dsInfo, tempoQuery, backendQuery, span)
defer func() {
if err := resp.Body.Close(); err != nil {
ctxLogger.Error("Failed to close response body", "error", err, "function", logEntrypoint())
if resp != nil && resp.Body != nil {
if err := resp.Body.Close(); err != nil {
ctxLogger.Error("Failed to close response body", "error", err, "function", logEntrypoint())
}
}
}()
if err != nil {
@@ -105,7 +107,7 @@ func (s *Service) runTraceQlQueryMetrics(ctx context.Context, pCtx backend.Plugi
return res, err
}
frames := traceql.TransformMetricsResponse(tempoQuery, queryResponse)
frames := traceql.TransformMetricsResponse(*tempoQuery.Query, queryResponse)
result.Frames = frames
}