Extract iobufpool

internal/iobufpool/iobufpool.go

@@ -0,0 +1,46 @@
+// Package iobufpool implements a global segregated-fit pool of buffers for IO.
+package iobufpool
+
+import "sync"
+
+const minPoolExpOf2 = 8
+
+var pools [18]*sync.Pool
+
+func init() {
+	for i := range pools {
+		bufLen := 1 << (minPoolExpOf2 + i)
+		pools[i] = &sync.Pool{New: func() any { return make([]byte, bufLen) }}
+	}
+}
+
+// Get gets a []byte with len >= size and len <= size*2.
+func Get(size int) []byte {
+	i := poolIdx(size)
+	if i >= len(pools) {
+		return make([]byte, size)
+	}
+	return pools[i].Get().([]byte)
+}
+
+// Put returns buf to the pool.
+func Put(buf []byte) {
+	i := poolIdx(len(buf))
+	if i >= len(pools) {
+		return
+	}
+
+	pools[i].Put(buf)
+}
+
+func poolIdx(size int) int {
+	size--
+	size >>= minPoolExpOf2
+	i := 0
+	for size > 0 {
+		size >>= 1
+		i++
+	}
+
+	return i
+}

internal/iobufpool/iobufpool_internal_test.go

@@ -0,0 +1,36 @@
+package iobufpool
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPoolIdx(t *testing.T) {
+	tests := []struct {
+		size     int
+		expected int
+	}{
+		{size: 0, expected: 0},
+		{size: 1, expected: 0},
+		{size: 255, expected: 0},
+		{size: 256, expected: 0},
+		{size: 257, expected: 1},
+		{size: 511, expected: 1},
+		{size: 512, expected: 1},
+		{size: 513, expected: 2},
+		{size: 1023, expected: 2},
+		{size: 1024, expected: 2},
+		{size: 1025, expected: 3},
+		{size: 2047, expected: 3},
+		{size: 2048, expected: 3},
+		{size: 2049, expected: 4},
+		{size: 8388607, expected: 15},
+		{size: 8388608, expected: 15},
+		{size: 8388609, expected: 16},
+	}
+	for _, tt := range tests {
+		idx := poolIdx(tt.size)
+		assert.Equalf(t, tt.expected, idx, "size: %d", tt.size)
+	}
+}

internal/iobufpool/iobufpool_test.go

@@ -0,0 +1,35 @@
+package iobufpool_test
+
+import (
+	"testing"
+
+	"github.com/jackc/pgx/v5/internal/iobufpool"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGet(t *testing.T) {
+	tests := []struct {
+		requestedLen int
+		expectedLen  int
+	}{
+		{requestedLen: 0, expectedLen: 256},
+		{requestedLen: 128, expectedLen: 256},
+		{requestedLen: 255, expectedLen: 256},
+		{requestedLen: 256, expectedLen: 256},
+		{requestedLen: 257, expectedLen: 512},
+		{requestedLen: 511, expectedLen: 512},
+		{requestedLen: 512, expectedLen: 512},
+		{requestedLen: 513, expectedLen: 1024},
+		{requestedLen: 1023, expectedLen: 1024},
+		{requestedLen: 1024, expectedLen: 1024},
+		{requestedLen: 33554431, expectedLen: 33554432},
+		{requestedLen: 33554432, expectedLen: 33554432},
+
+		// Above 32 MiB skip the pool and allocate exactly the requested size.
+		{requestedLen: 33554433, expectedLen: 33554433},
+	}
+	for _, tt := range tests {
+		buf := iobufpool.Get(tt.requestedLen)
+		assert.Equalf(t, tt.expectedLen, len(buf), "requestedLen: %d", tt.requestedLen)
+	}
+}