fix: SelectSeries not correct for single reports (#4617)

* fix: SelectSeries not correct for single reports

The query-backend will only call the aggregate functions, when there is
more than one report for particular type. In the time series case the
aggregate function is also the one that ensures order/correctness of the
shape of data.

I considered doing this inside every query, but eventually it felt better to
do this when we would normally "aggregate"

* Ensure we are holding the lock when initializing the aggreagator.

* Use error

Author: simonswine

pkg/querybackend/block_reader_test.go

@@ -239,18 +239,20 @@ func (s *testSuite) Test_SeriesLabels() {
 	s.Assert().JSONEq(string(expected), string(actual))
 }
 
+var startTime = time.Unix(1739263329, 0)
+
 func (s *testSuite) Test_QueryTimeSeries() {
 	query := &queryv1.Query{
 		QueryType: queryv1.QueryType_QUERY_TIME_SERIES,
 		TimeSeries: &queryv1.TimeSeriesQuery{
 			GroupBy: []string{"service_name"},
-			Step:    1.0, // 1 second step
+			Step:    30.0,
 		},
 	}
 
 	req := &queryv1.InvokeRequest{
-		StartTime:     time.Now().Add(-1 * time.Hour).UnixMilli(),
-		EndTime:       time.Now().UnixMilli(),
+		StartTime:     startTime.UnixMilli(),
+		EndTime:       startTime.Add(time.Hour).UnixMilli(),
 		Query:         []*queryv1.Query{query},
 		QueryPlan:     s.plan,
 		LabelSelector: "{}",
@@ -262,6 +264,47 @@ func (s *testSuite) Test_QueryTimeSeries() {
 	s.Require().NotNil(resp)
 	s.Require().Len(resp.Reports, 1)
 	s.Require().NotNil(resp.Reports[0].TimeSeries)
+
+	actual, _ := json.Marshal(resp.Reports[0].TimeSeries.TimeSeries)
+	expected, err := os.ReadFile("testdata/fixtures/time_series.json")
+	s.Require().NoError(err)
+	s.Assert().JSONEq(string(expected), string(actual))
+}
+
+// When there is only one report we don't run the aggregate method. This check ensures that the timeseries, is still correctly formatted.
+func (s *testSuite) Test_QueryTimeSeriesOneReport() {
+	query := &queryv1.Query{
+		QueryType: queryv1.QueryType_QUERY_TIME_SERIES,
+		TimeSeries: &queryv1.TimeSeriesQuery{
+			GroupBy: []string{"service_name"},
+			Step:    30.0,
+		},
+	}
+
+	// shorten plan so there is only one report
+	shorterPlan := s.plan.CloneVT()
+	shorterPlan.Root = s.plan.Root.CloneVT()
+	shorterPlan.Root.Blocks = s.plan.Root.Blocks[:1]
+
+	req := &queryv1.InvokeRequest{
+		StartTime:     startTime.UnixMilli(),
+		EndTime:       startTime.Add(time.Hour).UnixMilli(),
+		Query:         []*queryv1.Query{query},
+		QueryPlan:     shorterPlan,
+		LabelSelector: "{}",
+		Tenant:        s.tenant,
+	}
+
+	resp, err := s.reader.Invoke(s.ctx, req)
+	s.Require().NoError(err)
+	s.Require().NotNil(resp)
+	s.Require().Len(resp.Reports, 1)
+	s.Require().NotNil(resp.Reports[0].TimeSeries)
+
+	actual, _ := json.Marshal(resp.Reports[0].TimeSeries.TimeSeries)
+	expected, err := os.ReadFile("testdata/fixtures/time_series_first_block.json")
+	s.Require().NoError(err)
+	s.Assert().JSONEq(string(expected), string(actual))
 }
 
 func (s *testSuite) Test_QueryTree_All_Tenant_Isolation() {

pkg/querybackend/query.go

@@ -67,12 +67,13 @@ func registerQueryType(
 	rt queryv1.ReportType,
 	q queryHandler,
 	a aggregatorProvider,
+	alwaysAggregate bool, // this option will always call the aggregate method for this report type, so it will also run when there is only one report
 	deps ...block.Section,
 ) {
 	registerQueryReportType(qt, rt)
 	registerQueryHandler(qt, q)
 	registerQueryDependencies(qt, deps...)
-	registerAggregator(rt, a)
+	registerAggregator(rt, a, alwaysAggregate)
 }
 
 type blockContext struct {

pkg/querybackend/query_label_names.go

@@ -18,6 +18,7 @@ func init() {
 		queryv1.ReportType_REPORT_LABEL_NAMES,
 		queryLabelNames,
 		newLabelNameAggregator,
+		false,
 		[]block.Section{block.SectionTSDB}...,
 	)
 }

pkg/querybackend/query_label_values.go

@@ -20,6 +20,7 @@ func init() {
 		queryv1.ReportType_REPORT_LABEL_VALUES,
 		queryLabelValues,
 		newLabelValueAggregator,
+		false,
 		[]block.Section{block.SectionTSDB}...,
 	)
 }

pkg/querybackend/query_pprof.go

@@ -21,6 +21,7 @@ func init() {
 		queryv1.ReportType_REPORT_PPROF,
 		queryPprof,
 		newPprofAggregator,
+		false,
 		[]block.Section{
 			block.SectionTSDB,
 			block.SectionProfiles,

pkg/querybackend/query_series_labels.go

@@ -21,6 +21,7 @@ func init() {
 		queryv1.ReportType_REPORT_SERIES_LABELS,
 		querySeriesLabels,
 		newSeriesLabelsAggregator,
+		false,
 		[]block.Section{block.SectionTSDB}...,
 	)
 }

pkg/querybackend/query_time_series.go

@@ -22,6 +22,7 @@ func init() {
 		queryv1.ReportType_REPORT_TIME_SERIES,
 		queryTimeSeries,
 		newTimeSeriesAggregator,
+		true,
 		[]block.Section{
 			block.SectionTSDB,
 			block.SectionProfiles,

pkg/querybackend/query_tree.go

@@ -19,6 +19,7 @@ func init() {
 		queryv1.ReportType_REPORT_TREE,
 		queryTree,
 		newTreeAggregator,
+		false,
 		[]block.Section{
 			block.SectionTSDB,
 			block.SectionProfiles,

pkg/querybackend/report_aggregator.go

@@ -10,6 +10,7 @@ import (
 var (
 	aggregatorMutex = new(sync.RWMutex)
 	aggregators     = map[queryv1.ReportType]aggregatorProvider{}
+	alwaysAggregate = map[queryv1.ReportType]struct{}{}
 	queryReportType = map[queryv1.QueryType]queryv1.ReportType{}
 )
 
@@ -23,14 +24,29 @@ type aggregator interface {
 	build() *queryv1.Report
 }
 
-func registerAggregator(t queryv1.ReportType, ap aggregatorProvider) {
+func registerAggregator(t queryv1.ReportType, ap aggregatorProvider, always bool) {
 	aggregatorMutex.Lock()
 	defer aggregatorMutex.Unlock()
 	_, ok := aggregators[t]
 	if ok {
 		panic(fmt.Sprintf("%s: aggregator already registered", t))
 	}
 	aggregators[t] = ap
+
+	if always {
+		_, ok := alwaysAggregate[t]
+		if ok {
+			panic(fmt.Sprintf("%s: aggregator already registered to always aggregat", t))
+		}
+		alwaysAggregate[t] = struct{}{}
+	}
+}
+
+func isAlwaysAggregate(t queryv1.ReportType) bool {
+	aggregatorMutex.RLock()
+	defer aggregatorMutex.RUnlock()
+	_, result := alwaysAggregate[t]
+	return result
 }
 
 func getAggregator(r *queryv1.InvokeRequest, x *queryv1.Report) (aggregator, error) {
@@ -97,13 +113,22 @@ func (ra *reportAggregator) aggregateReport(r *queryv1.Report) (err error) {
 	ra.sm.Lock()
 	v, found := ra.staged[r.ReportType]
 	if !found {
-		// We delay aggregation until we have at least two
-		// reports of the same type. Otherwise, we just store
-		// the report and will return it as is, if it is the
-		// only one.
-		ra.staged[r.ReportType] = r
+		// For most ReportTypes we delay aggregation until we have at least two
+		// reports of the same type. In case there is only one we will
+		// return it as is.
+		if !isAlwaysAggregate(r.ReportType) {
+			ra.staged[r.ReportType] = r
+			ra.sm.Unlock()
+			return nil
+		}
+
+		// Some ReportTypes need to call the aggregator for correctness even when
+		// there is only single instance, in that case call the aggregator right
+		// away and mark the report type appropriately in the staged map.
+		err = ra.aggregateReportNoCheck(r)
+		ra.staged[r.ReportType] = nil
 		ra.sm.Unlock()
-		return nil
+		return err
 	}
 	// Found a staged report of the same type.
 	if v != nil {

pkg/querybackend/report_aggregator_test.go

@@ -50,7 +50,7 @@ func mockAggregatorProvider(req *queryv1.InvokeRequest) aggregator {
 
 func TestReportAggregator_SingleReport(t *testing.T) {
 	reportType := queryv1.ReportType(999) // use a high number that won't conflict with other registrations
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, reportType)
@@ -78,7 +78,7 @@ func TestReportAggregator_SingleReport(t *testing.T) {
 
 func TestReportAggregator_TwoReports(t *testing.T) {
 	reportType := queryv1.ReportType(999)
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, reportType)
@@ -115,8 +115,8 @@ func TestReportAggregator_MultipleTypes(t *testing.T) {
 	type1 := queryv1.ReportType(999)
 	type2 := queryv1.ReportType(998)
 
-	registerAggregator(type1, mockAggregatorProvider)
-	registerAggregator(type2, mockAggregatorProvider)
+	registerAggregator(type1, mockAggregatorProvider, false)
+	registerAggregator(type2, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, type1)
@@ -167,7 +167,7 @@ func TestReportAggregator_NilReport(t *testing.T) {
 
 func TestReportAggregator_AggregateResponse(t *testing.T) {
 	reportType := queryv1.ReportType(999)
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, reportType)
@@ -194,7 +194,7 @@ func TestReportAggregator_AggregateResponse(t *testing.T) {
 
 func TestReportAggregator_ConcurrentAccess(t *testing.T) {
 	reportType := queryv1.ReportType(999)
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, reportType)
@@ -230,7 +230,7 @@ func TestReportAggregator_ConcurrentAccess(t *testing.T) {
 
 func TestGetAggregator(t *testing.T) {
 	reportType := queryv1.ReportType(999)
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	defer func() {
 		aggregatorMutex.Lock()
 		delete(aggregators, reportType)
@@ -256,9 +256,9 @@ func TestGetAggregator_UnknownReportType(t *testing.T) {
 func TestRegisterAggregator_Duplicate(t *testing.T) {
 	reportType := queryv1.ReportType(999)
 
-	registerAggregator(reportType, mockAggregatorProvider)
+	registerAggregator(reportType, mockAggregatorProvider, false)
 	assert.Panics(t, func() {
-		registerAggregator(reportType, mockAggregatorProvider)
+		registerAggregator(reportType, mockAggregatorProvider, false)
 	})
 
 	aggregatorMutex.Lock()