github.com/cockroachdb/pebble@v0.0.0-20231214172447-ab4952c5f87b/sstable/buffer_pool.go

// Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.

package sstable

import (
	"github.com/cockroachdb/errors"
	"github.com/cockroachdb/pebble/internal/cache"
)

// A bufferHandle is a handle to manually-managed memory. The handle may point
// to a block in the block cache (h.Get() != nil), or a buffer that exists
// outside the block cache allocated from a BufferPool (b.Valid()).
type bufferHandle struct {
	h cache.Handle
	b Buf
}

// Get retrieves the underlying buffer referenced by the handle.
func (bh bufferHandle) Get() []byte {
	if v := bh.h.Get(); v != nil {
		return v
	} else if bh.b.p != nil {
		return bh.b.p.pool[bh.b.i].b
	}
	return nil
}

// Release releases the buffer, either back to the block cache or BufferPool.
func (bh bufferHandle) Release() {
	bh.h.Release()
	bh.b.Release()
}

// A BufferPool holds a pool of buffers for holding sstable blocks. An initial
// size of the pool is provided on Init, but a BufferPool will grow to meet the
// largest working set size. It'll never shrink. When a buffer is released, the
// BufferPool recycles the buffer for future allocations.
//
// A BufferPool should only be used for short-lived allocations with
// well-understood working set sizes to avoid excessive memory consumption.
//
// BufferPool is not thread-safe.
type BufferPool struct {
	// pool contains all the buffers held by the pool, including buffers that
	// are in-use. For every i < len(pool): pool[i].v is non-nil.
	pool []allocedBuffer
}

type allocedBuffer struct {
	v *cache.Value
	// b holds the current byte slice. It's backed by v, but may be a subslice
	// of v's memory while the buffer is in-use [ len(b) ≤ len(v.Buf()) ].
	//
	// If the buffer is not currently in-use, b is nil. When being recycled, the
	// BufferPool.Alloc will reset b to be a subslice of v.Buf().
	b []byte
}

// Init initializes the pool with an initial working set buffer size of
// `initialSize`.
func (p *BufferPool) Init(initialSize int) {
	*p = BufferPool{
		pool: make([]allocedBuffer, 0, initialSize),
	}
}

// initPreallocated is like Init but for internal sstable package use in
// instances where a pre-allocated slice of []allocedBuffer already exists. It's
// used to avoid an extra allocation initializing BufferPool.pool.
func (p *BufferPool) initPreallocated(pool []allocedBuffer) {
	*p = BufferPool{
		pool: pool[:0],
	}
}

// Release releases all buffers held by the pool and resets the pool to an
// uninitialized state.
func (p *BufferPool) Release() {
	for i := range p.pool {
		if p.pool[i].b != nil {
			panic(errors.AssertionFailedf("Release called on a BufferPool with in-use buffers"))
		}
		cache.Free(p.pool[i].v)
	}
	*p = BufferPool{}
}

// Alloc allocates a new buffer of size n. If the pool already holds a buffer at
// least as large as n, the pooled buffer is used instead.
//
// Alloc is O(MAX(N,M)) where N is the largest number of concurrently in-use
// buffers allocated and M is the initialSize passed to Init.
func (p *BufferPool) Alloc(n int) Buf {
	unusableBufferIdx := -1
	for i := 0; i < len(p.pool); i++ {
		if p.pool[i].b == nil {
			if len(p.pool[i].v.Buf()) >= n {
				p.pool[i].b = p.pool[i].v.Buf()[:n]
				return Buf{p: p, i: i}
			}
			unusableBufferIdx = i
		}
	}

	// If we would need to grow the size of the pool to allocate another buffer,
	// but there was a slot available occupied by a buffer that's just too
	// small, replace the too-small buffer.
	if len(p.pool) == cap(p.pool) && unusableBufferIdx >= 0 {
		i := unusableBufferIdx
		cache.Free(p.pool[i].v)
		p.pool[i].v = cache.Alloc(n)
		p.pool[i].b = p.pool[i].v.Buf()
		return Buf{p: p, i: i}
	}

	// Allocate a new buffer.
	v := cache.Alloc(n)
	p.pool = append(p.pool, allocedBuffer{v: v, b: v.Buf()[:n]})
	return Buf{p: p, i: len(p.pool) - 1}
}

// A Buf holds a reference to a manually-managed, pooled byte buffer.
type Buf struct {
	p *BufferPool
	// i holds the index into p.pool where the buffer may be found. This scheme
	// avoids needing to allocate the handle to the buffer on the heap at the
	// cost of copying two words instead of one.
	i int
}

// Valid returns true if the buf holds a valid buffer.
func (b Buf) Valid() bool {
	return b.p != nil
}

// Release releases the buffer back to the pool.
func (b *Buf) Release() {
	if b.p == nil {
		return
	}
	// Clear the allocedBuffer's byte slice. This signals the allocated buffer
	// is no longer in use and a future call to BufferPool.Alloc may reuse this
	// buffer.
	b.p.pool[b.i].b = nil
	b.p = nil
}