diff --git a/CHANGELOG.md b/CHANGELOG.md index a04a1784e99..ee9111abf12 100644 --- a/CHANGELOG.md +++ b/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 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 d90a0e1033e..adbaf8e43be 100644 --- a/pkg/querier/batch/batch_test.go +++ b/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()) + } + } + } + }) + } + }) + } +} diff --git a/pkg/querier/batch/merge.go b/pkg/querier/batch/merge.go index 33c0f91787e..3f617332012 100644 --- a/pkg/querier/batch/merge.go +++ b/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 } 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": {