From 374ca2159407c6f3ec786bc19e25da44ded62fcf Mon Sep 17 00:00:00 2001 From: Runxi Yu Date: Fri, 6 Mar 2026 11:22:29 +0800 Subject: internal/bufpool: Split files --- internal/bufpool/append.go | 16 +++ internal/bufpool/borrow.go | 31 ++++++ internal/bufpool/buffer.go | 24 +++++ internal/bufpool/buffers.go | 214 ----------------------------------------- internal/bufpool/bytes.go | 7 ++ internal/bufpool/capacity.go | 37 +++++++ internal/bufpool/class.go | 23 +++++ internal/bufpool/consts.go | 12 +++ internal/bufpool/doc.go | 3 + internal/bufpool/from_owned.go | 8 ++ internal/bufpool/index.go | 7 ++ internal/bufpool/pool.go | 17 ++++ internal/bufpool/release.go | 17 ++++ internal/bufpool/resize.go | 15 +++ internal/bufpool/return.go | 10 ++ 15 files changed, 227 insertions(+), 214 deletions(-) create mode 100644 internal/bufpool/append.go create mode 100644 internal/bufpool/borrow.go create mode 100644 internal/bufpool/buffer.go delete mode 100644 internal/bufpool/buffers.go create mode 100644 internal/bufpool/bytes.go create mode 100644 internal/bufpool/capacity.go create mode 100644 internal/bufpool/class.go create mode 100644 internal/bufpool/consts.go create mode 100644 internal/bufpool/doc.go create mode 100644 internal/bufpool/from_owned.go create mode 100644 internal/bufpool/index.go create mode 100644 internal/bufpool/pool.go create mode 100644 internal/bufpool/release.go create mode 100644 internal/bufpool/resize.go create mode 100644 internal/bufpool/return.go diff --git a/internal/bufpool/append.go b/internal/bufpool/append.go new file mode 100644 index 00000000..f19dbc78 --- /dev/null +++ b/internal/bufpool/append.go @@ -0,0 +1,16 @@ +package bufpool + +// 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) +} diff --git a/internal/bufpool/borrow.go b/internal/bufpool/borrow.go new file mode 100644 index 00000000..ff212a9b --- /dev/null +++ b/internal/bufpool/borrow.go @@ -0,0 +1,31 @@ +package bufpool + +// 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 +} diff --git a/internal/bufpool/buffer.go b/internal/bufpool/buffer.go new file mode 100644 index 00000000..b2d648a1 --- /dev/null +++ b/internal/bufpool/buffer.go @@ -0,0 +1,24 @@ +package bufpool + +// 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 +} diff --git a/internal/bufpool/buffers.go b/internal/bufpool/buffers.go deleted file mode 100644 index 91e30a31..00000000 --- a/internal/bufpool/buffers.go +++ /dev/null @@ -1,214 +0,0 @@ -// 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) -} diff --git a/internal/bufpool/bytes.go b/internal/bufpool/bytes.go new file mode 100644 index 00000000..bcefbdfd --- /dev/null +++ b/internal/bufpool/bytes.go @@ -0,0 +1,7 @@ +package bufpool + +// 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 +} diff --git a/internal/bufpool/capacity.go b/internal/bufpool/capacity.go new file mode 100644 index 00000000..ecbd7d76 --- /dev/null +++ b/internal/bufpool/capacity.go @@ -0,0 +1,37 @@ +package bufpool + +// 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 + } +} diff --git a/internal/bufpool/class.go b/internal/bufpool/class.go new file mode 100644 index 00000000..60842d5e --- /dev/null +++ b/internal/bufpool/class.go @@ -0,0 +1,23 @@ +package bufpool + +var sizeClasses = [...]int{ + DefaultBufferCap, + 64 << 10, + 128 << 10, + 256 << 10, + 512 << 10, + 1 << 20, + 2 << 20, + 4 << 20, + maxPooledBuffer, +} + +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 +} diff --git a/internal/bufpool/consts.go b/internal/bufpool/consts.go new file mode 100644 index 00000000..4c205879 --- /dev/null +++ b/internal/bufpool/consts.go @@ -0,0 +1,12 @@ +package bufpool + +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 +) diff --git a/internal/bufpool/doc.go b/internal/bufpool/doc.go new file mode 100644 index 00000000..cadfe26e --- /dev/null +++ b/internal/bufpool/doc.go @@ -0,0 +1,3 @@ +// Package bufpool provides a lightweight byte-buffer type with optional +// pooling. +package bufpool diff --git a/internal/bufpool/from_owned.go b/internal/bufpool/from_owned.go new file mode 100644 index 00000000..65c5f471 --- /dev/null +++ b/internal/bufpool/from_owned.go @@ -0,0 +1,8 @@ +package bufpool + +// 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} +} diff --git a/internal/bufpool/index.go b/internal/bufpool/index.go new file mode 100644 index 00000000..5f59b0ed --- /dev/null +++ b/internal/bufpool/index.go @@ -0,0 +1,7 @@ +package bufpool + +type poolIndex int8 + +const ( + unpooled poolIndex = -1 +) diff --git a/internal/bufpool/pool.go b/internal/bufpool/pool.go new file mode 100644 index 00000000..30a4a2fb --- /dev/null +++ b/internal/bufpool/pool.go @@ -0,0 +1,17 @@ +package bufpool + +import "sync" + +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 +}() diff --git a/internal/bufpool/release.go b/internal/bufpool/release.go new file mode 100644 index 00000000..d8a52061 --- /dev/null +++ b/internal/bufpool/release.go @@ -0,0 +1,17 @@ +package bufpool + +// 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 +} diff --git a/internal/bufpool/resize.go b/internal/bufpool/resize.go new file mode 100644 index 00000000..78dc1dd7 --- /dev/null +++ b/internal/bufpool/resize.go @@ -0,0 +1,15 @@ +package bufpool + +// 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] +} diff --git a/internal/bufpool/return.go b/internal/bufpool/return.go new file mode 100644 index 00000000..fd08c121 --- /dev/null +++ b/internal/bufpool/return.go @@ -0,0 +1,10 @@ +package bufpool + +func (buf *Buffer) returnToPool() { + if buf.pool == unpooled { + return + } + + tmp := buf.buf[:0] + bufferPools[int(buf.pool)].Put(&tmp) +} -- cgit v1.3.1-10-gc9f91