71 lines
1.4 KiB
Go
71 lines
1.4 KiB
Go
package pdu
|
|
|
|
import (
|
|
"bytes"
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
type BufferPoolManager struct {
|
|
pools map[int]*sync.Pool
|
|
mu sync.RWMutex
|
|
debug bool
|
|
}
|
|
|
|
func (bpm *BufferPoolManager) logf(format string, args ...interface{}) {
|
|
if bpm.debug {
|
|
log.Printf(format, args...)
|
|
}
|
|
}
|
|
|
|
func NewBufferPoolManager() *BufferPoolManager {
|
|
return &BufferPoolManager{
|
|
pools: make(map[int]*sync.Pool),
|
|
}
|
|
}
|
|
|
|
func (bpm *BufferPoolManager) Get(size int) *bytes.Buffer {
|
|
bpm.mu.RLock()
|
|
pool, exists := bpm.pools[size]
|
|
bpm.mu.RUnlock()
|
|
|
|
if !exists {
|
|
bpm.mu.Lock()
|
|
// Double-check if another goroutine added the pool while we were waiting
|
|
pool, exists = bpm.pools[size]
|
|
if !exists {
|
|
bpm.logf("Creating new pool for size %d\n", size)
|
|
pool = &sync.Pool{
|
|
New: func() interface{} {
|
|
bpm.logf("Creating new buffer of size %d\n", size)
|
|
buf := bytes.NewBuffer(make([]byte, size))
|
|
buf.Reset()
|
|
return buf
|
|
},
|
|
}
|
|
bpm.pools[size] = pool
|
|
}
|
|
bpm.mu.Unlock()
|
|
}
|
|
|
|
buf := pool.Get().(*bytes.Buffer)
|
|
bpm.logf("Returning buffer of size %d: %p\n", buf.Cap(), buf)
|
|
return buf
|
|
}
|
|
|
|
func (bpm *BufferPoolManager) Put(buf *bytes.Buffer) {
|
|
size := buf.Cap()
|
|
bpm.mu.RLock()
|
|
pool, exists := bpm.pools[size]
|
|
bpm.mu.RUnlock()
|
|
|
|
if !exists {
|
|
bpm.logf("Cannot return %p, No pool for size %d\n", buf, size)
|
|
return
|
|
}
|
|
|
|
buf.Reset()
|
|
pool.Put(buf)
|
|
bpm.logf("Returned buffer of size %d: %p\n", size, buf)
|
|
}
|