From aed354b08c468847563ad12941cf56a27dd2d801 Mon Sep 17 00:00:00 2001
From: PapaCharlie <paul.chesnais@gmail.com>
Date: Fri, 4 Oct 2024 11:40:06 -0700
Subject: [PATCH] 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.
---
 mem/buffers.go      |  5 ++++-
 mem/buffers_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/mem/buffers.go b/mem/buffers.go
index 565ec0a13949..5144b7c86536 100644
--- a/mem/buffers.go
+++ b/mem/buffers.go
@@ -92,7 +92,10 @@ 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..41f0c9bfacd8 100644
--- a/mem/buffers_test.go
+++ b/mem/buffers_test.go
@@ -98,6 +98,32 @@ 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 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 {