From e42270a07c99b0b51b5d0a7ade7d326c39bafb9c Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Mon, 17 Aug 2026 12:34:43 -0700 Subject: [PATCH] perf(batch): pool batchesBuf in buildNextBatch via sync.Pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace per-iterator batchesBuf with a sync.Pool. Each buildNextBatch call borrows a scratch buffer from the pool, merges into it, copies the result into c.batches, and returns the buffer immediately. Idle iterators hold no buffer at all. This reduces heap usage from O(series) to O(concurrent_evaluations) — saving ~1200 bytes per series (~11 GB at 9M series). The pool approach is engine-agnostic: safe for both sequential and parallel series processing within a single Select call. Signed-off-by: Paurush Garg --- CHANGELOG.md | 1 + pkg/querier/batch/batch_test.go | 30 +++++++++++++++++++ pkg/querier/batch/merge.go | 51 ++++++++++++++++++++++----------- 3 files changed, 66 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index abd8bcc7927..993c725fd02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -53,6 +53,7 @@ * [ENHANCEMENT] Compactor: Reduce object storage GET calls when updating the bucket index by skipping re-reading parquet converter markers for blocks that already have a valid-version parquet entry in the previous index. #7669 * [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] Querier: Use sync.Pool for mergeIterator batchesBuf to reduce memory allocations during series iteration. #7765 * [BUGFIX] Querier: Fix queryWithRetry and labelsWithRetry returning (nil, nil) on cancelled context by propagating ctx.Err(). #7370 * [BUGFIX] Metrics Helper: Fix non-deterministic bucket order in merged histograms by sorting buckets after map iteration, matching Prometheus client library behavior. #7380 * [BUGFIX] Distributor: Return HTTP 401 Unauthorized when tenant ID resolution fails in the Prometheus Remote Write 2.0 path. #7389 diff --git a/pkg/querier/batch/batch_test.go b/pkg/querier/batch/batch_test.go index d90a0e1033e..8c44ab4c37c 100644 --- a/pkg/querier/batch/batch_test.go +++ b/pkg/querier/batch/batch_test.go @@ -171,3 +171,33 @@ func createChunks(b *testing.B, step time.Duration, numChunks, numSamplesPerChun return result } + +func BenchmarkNewChunkMergeIterator_ManyIterators(b *testing.B) { + const numSeries = 10000 + chunks := createChunks(b, step, 10, 100, 3, promchunk.PrometheusXorChunk) + + b.Run("create_only", 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) + } + _ = iters + } + }) + + b.Run("create_and_iterate_sequential", 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 { + } + } + } + }) +} diff --git a/pkg/querier/batch/merge.go b/pkg/querier/batch/merge.go index 33c0f91787e..be6a44ca132 100644 --- a/pkg/querier/batch/merge.go +++ b/pkg/querier/batch/merge.go @@ -3,12 +3,20 @@ package batch import ( "container/heap" "sort" + "sync" "github.com/prometheus/prometheus/tsdb/chunkenc" promchunk "github.com/cortexproject/cortex/pkg/chunk" ) +var batchesBufPool = sync.Pool{ + New: func() any { + buf := make(batchStream, 3) + return &buf + }, +} + type mergeIterator struct { its []*nonOverlappingIterator h iteratorHeap @@ -16,10 +24,10 @@ type mergeIterator struct { // Store the current sorted batchStream batches batchStream - // Buffers to merge in. - batchesBuf batchStream nextBatchBuf [1]promchunk.Batch + numPartitions int + currErr error } @@ -32,9 +40,9 @@ 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), } } @@ -65,15 +73,7 @@ func (c *mergeIterator) Reset(size int) *mergeIterator { c.its = c.its[:0] c.h = c.h[:0] c.batches = c.batches[:0] - - 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{} - } - } + c.numPartitions = size for i := range len(c.nextBatchBuf) { c.nextBatchBuf[i] = promchunk.Batch{} @@ -141,12 +141,31 @@ 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 + } + + bp := batchesBufPool.Get().(*batchStream) + batchesBuf := *bp + if cap(batchesBuf) < c.numPartitions { + batchesBuf = make(batchStream, c.numPartitions) + } else { + batchesBuf = batchesBuf[:c.numPartitions] + for i := range batchesBuf { + batchesBuf[i] = promchunk.Batch{} + } + } + defer batchesBufPool.Put(&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...) + batchesBuf = mergeStreams(c.batches, c.nextBatchBuf[:], batchesBuf, size) + c.batches = append(c.batches[:0], batchesBuf...) if valType := c.h[0].Next(size); valType != chunkenc.ValNone { heap.Fix(&c.h, 0)