From 1c35ade12488f55e73365cb1d203c337d2b61fb2 Mon Sep 17 00:00:00 2001 From: Matheus Degiovani Date: Tue, 17 Sep 2024 18:33:44 -0300 Subject: [PATCH] bufferpool: Improve performance This improves the performance of bufferpool by removing the mutex called on every operation (via init()) and improving bucket indexing by removing the loop and using direct index calculation. Warning: this is a breaking API change. Zero valued Pool objects are no longer safe to be used and MUST be initialized by NewPool(). --- exp/bufferpool/pool.go | 180 ++++++++++++++++----------- exp/bufferpool/pool_internal_test.go | 96 ++++++++++++++ exp/bufferpool/pool_test.go | 42 +++---- 3 files changed, 223 insertions(+), 95 deletions(-) create mode 100644 exp/bufferpool/pool_internal_test.go diff --git a/exp/bufferpool/pool.go b/exp/bufferpool/pool.go index 170917d8..4b588854 100644 --- a/exp/bufferpool/pool.go +++ b/exp/bufferpool/pool.go @@ -2,7 +2,7 @@ package bufferpool import ( - "sync" + "math/bits" "github.com/colega/zeropool" ) @@ -13,35 +13,21 @@ const ( ) // A default global pool. -var Default Pool - -// Pool maintains a list of BucketCount buckets that contain buffers -// of exponentially-increasing capacity, 1 << 0 to 1 << BucketCount. -// -// The MinAlloc field specifies the minimum capacity of new buffers -// allocated by Pool, which improves reuse of small buffers. For the -// avoidance of doubt: calls to Get() with size < MinAlloc return a -// buffer of len(buf) = size and cap(buf) >= MinAlloc. MinAlloc MUST -// NOT exceed 1 << BucketCount, or method calls to Pool will panic. // -// The zero-value Pool is ready to use, defaulting to BucketCount=20 -// and MinAlloc=1024 (max size = ~1MiB). Most applications will not -// benefit from tuning these parameters. +// This pool defaults to bucketCount=20 and minAlloc=1024 (max size = ~1MiB). +var Default Pool = *NewPool(defaultMinSize, defaultBucketCount) + +// Pool maintains a list of buffers, exponentially increasing in size. Values +// MUST be initialized by NewPool(). // -// As a general rule, increasing MinAlloc reduces GC latency at the -// expense of increased memory usage. Increasing BucketCount can -// reduce GC latency in applications that frequently allocate large -// buffers. +// Buffer instances are safe for concurrent access. type Pool struct { - once sync.Once - MinAlloc, BucketCount int - buckets bucketSlice + minAlloc int + buckets bucketSlice } // Get a buffer of len(buf) == size and cap >= size. func (p *Pool) Get(size int) []byte { - p.init() - if buf := p.buckets.Get(size); buf != nil { return buf[:size] } @@ -52,76 +38,122 @@ func (p *Pool) Get(size int) []byte { // Put returns the buffer to the pool. The first len(buf) bytes // of the buffer are zeroed. func (p *Pool) Put(buf []byte) { - p.init() - for i := range buf { buf[i] = 0 } + // Do not store buffers less than the min alloc size (prevents storing + // buffers that do not conform to the min alloc policy of this pool). + if cap(buf) < p.minAlloc { + return + } + p.buckets.Put(buf[:cap(buf)]) } -func (p *Pool) init() { - p.once.Do(func() { - if p.MinAlloc <= 0 { - p.MinAlloc = defaultMinSize - } - - if p.BucketCount <= 0 { - p.BucketCount = defaultBucketCount - } - - if p.MinAlloc > (1 << p.BucketCount) { - panic("MinAlloc greater than largest bucket") - } - - // Get the index of the bucket responsible for MinAlloc. - var idx int - for idx = range p.buckets { - if 1<= p.MinAlloc { - break - } - } - - p.buckets = make(bucketSlice, p.BucketCount) - for i := range p.buckets { - if i < idx { - // Set the 'New' function for all "small" buckets to - // n.buckets[idx].Get, so as to allow reuse of buffers - // smaller than MinAlloc that are passed to Put, while - // still maximizing reuse of buffers allocated by Get. - // Note that we cannot simply use n.buckets[idx].New, - // as this would side-step pooling. - p.buckets[i] = zeropool.New(p.buckets[idx].Get) - } else { - p.buckets[i] = zeropool.New(newAllocFunc(i)) - } - } - }) +// NewPool creates a list of BucketCount buckets that contain buffers +// of exponentially-increasing capacity, 1 << 0 to 1 << BucketCount. +// +// The minAlloc field specifies the minimum capacity of new buffers +// allocated by Pool, which improves reuse of small buffers. For the +// avoidance of doubt: calls to Get() with size < minAlloc return a +// buffer of len(buf) = size and cap(buf) >= minAlloc. MinAlloc MUST +// NOT exceed 1 << BucketCount, or method calls to Pool will panic. +// +// Passing zero to the parameters will default bucketCount to 20 +// and minAlloc to 1024 (max size = ~1MiB). +// +// As a general rule, increasing MinAlloc reduces GC latency at the +// expense of increased memory usage. Increasing BucketCount can +// reduce GC latency in applications that frequently allocate large +// buffers. +func NewPool(minAlloc, bucketCount int) *Pool { + if minAlloc <= 0 { + minAlloc = defaultMinSize + } + + if bucketCount <= 0 { + bucketCount = defaultBucketCount + } + + if minAlloc > (1 << bucketCount) { + panic("MinAlloc greater than largest bucket") + } + + if !isPowerOf2(minAlloc) { + panic("MinAlloc not a power of two") + } + + return &Pool{ + minAlloc: minAlloc, + buckets: makeBucketSlice(minAlloc, bucketCount), + } +} + +type bucketSlice []*zeropool.Pool[[]byte] + +func isPowerOf2(i int) bool { + return i&(i-1) == 0 +} + +func bucketToGet(size int) int { + i := bits.Len(uint(size)) + if isPowerOf2(size) && size > 0 { + // When the size is a power of two, reduce by one (because + // bucket i is for sizes <= 1<< i). + i -= 1 + } + return i } -type bucketSlice []zeropool.Pool[[]byte] +func bucketToPut(size int) int { + i := bits.Len(uint(size)) + + // Always put on the bucket whose upper bound is size == 1<= size { - return bs[i].Get() - } + i := bucketToGet(size) + if i < len(bs) { + r := bs[i].Get() + return r } - return nil } func (bs bucketSlice) Put(buf []byte) { - for i := range bs { - if cap(buf) >= 1<= the bucket that stores the min + // allocation size. + minBucket := bucketToGet(minAlloc) + buckets := make(bucketSlice, bucketCount) + for i := minBucket; i < bucketCount; i++ { + bp := zeropool.New(newAllocFuncForBucket(i)) + buckets[i] = &bp + } + + // Buckets smaller than the min bucket size all get/put buffers in the + // minimum bucket size. + for i := 0; i < minBucket; i++ { + buckets[i] = buckets[minBucket] + } + + return buckets +} + +// newAllocFuncForBucket returns a function to allocate a byte slice of size +// 2^i. +func newAllocFuncForBucket(i int) func() []byte { return func() []byte { return make([]byte, 1< last bucket size is not allocated. + wantLen: 0, + wantCap: 0, + }} + + for _, tc := range tests { + t.Run(fmt.Sprintf("%d", tc.size), func(t *testing.T) { + bs := makeBucketSlice(minAlloc, bucketCount) + require.Len(t, bs, bucketCount) + buf := bs.Get(tc.size) + require.Len(t, buf, tc.wantLen) + require.Equal(t, tc.wantCap, cap(buf)) + }) + } +} diff --git a/exp/bufferpool/pool_test.go b/exp/bufferpool/pool_test.go index cac48c75..37d594ee 100644 --- a/exp/bufferpool/pool_test.go +++ b/exp/bufferpool/pool_test.go @@ -10,13 +10,8 @@ import ( func TestInvariants(t *testing.T) { t.Parallel() - pool := bufferpool.Pool{ - BucketCount: 1, - MinAlloc: 1024, - } - assert.Panics(t, func() { - pool.Get(32) + bufferpool.NewPool(1024, 1) }, "should panic when BucketCount cannot satisfy MinAlloc.") } @@ -24,38 +19,40 @@ func TestGet(t *testing.T) { t.Parallel() t.Helper() - pool := bufferpool.Pool{ - BucketCount: 8, - MinAlloc: 64, - } + minAlloc, bucketCount := 64, 8 + pool := bufferpool.NewPool(minAlloc, bucketCount) t.Run("SizeMinAlloc", func(t *testing.T) { - assert.Len(t, pool.Get(33), 33, "should return buffer with len=33") - assert.Equal(t, 128, cap(pool.Get(128)), "should return buffer with cap=128") + size := minAlloc + 1 + nextAlloc := minAlloc << 1 + assert.Len(t, pool.Get(size), size, "should return buffer with len=%d", size) + assert.Equal(t, nextAlloc, cap(pool.Get(size)), "should return buffer with cap=%d", nextAlloc) }) t.Run("Size>MaxSize", func(t *testing.T) { - assert.Len(t, pool.Get(512), 512, "should return buffer with len=512") - assert.GreaterOrEqual(t, 512, cap(pool.Get(512)), "should return buffer with cap>=512") + size := 1<<(bucketCount+1) + 1 + assert.Len(t, pool.Get(size), size, "should return buffer with len=%d", size) + assert.GreaterOrEqual(t, size, cap(pool.Get(size)), "should return buffer with cap>=%d", size) }) } func TestPut(t *testing.T) { t.Parallel() - var pool bufferpool.Pool + pool := bufferpool.NewPool(1024, 20) - buf := make([]byte, 1024) + buf := make([]byte, 1024, 1024) for i := range buf { buf[i] = byte(i) } @@ -63,12 +60,15 @@ func TestPut(t *testing.T) { pool.Put(buf) buf = pool.Get(8) + fullBuf := buf[:cap(buf)] + assert.Len(t, fullBuf, 1024, "buf should have len and cap 1024") assert.Equal(t, make([]byte, 8), buf, "should zero first 8 bytes") + assert.NotEqual(t, 0, fullBuf[9], "first byte after clearing should not be zero") } func BenchmarkPool(b *testing.B) { - var pool bufferpool.Pool + pool := bufferpool.NewPool(0, 0) const size = 32 // Make cache hot.