pool.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/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 = 1 << 32

// 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
}

// 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
	}
	if length > MaxLength {
		return make([]byte, length)
	}
	idx := nextLogBase2(uint32(length))
	if buf := p.pools[idx].Get(); buf != nil {
		return buf.([]byte)[:uint32(length)]
	}
	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))
	p.pools[idx].Put(buf)
}

// 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
}