chore(blooms): allocator is best effort; uses pool if available otherwise allocs

pkg/bloomgateway/processor.go

@@ -10,6 +10,7 @@ import (
 	"github.com/pkg/errors"
 
 	"github.com/grafana/dskit/concurrency"
+	"github.com/grafana/dskit/multierror"
 
 	v1 "github.com/grafana/loki/v3/pkg/storage/bloom/v1"
 	"github.com/grafana/loki/v3/pkg/storage/config"
@@ -113,16 +114,6 @@ func (p *processor) processTasks(ctx context.Context, tenant string, day config.
 }
 
 func (p *processor) processBlocks(ctx context.Context, bqs []*bloomshipper.CloseableBlockQuerier, data []blockWithTasks) error {
-
-	defer func() {
-		for i := range bqs {
-			if bqs[i] == nil {
-				continue
-			}
-			bqs[i].Close()
-		}
-	}()
-
 	return concurrency.ForEachJob(ctx, len(bqs), p.concurrency, func(ctx context.Context, i int) error {
 		bq := bqs[i]
 		if bq == nil {
@@ -136,15 +127,23 @@ func (p *processor) processBlocks(ctx context.Context, bqs []*bloomshipper.Close
 			return errors.Errorf("block and querier bounds differ: %s vs %s", block.ref.Bounds, bq.Bounds)
 		}
 
-		err := p.processBlock(ctx, bq.BlockQuerier, block.tasks)
+		err := p.processBlock(ctx, bq, block.tasks)
 		if err != nil {
 			return errors.Wrap(err, "processing block")
 		}
 		return nil
 	})
 }
 
-func (p *processor) processBlock(_ context.Context, blockQuerier *v1.BlockQuerier, tasks []Task) error {
+func (p *processor) processBlock(_ context.Context, bq *bloomshipper.CloseableBlockQuerier, tasks []Task) (err error) {
+	defer func() {
+		var errs multierror.MultiError
+		errs.Add(err)
+		errs.Add(bq.Close())
+		err = errs.Err()
+	}()
+
+	blockQuerier := bq.BlockQuerier
 	schema, err := blockQuerier.Schema()
 	if err != nil {
 		return err

pkg/storage/bloom/v1/block.go

@@ -136,6 +136,7 @@ func (bq *BlockQuerier) Schema() (Schema, error) {
 }
 
 func (bq *BlockQuerier) Reset() error {
+	bq.blooms.Reset()
 	return bq.LazySeriesIter.Seek(0)
 }
 
@@ -147,10 +148,6 @@ func (bq *BlockQuerier) Err() error {
 	return bq.blooms.Err()
 }
 
-func (bq *BlockQuerier) Close() {
-	bq.blooms.Close()
-}
-
 type BlockQuerierIter struct {
 	*BlockQuerier
 }
@@ -163,7 +160,11 @@ func (bq *BlockQuerier) Iter() *BlockQuerierIter {
 }
 
 func (b *BlockQuerierIter) Next() bool {
-	return b.LazySeriesIter.Next()
+	next := b.LazySeriesIter.Next()
+	if !next {
+		b.blooms.Reset()
+	}
+	return next
 }
 
 func (b *BlockQuerierIter) At() *SeriesWithBlooms {

pkg/storage/bloom/v1/bloom.go

@@ -65,13 +65,10 @@ func (b *Bloom) Decode(dec *encoding.Decbuf) error {
 }
 
 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")
-	}
+	data := alloc.Get(page.Len)
 	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")
 	}
@@ -87,12 +84,9 @@ func LazyDecodeBloomPage(r io.Reader, alloc mempool.Allocator, pool chunkenc.Rea
 	}
 	defer pool.PutReader(decompressor)
 
-	b, err := alloc.Get(page.DecompressedLen)
-	if err != nil {
-		return nil, errors.Wrap(err, "allocating buffer")
-	}
+	b := alloc.Get(page.DecompressedLen)
 
-	if _, err = io.ReadFull(decompressor, b); err != nil {
+	if _, err := io.ReadFull(decompressor, b); err != nil {
 		return nil, errors.Wrap(err, "decompressing bloom page")
 	}
 
@@ -108,12 +102,9 @@ func LazyDecodeBloomPageNoCompression(r io.Reader, alloc mempool.Allocator, page
 		return nil, errors.New("the Len and DecompressedLen of the page do not match")
 	}
 
-	data, err := alloc.Get(page.Len)
-	if err != nil {
-		return nil, errors.Wrap(err, "allocating buffer")
-	}
+	data := alloc.Get(page.Len)
 
-	_, err = io.ReadFull(r, data)
+	_, err := io.ReadFull(r, data)
 	if err != nil {
 		return nil, errors.Wrap(err, "reading bloom page")
 	}
@@ -175,10 +166,12 @@ func (d *BloomPageDecoder) Relinquish(alloc mempool.Allocator) {
 
 	data := d.data
 	d.data = nil
+	d.Reset() // Reset for cleaning up residual references to data via `dec`
 
 	if cap(data) > 0 {
 		_ = alloc.Put(data)
 	}
+
 }
 
 func (d *BloomPageDecoder) Reset() {

pkg/storage/bloom/v1/bloom_querier.go

@@ -162,6 +162,11 @@ func (it *LazyBloomIter) Err() error {
 	}
 }
 
-func (it *LazyBloomIter) Close() {
-	it.curPage.Relinquish(it.alloc)
+func (it *LazyBloomIter) Reset() {
+	it.err = nil
+	it.curPageIndex = 0
+	if it.curPage != nil {
+		it.curPage.Relinquish(it.alloc)
+	}
+	it.curPage = nil
 }

pkg/storage/bloom/v1/index.go

@@ -166,7 +166,7 @@ func (b *BlockIndex) NewSeriesPageDecoder(r io.ReadSeeker, header SeriesPageHead
 		return nil, errors.Wrap(err, "seeking to series page")
 	}
 
-	data, _ := SeriesPagePool.Get(header.Len)
+	data := SeriesPagePool.Get(header.Len)
 	defer SeriesPagePool.Put(data)
 	_, err = io.ReadFull(r, data)
 	if err != nil {

pkg/storage/stores/shipper/bloomshipper/cache.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/multierror"
 	"github.com/pkg/errors"
 
 	v1 "github.com/grafana/loki/v3/pkg/storage/bloom/v1"
@@ -23,11 +24,12 @@ type CloseableBlockQuerier struct {
 }
 
 func (c *CloseableBlockQuerier) Close() error {
-	c.BlockQuerier.Close()
+	var err multierror.MultiError
+	err.Add(c.BlockQuerier.Reset())
 	if c.close != nil {
-		return c.close()
+		err.Add(c.close())
 	}
-	return nil
+	return err.Err()
 }
 
 func (c *CloseableBlockQuerier) SeriesIter() (v1.PeekingIterator[*v1.SeriesWithBlooms], error) {

pkg/util/mempool/allocator.go

@@ -7,15 +7,15 @@ import (
 // Allocator handles byte slices for bloom queriers.
 // It exists to reduce the cost of allocations and allows to re-use already allocated memory.
 type Allocator interface {
-	Get(size int) ([]byte, error)
-	Put([]byte) bool
+	Get(size int) []byte
+	Put([]byte) (returned bool)
 }
 
 // SimpleHeapAllocator allocates a new byte slice every time and does not re-cycle buffers.
 type SimpleHeapAllocator struct{}
 
-func (a *SimpleHeapAllocator) Get(size int) ([]byte, error) {
-	return make([]byte, size), nil
+func (a *SimpleHeapAllocator) Get(size int) []byte {
+	return make([]byte, size)
 }
 
 func (a *SimpleHeapAllocator) Put([]byte) bool {
@@ -38,8 +38,8 @@ func NewBytePoolAllocator(minSize, maxSize int, factor float64) *BytePool {
 }
 
 // Get implements Allocator
-func (p *BytePool) Get(size int) ([]byte, error) {
-	return p.pool.Get(size).([]byte)[:size], nil
+func (p *BytePool) Get(size int) []byte {
+	return p.pool.Get(size).([]byte)[:size]
 }
 
 // Put implements Allocator

pkg/util/mempool/pool.go

@@ -2,7 +2,6 @@ package mempool
 
 import (
 	"errors"
-	"fmt"
 	"sync"
 	"unsafe"
 
@@ -18,9 +17,9 @@ var (
 )
 
 type slab struct {
+	once        sync.Once
 	buffer      chan unsafe.Pointer
 	size, count int
-	mtx         sync.Mutex
 	metrics     *metrics
 	name        string
 }
@@ -47,13 +46,9 @@ func (s *slab) init() {
 	s.metrics.availableBuffersPerSlab.WithLabelValues(s.name).Set(float64(s.count))
 }
 
-func (s *slab) get(size int) ([]byte, error) {
+func (s *slab) get(size int) []byte {
 	s.metrics.accesses.WithLabelValues(s.name, opTypeGet).Inc()
-	s.mtx.Lock()
-	if s.buffer == nil {
-		s.init()
-	}
-	defer s.mtx.Unlock()
+	s.once.Do(s.init)
 
 	// wait for available buffer on channel
 	var buf []byte
@@ -62,20 +57,27 @@ func (s *slab) get(size int) ([]byte, error) {
 		buf = unsafe.Slice((*byte)(ptr), s.size)
 	default:
 		s.metrics.errorsCounter.WithLabelValues(s.name, reasonSlabExhausted).Inc()
-		return nil, errSlabExhausted
+		buf = make([]byte, s.size)
 	}
 
-	return buf[:size], nil
+	return buf[:size]
 }
 
-func (s *slab) put(buf []byte) {
+func (s *slab) put(buf []byte) (returned bool) {
 	s.metrics.accesses.WithLabelValues(s.name, opTypePut).Inc()
 	if s.buffer == nil {
 		panic("slab is not initialized")
 	}
 
 	ptr := unsafe.Pointer(unsafe.SliceData(buf))
-	s.buffer <- ptr
+
+	// try to put buffer back on channel; if channel is full, discard buffer
+	select {
+	case s.buffer <- ptr:
+		return true
+	default:
+		return false
+	}
 }
 
 // MemPool is an Allocator implementation that uses a fixed size memory pool
@@ -101,28 +103,35 @@ func New(name string, buckets []Bucket, r prometheus.Registerer) *MemPool {
 // Get satisfies Allocator interface
 // Allocating a buffer from an exhausted pool/slab, or allocating a buffer that
 // exceeds the largest slab size will return an error.
-func (a *MemPool) Get(size int) ([]byte, error) {
+func (a *MemPool) Get(size int) []byte {
 	for i := 0; i < len(a.slabs); i++ {
 		if a.slabs[i].size < size {
 			continue
 		}
 		return a.slabs[i].get(size)
 	}
 	a.metrics.errorsCounter.WithLabelValues("pool", reasonSizeExceeded).Inc()
-	return nil, fmt.Errorf("no slab found for size: %d", size)
+	return make([]byte, size)
 }
 
 // Put satisfies Allocator interface
 // Every buffer allocated with Get(size int) needs to be returned to the pool
 // using Put(buffer []byte) so it can be re-cycled.
-func (a *MemPool) Put(buffer []byte) bool {
+// NB(owen-d): MemPool ensures that buffer capacities are _exactly_ the same
+// as individual slab sizes before returning them to the pool.
+func (a *MemPool) Put(buffer []byte) (returned bool) {
 	size := cap(buffer)
 	for i := 0; i < len(a.slabs); i++ {
-		if a.slabs[i].size < size {
-			continue
+		if a.slabs[i].size == size {
+			return a.slabs[i].put(buffer)
+		}
+
+		// if the current slab is too large, exit early
+		// as slabs are sorted by size and we won't find a smaller slab
+		if a.slabs[i].size > size {
+			break
 		}
-		a.slabs[i].put(buffer)
-		return true
 	}
+
 	return false
 }