Introduce mem package

This package is required for #7356 but is introduced early here to allow additional features to be built on top of it. It provides a mechanism to reuse and recycle buffers to reduce GC pressure.
grpc · easwars · Aug 1, 2024 · Jul 22, 2024 · Jul 23, 2024 · Jul 23, 2024
commit be4d4ab8d90af3d547096fd09e7b910d0326cc3e
diff --git a/mem/buffer_pool.go b/mem/buffer_pool.go
@@ -0,0 +1,184 @@
+/*
+ *
+ * Copyright 2024 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package mem
+
+import (
+	"sort"
+	"sync"
+	"sync/atomic"
+)
+
+// BufferPool is a pool of buffers that can be shared, resulting in
+// decreased memory allocation.
+type BufferPool interface {
+	// Get returns a buffer with specified length from the pool.
+	Get(length int) []byte
+
+	// Put returns a buffer to the pool.
+	Put([]byte)
+}
+
+var defaultBufferPoolSizes = []int{
+	256,
+	4 << 10,  // 4KB (go page size)
+	16 << 10, // 16KB (max HTTP/2 frame size used by gRPC)
+	32 << 10, // 32KB (default buffer size for io.Copy)
+	// TODO: Report the buffer sizes requested with Get to tune this properly.
+	1 << 20, // 1MB
+}
+
+var defaultBufferPool = func() *atomic.Pointer[BufferPool] {
+	pool := NewBufferPool(defaultBufferPoolSizes...)
+	ptr := new(atomic.Pointer[BufferPool])
+	ptr.Store(&pool)
+	return ptr
+}()
+
+// DefaultBufferPool returns the current default buffer pool. It is a BufferPool
+// created with NewBufferPool that uses a set of default sizes optimized for
+// expected workflows.
+func DefaultBufferPool() BufferPool {
+	return *defaultBufferPool.Load()
+}
+
+// SetDefaultBufferPoolForTesting updates the default buffer pool, for testing
+// purposes.
+func SetDefaultBufferPoolForTesting(pool BufferPool) {
+	defaultBufferPool.Store(&pool)
+}
+
+// NewBufferPool returns a BufferPool implementation that uses multiple
+// underlying pools of the given pool sizes. When a buffer is requested from the
+// returned pool, it will have a
+func NewBufferPool(poolSizes ...int) BufferPool {
+	sort.Ints(poolSizes)
+	pools := make([]*sizedBufferPool, len(poolSizes))
+	for i, s := range poolSizes {
+		pools[i] = newBufferPool(s)
+	}
+	return &tieredBufferPool{
+		sizedPools: pools,
+	}
+}
+
+// tieredBufferPool implements the BufferPool interface with multiple tiers of
+// buffer pools for different sizes of buffers.
+type tieredBufferPool struct {
+	sizedPools   []*sizedBufferPool
+	fallbackPool simpleBufferPool
+}
+
+func (p *tieredBufferPool) Get(size int) []byte {
+	return p.getPool(size).Get(size)
+}
+
+func (p *tieredBufferPool) Put(buf []byte) {
+	p.getPool(cap(buf)).Put(buf)
+}
+
+func (p *tieredBufferPool) getPool(size int) BufferPool {
+	poolIdx := sort.Search(len(p.sizedPools), func(i int) bool {
+		return p.sizedPools[i].defaultSize >= size
+	})
+
+	if poolIdx == len(p.sizedPools) {
+		return &p.fallbackPool
+	}
+
+	return p.sizedPools[poolIdx]
+}
+
+// sizedBufferPool is a BufferPool implementation that is optimized for specific
+// buffer sizes. For example, HTTP/2 frames within grpc are always 16kb and a
+// sizedBufferPool can be configured to only return buffers with a capacity of
+// 16kb. Note that however it does not support returning larger buffers and in
+// fact panics if such a buffer is requested.
+type sizedBufferPool struct {
+	pool        sync.Pool
+	defaultSize int
+}
+
+func (p *sizedBufferPool) Get(size int) []byte {
+	bs := *p.pool.Get().(*[]byte)
+	return bs[:size]
+}
+
+func (p *sizedBufferPool) Put(buf []byte) {
+	if cap(buf) < p.defaultSize {
+		// Ignore buffers that are too small to fit in the pool. Otherwise, when Get is
+		// called it will panic as it tries to index outside the bounds of the buffer.
+		return
+	}
+	buf = buf[:cap(buf)]
+	clear(buf)
+	p.pool.Put(&buf)
+}
+
+func newBufferPool(size int) *sizedBufferPool {
+	return &sizedBufferPool{
+		pool: sync.Pool{
+			New: func() any {
+				buf := make([]byte, size)
+				return &buf
+			},
+		},
+		defaultSize: size,
+	}
+}
+
+var _ BufferPool = (*simpleBufferPool)(nil)
+
+// simpleBufferPool is an implementation of the BufferPool interface that
+// attempts to pool buffers with a sync.Pool. When Get is invoked, it tries to
+// acquire a buffer from the pool but if that buffer is too small, it returns it
+// to the pool and creates a new one.
+type simpleBufferPool struct {
+	pool sync.Pool
+}
+
+func (p *simpleBufferPool) Get(size int) []byte {
+	bs, ok := p.pool.Get().(*[]byte)
+	if ok && cap(*bs) >= size {
+		return (*bs)[:size]
+	}
+
+	if ok {
+		p.pool.Put(bs)
+	}
+
+	return make([]byte, size)
+}
+
+func (p *simpleBufferPool) Put(buf []byte) {
+	buf = buf[:cap(buf)]
+	clear(buf)
+	p.pool.Put(&buf)
+}
+
+var _ BufferPool = NopBufferPool{}
+
+// NopBufferPool is a buffer pool just makes new buffer without pooling.
+type NopBufferPool struct{}
+
+func (NopBufferPool) Get(length int) []byte {
+	return make([]byte, length)
+}
+
+func (NopBufferPool) Put([]byte) {
+}
diff --git a/mem/buffer_pool_test.go b/mem/buffer_pool_test.go
@@ -0,0 +1,61 @@
+/*
+ *
+ * Copyright 2023 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package mem
+
+import "testing"
+
+func TestSharedBufferPool(t *testing.T) {
+	pools := []BufferPool{
+		NopBufferPool{},
+		NewBufferPool(defaultBufferPoolSizes...),
+	}
+
+	testSizes := append(defaultBufferPoolSizes, 1<<20+1)
+
+	for _, p := range pools {
+		for _, l := range testSizes {
+			bs := p.Get(l)
+			if len(bs) != l {
+				t.Fatalf("Expected buffer of length %d, got %d", l, len(bs))
+			}
+
+			p.Put(bs)
+		}
+	}
+}
+
+func TestTieredBufferPool(t *testing.T) {
+	pool := &tieredBufferPool{
+		sizedPools: []*sizedBufferPool{
+			newBufferPool(10),
+			newBufferPool(20),
+		},
+	}
+	buf := pool.Get(1)
+	if cap(buf) != 10 {
+		t.Fatalf("Unexpected buffer capacity: %d", cap(buf))
+	}
+
+	// Insert a short buffer into the pool, which is currently empty.
+	pool.Put(make([]byte, 1))
+	// Then immediately request a buffer that would be pulled from the pool where the
+	// short buffer would have been returned. If the short buffer is pulled from the
+	// pool, it could cause a panic.
+	pool.Get(10)
+}