85 lines
2.1 KiB
Go
Raw Normal View History

2021-12-04 16:42:11 +00:00
// Copyright 2020 lesismal. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
package bytes
import (
"sync"
)
// maxAppendSize represents the max size to append to a slice.
const maxAppendSize = 1024 * 1024 * 4
// Pool is the default instance of []byte pool.
// User can customize a Pool implementation and reset this instance if needed.
var Pool interface {
Get() []byte
GetN(size int) []byte
Put(b []byte)
} = NewPool(64)
// bufferPool is a default implementatiion of []byte Pool.
type bufferPool struct {
sync.Pool
MinSize int
}
// NewPool creates and returns a bufferPool instance.
// All slice created by this instance has an initial cap of minSize.
func NewPool(minSize int) *bufferPool {
if minSize <= 0 {
minSize = 64
}
bp := &bufferPool{
MinSize: minSize,
}
bp.Pool.New = func() interface{} {
buf := make([]byte, bp.MinSize)
return &buf
}
return bp
}
// Get gets a slice from the pool and returns it with length 0.
// User can append the slice and should Put it back to the pool after being used over.
func (bp *bufferPool) Get() []byte {
pbuf := bp.Pool.Get().(*[]byte)
return (*pbuf)[0:0]
}
// GetN returns a slice with length size.
// To reuse slices as possible,
// if the cap of the slice got from the pool is not enough,
// It will append the slice,
// or put the slice back to the pool and create a new slice with cap of size.
//
// User can use the slice both by the size or append it,
// and should Put it back to the pool after being used over.
func (bp *bufferPool) GetN(size int) []byte {
pbuf := bp.Pool.Get().(*[]byte)
need := size - cap(*pbuf)
if need > 0 {
if need <= maxAppendSize {
*pbuf = (*pbuf)[:cap(*pbuf)]
*pbuf = append(*pbuf, make([]byte, need)...)
} else {
bp.Pool.Put(pbuf)
newBuf := make([]byte, size)
pbuf = &newBuf
}
}
return (*pbuf)[:size]
}
// Put puts a slice back to the pool.
// If the slice's cap is smaller than MinSize,
// it will not be put back to the pool but dropped.
func (bp *bufferPool) Put(b []byte) {
if cap(b) < bp.MinSize {
return
}
bp.Pool.Put(&b)
}