perf: Re-introduce fixed size memory pool for bloom querier

pkg/bloombuild/builder/spec_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/chunkenc"
 	v1 "github.com/grafana/loki/v3/pkg/storage/bloom/v1"
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 func blocksFromSchema(t *testing.T, n int, options v1.BlockOptions) (res []*v1.Block, data []v1.SeriesWithBlooms, refs []bloomshipper.BlockRef) {
@@ -74,7 +75,7 @@ func dummyBloomGen(t *testing.T, opts v1.BlockOptions, store v1.Iterator[*v1.Ser
 	for i, b := range blocks {
 		bqs = append(bqs, &bloomshipper.CloseableBlockQuerier{
 			BlockRef:     refs[i],
-			BlockQuerier: v1.NewBlockQuerier(b, false, v1.DefaultMaxPageSize),
+			BlockQuerier: v1.NewBlockQuerier(b, &mempool.SimpleHeapAllocator{}, v1.DefaultMaxPageSize),
 		})
 	}
 
@@ -152,7 +153,7 @@ func TestSimpleBloomGenerator(t *testing.T) {
 				expectedRefs := v1.PointerSlice(data)
 				outputRefs := make([]*v1.SeriesWithBlooms, 0, len(data))
 				for _, block := range outputBlocks {
-					bq := v1.NewBlockQuerier(block, false, v1.DefaultMaxPageSize).Iter()
+					bq := v1.NewBlockQuerier(block, &mempool.SimpleHeapAllocator{}, v1.DefaultMaxPageSize).Iter()
 					for bq.Next() {
 						outputRefs = append(outputRefs, bq.At())
 					}

pkg/bloombuild/planner/planner_test.go

@@ -29,6 +29,7 @@ import (
 	bloomshipperconfig "github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper/config"
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/indexshipper/tsdb"
 	"github.com/grafana/loki/v3/pkg/storage/types"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 var testDay = parseDayTime("2023-09-01")
@@ -411,7 +412,7 @@ func createPlanner(
 	reg := prometheus.NewPedanticRegistry()
 	metasCache := cache.NewNoopCache()
 	blocksCache := bloomshipper.NewFsBlocksCache(storageCfg.BloomShipperConfig.BlocksCache, reg, logger)
-	bloomStore, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageCfg, storage.ClientMetrics{}, metasCache, blocksCache, reg, logger)
+	bloomStore, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageCfg, storage.ClientMetrics{}, metasCache, blocksCache, &mempool.SimpleHeapAllocator{}, reg, logger)
 	require.NoError(t, err)
 
 	planner, err := New(cfg, limits, schemaCfg, storageCfg, storage.ClientMetrics{}, bloomStore, logger, reg)

pkg/bloomcompactor/retention_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper/config"
 	"github.com/grafana/loki/v3/pkg/storage/types"
 	util_log "github.com/grafana/loki/v3/pkg/util/log"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 	lokiring "github.com/grafana/loki/v3/pkg/util/ring"
 	"github.com/grafana/loki/v3/pkg/validation"
 )
@@ -822,7 +823,7 @@ func NewMockBloomStoreWithWorkDir(t *testing.T, workDir string) (*bloomshipper.B
 	metasCache := cache.NewMockCache()
 	blocksCache := bloomshipper.NewFsBlocksCache(storageConfig.BloomShipperConfig.BlocksCache, prometheus.NewPedanticRegistry(), logger)
 
-	store, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageConfig, metrics, metasCache, blocksCache, reg, logger)
+	store, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageConfig, metrics, metasCache, blocksCache, &mempool.SimpleHeapAllocator{}, reg, logger)
 	if err == nil {
 		t.Cleanup(store.Stop)
 	}

pkg/bloomcompactor/spec_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/chunkenc"
 	v1 "github.com/grafana/loki/v3/pkg/storage/bloom/v1"
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 func blocksFromSchema(t *testing.T, n int, options v1.BlockOptions) (res []*v1.Block, data []v1.SeriesWithBlooms, refs []bloomshipper.BlockRef) {
@@ -74,7 +75,7 @@ func dummyBloomGen(t *testing.T, opts v1.BlockOptions, store v1.Iterator[*v1.Ser
 	for i, b := range blocks {
 		bqs = append(bqs, &bloomshipper.CloseableBlockQuerier{
 			BlockRef:     refs[i],
-			BlockQuerier: v1.NewBlockQuerier(b, false, v1.DefaultMaxPageSize),
+			BlockQuerier: v1.NewBlockQuerier(b, &mempool.SimpleHeapAllocator{}, v1.DefaultMaxPageSize),
 		})
 	}
 
@@ -152,7 +153,7 @@ func TestSimpleBloomGenerator(t *testing.T) {
 				expectedRefs := v1.PointerSlice(data)
 				outputRefs := make([]*v1.SeriesWithBlooms, 0, len(data))
 				for _, block := range outputBlocks {
-					bq := v1.NewBlockQuerier(block, false, v1.DefaultMaxPageSize).Iter()
+					bq := v1.NewBlockQuerier(block, &mempool.SimpleHeapAllocator{}, v1.DefaultMaxPageSize).Iter()
 					for bq.Next() {
 						outputRefs = append(outputRefs, bq.At())
 					}

pkg/bloomgateway/bloomgateway_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper"
 	bloomshipperconfig "github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper/config"
 	"github.com/grafana/loki/v3/pkg/storage/types"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 	"github.com/grafana/loki/v3/pkg/validation"
 )
 
@@ -92,7 +93,7 @@ func setupBloomStore(t *testing.T) *bloomshipper.BloomStore {
 
 	reg := prometheus.NewRegistry()
 	blocksCache := bloomshipper.NewFsBlocksCache(storageCfg.BloomShipperConfig.BlocksCache, nil, logger)
-	store, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageCfg, cm, nil, blocksCache, reg, logger)
+	store, err := bloomshipper.NewBloomStore(schemaCfg.Configs, storageCfg, cm, nil, blocksCache, &mempool.SimpleHeapAllocator{}, reg, logger)
 	require.NoError(t, err)
 	t.Cleanup(store.Stop)
 

pkg/bloomgateway/processor.go

@@ -88,7 +88,7 @@ func (p *processor) processTasks(ctx context.Context, tenant string, day config.
 		// after iteration for performance (alloc reduction).
 		// This is safe to do here because we do not capture
 		// the underlying bloom []byte outside of iteration
-		bloomshipper.WithPool(true),
+		bloomshipper.WithPool(p.store.Allocator()),
 	)
 	duration = time.Since(startBlocks)
 	level.Debug(p.logger).Log("msg", "fetched blocks", "count", len(refs), "duration", duration, "err", err)

pkg/bloomgateway/processor_test.go

@@ -20,16 +20,21 @@ import (
 	"github.com/grafana/loki/v3/pkg/storage/config"
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/bloomshipper"
 	"github.com/grafana/loki/v3/pkg/util/constants"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 var _ bloomshipper.Store = &dummyStore{}
 
 // refs and blocks must be in 1-1 correspondence.
 func newMockBloomStore(refs []bloomshipper.BlockRef, blocks []*v1.Block, metas []bloomshipper.Meta) *dummyStore {
+	allocator := mempool.New("bloompages", mempool.Buckets{
+		{Size: 32, Capacity: 512 << 10},
+	}, nil)
 	return &dummyStore{
-		refs:   refs,
-		blocks: blocks,
-		metas:  metas,
+		refs:      refs,
+		blocks:    blocks,
+		metas:     metas,
+		allocator: allocator,
 	}
 }
 
@@ -38,6 +43,8 @@ type dummyStore struct {
 	refs   []bloomshipper.BlockRef
 	blocks []*v1.Block
 
+	allocator mempool.Allocator
+
 	// mock how long it takes to serve block queriers
 	delay time.Duration
 	// mock response error when serving block queriers in ForEach
@@ -76,6 +83,10 @@ func (s *dummyStore) Client(_ model.Time) (bloomshipper.Client, error) {
 	return nil, nil
 }
 
+func (s *dummyStore) Allocator() mempool.Allocator {
+	return s.allocator
+}
+
 func (s *dummyStore) Stop() {
 }
 
@@ -92,7 +103,7 @@ func (s *dummyStore) FetchBlocks(_ context.Context, refs []bloomshipper.BlockRef
 			if ref.Bounds.Equal(s.refs[i].Bounds) {
 				blockCopy := *block
 				bq := &bloomshipper.CloseableBlockQuerier{
-					BlockQuerier: v1.NewBlockQuerier(&blockCopy, false, v1.DefaultMaxPageSize),
+					BlockQuerier: v1.NewBlockQuerier(&blockCopy, s.Allocator(), v1.DefaultMaxPageSize),
 					BlockRef:     s.refs[i],
 				}
 				result = append(result, bq)

pkg/bloomgateway/util_test.go

@@ -432,6 +432,7 @@ func createBlocks(t *testing.T, tenant string, n int, from, through model.Time,
 		// 		t.Log(i, j, string(keys[i][j]))
 		// 	}
 		// }
+
 		blocks = append(blocks, block)
 		metas = append(metas, meta)
 		blockRefs = append(blockRefs, blockRef)

pkg/loki/modules.go

@@ -79,6 +79,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/util/httpreq"
 	"github.com/grafana/loki/v3/pkg/util/limiter"
 	util_log "github.com/grafana/loki/v3/pkg/util/log"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 	"github.com/grafana/loki/v3/pkg/util/querylimits"
 	lokiring "github.com/grafana/loki/v3/pkg/util/ring"
 	serverutil "github.com/grafana/loki/v3/pkg/util/server"
@@ -753,7 +754,24 @@ func (t *Loki) initBloomStore() (services.Service, error) {
 		level.Warn(logger).Log("msg", "failed to preload blocks cache", "err", err)
 	}
 
-	t.BloomStore, err = bloomshipper.NewBloomStore(t.Cfg.SchemaConfig.Configs, t.Cfg.StorageConfig, t.ClientMetrics, metasCache, blocksCache, reg, logger)
+	var pageAllocator mempool.Allocator
+
+	// Set global BloomPageAllocator variable
+	switch bsCfg.MemoryManagement.BloomPageAllocationType {
+	case "simple":
+		pageAllocator = &mempool.SimpleHeapAllocator{}
+	case "dynamic":
+		// sync buffer pool for bloom pages
+		// 128KB 256KB 512KB 1MB 2MB 4MB 8MB 16MB 32MB 64MB 128MB
+		pageAllocator = mempool.NewBytePoolAllocator(128<<10, 128<<20, 2)
+	case "fixed":
+		pageAllocator = mempool.New("bloom-page-pool", bsCfg.MemoryManagement.BloomPageMemPoolBuckets, reg)
+	default:
+		// should not happen as the type is validated upfront
+		return nil, fmt.Errorf("failed to create bloom store: invalid allocator type")
+	}
+
+	t.BloomStore, err = bloomshipper.NewBloomStore(t.Cfg.SchemaConfig.Configs, t.Cfg.StorageConfig, t.ClientMetrics, metasCache, blocksCache, pageAllocator, reg, logger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create bloom store: %w", err)
 	}

pkg/storage/bloom/v1/block.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 
 	"github.com/pkg/errors"
+
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 type BlockMetadata struct {
@@ -110,17 +112,18 @@ type BlockQuerier struct {
 }
 
 // NewBlockQuerier returns a new BlockQuerier for the given block.
-// WARNING: If noCapture is true, the underlying byte slice of the bloom page
-// will be returned to the pool for efficiency. This can only safely be used
-// when the underlying bloom bytes don't escape the decoder, i.e.
-// when loading blooms for querying (bloom-gw) but not for writing (bloom-compactor).
-// When usePool is true, the bloom MUST NOT be captured by the caller. Rather,
-// it should be discarded before another call to Next().
-func NewBlockQuerier(b *Block, usePool bool, maxPageSize int) *BlockQuerier {
+// WARNING: You can pass an implementation of Allocator that is responsible for
+// whether the underlying byte slice of the bloom page will be returned to the
+// pool for efficiency or not. Returning to the pool can only safely be used
+// when the underlying bloom bytes don't escape the decoder, i.e. when loading
+// blooms for querying (bloom-gateway), but not for writing (bloom-compactor).
+// Therefore, when calling NewBlockQuerier on the write path, you should always
+// pass the SimpleHeapAllocator implementation of the Allocator interface.
+func NewBlockQuerier(b *Block, alloc mempool.Allocator, maxPageSize int) *BlockQuerier {
 	return &BlockQuerier{
 		block:          b,
 		LazySeriesIter: NewLazySeriesIter(b),
-		blooms:         NewLazyBloomIter(b, usePool, maxPageSize),
+		blooms:         NewLazyBloomIter(b, alloc, maxPageSize),
 	}
 }
 
@@ -144,6 +147,10 @@ func (bq *BlockQuerier) Err() error {
 	return bq.blooms.Err()
 }
 
+func (bq *BlockQuerier) Close() {
+	bq.blooms.Close()
+}
+
 type BlockQuerierIter struct {
 	*BlockQuerier
 }

pkg/storage/bloom/v1/bloom.go

@@ -10,6 +10,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/chunkenc"
 	"github.com/grafana/loki/v3/pkg/storage/bloom/v1/filter"
 	"github.com/grafana/loki/v3/pkg/util/encoding"
+	"github.com/grafana/loki/v3/pkg/util/mempool"
 )
 
 // NB(chaudum): Some block pages are way bigger than others (400MiB and
@@ -24,7 +25,7 @@ type Bloom struct {
 
 func (b *Bloom) Encode(enc *encoding.Encbuf) error {
 	// divide by 8 b/c bloom capacity is measured in bits, but we want bytes
-	buf := bytes.NewBuffer(BloomPagePool.Get(int(b.Capacity() / 8)))
+	buf := bytes.NewBuffer(make([]byte, 0, int(b.Capacity()/8)))
 
 	// TODO(owen-d): have encoder implement writer directly so we don't need
 	// to indirect via a buffer
@@ -36,7 +37,6 @@ func (b *Bloom) Encode(enc *encoding.Encbuf) error {
 	data := buf.Bytes()
 	enc.PutUvarint(len(data)) // length of bloom filter
 	enc.PutBytes(data)
-	BloomPagePool.Put(data[:0]) // release to pool
 	return nil
 }
 
@@ -64,11 +64,14 @@ func (b *Bloom) Decode(dec *encoding.Decbuf) error {
 	return nil
 }
 
-func LazyDecodeBloomPage(r io.Reader, pool chunkenc.ReaderPool, page BloomPageHeader) (*BloomPageDecoder, error) {
-	data := BloomPagePool.Get(page.Len)[:page.Len]
-	defer BloomPagePool.Put(data)
+func LazyDecodeBloomPage(r io.Reader, alloc mempool.Allocator, pool chunkenc.ReaderPool, page BloomPageHeader) (*BloomPageDecoder, error) {
+	data, err := alloc.Get(page.Len)
+	if err != nil {
+		return nil, errors.Wrap(err, "allocating buffer")
+	}
+	defer alloc.Put(data)
 
-	_, err := io.ReadFull(r, data)
+	_, err = io.ReadFull(r, data)
 	if err != nil {
 		return nil, errors.Wrap(err, "reading bloom page")
 	}
@@ -84,7 +87,10 @@ func LazyDecodeBloomPage(r io.Reader, pool chunkenc.ReaderPool, page BloomPageHe
 	}
 	defer pool.PutReader(decompressor)
 
-	b := BloomPagePool.Get(page.DecompressedLen)[:page.DecompressedLen]
+	b, err := alloc.Get(page.DecompressedLen)
+	if err != nil {
+		return nil, errors.Wrap(err, "allocating buffer")
+	}
 
 	if _, err = io.ReadFull(decompressor, b); err != nil {
 		return nil, errors.Wrap(err, "decompressing bloom page")
@@ -96,14 +102,18 @@ func LazyDecodeBloomPage(r io.Reader, pool chunkenc.ReaderPool, page BloomPageHe
 }
 
 // shortcut to skip allocations when we know the page is not compressed
-func LazyDecodeBloomPageNoCompression(r io.Reader, page BloomPageHeader) (*BloomPageDecoder, error) {
+func LazyDecodeBloomPageNoCompression(r io.Reader, alloc mempool.Allocator, page BloomPageHeader) (*BloomPageDecoder, error) {
 	// data + checksum
 	if page.Len != page.DecompressedLen+4 {
 		return nil, errors.New("the Len and DecompressedLen of the page do not match")
 	}
-	data := BloomPagePool.Get(page.Len)[:page.Len]
 
-	_, err := io.ReadFull(r, data)
+	data, err := alloc.Get(page.Len)
+	if err != nil {
+		return nil, errors.Wrap(err, "allocating buffer")
+	}
+
+	_, err = io.ReadFull(r, data)
 	if err != nil {
 		return nil, errors.Wrap(err, "reading bloom page")
 	}
@@ -158,12 +168,16 @@ type BloomPageDecoder struct {
 // This can only safely be used when the underlying bloom
 // bytes don't escape the decoder:
 // on reads in the bloom-gw but not in the bloom-compactor
-func (d *BloomPageDecoder) Relinquish() {
+func (d *BloomPageDecoder) Relinquish(alloc mempool.Allocator) {
+	if d == nil {
+		return
+	}
+
 	data := d.data
 	d.data = nil
 
 	if cap(data) > 0 {
-		BloomPagePool.Put(data)
+		_ = alloc.Put(data)
 	}
 }
 
@@ -271,7 +285,7 @@ func (b *BloomBlock) DecodeHeaders(r io.ReadSeeker) (uint32, error) {
 // BloomPageDecoder returns a decoder for the given page index.
 // It may skip the page if it's too large.
 // NB(owen-d): if `skip` is true, err _must_ be nil.
-func (b *BloomBlock) BloomPageDecoder(r io.ReadSeeker, pageIdx int, maxPageSize int, metrics *Metrics) (res *BloomPageDecoder, skip bool, err error) {
+func (b *BloomBlock) BloomPageDecoder(r io.ReadSeeker, alloc mempool.Allocator, pageIdx int, maxPageSize int, metrics *Metrics) (res *BloomPageDecoder, skip bool, err error) {
 	if pageIdx < 0 || pageIdx >= len(b.pageHeaders) {
 		metrics.pagesSkipped.WithLabelValues(pageTypeBloom, skipReasonOOB).Inc()
 		metrics.bytesSkipped.WithLabelValues(pageTypeBloom, skipReasonOOB).Add(float64(b.pageHeaders[pageIdx].DecompressedLen))
@@ -294,9 +308,9 @@ func (b *BloomBlock) BloomPageDecoder(r io.ReadSeeker, pageIdx int, maxPageSize
 	}
 
 	if b.schema.encoding == chunkenc.EncNone {
-		res, err = LazyDecodeBloomPageNoCompression(r, page)
+		res, err = LazyDecodeBloomPageNoCompression(r, alloc, page)
 	} else {
-		res, err = LazyDecodeBloomPage(r, b.schema.DecompressorPool(), page)
+		res, err = LazyDecodeBloomPage(r, alloc, b.schema.DecompressorPool(), page)
 	}
 
 	if err != nil {