status-go/vendor/github.com/libp2p/go-buffer-pool/pool.go

118 lines
2.9 KiB
Go

// Package pool provides a sync.Pool equivalent that buckets incoming
// requests to one of 32 sub-pools, one for each power of 2, 0-32.
//
// import (pool "github.com/libp2p/go-buffer-pool")
// var p pool.BufferPool
//
// small := make([]byte, 1024)
// large := make([]byte, 4194304)
// p.Put(small)
// p.Put(large)
//
// small2 := p.Get(1024)
// large2 := p.Get(4194304)
// fmt.Println("small2 len:", len(small2))
// fmt.Println("large2 len:", len(large2))
//
// // Output:
// // small2 len: 1024
// // large2 len: 4194304
//
package pool
import (
"math"
"math/bits"
"sync"
)
// GlobalPool is a static Pool for reusing byteslices of various sizes.
var GlobalPool = new(BufferPool)
// MaxLength is the maximum length of an element that can be added to the Pool.
const MaxLength = math.MaxInt32
// BufferPool is a pool to handle cases of reusing elements of varying sizes. It
// maintains 32 internal pools, for each power of 2 in 0-32.
//
// You should generally just call the package level Get and Put methods or use
// the GlobalPool BufferPool instead of constructing your own.
//
// You MUST NOT copy Pool after using.
type BufferPool struct {
pools [32]sync.Pool // a list of singlePools
ptrs sync.Pool
}
type bufp struct {
buf []byte
}
// Get retrieves a buffer of the appropriate length from the buffer pool or
// allocates a new one. Get may choose to ignore the pool and treat it as empty.
// Callers should not assume any relation between values passed to Put and the
// values returned by Get.
//
// If no suitable buffer exists in the pool, Get creates one.
func (p *BufferPool) Get(length int) []byte {
if length == 0 {
return nil
}
// Calling this function with a negative length is invalid.
// make will panic if length is negative, so we don't have to.
if length > MaxLength || length < 0 {
return make([]byte, length)
}
idx := nextLogBase2(uint32(length))
if ptr := p.pools[idx].Get(); ptr != nil {
bp := ptr.(*bufp)
buf := bp.buf[:uint32(length)]
bp.buf = nil
p.ptrs.Put(ptr)
return buf
}
return make([]byte, 1<<idx)[:uint32(length)]
}
// Put adds x to the pool.
func (p *BufferPool) Put(buf []byte) {
capacity := cap(buf)
if capacity == 0 || capacity > MaxLength {
return // drop it
}
idx := prevLogBase2(uint32(capacity))
var bp *bufp
if ptr := p.ptrs.Get(); ptr != nil {
bp = ptr.(*bufp)
} else {
bp = new(bufp)
}
bp.buf = buf
p.pools[idx].Put(bp)
}
// Get retrieves a buffer of the appropriate length from the global buffer pool
// (or allocates a new one).
func Get(length int) []byte {
return GlobalPool.Get(length)
}
// Put returns a buffer to the global buffer pool.
func Put(slice []byte) {
GlobalPool.Put(slice)
}
// Log of base two, round up (for v > 0).
func nextLogBase2(v uint32) uint32 {
return uint32(bits.Len32(v - 1))
}
// Log of base two, round down (for v > 0)
func prevLogBase2(num uint32) uint32 {
next := nextLogBase2(num)
if num == (1 << uint32(next)) {
return next
}
return next - 1
}