From 4fa8b8972844da5b2ae7e3946d67f8d0596f829b Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Mon, 17 Aug 2026 12:34:43 -0700 Subject: [PATCH 1/4] 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 | 53 +++++++++++++++++++++++---------- 3 files changed, 68 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..fc5a75cdb09 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,33 @@ 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) + defer func() { + batchesBufPool.Put(bp) + }() + + if cap(*bp) < c.numPartitions { + *bp = make(batchStream, c.numPartitions) + } else { + *bp = (*bp)[:c.numPartitions] + for i := range *bp { + (*bp)[i] = promchunk.Batch{} + } + } + // 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...) + *bp = mergeStreams(c.batches, c.nextBatchBuf[:], *bp, size) + c.batches = append(c.batches[:0], *bp...) if valType := c.h[0].Next(size); valType != chunkenc.ValNone { heap.Fix(&c.h, 0) From 54cbba87b5ea244868d23aa77bafc55543f3d877 Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Mon, 24 Aug 2026 13:57:21 -0700 Subject: [PATCH 2/4] zero batchesBuf before pool return and cap pool entry size Signed-off-by: Paurush Garg --- pkg/querier/batch/merge.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pkg/querier/batch/merge.go b/pkg/querier/batch/merge.go index fc5a75cdb09..be5ea7a41ba 100644 --- a/pkg/querier/batch/merge.go +++ b/pkg/querier/batch/merge.go @@ -10,6 +10,8 @@ import ( promchunk "github.com/cortexproject/cortex/pkg/chunk" ) +const maxPooledBatchesBufCap = 32 + var batchesBufPool = sync.Pool{ New: func() any { buf := make(batchStream, 3) @@ -150,16 +152,18 @@ func (c *mergeIterator) buildNextBatch(size int) chunkenc.ValueType { bp := batchesBufPool.Get().(*batchStream) defer func() { - batchesBufPool.Put(bp) + for i := range *bp { + (*bp)[i] = promchunk.Batch{} + } + if cap(*bp) <= maxPooledBatchesBufCap { + batchesBufPool.Put(bp) + } }() if cap(*bp) < c.numPartitions { *bp = make(batchStream, c.numPartitions) } else { *bp = (*bp)[:c.numPartitions] - for i := range *bp { - (*bp)[i] = promchunk.Batch{} - } } // All we need to do is get enough batches that our first batch's last entry From c3560314e5b04186114be04fb95ae7dcd760267a Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Mon, 24 Aug 2026 16:06:10 -0700 Subject: [PATCH 3/4] Updating Benchmark tests Signed-off-by: Paurush Garg --- ...5-go1.26.5-linux-arm64-2026-08-13.v1.count | Bin 0 -> 16384 bytes ...5-go1.26.5-linux-arm64-2026-08-13.v1.count | Bin 0 -> 16384 bytes ...5-go1.26.5-linux-arm64-2026-08-13.v1.count | Bin 0 -> 16384 bytes ...5-go1.26.5-linux-arm64-2026-08-13.v1.count | Bin 0 -> 16384 bytes ...5-go1.26.5-linux-arm64-2026-08-13.v1.count | Bin 0 -> 16384 bytes .config/go/telemetry/local/upload.token | 0 .config/go/telemetry/local/weekends | 1 + pkg/querier/batch/batch_test.go | 71 ++++++++++++------ 8 files changed, 49 insertions(+), 23 deletions(-) create mode 100644 .config/go/telemetry/local/asm@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count create mode 100644 .config/go/telemetry/local/cgo@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count create mode 100644 .config/go/telemetry/local/compile@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count create mode 100644 .config/go/telemetry/local/go@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count create mode 100644 .config/go/telemetry/local/link@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count create mode 100644 .config/go/telemetry/local/upload.token create mode 100644 .config/go/telemetry/local/weekends diff --git a/.config/go/telemetry/local/asm@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count b/.config/go/telemetry/local/asm@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count new file mode 100644 index 0000000000000000000000000000000000000000..5388060db2b0c10c788f08f665ff7f1b11af6e43 GIT binary patch literal 16384 zcmeI!Pfvp|9L8~u%EcJpKoif~*wkq_`8U(kj4mF$;nrajN*7>@JDB)Vd?~&Z--$Ar zW*}tAvV$>SLQ>j?Unw6rbXO!YkXk0uw9z-?Fp-fMsz8cK(|)rotEaR)mA(pH(Q;bt znsZcZ?)4nU{W|aMS$+`?7Sh9o^g1$r=L4yL5#Y>Gr;Dr%TEozen;pz5ge$x!zZLTl0MDT^+M66L>9rzxR!0 z732uy1&#~)=_aJ1!inj7NxBb-t?xYDM*QT% z!1LYntgte_%JN#)SUbgfUxx8idk-p}-%;LF&Tq2&EAvr4-!48siIg6B$z%Cr8ASjA z1Q0*~0R#|0009ILKmY**5I_I{1Q0*~0R#|0009ILKmY**5I_I{1Q0*~0R#|0009IL RKmY**5I_I{1Q7V+0-yCfV?F=? literal 0 HcmV?d00001 diff --git a/.config/go/telemetry/local/cgo@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count b/.config/go/telemetry/local/cgo@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count new file mode 100644 index 0000000000000000000000000000000000000000..e189e27dc18edd58898ebeb18fcdda3a218c68b4 GIT binary patch literal 16384 zcmeI!&rZTH90%|k6QjFm^x(nE63@UuMYEHCV(^Ah;=v24TbY$@CEXx$<57GR9)&m2 zg3NRb5+oAEdARRqp|>1_QqpL3j@rT!&@&F_`(;k_hpD`GJCp|i6wenqwkc^@{_C8WT^j&NwK z?jaFb8^;f{c*Ccj&4ZSZab2n_k-Fm4Ro-}H`X|ZXpOGJRBZvBaI!Sw4AGi0W=eKD5 ztc}h#ll+~@eBT{zBTC}>{zv;0PcJT)B#G}%=DXeSlOQ{%-jnQqTE4o#HJQ{$XWPmA znf0mpMLYglU&cqpThwV%k7bLZ2Lcd)00bZa0SG_<0uX=z1Rwwb2tWV=5P$##AOHaf uKmY;|fB*y_009U<00Izz00bZa0SG_<0uX=z1Rwwb2tWV=5P-lx5cmLlSBxJ3 literal 0 HcmV?d00001 diff --git a/.config/go/telemetry/local/compile@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count b/.config/go/telemetry/local/compile@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count new file mode 100644 index 0000000000000000000000000000000000000000..15a44fde7f941778d0619bad029f4fea1df2bbc1 GIT binary patch literal 16384 zcmeI!%We}f6b4|k0KtF|_YE6MtkWg|0-99;p|T5tiUq4UGtSh)9!DN0&|M#;j}qP_ zkHC@@FiAY2PMxVxS}8(*R4KmX9H04oq|Os+x#WslbF`c2QEs`R10gvbx1tFl37pHdL8=V=ZVw$M5Os|Gz=I9eQh0tZf zOt*I0ot@{=z7DK)yB}jJMLzl#?RQ^&c=Il%%&5+@C@RZZeok&Jhu(m$S@(W~2?Q=# z;MnWdOFnE%)1Ta|FO1%(-`|$qiMz|mruClVu`SR0&vjg2ukZb@Mr%;_F3@hx5#I9j zx;G6@dalopH|_V8>{Iuw)AEjc8~&f}ozo}3t}PC~Pss1t$i0N@ihQgSX3O`GVx>BN z`qRZfScxAne6D4ANpev}_@1ixpsyyRc7OPHKUn{e zC11R!Jv;f~mhOb=FKk+!AG}{{%12}N@NC-UO2ds-`*+@dtWbxkS^hROsN#eDTO29Y z7hzbJ?5g_?*3VU+r?u@tuDJMnA)oXsGlNjoSc3orAOHafKmY;|fB*y_009U<00Izz z00bZa0SG_<0uX=z1Rwwb2tWV=5P$##AOHafKmY;|fB*y_009U<00Izz00d?Ne*iL0 B+RFd{ literal 0 HcmV?d00001 diff --git a/.config/go/telemetry/local/go@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count b/.config/go/telemetry/local/go@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count new file mode 100644 index 0000000000000000000000000000000000000000..c9fcb85d4506d1741be3bbd2a5a4f37604d4b24a GIT binary patch literal 16384 zcmeI!L2uJA6bJBdK_f^LaNvfBn6z7(v_P4B8cb;1ieTcv2{&`HGdvP?3?toe~-x+*L)(MU)ZO}gPU2!aEl z*gM8WL84yT>&NL%+Ngd99zWF0Worbjxcz-B@46 z;|F(K`-W=UytDTc9OZo>~<*X%ar zuj|)rf+y|xUcM5gJA=!vpD!+=-ud$K$HFd4MhpTFfB*y_009U<00Izz00bZa0SG_< z0uX=z1Rwwb2tWV=5P$##AOHafKmY;|fB*y_009U<00Izz00bZa0SG_<0uWehfioIx B$C&^C literal 0 HcmV?d00001 diff --git a/.config/go/telemetry/local/link@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count b/.config/go/telemetry/local/link@go1.26.5-go1.26.5-linux-arm64-2026-08-13.v1.count new file mode 100644 index 0000000000000000000000000000000000000000..4ffee2db735a28f96f914a5f1a018a8e2dbb284a GIT binary patch literal 16384 zcmeI!K}y3w7>41vQsY9!3kQktviq6Q9Md-qG@R% zHmwR02=)zug!%aYnR!S;_C+EC8OkJnsk*u!B{CLW6-Y6tneT>SbW|wMrLQ7e)UA4> zVx3lMhaJnZC+EF6&MzY`lb&SK>sb4-6WYQJy=tJMNApg`J*CrbzOEhA8wbZ`OD`B} zw{LBc7W9W^t9^EJab*i94jV_NnXV~&H0KAhVZryK&4MSO{r?gu_*orol#5Fgyx(%| zw&2yQ;J2o^f_r`~`JUfOzP|alaD6(kV;CRf+wo6GRYe2sI!XE;(p#J6yV<-OIKKVR zSAnOzX};_-tKVB(A8Jo7E1itb>bGc5riufNQn&wD@Q z&)?2VhyVfzAb literal 0 HcmV?d00001 diff --git a/.config/go/telemetry/local/upload.token b/.config/go/telemetry/local/upload.token new file mode 100644 index 00000000000..e69de29bb2d diff --git a/.config/go/telemetry/local/weekends b/.config/go/telemetry/local/weekends new file mode 100644 index 00000000000..d00491fd7e5 --- /dev/null +++ b/.config/go/telemetry/local/weekends @@ -0,0 +1 @@ +1 diff --git a/pkg/querier/batch/batch_test.go b/pkg/querier/batch/batch_test.go index 8c44ab4c37c..9ec9dc908a4 100644 --- a/pkg/querier/batch/batch_test.go +++ b/pkg/querier/batch/batch_test.go @@ -174,30 +174,55 @@ func createChunks(b *testing.B, step time.Duration, numChunks, numSamplesPerChun 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) + + 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 _, 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+"/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 } - for _, it := range iters { - for it.Next() != chunkenc.ValNone { + }) + + b.Run(name+"/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 { + it.At() + } + if it.Err() != nil { + b.Fatal(it.Err().Error()) + } } } - } - }) + }) + } } From a25906261bb15901f8237c02636f8d6508460ca2 Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Fri, 11 Sep 2026 16:21:56 -0700 Subject: [PATCH 4/4] gate batchesBuf pool behind config flag and remove pool cap Signed-off-by: Paurush Garg --- CHANGELOG.md | 2 +- docs/blocks-storage/querier.md | 5 ++ docs/configuration/config-file-reference.md | 5 ++ pkg/querier/batch/batch_test.go | 55 ++++++++++--------- pkg/querier/batch/merge.go | 60 ++++++++++++++++----- pkg/querier/batch/merge_test.go | 43 +++++++++++++++ pkg/querier/querier.go | 6 +++ schemas/cortex-config-schema.json | 6 +++ 8 files changed, 143 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 40b695e2d2f..ee9111abf12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,7 +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: Use sync.Pool for mergeIterator batchesBuf to reduce memory allocations during series iteration. #7765 +* [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 diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 3dc3c704b21..e1b53c5ba63 100644 --- a/docs/blocks-storage/querier.md +++ b/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: | 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: | default = false] ``` ### `blocks_storage_config` diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index 1321bb42b75..a98b59755e4 100644 --- a/docs/configuration/config-file-reference.md +++ b/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: | 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: | default = false] ``` ### `query_frontend_config` diff --git a/pkg/querier/batch/batch_test.go b/pkg/querier/batch/batch_test.go index 6dcc204d482..adbaf8e43be 100644 --- a/pkg/querier/batch/batch_test.go +++ b/pkg/querier/batch/batch_test.go @@ -187,31 +187,38 @@ func BenchmarkNewChunkMergeIterator_NoReuse(b *testing.B) { {numChunks: 10, numSamplesPerChunk: 100, duplicationFactor: 3, enc: promchunk.PrometheusHistogramChunk}, } - 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() + 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()) + } + } } - if it.Err() != nil { - b.Fatal(it.Err().Error()) - } - } + }) } }) } diff --git a/pkg/querier/batch/merge.go b/pkg/querier/batch/merge.go index be5ea7a41ba..3f617332012 100644 --- a/pkg/querier/batch/merge.go +++ b/pkg/querier/batch/merge.go @@ -6,11 +6,17 @@ import ( "sync" "github.com/prometheus/prometheus/tsdb/chunkenc" + "go.uber.org/atomic" promchunk "github.com/cortexproject/cortex/pkg/chunk" ) -const maxPooledBatchesBufCap = 32 +// 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 { @@ -26,9 +32,12 @@ type mergeIterator struct { // Store the current sorted batchStream batches batchStream + // Buffers to merge in. + batchesBuf batchStream nextBatchBuf [1]promchunk.Batch numPartitions int + usePool bool currErr error } @@ -45,6 +54,11 @@ func newMergeIterator(it iterator, cs []GenericChunk) *mergeIterator { 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)) } } @@ -77,6 +91,17 @@ func (c *mergeIterator) Reset(size int) *mergeIterator { c.batches = c.batches[:0] c.numPartitions = size + 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{} + } + } + } + for i := range len(c.nextBatchBuf) { c.nextBatchBuf[i] = promchunk.Batch{} } @@ -150,28 +175,31 @@ func (c *mergeIterator) buildNextBatch(size int) chunkenc.ValueType { return chunkenc.ValNone } - bp := batchesBufPool.Get().(*batchStream) - defer func() { - for i := range *bp { - (*bp)[i] = promchunk.Batch{} - } - if cap(*bp) <= maxPooledBatchesBufCap { + 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] } - }() - - if cap(*bp) < c.numPartitions { - *bp = make(batchStream, c.numPartitions) + buf = *bp } else { - *bp = (*bp)[:c.numPartitions] + 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() - *bp = mergeStreams(c.batches, c.nextBatchBuf[:], *bp, size) - c.batches = append(c.batches[:0], *bp...) + 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) @@ -180,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 } diff --git a/pkg/querier/batch/merge_test.go b/pkg/querier/batch/merge_test.go index a7ab54b94b8..535f3a9b9a5 100644 --- a/pkg/querier/batch/merge_test.go +++ b/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) + }) +} diff --git a/pkg/querier/querier.go b/pkg/querier/querier.go index bef94e47b17..5284240c7da 100644 --- a/pkg/querier/querier.go +++ b/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 diff --git a/schemas/cortex-config-schema.json b/schemas/cortex-config-schema.json index 2fa4d4eaf97..0ba3aabfbca 100644 --- a/schemas/cortex-config-schema.json +++ b/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": {