Merge pull request #2432 from analytically/iobufpool_putoptim

perf(iobufpool): optimize pool index calculations with O(1) bit manipulation

Author: jackc

internal/iobufpool/iobufpool.go

@@ -4,7 +4,10 @@
 // an allocation is purposely not documented. https://github.com/golang/go/issues/16323
 package iobufpool
 
-import "sync"
+import (
+	"math/bits"
+	"sync"
+)
 
 const minPoolExpOf2 = 8
 
@@ -37,15 +40,14 @@ func Get(size int) *[]byte {
 }
 
 func getPoolIdx(size int) int {
-	size--
-	size >>= minPoolExpOf2
-	i := 0
-	for size > 0 {
-		size >>= 1
-		i++
+	if size < 2 {
+		return 0
 	}
-
-	return i
+	idx := bits.Len(uint(size-1)) - minPoolExpOf2
+	if idx < 0 {
+		return 0
+	}
+	return idx
 }
 
 // Put returns buf to the pool.
@@ -59,12 +61,18 @@ func Put(buf *[]byte) {
 }
 
 func putPoolIdx(size int) int {
-	minPoolSize := 1 << minPoolExpOf2
-	for i := range pools {
-		if size == minPoolSize<<i {
-			return i
-		}
+	// Only exact power-of-2 sizes match pool buckets
+	if size&(size-1) != 0 {
+		return -1
+	}
+
+	// Calculate log2(size) using trailing zeros count
+	exp := bits.TrailingZeros(uint(size))
+	idx := exp - minPoolExpOf2
+
+	if idx < 0 || idx >= len(pools) {
+		return -1
 	}
 
-	return -1
+	return idx
 }