Files
go-mqtt/mempool/bufpool.go
2024-08-21 15:32:05 +08:00

84 lines
1.8 KiB
Go

package mempool
import (
"bytes"
"sync"
)
var bufPool = NewBuffer(0)
// GetBuffer 从默认缓冲池中获取一个缓冲区
// GetBuffer takes a Buffer from the default buffer pool
func GetBuffer() *bytes.Buffer { return bufPool.Get() }
// PutBuffer 将 Buffer 返回到默认缓冲池
// PutBuffer returns Buffer to the default buffer pool
func PutBuffer(x *bytes.Buffer) { bufPool.Put(x) }
// BufferPool 缓冲池接口
type BufferPool interface {
Get() *bytes.Buffer
Put(x *bytes.Buffer)
}
// NewBuffer returns a buffer pool. The max specify the max capacity of the Buffer the pool will
// return. If the Buffer becoomes large than max, it will no longer be returned to the pool. If
// max <= 0, no limit will be enforced.
func NewBuffer(max int) BufferPool {
if max > 0 {
return newBufferWithCap(max)
}
return newBuffer()
}
// Buffer is a Buffer pool.
type Buffer struct {
pool *sync.Pool
}
func newBuffer() *Buffer {
return &Buffer{
pool: &sync.Pool{
New: func() any { return new(bytes.Buffer) },
},
}
}
// Get a Buffer from the pool.
func (b *Buffer) Get() *bytes.Buffer {
return b.pool.Get().(*bytes.Buffer)
}
// Put the Buffer back into pool. It resets the Buffer for reuse.
func (b *Buffer) Put(x *bytes.Buffer) {
x.Reset()
b.pool.Put(x)
}
// BufferWithCap is a Buffer pool that
type BufferWithCap struct {
bp *Buffer
max int
}
func newBufferWithCap(max int) *BufferWithCap {
return &BufferWithCap{
bp: newBuffer(),
max: max,
}
}
// Get a Buffer from the pool.
func (b *BufferWithCap) Get() *bytes.Buffer {
return b.bp.Get()
}
// Put the Buffer back into the pool if the capacity doesn't exceed the limit. It resets the Buffer
// for reuse.
func (b *BufferWithCap) Put(x *bytes.Buffer) {
if x.Cap() > b.max {
return
}
b.bp.Put(x)
}