Prometheus: Add Exemplar sampling for streaming parser (#56049) (#56571)

(cherry picked from commit 152c7f149a)

Co-authored-by: Todd Treece <360020+toddtreece@users.noreply.github.com>

pkg/tsdb/prometheus/buffered/framing_test.go

@@ -26,16 +26,17 @@ import (
 
 var update = true
 
-func TestMatrixResponses(t *testing.T) {
+func TestResponses(t *testing.T) {
 	tt := []struct {
 		name     string
 		filepath string
 	}{
 		{name: "parse a simple matrix response", filepath: "range_simple"},
 		{name: "parse a simple matrix response with value missing steps", filepath: "range_missing"},
-		{name: "parse a response with Infinity", filepath: "range_infinity"},
-		{name: "parse a response with NaN", filepath: "range_nan"},
+		{name: "parse a matrix response with Infinity", filepath: "range_infinity"},
+		{name: "parse a matrix response with NaN", filepath: "range_nan"},
 		{name: "parse a response with legendFormat __auto", filepath: "range_auto"},
+		{name: "parse an exemplar response", filepath: "exemplar"},
 	}
 
 	for _, test := range tt {
@@ -96,13 +97,14 @@ func makeMockedApi(responseBytes []byte) (apiv1.API, error) {
 // struct here, because it has `time.time` and `time.duration` fields that
 // cannot be unmarshalled from JSON automatically.
 type storedPrometheusQuery struct {
-	RefId        string
-	RangeQuery   bool
-	Start        int64
-	End          int64
-	Step         int64
-	Expr         string
-	LegendFormat string
+	RefId         string
+	RangeQuery    bool
+	ExemplarQuery bool
+	Start         int64
+	End           int64
+	Step          int64
+	Expr          string
+	LegendFormat  string
 }
 
 func loadStoredPrometheusQuery(fileName string) (storedPrometheusQuery, error) {
@@ -126,11 +128,12 @@ func runQuery(response []byte, sq storedPrometheusQuery) (*backend.QueryDataResp
 	tracer := tracing.InitializeTracerForTest()
 
 	qm := QueryModel{
-		RangeQuery:   sq.RangeQuery,
-		Expr:         sq.Expr,
-		Interval:     fmt.Sprintf("%ds", sq.Step),
-		IntervalMS:   sq.Step * 1000,
-		LegendFormat: sq.LegendFormat,
+		RangeQuery:    sq.RangeQuery,
+		ExemplarQuery: sq.ExemplarQuery,
+		Expr:          sq.Expr,
+		Interval:      fmt.Sprintf("%ds", sq.Step),
+		IntervalMS:    sq.Step * 1000,
+		LegendFormat:  sq.LegendFormat,
 	}
 
 	b := Buffered{
@@ -165,6 +168,14 @@ func runQuery(response []byte, sq storedPrometheusQuery) (*backend.QueryDataResp
 		return nil, err
 	}
 
+	// parseTimeSeriesQuery forces range queries if the only query is an exemplar query
+	// so we need to set it back to false
+	if qm.ExemplarQuery {
+		for i := range queries {
+			queries[i].RangeQuery = false
+		}
+	}
+
 	return b.runQueries(context.Background(), queries)
 }
 

pkg/tsdb/prometheus/buffered/prometeus_bench_test.go

@@ -2,8 +2,11 @@ package buffered
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"math/rand"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
@@ -13,9 +16,37 @@ import (
 )
 
 // when memory-profiling this benchmark, these commands are recommended:
-// - go test -benchmem -run=^$ -benchtime 1x -memprofile memprofile.out -memprofilerate 1 -bench ^BenchmarkJson$ github.com/grafana/grafana/pkg/tsdb/prometheus
+// - go test -benchmem -run=^$ -benchtime 1x -memprofile memprofile.out -memprofilerate 1 -bench ^BenchmarkExemplarJson$ github.com/grafana/grafana/pkg/tsdb/prometheus/buffered
 // - go tool pprof -http=localhost:6061 memprofile.out
-func BenchmarkJson(b *testing.B) {
+func BenchmarkExemplarJson(b *testing.B) {
+	queryFileName := filepath.Join("../testdata", "exemplar.query.json")
+	query, err := loadStoredQuery(queryFileName)
+	require.NoError(b, err)
+
+	responseFileName := filepath.Join("../testdata", "exemplar.result.json")
+	// This is a test, so it's safe to ignore gosec warning G304.
+	// nolint:gosec
+	responseBytes, err := os.ReadFile(responseFileName)
+	require.NoError(b, err)
+
+	api, err := makeMockedApi(responseBytes)
+	require.NoError(b, err)
+
+	tracer := tracing.InitializeTracerForTest()
+
+	s := Buffered{tracer: tracer, log: &fakeLogger{}, client: api}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, err := s.runQueries(context.Background(), []*PrometheusQuery{query})
+		require.NoError(b, err)
+	}
+}
+
+// when memory-profiling this benchmark, these commands are recommended:
+// - go test -benchmem -run=^$ -benchtime 1x -memprofile memprofile.out -memprofilerate 1 -bench ^BenchmarkRangeJson$ github.com/grafana/grafana/pkg/tsdb/prometheus/buffered
+// - go tool pprof -http=localhost:6061 memprofile.out
+func BenchmarkRangeJson(b *testing.B) {
 	resp, query := createJsonTestData(1642000000, 1, 300, 400)
 
 	api, err := makeMockedApi(resp)
@@ -82,3 +113,28 @@ func createJsonTestData(start int64, step int64, timestampCount int, seriesCount
 
 	return bytes, query
 }
+
+func loadStoredQuery(fileName string) (*PrometheusQuery, error) {
+	// This is a test, so it's safe to ignore gosec warning G304.
+	// nolint:gosec
+	bytes, err := os.ReadFile(fileName)
+	if err != nil {
+		return nil, err
+	}
+
+	var sq storedPrometheusQuery
+
+	err = json.Unmarshal(bytes, &sq)
+	if err != nil {
+		return nil, err
+	}
+
+	return &PrometheusQuery{
+		RefId:         "A",
+		ExemplarQuery: sq.ExemplarQuery,
+		Start:         time.Unix(sq.Start, 0),
+		End:           time.Unix(sq.End, 0),
+		Step:          time.Second * time.Duration(sq.Step),
+		Expr:          sq.Expr,
+	}, nil
+}

pkg/tsdb/prometheus/buffered/time_series_query.go

@@ -568,6 +568,10 @@ func exemplarToDataFrames(response []apiv1.ExemplarQueryResult, query *Prometheu
 		}
 	}
 
+	sort.SliceStable(sampleExemplars, func(i, j int) bool {
+		return sampleExemplars[i].Time.Before(sampleExemplars[j].Time)
+	})
+
 	// Create DF from sampled exemplars
 	timeField := data.NewFieldFromFieldType(data.FieldTypeTime, len(sampleExemplars))
 	timeField.Name = "Time"