87 lines
1.9 KiB
Go
87 lines
1.9 KiB
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 BufferPool 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 *BufferPool
|
||
|
|
}
|
||
|
|
|
||
|
|
// 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 *BufferPool) Get() *Buffer {
|
||
|
|
buf := b.pool.Get().(*Buffer)
|
||
|
|
b.currentCapacity.Add(-uint64(buf.Cap()))
|
||
|
|
return buf
|
||
|
|
}
|
||
|
|
|
||
|
|
func (b *BufferPool) 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) *BufferPool {
|
||
|
|
b := &BufferPool{
|
||
|
|
maxSharedCapacity: sharedCapacity,
|
||
|
|
maxBufferCapacity: bufferCapacity,
|
||
|
|
}
|
||
|
|
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)
|
||
|
|
}
|