add context cancellation checks on GetSeries

Signed-off-by: Erlan Zholdubai uulu <[email protected]>

Author: erlan-z

pkg/querier/distributor_queryable.go

@@ -136,7 +136,7 @@ func (q *distributorQuerier) Select(ctx context.Context, sortSeries bool, sp *st
 		if err != nil {
 			return storage.ErrSeriesSet(err)
 		}
-		return series.MetricsToSeriesSet(sortSeries, ms)
+		return series.MetricsToSeriesSet(ctx, sortSeries, ms)
 	}
 
 	return q.streamingSelect(ctx, sortSeries, minT, maxT, matchers)

pkg/querier/querier.go

@@ -437,7 +437,7 @@ func (q querier) Select(ctx context.Context, sortSeries bool, sp *storage.Select
 	// we have all the sets from different sources (chunk from store, chunks from ingesters,
 	// time series from store and time series from ingesters).
 	// mergeSeriesSets will return sorted set.
-	return q.mergeSeriesSets(result)
+	return q.mergeSeriesSets(ctx, result)
 }
 
 // LabelValues implements storage.Querier.
@@ -553,7 +553,7 @@ func (querier) Close() error {
 	return nil
 }
 
-func (q querier) mergeSeriesSets(sets []storage.SeriesSet) storage.SeriesSet {
+func (q querier) mergeSeriesSets(ctx context.Context, sets []storage.SeriesSet) storage.SeriesSet {
 	// Here we deal with sets that are based on chunks and build single set from them.
 	// Remaining sets are merged with chunks-based one using storage.NewMergeSeriesSet
 
@@ -565,6 +565,9 @@ func (q querier) mergeSeriesSets(sets []storage.SeriesSet) storage.SeriesSet {
 
 		// SeriesSet may have some series backed up by chunks, and some not.
 		for set.Next() {
+			if ctx.Err() != nil {
+				return storage.ErrSeriesSet(ctx.Err())
+			}
 			s := set.At()
 
 			if sc, ok := s.(SeriesWithChunks); ok {

pkg/querier/series/series_set.go

@@ -17,6 +17,7 @@
 package series
 
 import (
+	"context"
 	"sort"
 
 	"github.com/prometheus/common/model"
@@ -167,9 +168,12 @@ func MatrixToSeriesSet(sortSeries bool, m model.Matrix) storage.SeriesSet {
 }
 
 // MetricsToSeriesSet creates a storage.SeriesSet from a []metric.Metric
-func MetricsToSeriesSet(sortSeries bool, ms []metric.Metric) storage.SeriesSet {
+func MetricsToSeriesSet(ctx context.Context, sortSeries bool, ms []metric.Metric) storage.SeriesSet {
 	series := make([]storage.Series, 0, len(ms))
 	for _, m := range ms {
+		if ctx.Err() != nil {
+			return storage.ErrSeriesSet(ctx.Err())
+		}
 		series = append(series, &ConcreteSeries{
 			labels:  metricToLabels(m.Metric),
 			samples: nil,