From f8e5d8f704180341e3f1034bc503a324d05d91dc Mon Sep 17 00:00:00 2001
From: Paul Chesnais <paul.chesnais@gmail.com>
Date: Thu, 24 Oct 2024 13:14:00 -0400
Subject: [PATCH] mem: use slice capacity instead of length, to determine
 whether to pool buffers or directly allocate them (#7702)

* Address #7631 by correctly pooling large-capacity buffers

As the issue states, `mem.NewBuffer` would not pool buffers with a length below
the pooling threshold but whose capacity is actually larger than the pooling
threshold. This can lead to buffers being leaked.

---------

Co-authored-by: Purnesh Dixit <purneshdixit@google.com>
Co-authored-by: Easwar Swaminathan <easwars@google.com>
---
 mem/buffers.go      |  6 +++++-
 mem/buffers_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/mem/buffers.go b/mem/buffers.go
index 565ec0a13949..ecbf0b9a73ea 100644
--- a/mem/buffers.go
+++ b/mem/buffers.go
@@ -92,7 +92,11 @@ func newBuffer() *buffer {
 //
 // Note that the backing array of the given data is not copied.
 func NewBuffer(data *[]byte, pool BufferPool) Buffer {
-	if pool == nil || IsBelowBufferPoolingThreshold(len(*data)) {
+	// Use the buffer's capacity instead of the length, otherwise buffers may
+	// not be reused under certain conditions. For example, if a large buffer
+	// is acquired from the pool, but fewer bytes than the buffering threshold
+	// are written to it, the buffer will not be returned to the pool.
+	if pool == nil || IsBelowBufferPoolingThreshold(cap(*data)) {
 		return (SliceBuffer)(*data)
 	}
 	b := newBuffer()
diff --git a/mem/buffers_test.go b/mem/buffers_test.go
index c17995745209..2b0627da159d 100644
--- a/mem/buffers_test.go
+++ b/mem/buffers_test.go
@@ -98,6 +98,33 @@ func (s) TestBuffer_NewBufferRefAndFree(t *testing.T) {
 	}
 }
 
+func (s) TestBuffer_NewBufferHandlesShortBuffers(t *testing.T) {
+	const threshold = 100
+
+	// Update the pooling threshold, since that's what's being tested.
+	internal.SetBufferPoolingThresholdForTesting.(func(int))(threshold)
+	t.Cleanup(func() {
+		internal.SetBufferPoolingThresholdForTesting.(func(int))(0)
+	})
+
+	// Make a pool with a buffer whose capacity is larger than the pooling
+	// threshold, but whose length is less than the threshold.
+	b := make([]byte, threshold/2, threshold*2)
+	pool := &singleBufferPool{
+		t:    t,
+		data: &b,
+	}
+
+	// Get a Buffer, then free it. If NewBuffer decided that the Buffer
+	// shouldn't get pooled, Free will be a noop and singleBufferPool will not
+	// have been updated.
+	mem.NewBuffer(&b, pool).Free()
+
+	if pool.data != nil {
+		t.Fatalf("Buffer not returned to pool")
+	}
+}
+
 func (s) TestBuffer_FreeAfterFree(t *testing.T) {
 	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
 	if buf.Len() != 4 {