rpc_util: use preconfigured size buffers for each pool

grpc · Mar 17, 2023 · 76caf74 · 76caf74
1 parent 57b9c67
commit 76caf74
Showing 1 changed file with 33 additions and 8 deletions.
diff --git a/shared_buffer_pool.go b/shared_buffer_pool.go
@@ -46,12 +46,12 @@ type SharedBufferPool interface {
 // later release.
 func NewSimpleSharedBufferPool() SharedBufferPool {
 	return &simpleSharedBufferPool{
-		pool0:   makeBytesPool(),
-		pool1:   makeBytesPool(),
-		pool2:   makeBytesPool(),
-		pool3:   makeBytesPool(),
-		pool4:   makeBytesPool(),
-		poolMax: makeBytesPool(),
+		pool0:   makeBytesPool(level0PoolMaxSize),
+		pool1:   makeBytesPool(level1PoolMaxSize),
+		pool2:   makeBytesPool(level2PoolMaxSize),
+		pool3:   makeBytesPool(level3PoolMaxSize),
+		pool4:   makeBytesPool(level4PoolMaxSize),
+		poolMax: makeFallbackBytesPool(),
 	}
 }
 
@@ -113,6 +113,31 @@ type bufferPool struct {
 
 func (p *bufferPool) Get(size int) []byte {
 	bs := p.Pool.Get().(*[]byte)
+
+	return (*bs)[:size]
+}
+
+func (p *bufferPool) Put(bs *[]byte) {
+	p.Pool.Put(bs)
+}
+
+func makeBytesPool(size int) bufferPool {
+	return bufferPool{
+		sync.Pool{
+			New: func() interface{} {
+				bs := make([]byte, size)
+				return &bs
+			},
+		},
+	}
+}
+
+type fallbackBufferPool struct {
+	sync.Pool
+}
+
+func (p *fallbackBufferPool) Get(size int) []byte {
+	bs := p.Pool.Get().(*[]byte)
 	if cap(*bs) < size {
 		*bs = make([]byte, size)
 		return *bs
@@ -121,11 +146,11 @@ func (p *bufferPool) Get(size int) []byte {
 	return (*bs)[:size]
 }
 
-func (p *bufferPool) Put(bs *[]byte) {
+func (p *fallbackBufferPool) Put(bs *[]byte) {
 	p.Pool.Put(bs)
 }
 
-func makeBytesPool() bufferPool {
+func makeFallbackBytesPool() bufferPool {
 	return bufferPool{
 		sync.Pool{
 			New: func() interface{} {