From a4f913cdfcaee98dde9f1896234ec20871743909 Mon Sep 17 00:00:00 2001
From: Ben Ye <benye@amazon.com>
Date: Tue, 13 Aug 2024 10:59:22 -0700
Subject: [PATCH] Check context cancellation every 128 iterations (#6159)

---
 pkg/ingester/ingester.go                 |  4 +++-
 pkg/querier/series/series_set.go         |  5 +++--
 pkg/storage/tsdb/inmemory_index_cache.go | 10 ++++++----
 pkg/util/context.go                      |  6 ++++++
 4 files changed, 18 insertions(+), 7 deletions(-)
 create mode 100644 pkg/util/context.go

diff --git a/pkg/ingester/ingester.go b/pkg/ingester/ingester.go
index e42180bd54..59ede90f83 100644
--- a/pkg/ingester/ingester.go
+++ b/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()
 		}
 
diff --git a/pkg/querier/series/series_set.go b/pkg/querier/series/series_set.go
index 0605ac9437..f246d73816 100644
--- a/pkg/querier/series/series_set.go
+++ b/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{
diff --git a/pkg/storage/tsdb/inmemory_index_cache.go b/pkg/storage/tsdb/inmemory_index_cache.go
index 1530bf99f8..95e10803cd 100644
--- a/pkg/storage/tsdb/inmemory_index_cache.go
+++ b/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
diff --git a/pkg/util/context.go b/pkg/util/context.go
new file mode 100644
index 0000000000..6cd5766fc6
--- /dev/null
+++ b/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
+)