From ec6836162234e72ea740c9cbb1730042c22d0c36 Mon Sep 17 00:00:00 2001
From: Aaron Beitch <aaronb@arista.com>
Date: Mon, 23 Dec 2024 16:49:48 -0800
Subject: [PATCH] region: Use bufferPool for read buffers

Reuse buffers for reading responses from HBase, both for the
compressed response and decompressed response. This should reduce
memory pressure when reading lots of data.
---
 region/client.go     | 12 +++++-------
 region/compressor.go |  2 +-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/region/client.go b/region/client.go
index 0c334978..ed4e4ac9 100644
--- a/region/client.go
+++ b/region/client.go
@@ -103,12 +103,8 @@ const (
 var bufferPool sync.Pool
 
 func newBuffer(size int) []byte {
-	v := bufferPool.Get()
-	var b []byte
-	if v != nil {
-		b = v.([]byte)
-	}
-	return append(b[:0], make([]byte, size)...)
+	b, _ := bufferPool.Get().([]byte)
+	return append(b, make([]byte, size)...)
 }
 
 func freeBuffer(b []byte) {
@@ -477,7 +473,8 @@ func (c *client) receive(r io.Reader) (err error) {
 	}
 
 	size := binary.BigEndian.Uint32(sz[:])
-	b := make([]byte, size)
+	b := newBuffer(int(size))
+	defer freeBuffer(b)
 
 	_, err = io.ReadFull(r, b)
 	if err != nil {
@@ -546,6 +543,7 @@ func (c *client) receive(r io.Reader) (err error) {
 		b := b[size-cellsLen:]
 		if c.compressor != nil {
 			b, err = c.compressor.decompressCellblocks(b)
+			defer freeBuffer(b)
 			if err != nil {
 				err = RetryableError{fmt.Errorf("failed to decompress the response: %s", err)}
 				return
diff --git a/region/compressor.go b/region/compressor.go
index a71b9ffa..e22e0b44 100644
--- a/region/compressor.go
+++ b/region/compressor.go
@@ -92,7 +92,7 @@ func readUint32(b []byte) (uint32, []byte, error) {
 func (c *compressor) decompressCellblocks(b []byte) ([]byte, error) {
 	var (
 		err                  error
-		out                  []byte
+		out                  = newBuffer(0) // gets freed in receive
 		compressedChunk      []byte
 		compressedChunkLen   uint32
 		uncompressedBlockLen uint32