diff --git a/common/buf/buffer.go b/common/buf/buffer.go index 26c6c74f..a9e545a2 100644 --- a/common/buf/buffer.go +++ b/common/buf/buffer.go @@ -13,59 +13,65 @@ import ( ) type Buffer struct { - data []byte - start int - end int - capacity int - refs atomic.Int32 - managed bool + data []byte + start int + end int + capacity int + refs atomic.Int32 + managed bool + dataManaged bool } func New() *Buffer { - return &Buffer{ - data: Get(BufferSize), - capacity: BufferSize, - managed: true, - } + return NewSize(BufferSize) } func NewPacket() *Buffer { - return &Buffer{ - data: Get(UDPBufferSize), - capacity: UDPBufferSize, - managed: true, - } + return NewSize(UDPBufferSize) } func NewSize(size int) *Buffer { + buffer := getBuffer() if size == 0 { - return &Buffer{} + *buffer = Buffer{ + managed: true, + } } else if size > 65535 { - return &Buffer{ + *buffer = Buffer{ data: make([]byte, size), capacity: size, + managed: true, + } + } else { + *buffer = Buffer{ + data: Get(size), + capacity: size, + managed: true, + dataManaged: true, } } - return &Buffer{ - data: Get(size), - capacity: size, - managed: true, - } + return buffer } func As(data []byte) *Buffer { - return &Buffer{ + buffer := getBuffer() + *buffer = Buffer{ data: data, end: len(data), capacity: len(data), + managed: true, } + return buffer } func With(data []byte) *Buffer { - return &Buffer{ + buffer := getBuffer() + *buffer = Buffer{ data: data, capacity: len(data), + managed: true, } + return buffer } func (b *Buffer) Byte(index int) byte { @@ -293,19 +299,28 @@ func (b *Buffer) DecRef() { } func (b *Buffer) Release() { - if b == nil || !b.managed { + if b == nil { + return + } + managed, dataManaged := b.managed, b.dataManaged + if !(managed || dataManaged) { return } if b.refs.Load() > 0 { return } - common.Must(Put(b.data)) + if dataManaged { + common.Must(Put(b.data)) + } *b = Buffer{} + if managed { + putBuffer(b) + } } func (b *Buffer) Leak() { if debug.Enabled { - if b == nil || !b.managed { + if b == nil || !(b.managed || b.dataManaged) { return } refs := b.refs.Load() diff --git a/common/buf/pool.go b/common/buf/pool.go index 37f1232c..e33c7576 100644 --- a/common/buf/pool.go +++ b/common/buf/pool.go @@ -1,5 +1,7 @@ package buf +import "sync" + func Get(size int) []byte { if size == 0 { return nil @@ -11,6 +13,20 @@ func Put(buf []byte) error { return DefaultAllocator.Put(buf) } +var bufferPool = sync.Pool{ + New: func() interface{} { + return new(Buffer) + }, +} + +func getBuffer() *Buffer { + return bufferPool.Get().(*Buffer) +} + +func putBuffer(b *Buffer) { + bufferPool.Put(b) +} + // Deprecated: use array instead. func Make(size int) []byte { return make([]byte, size)