ceremonyclient/pebble/sstable/buffer_pool.go

149 lines
4.5 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
// 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
}