fixed putting wrong size bufs

internal/iobufpool/iobufpool.go

@@ -14,26 +14,16 @@ func init() {
 	}
 }
 
-// Get gets a []byte with len >= size and len <= size*2.
+// Get gets a []byte of len size with cap <= size*2.
 func Get(size int) []byte {
-	i := poolIdx(size)
+	i := getPoolIdx(size)
 	if i >= len(pools) {
 		return make([]byte, size)
 	}
-	return pools[i].Get().([]byte)
+	return pools[i].Get().([]byte)[:size]
 }
 
-// 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 {
+func getPoolIdx(size int) int {
 	size--
 	size >>= minPoolExpOf2
 	i := 0
@@ -44,3 +34,26 @@ func poolIdx(size int) int {
 
 	return i
 }
+
+// Put returns buf to the pool.
+func Put(buf []byte) {
+	buf = buf[:cap(buf)]
+
+	i := putPoolIdx(len(buf))
+	if i < 0 {
+		return
+	}
+
+	pools[i].Put(buf)
+}
+
+func putPoolIdx(size int) int {
+	minPoolSize := 1 << minPoolExpOf2
+	for i := range pools {
+		if size == minPoolSize<<i {
+			return i
+		}
+	}
+
+	return -1
+}

internal/iobufpool/iobufpool_internal_test.go

@@ -30,7 +30,7 @@ func TestPoolIdx(t *testing.T) {
 		{size: 8388609, expected: 16},
 	}
 	for _, tt := range tests {
-		idx := poolIdx(tt.size)
+		idx := getPoolIdx(tt.size)
 		assert.Equalf(t, tt.expected, idx, "size: %d", tt.size)
 	}
 }

internal/iobufpool/iobufpool_test.go

@@ -5,31 +5,74 @@ import (
 
 	"github.com/jackc/pgx/v5/internal/iobufpool"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-func TestGet(t *testing.T) {
+func TestGetCap(t *testing.T) {
 	tests := []struct {
 		requestedLen int
-		expectedLen  int
+		expectedCap  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},
+		{requestedLen: 0, expectedCap: 256},
+		{requestedLen: 128, expectedCap: 256},
+		{requestedLen: 255, expectedCap: 256},
+		{requestedLen: 256, expectedCap: 256},
+		{requestedLen: 257, expectedCap: 512},
+		{requestedLen: 511, expectedCap: 512},
+		{requestedLen: 512, expectedCap: 512},
+		{requestedLen: 513, expectedCap: 1024},
+		{requestedLen: 1023, expectedCap: 1024},
+		{requestedLen: 1024, expectedCap: 1024},
+		{requestedLen: 33554431, expectedCap: 33554432},
+		{requestedLen: 33554432, expectedCap: 33554432},
 
 		// Above 32 MiB skip the pool and allocate exactly the requested size.
-		{requestedLen: 33554433, expectedLen: 33554433},
+		{requestedLen: 33554433, expectedCap: 33554433},
 	}
 	for _, tt := range tests {
 		buf := iobufpool.Get(tt.requestedLen)
-		assert.Equalf(t, tt.expectedLen, len(buf), "requestedLen: %d", tt.requestedLen)
+		assert.Equalf(t, tt.requestedLen, len(buf), "bad len for requestedLen: %d", len(buf), tt.requestedLen)
+		assert.Equalf(t, tt.expectedCap, cap(buf), "bad cap for requestedLen: %d", tt.requestedLen)
 	}
 }
+
+func TestPutHandlesWrongSizedBuffers(t *testing.T) {
+	for putBufSize := range []int{0, 1, 128, 250, 256, 257, 1023, 1024, 1025, 1 << 28} {
+		putBuf := make([]byte, putBufSize)
+		iobufpool.Put(putBuf)
+
+		tests := []struct {
+			requestedLen int
+			expectedCap  int
+		}{
+			{requestedLen: 0, expectedCap: 256},
+			{requestedLen: 128, expectedCap: 256},
+			{requestedLen: 255, expectedCap: 256},
+			{requestedLen: 256, expectedCap: 256},
+			{requestedLen: 257, expectedCap: 512},
+			{requestedLen: 511, expectedCap: 512},
+			{requestedLen: 512, expectedCap: 512},
+			{requestedLen: 513, expectedCap: 1024},
+			{requestedLen: 1023, expectedCap: 1024},
+			{requestedLen: 1024, expectedCap: 1024},
+			{requestedLen: 33554431, expectedCap: 33554432},
+			{requestedLen: 33554432, expectedCap: 33554432},
+
+			// Above 32 MiB skip the pool and allocate exactly the requested size.
+			{requestedLen: 33554433, expectedCap: 33554433},
+		}
+		for _, tt := range tests {
+			getBuf := iobufpool.Get(tt.requestedLen)
+			assert.Equalf(t, tt.requestedLen, len(getBuf), "len(putBuf): %d, requestedLen: %d", len(putBuf), tt.requestedLen)
+			assert.Equalf(t, tt.expectedCap, cap(getBuf), "cap(putBuf): %d, requestedLen: %d", cap(putBuf), tt.requestedLen)
+		}
+	}
+}
+
+func TestPutGetBufferReuse(t *testing.T) {
+	buf := iobufpool.Get(4)
+	buf[0] = 1
+	iobufpool.Put(buf)
+	buf = iobufpool.Get(4)
+	require.Equal(t, byte(1), buf[0])
+}

internal/nbbconn/nbbconn.go

@@ -83,7 +83,6 @@ func (c *Conn) Write(b []byte) (n int, err error) {
 	}
 
 	buf := iobufpool.Get(len(b))
-	buf = buf[:len(b)]
 	copy(buf, b)
 	c.writeQueue.pushBack(buf)
 	return len(b), nil