zen/internal/bufferpool/bufferpool.go

88 lines
1.9 KiB
Go
Raw Normal View History

// 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"
)
2024-08-27 15:22:42 +07:00
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
2024-08-27 15:22:42 +07:00
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
}
2024-08-27 15:22:42 +07:00
func (b *Pool) Get() *Buffer {
buf := b.pool.Get().(*Buffer)
b.currentCapacity.Add(-uint64(buf.Cap()))
return buf
}
2024-08-27 15:22:42 +07:00
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)
}
}
2024-08-27 15:22:42 +07:00
func New(sharedCapacity uint64, bufferCapacity int) *Pool {
b := &Pool{
maxSharedCapacity: sharedCapacity,
maxBufferCapacity: bufferCapacity,
2024-09-03 13:38:45 +07:00
pool: &sync.Pool{},
}
b.pool.New = func() any {
return &Buffer{&bytes.Buffer{}, b}
}
return b
}
const (
2024-08-27 15:22:42 +07:00
SharedCap uint64 = 64 * 1024 * 1024 // 64MB
BufCap int = 8 * 1024 * 1024 // 8MB
)
2024-08-27 15:22:42 +07:00
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)
}