zen/internal/bufferpool/bufferpool.go

// bufferpool is a simple buffer pool for bytes.Buffer.
//
// It is useful for reducing memory allocations when working with bytes.Buffer.
//
// Maximum buffer size is 4MB. Anything bigger will be discarded to be garbage collected
// and avoid huge memory usage.

package bufferpool

import (
	"bytes"
	"sync"
	"sync/atomic"
)

type Pool struct {
	pool              *sync.Pool
	maxSharedCapacity uint64
	currentCapacity   atomic.Uint64
	maxBufferCapacity int
}

// Buffer is a wrapper around bytes.Buffer
// that contains a reference to the pool.
//
// When Buffer Close method is called, it will be put back to the pool.
//
// Close never returns an error. The signature is to implement io.Closer.
type Buffer struct {
	*bytes.Buffer
	pool *Pool
}

// Close puts the buffer back to the pool.
// It never returns an error.
func (buf *Buffer) Close() error {
	buf.pool.Put(buf)
	return nil
}

func (b *Pool) Get() *Buffer {
	buf := b.pool.Get().(*Buffer)
	b.currentCapacity.Add(-uint64(buf.Cap()))
	return buf
}

func (b *Pool) Put(buf *Buffer) {
	bufCap := uint64(buf.Cap())
	overloaded := b.currentCapacity.Add(bufCap) > b.maxSharedCapacity
	if buf.Cap() < b.maxBufferCapacity && !overloaded {
		buf.Reset()
		b.pool.Put(buf)
	} else {
		b.currentCapacity.Add(-bufCap)
	}
}

func New(sharedCapacity uint64, bufferCapacity int) *Pool {
	b := &Pool{
		maxSharedCapacity: sharedCapacity,
		maxBufferCapacity: bufferCapacity,
		pool:              &sync.Pool{},
	}
	b.pool.New = func() any {
		return &Buffer{&bytes.Buffer{}, b}
	}
	return b
}

const (
	SharedCap uint64 = 64 * 1024 * 1024 // 64MB
	BufCap    int    = 8 * 1024 * 1024  // 8MB
)

var pool = New(SharedCap, BufCap)

// Get returns a bytes.Buffer from the global pool.
func Get() *Buffer {
	return pool.Get()
}

// Put puts the bytes.Buffer back to the the pool it origins from.
//
// Same method as calling Close on the Buffer.
func Put(buf *Buffer) {
	buf.pool.Put(buf)
}