ref: 7f84e2e246aa9a9b5adb28ccd86fc61635d5c0ea
dir: /internal/bufpool/buffers.go/
// Package bufpool provides a lightweight byte-buffer type with optional
// pooling.
package bufpool
import "sync"
const (
// DefaultBufferCap is the minimum capacity a borrowed buffer will have.
// Borrow() will allocate or retrieve a buffer with at least this capacity.
DefaultBufferCap = 32 * 1024
// maxPooledBuffer defines the maximum capacity of a buffer that may be
// returned to the pool. Buffers larger than this will not be pooled to
// avoid unbounded memory usage.
maxPooledBuffer = 8 << 20
)
// Buffer is a growable byte container that optionally participates in a
// memory pool. A Buffer may be obtained through Borrow() or constructed
// directly from owned data via FromOwned().
//
// A Buffer's underlying slice may grow as needed. When finished with a
// pooled buffer, the caller should invoke Release() to return it to the pool.
//
// Buffers must not be copied after first use; doing so can cause double-returns
// to the pool and data races.
//
// In general, pass Buffer around when used internally, and directly .Bytes() when
// returning output across our API boundary. It is neither necessary nor efficient
// to copy/append the .Bytes() to a newly-allocated slice; in cases where we do
// want the raw byte slice out of our API boundary, it is perfectly acceptable to
// simply not call Release().
//
//go:nocopy
type Buffer struct {
_ struct{} // for nocopy
buf []byte
pool poolIndex
}
type poolIndex int8
const (
unpooled poolIndex = -1
)
var sizeClasses = [...]int{
DefaultBufferCap,
64 << 10,
128 << 10,
256 << 10,
512 << 10,
1 << 20,
2 << 20,
4 << 20,
maxPooledBuffer,
}
var bufferPools = func() []sync.Pool {
pools := make([]sync.Pool, len(sizeClasses))
for i, classCap := range sizeClasses {
capCopy := classCap
pools[i].New = func() any {
buf := make([]byte, 0, capCopy)
return &buf
}
}
return pools
}()
// Borrow retrieves a Buffer suitable for storing up to capHint bytes.
// The returned Buffer may come from an internal sync.Pool.
//
// If capHint is smaller than DefaultBufferCap, it is automatically raised
// to DefaultBufferCap. If no pooled buffer has sufficient capacity, a new
// unpooled buffer is allocated.
//
// The caller must call Release() when finished using the returned Buffer.
func Borrow(capHint int) Buffer {
if capHint < DefaultBufferCap {
capHint = DefaultBufferCap
}
classIdx, classCap, pooled := classFor(capHint)
if !pooled {
newBuf := make([]byte, 0, capHint)
return Buffer{buf: newBuf, pool: unpooled}
}
//nolint:forcetypeassert
buf := bufferPools[classIdx].Get().(*[]byte)
if cap(*buf) < classCap {
*buf = make([]byte, 0, classCap)
}
slice := (*buf)[:0]
return Buffer{buf: slice, pool: poolIndex(classIdx)} //#nosec G115
}
// FromOwned constructs a Buffer from a caller-owned byte slice. The resulting
// Buffer does not participate in pooling and will never be returned to the
// internal pool when released.
func FromOwned(buf []byte) Buffer {
return Buffer{buf: buf, pool: unpooled}
}
// Resize adjusts the length of the buffer to n bytes. If n exceeds the current
// capacity, the underlying storage is grown. If n is negative, it is treated
// as zero.
//
// The buffer's new contents beyond the previous length are undefined.
func (buf *Buffer) Resize(n int) {
if n < 0 {
n = 0
}
buf.ensureCapacity(n)
buf.buf = buf.buf[:n]
}
// Append copies the provided bytes onto the end of the buffer, growing its
// capacity if required. If src is empty, the method does nothing.
//
// The receiver retains ownership of the data; the caller may reuse src freely.
func (buf *Buffer) Append(src []byte) {
if len(src) == 0 {
return
}
start := len(buf.buf)
buf.ensureCapacity(start + len(src))
buf.buf = buf.buf[:start+len(src)]
copy(buf.buf[start:], src)
}
// Bytes returns the underlying byte slice that represents the current contents
// of the buffer. Modifying the returned slice modifies the Buffer itself.
func (buf *Buffer) Bytes() []byte {
return buf.buf
}
// Release returns the buffer to the global pool if it originated from the
// pool and its capacity is no larger than maxPooledBuffer. After release, the
// Buffer becomes invalid and should not be used further.
//
// Releasing a non-pooled buffer has no effect beyond clearing its internal
// storage.
func (buf *Buffer) Release() {
if buf.buf == nil {
return
}
buf.returnToPool()
buf.buf = nil
buf.pool = unpooled
}
// ensureCapacity grows the underlying buffer to accommodate the requested
// number of bytes. Growth doubles the capacity by default unless a larger
// expansion is needed. If the previous storage was pooled and not oversized,
// it is returned to the pool.
func (buf *Buffer) ensureCapacity(needed int) {
if cap(buf.buf) >= needed {
return
}
classIdx, classCap, pooled := classFor(needed)
var newBuf []byte
if pooled {
//nolint:forcetypeassert
raw := bufferPools[classIdx].Get().(*[]byte)
if cap(*raw) < classCap {
*raw = make([]byte, 0, classCap)
}
newBuf = (*raw)[:len(buf.buf)]
} else {
newBuf = make([]byte, len(buf.buf), classCap)
}
copy(newBuf, buf.buf)
buf.returnToPool()
buf.buf = newBuf
if pooled {
buf.pool = poolIndex(classIdx) //#nosec G115
} else {
buf.pool = unpooled
}
}
func classFor(size int) (idx, classCap int, ok bool) {
for i, class := range sizeClasses {
if size <= class {
return i, class, true
}
}
return -1, size, false
}
func (buf *Buffer) returnToPool() {
if buf.pool == unpooled {
return
}
tmp := buf.buf[:0]
bufferPools[int(buf.pool)].Put(&tmp)
}