113 lines
2.8 KiB
Go
113 lines
2.8 KiB
Go
// Copyright (C) 2016 The Protocol Authors.
|
|
|
|
package protocol
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// Global pool to get buffers from. Requires Blocksizes to be initialised,
|
|
// therefore it is initialized in the same init() as BlockSizes
|
|
var BufferPool bufferPool
|
|
|
|
type bufferPool struct {
|
|
puts atomic.Int64
|
|
skips atomic.Int64
|
|
misses atomic.Int64
|
|
pools []sync.Pool
|
|
hits []atomic.Int64
|
|
}
|
|
|
|
func newBufferPool() bufferPool {
|
|
return bufferPool{
|
|
pools: make([]sync.Pool, len(BlockSizes)),
|
|
hits: make([]atomic.Int64, len(BlockSizes)),
|
|
}
|
|
}
|
|
|
|
func (p *bufferPool) Get(size int) []byte {
|
|
// Too big, isn't pooled
|
|
if size > MaxBlockSize {
|
|
p.skips.Add(1)
|
|
return make([]byte, size)
|
|
}
|
|
|
|
// Try the fitting and all bigger pools
|
|
bkt := getBucketForLen(size)
|
|
for j := bkt; j < len(BlockSizes); j++ {
|
|
if intf := p.pools[j].Get(); intf != nil {
|
|
p.hits[j].Add(1)
|
|
bs := *intf.(*[]byte)
|
|
return bs[:size]
|
|
}
|
|
}
|
|
|
|
p.misses.Add(1)
|
|
|
|
// All pools are empty, must allocate. For very small slices where we
|
|
// didn't have a block to reuse, just allocate a small slice instead of
|
|
// a large one. We won't be able to reuse it, but avoid some overhead.
|
|
if size < MinBlockSize/64 {
|
|
return make([]byte, size)
|
|
}
|
|
return make([]byte, BlockSizes[bkt])[:size]
|
|
}
|
|
|
|
// Put makes the given byte slice available again in the global pool.
|
|
// You must only Put() slices that were returned by Get() or Upgrade().
|
|
func (p *bufferPool) Put(bs []byte) {
|
|
// Don't buffer slices outside of our pool range
|
|
if cap(bs) > MaxBlockSize || cap(bs) < MinBlockSize {
|
|
p.skips.Add(1)
|
|
return
|
|
}
|
|
|
|
p.puts.Add(1)
|
|
bkt := putBucketForCap(cap(bs))
|
|
p.pools[bkt].Put(&bs)
|
|
}
|
|
|
|
// Upgrade grows the buffer to the requested size, while attempting to reuse
|
|
// it if possible.
|
|
func (p *bufferPool) Upgrade(bs []byte, size int) []byte {
|
|
if cap(bs) >= size {
|
|
// Reslicing is enough, lets go!
|
|
return bs[:size]
|
|
}
|
|
|
|
// It was too small. But it pack into the pool and try to get another
|
|
// buffer.
|
|
p.Put(bs)
|
|
return p.Get(size)
|
|
}
|
|
|
|
// getBucketForLen returns the bucket where we should get a slice of a
|
|
// certain length. Each bucket is guaranteed to hold slices that are
|
|
// precisely the block size for that bucket, so if the block size is larger
|
|
// than our size we are good.
|
|
func getBucketForLen(len int) int {
|
|
for i, blockSize := range BlockSizes {
|
|
if len <= blockSize {
|
|
return i
|
|
}
|
|
}
|
|
|
|
panic(fmt.Sprintf("bug: tried to get impossible block len %d", len))
|
|
}
|
|
|
|
// putBucketForCap returns the bucket where we should put a slice of a
|
|
// certain capacity. Each bucket is guaranteed to hold slices that are
|
|
// precisely the block size for that bucket, so we just find the matching
|
|
// one.
|
|
func putBucketForCap(cap int) int {
|
|
for i, blockSize := range BlockSizes {
|
|
if cap == blockSize {
|
|
return i
|
|
}
|
|
}
|
|
|
|
panic(fmt.Sprintf("bug: tried to put impossible block cap %d", cap))
|
|
}
|