From d2c5fcbdfb533b9053f8930a050f60ade93727db Mon Sep 17 00:00:00 2001
From: Matheus Degiovani <opensource@matheusd.com>
Date: Tue, 5 Mar 2024 17:28:32 -0300
Subject: [PATCH] bufferpool: Switch to zeropool

This commit switches the bufferpool to use the zeropool implementation
for sync pools.

The stdlib sync.Pool implementation has an issue where it causes an
additional heap allocation per Put() call when used with byte slices.

github.com/colega/zeropool package has been specifically designed to
work around this issue, which reduces GC pressure and improves
performance.

This also fixes the bufferpool's pkg benchmark to use a new pool per
test, to avoid other tests influencing the behavior of the benchmark and
sets it to report the allocations.
---
 exp/bufferpool/pool.go      | 14 ++++++++------
 exp/bufferpool/pool_test.go | 13 +++++++++++--
 go.mod                      |  1 +
 go.sum                      |  2 ++
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/exp/bufferpool/pool.go b/exp/bufferpool/pool.go
index 22112687..170917d8 100644
--- a/exp/bufferpool/pool.go
+++ b/exp/bufferpool/pool.go
@@ -3,6 +3,8 @@ package bufferpool
 
 import (
 	"sync"
+
+	"github.com/colega/zeropool"
 )
 
 const (
@@ -90,20 +92,20 @@ func (p *Pool) init() {
 				// still maximizing reuse of buffers allocated by Get.
 				// Note that we cannot simply use n.buckets[idx].New,
 				// as this would side-step pooling.
-				p.buckets[i].New = p.buckets[idx].Get
+				p.buckets[i] = zeropool.New(p.buckets[idx].Get)
 			} else {
-				p.buckets[i].New = newAllocFunc(i)
+				p.buckets[i] = zeropool.New(newAllocFunc(i))
 			}
 		}
 	})
 }
 
-type bucketSlice []sync.Pool
+type bucketSlice []zeropool.Pool[[]byte]
 
 func (bs bucketSlice) Get(size int) []byte {
 	for i := range bs {
 		if 1<<i >= size {
-			return bs[i].Get().([]byte)
+			return bs[i].Get()
 		}
 	}
 
@@ -119,8 +121,8 @@ func (bs bucketSlice) Put(buf []byte) {
 	}
 }
 
-func newAllocFunc(i int) func() any {
-	return func() any {
+func newAllocFunc(i int) func() []byte {
+	return func() []byte {
 		return make([]byte, 1<<i)
 	}
 }
diff --git a/exp/bufferpool/pool_test.go b/exp/bufferpool/pool_test.go
index e497db0f..cac48c75 100644
--- a/exp/bufferpool/pool_test.go
+++ b/exp/bufferpool/pool_test.go
@@ -68,8 +68,17 @@ func TestPut(t *testing.T) {
 }
 
 func BenchmarkPool(b *testing.B) {
+	var pool bufferpool.Pool
+	const size = 32
+
+	// Make cache hot.
+	buf := pool.Get(size)
+	pool.Put(buf)
+
+	b.ResetTimer()
+	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
-		buf := bufferpool.Default.Get(32)
-		bufferpool.Default.Put(buf)
+		buf := pool.Get(size)
+		pool.Put(buf)
 	}
 }
diff --git a/go.mod b/go.mod
index 94c24610..fbd8f94d 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module capnproto.org/go/capnp/v3
 go 1.19
 
 require (
+	github.com/colega/zeropool v0.0.0-20230505084239-6fb4a4f75381
 	github.com/kylelemons/godebug v1.1.0
 	github.com/stretchr/testify v1.8.2
 	github.com/tinylib/msgp v1.1.5
diff --git a/go.sum b/go.sum
index 37eaf814..1076665e 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/colega/zeropool v0.0.0-20230505084239-6fb4a4f75381 h1:d5EKgQfRQvO97jnISfR89AiCCCJMwMFoSxUiU0OGCRU=
+github.com/colega/zeropool v0.0.0-20230505084239-6fb4a4f75381/go.mod h1:OU76gHeRo8xrzGJU3F3I1CqX1ekM8dfJw0+wPeMwnp0=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=