Use sync.Pool for mergeIterator batchesBuf

CHANGELOG.md

@@ -67,6 +67,7 @@
 * [ENHANCEMENT] Upgrade Thanos and promql-engine to latest. #7740
 * [ENHANCEMENT] Ruler: Adjust ruler frontend decoder to not wrap query error messages with execution prefix, this makes error responses consistent between internal and external ruler paths. #7741
 * [ENHANCEMENT] Distributor: Deduplicate metric metadata when converting PRW 2.0 requests. PRW 2.0 attaches metadata to every series, so a metric family was previously expanded into one `MetricMetadata` per series. #7760
+* [ENHANCEMENT] Querier: Add `-querier.pool-iterator-batches-buf` flag to pool mergeIterator scratch buffers via sync.Pool, reducing per-iterator memory allocation. #7765
 * [ENHANCEMENT] Querier: Use non-pointer HistogramBucket slice in response codec. #7809
 * [ENHANCEMENT] Update build image and Go version to 1.27.0. #7814
 * [BUGFIX] Querier: Fix queryWithRetry and labelsWithRetry returning (nil, nil) on cancelled context by propagating ctx.Err(). #7370

docs/blocks-storage/querier.md

@@ -377,6 +377,11 @@ querier:
       # when resource thresholds are breached.
       # CLI flag: -querier.query-protection.eviction.max-evictions-per-cycle
       [max_evictions_per_cycle: <int> | default = 1]
+
+  # Pool the merge iterator scratch buffer (batchesBuf) via sync.Pool instead of
+  # allocating one per iterator.
+  # CLI flag: -querier.pool-iterator-batches-buf
+  [pool_iterator_batches_buf: <boolean> | default = false]
 ```
 
 ### `blocks_storage_config`

docs/configuration/config-file-reference.md

@@ -5548,6 +5548,11 @@ query_protection:
     # when resource thresholds are breached.
     # CLI flag: -querier.query-protection.eviction.max-evictions-per-cycle
     [max_evictions_per_cycle: <int> | default = 1]
+
+# Pool the merge iterator scratch buffer (batchesBuf) via sync.Pool instead of
+# allocating one per iterator.
+# CLI flag: -querier.pool-iterator-batches-buf
+[pool_iterator_batches_buf: <boolean> | default = false]
 ```
 
 ### `query_frontend_config`

pkg/querier/batch/batch_test.go

@@ -171,3 +171,55 @@ func createChunks(b *testing.B, step time.Duration, numChunks, numSamplesPerChun
 
 	return result
 }
+
+func BenchmarkNewChunkMergeIterator_NoReuse(b *testing.B) {
+	const numSeries = 10000
+
+	scenarios := []struct {
+		numChunks          int
+		numSamplesPerChunk int
+		duplicationFactor  int
+		enc                promchunk.Encoding
+	}{
+		{numChunks: 10, numSamplesPerChunk: 100, duplicationFactor: 1, enc: promchunk.PrometheusXorChunk},
+		{numChunks: 10, numSamplesPerChunk: 100, duplicationFactor: 3, enc: promchunk.PrometheusXorChunk},
+		{numChunks: 10, numSamplesPerChunk: 100, duplicationFactor: 1, enc: promchunk.PrometheusHistogramChunk},
+		{numChunks: 10, numSamplesPerChunk: 100, duplicationFactor: 3, enc: promchunk.PrometheusHistogramChunk},
+	}
+
+	for _, usePool := range []bool{false, true} {
+		b.Run(fmt.Sprintf("pool=%t", usePool), func(b *testing.B) {
+			poolBatchesBuf.Store(usePool)
+			defer func() { poolBatchesBuf.Store(false) }()
+
+			for _, scenario := range scenarios {
+				name := fmt.Sprintf("chunks: %d samples per chunk: %d duplication factor: %d encoding: %s",
+					scenario.numChunks,
+					scenario.numSamplesPerChunk,
+					scenario.duplicationFactor,
+					scenario.enc.String())
+
+				chunks := createChunks(b, step, scenario.numChunks, scenario.numSamplesPerChunk, scenario.duplicationFactor, scenario.enc)
+
+				b.Run(name, func(b *testing.B) {
+					b.ReportAllocs()
+
+					for b.Loop() {
+						iters := make([]chunkenc.Iterator, numSeries)
+						for i := range numSeries {
+							iters[i] = NewChunkMergeIterator(nil, chunks, 0, 0)
+						}
+						for _, it := range iters {
+							for it.Next() != chunkenc.ValNone {
+								it.At()
+							}
+							if it.Err() != nil {
+								b.Fatal(it.Err().Error())
+							}
+						}
+					}
+				})
+			}
+		})
+	}
+}

pkg/querier/batch/merge.go

@@ -3,12 +3,28 @@ package batch
 import (
 	"container/heap"
 	"sort"
+	"sync"
 
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"go.uber.org/atomic"
 
 	promchunk "github.com/cortexproject/cortex/pkg/chunk"
 )
 
+// poolBatchesBuf controls whether buildNextBatch uses sync.Pool or per-iterator batchesBuf.
+var poolBatchesBuf atomic.Bool
+
+func SetPoolBatchesBuf(enabled bool) {
+	poolBatchesBuf.Store(enabled)
+}
+
+var batchesBufPool = sync.Pool{
+	New: func() any {
+		buf := make(batchStream, 3)
+		return &buf
+	},
+}
+
 type mergeIterator struct {
 	its []*nonOverlappingIterator
 	h   iteratorHeap
@@ -20,6 +36,9 @@ type mergeIterator struct {
 	batchesBuf   batchStream
 	nextBatchBuf [1]promchunk.Batch
 
+	numPartitions int
+	usePool       bool
+
 	currErr error
 }
 
@@ -32,9 +51,14 @@ func newMergeIterator(it iterator, cs []GenericChunk) *mergeIterator {
 		c = mIterator.Reset(len(css))
 	} else {
 		c = &mergeIterator{
-			h:          make(iteratorHeap, 0, len(css)),
-			batches:    make(batchStream, 0, len(css)),
-			batchesBuf: make(batchStream, len(css)),
+			h:             make(iteratorHeap, 0, len(css)),
+			batches:       make(batchStream, 0, len(css)),
+			numPartitions: len(css),
+			usePool:       poolBatchesBuf.Load(),
+		}
+
+		if !poolBatchesBuf.Load() {
+			c.batchesBuf = make(batchStream, len(css))
 		}
 	}
 
@@ -65,13 +89,16 @@ func (c *mergeIterator) Reset(size int) *mergeIterator {
 	c.its = c.its[:0]
 	c.h = c.h[:0]
 	c.batches = c.batches[:0]
+	c.numPartitions = size
 
-	if size > cap(c.batchesBuf) {
-		c.batchesBuf = make(batchStream, len(c.its))
-	} else {
-		c.batchesBuf = c.batchesBuf[:size]
-		for i := range size {
-			c.batchesBuf[i] = promchunk.Batch{}
+	if !c.usePool {
+		if size > cap(c.batchesBuf) {
+			c.batchesBuf = make(batchStream, size)
+		} else {
+			c.batchesBuf = c.batchesBuf[:size]
+			for i := range size {
+				c.batchesBuf[i] = promchunk.Batch{}
+			}
 		}
 	}
 
@@ -141,12 +168,38 @@ func (c *mergeIterator) nextBatchEndTime() int64 {
 }
 
 func (c *mergeIterator) buildNextBatch(size int) chunkenc.ValueType {
+	if len(c.h) == 0 && len(c.batches) > 0 {
+		return c.batches[0].ValType
+	}
+	if len(c.h) == 0 {
+		return chunkenc.ValNone
+	}
+
+	var buf batchStream
+	if c.usePool {
+		bp := batchesBufPool.Get().(*batchStream)
+		defer func() {
+			for i := range *bp {
+				(*bp)[i] = promchunk.Batch{}
+			}
+			batchesBufPool.Put(bp)
+		}()
+		if cap(*bp) < c.numPartitions {
+			*bp = make(batchStream, c.numPartitions)
+		} else {
+			*bp = (*bp)[:c.numPartitions]
+		}
+		buf = *bp
+	} else {
+		buf = c.batchesBuf
+	}
+
 	// All we need to do is get enough batches that our first batch's last entry
 	// is before all iterators next entry.
 	for len(c.h) > 0 && (len(c.batches) == 0 || c.nextBatchEndTime() >= c.h[0].AtTime()) {
 		c.nextBatchBuf[0] = c.h[0].Batch()
-		c.batchesBuf = mergeStreams(c.batches, c.nextBatchBuf[:], c.batchesBuf, size)
-		c.batches = append(c.batches[:0], c.batchesBuf...)
+		buf = mergeStreams(c.batches, c.nextBatchBuf[:], buf, size)
+		c.batches = append(c.batches[:0], buf...)
 
 		if valType := c.h[0].Next(size); valType != chunkenc.ValNone {
 			heap.Fix(&c.h, 0)
@@ -155,6 +208,10 @@ func (c *mergeIterator) buildNextBatch(size int) chunkenc.ValueType {
 		}
 	}
 
+	if !c.usePool {
+		c.batchesBuf = buf
+	}
+
 	if len(c.batches) > 0 {
 		return c.batches[0].ValType
 	}

pkg/querier/batch/merge_test.go

@@ -28,6 +28,25 @@ func TestMergeIter(t *testing.T) {
 	})
 }
 
+func TestMergeIterPooled(t *testing.T) {
+	poolBatchesBuf.Store(true)
+	defer func() { poolBatchesBuf.Store(false) }()
+
+	forEncodings(t, func(t *testing.T, enc encoding.Encoding) {
+		chunk1 := mkGenericChunk(t, 0, 100, enc)
+		chunk2 := mkGenericChunk(t, model.TimeFromUnix(25), 100, enc)
+		chunk3 := mkGenericChunk(t, model.TimeFromUnix(50), 100, enc)
+		chunk4 := mkGenericChunk(t, model.TimeFromUnix(75), 100, enc)
+		chunk5 := mkGenericChunk(t, model.TimeFromUnix(100), 100, enc)
+
+		iter := newMergeIterator(nil, []GenericChunk{chunk1, chunk2, chunk3, chunk4, chunk5})
+		testIter(t, 200, newIteratorAdapter(iter), enc)
+
+		iter = newMergeIterator(iter, []GenericChunk{chunk1, chunk2, chunk3, chunk4, chunk5})
+		testSeek(t, 200, newIteratorAdapter(iter), enc)
+	})
+}
+
 func BenchmarkMergeIterator(b *testing.B) {
 	chunks := make([]GenericChunk, 0, 10)
 	for i := range 10 {
@@ -74,3 +93,27 @@ func TestMergeHarder(t *testing.T) {
 		testSeek(t, offset*numChunks+samples-offset, newIteratorAdapter(iter), enc)
 	})
 }
+
+func TestMergeHarderPooled(t *testing.T) {
+	poolBatchesBuf.Store(true)
+	defer func() { poolBatchesBuf.Store(false) }()
+
+	forEncodings(t, func(t *testing.T, enc encoding.Encoding) {
+		var (
+			numChunks = 24 * 15
+			chunks    = make([]GenericChunk, 0, numChunks)
+			from      = model.Time(0)
+			offset    = 30
+			samples   = 100
+		)
+		for range numChunks {
+			chunks = append(chunks, mkGenericChunk(t, from, samples, enc))
+			from = from.Add(time.Duration(offset) * time.Second)
+		}
+		iter := newMergeIterator(nil, chunks)
+		testIter(t, offset*numChunks+samples-offset, newIteratorAdapter(iter), enc)
+
+		iter = newMergeIterator(iter, chunks)
+		testSeek(t, offset*numChunks+samples-offset, newIteratorAdapter(iter), enc)
+	})
+}

pkg/querier/querier.go

@@ -110,6 +110,10 @@ type Config struct {
 
 	// Query protection: resource-based rejection.
 	QueryProtection configs.QueryProtection `yaml:"query_protection"`
+
+	// Pool the merge iterator scratch buffer (batchesBuf) via sync.Pool instead of
+	// allocating one per iterator.
+	PoolIteratorBatchesBuf bool `yaml:"pool_iterator_batches_buf"`
 }
 
 var (
@@ -160,6 +164,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.TimeoutClassificationDeadline, "querier.timeout-classification-deadline", time.Minute+59*time.Second, "The total time before the querier proactively cancels a query for timeout classification. Set this a few seconds less than the querier timeout.")
 	f.DurationVar(&cfg.TimeoutClassificationEvalThreshold, "querier.timeout-classification-eval-threshold", time.Minute+30*time.Second, "Eval time threshold above which a timeout is classified as user error (4XX).")
 	cfg.QueryProtection.RegisterFlagsWithPrefix(f, "querier.")
+	f.BoolVar(&cfg.PoolIteratorBatchesBuf, "querier.pool-iterator-batches-buf", false, "Pool the merge iterator scratch buffer (batchesBuf) via sync.Pool instead of allocating one per iterator.")
 }
 
 // Validate the config
@@ -228,6 +233,7 @@ func getChunksIteratorFunction(_ Config) chunkIteratorFunc {
 // New builds a queryable and promql engine.
 func New(cfg Config, limits *validation.Overrides, distributor Distributor, stores []QueryableWithFilter, reg prometheus.Registerer, logger log.Logger, isPartialDataEnabled partialdata.IsCfgEnabledFunc, resourceMonitor resource.IMonitor) (storage.SampleAndChunkQueryable, storage.ExemplarQueryable, engine.QueryEngine, services.Service) {
 	iteratorFunc := getChunksIteratorFunction(cfg)
+	batch.SetPoolBatchesBuf(cfg.PoolIteratorBatchesBuf)
 
 	// Create resource-based limiter if resource monitor is available and thresholds are configured.
 	var resourceBasedLimiter *limiter.ResourceBasedLimiter

schemas/cortex-config-schema.json

@@ -6665,6 +6665,12 @@
           "type": "boolean",
           "x-cli-flag": "querier.per-step-stats-enabled"
         },
+        "pool_iterator_batches_buf": {
+          "default": false,
+          "description": "Pool the merge iterator scratch buffer (batchesBuf) via sync.Pool instead of allocating one per iterator.",
+          "type": "boolean",
+          "x-cli-flag": "querier.pool-iterator-batches-buf"
+        },
         "query_protection": {
           "properties": {
             "eviction": {