mirror of
https://github.com/mochi-mqtt/server.git
synced 2025-09-26 20:21:12 +08:00
c6c7c296f6
* Dynamically allocate buffer for writes if needed * Remove unused net.Buffer * Return bytes written to buffer instead of conn * Dynamic write buffer * Reduce double write of pk.Payload * Use memory pool for packet encode * Pool doesn't guarantee value between Put and Get * Add benchmark for bufpool * Fix issue #346 * Change default pool not to have size cap --------- Co-authored-by: JB <28275108+mochi-co@users.noreply.github.com>
82 lines
1.6 KiB
Go
82 lines
1.6 KiB
Go
package mempool
|
|
|
|
import (
|
|
"bytes"
|
|
"sync"
|
|
)
|
|
|
|
var bufPool = NewBuffer(0)
|
|
|
|
// GetBuffer takes a Buffer from the default buffer pool
|
|
func GetBuffer() *bytes.Buffer { return bufPool.Get() }
|
|
|
|
// PutBuffer returns Buffer to the default buffer pool
|
|
func PutBuffer(x *bytes.Buffer) { bufPool.Put(x) }
|
|
|
|
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)
|
|
}
|