Check context cancellation every 128 iterations (#6159)

Author: yeya24

pkg/ingester/ingester.go

@@ -1737,9 +1737,11 @@ func (i *Ingester) metricsForLabelMatchersCommon(ctx context.Context, req *clien
 		Metric: make([]*cortexpb.Metric, 0),
 	}
 
+	cnt := 0
 	for mergedSet.Next() {
+		cnt++
 		// Interrupt if the context has been canceled.
-		if ctx.Err() != nil {
+		if cnt%util.CheckContextEveryNIterations == 0 && ctx.Err() != nil {
 			return nil, cleanup, ctx.Err()
 		}
 

pkg/querier/series/series_set.go

@@ -27,6 +27,7 @@ import (
 	"github.com/prometheus/prometheus/util/annotations"
 
 	"github.com/cortexproject/cortex/pkg/querier/iterators"
+	"github.com/cortexproject/cortex/pkg/util"
 )
 
 // ConcreteSeriesSet implements storage.SeriesSet.
@@ -143,8 +144,8 @@ func MatrixToSeriesSet(sortSeries bool, m model.Matrix) storage.SeriesSet {
 // MetricsToSeriesSet creates a storage.SeriesSet from a []metric.Metric
 func MetricsToSeriesSet(ctx context.Context, sortSeries bool, ms []model.Metric) storage.SeriesSet {
 	series := make([]storage.Series, 0, len(ms))
-	for _, m := range ms {
-		if ctx.Err() != nil {
+	for i, m := range ms {
+		if (i+1)%util.CheckContextEveryNIterations == 0 && ctx.Err() != nil {
 			return storage.ErrSeriesSet(ctx.Err())
 		}
 		series = append(series, &ConcreteSeries{

pkg/storage/tsdb/inmemory_index_cache.go

@@ -16,6 +16,8 @@ import (
 
 	storecache "github.com/thanos-io/thanos/pkg/store/cache"
 	"github.com/thanos-io/thanos/pkg/tenancy"
+
+	"github.com/cortexproject/cortex/pkg/util"
 )
 
 type InMemoryIndexCache struct {
@@ -147,8 +149,8 @@ func (c *InMemoryIndexCache) FetchMultiPostings(ctx context.Context, blockID uli
 	blockIDKey := blockID.String()
 	requests := 0
 	hit := 0
-	for _, key := range keys {
-		if ctx.Err() != nil {
+	for i, key := range keys {
+		if (i+1)%util.CheckContextEveryNIterations == 0 && ctx.Err() != nil {
 			c.commonMetrics.RequestTotal.WithLabelValues(storecache.CacheTypePostings, tenant).Add(float64(requests))
 			c.commonMetrics.HitsTotal.WithLabelValues(storecache.CacheTypePostings, tenant).Add(float64(hit))
 			return hits, misses
@@ -208,8 +210,8 @@ func (c *InMemoryIndexCache) FetchMultiSeries(ctx context.Context, blockID ulid.
 	blockIDKey := blockID.String()
 	requests := 0
 	hit := 0
-	for _, id := range ids {
-		if ctx.Err() != nil {
+	for i, id := range ids {
+		if (i+1)%util.CheckContextEveryNIterations == 0 && ctx.Err() != nil {
 			c.commonMetrics.RequestTotal.WithLabelValues(storecache.CacheTypeSeries, tenant).Add(float64(requests))
 			c.commonMetrics.HitsTotal.WithLabelValues(storecache.CacheTypeSeries, tenant).Add(float64(hit))
 			return hits, misses

pkg/util/context.go

@@ -0,0 +1,6 @@
+package util
+
+const (
+	// CheckContextEveryNIterations is used in some tight loops to check if the context is done.
+	CheckContextEveryNIterations = 128
+)