55 lines
1.3 KiB
Go
55 lines
1.3 KiB
Go
|
package pool
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Buffer contains a byte slice.
|
||
|
|
type Buffer struct {
|
||
|
|
Data []byte
|
||
|
|
}
|
||
|
|
|
||
|
|
// BufferPool manages a pool of Buffers.
|
||
|
|
type BufferPool struct {
|
||
|
|
poolSliceSize uint32 // Size for the buffer slices in the pool.
|
||
|
|
buffersPool *sync.Pool
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewBufferPool creates a BufferPool with a specified size.
|
||
|
|
func NewBufferPool(poolSliceSize uint32) BufferPool {
|
||
|
|
pool := sync.Pool{
|
||
|
|
New: func() any {
|
||
|
|
return new(Buffer)
|
||
|
|
},
|
||
|
|
}
|
||
|
|
return BufferPool{poolSliceSize: poolSliceSize, buffersPool: &pool}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get retrieves a Buffer from the pool or creates a new one if necessary.
|
||
|
|
// It ensures the buffer's capacity is at least the specified size.
|
||
|
|
func (pool BufferPool) Get(size uint32) *Buffer {
|
||
|
|
result := pool.buffersPool.Get().(*Buffer)
|
||
|
|
|
||
|
|
if cap(result.Data) < int(size) {
|
||
|
|
result.Data = make([]byte, size)
|
||
|
|
} else {
|
||
|
|
result.Data = result.Data[:size]
|
||
|
|
}
|
||
|
|
return result
|
||
|
|
}
|
||
|
|
|
||
|
|
// Put returns a Buffer to the pool if its capacity does not exceed poolSliceSize.
|
||
|
|
func (pool BufferPool) Put(buf *Buffer) {
|
||
|
|
if cap(buf.Data) > int(pool.poolSliceSize) {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
buf.Data = buf.Data[:0]
|
||
|
|
pool.buffersPool.Put(buf)
|
||
|
|
}
|
||
|
|
|
||
|
|
// PoolSliceSize returns the size for buffer slices in the pool.
|
||
|
|
func (pool BufferPool) PoolSliceSize() uint32 {
|
||
|
|
return uint32(pool.poolSliceSize)
|
||
|
|
}
|