Handle word size, add benchmark, return construction error

Author: arjan-bal

mem/buffer_pool.go

@@ -19,6 +19,7 @@
 package mem
 
 import (
+	"fmt"
 	"math/bits"
 	"slices"
 	"sort"
@@ -51,10 +52,17 @@ var defaultBufferPoolSizeExponents = []uint8{
 	20, // 1MB
 }
 
-var defaultBufferPool BufferPool
+var (
+	defaultBufferPool BufferPool
+	uintSize          = bits.UintSize // use a variable for mocking during tests.
+)
 
 func init() {
-	defaultBufferPool = NewBinaryTieredBufferPool(defaultBufferPoolSizeExponents...)
+	var err error
+	defaultBufferPool, err = NewBinaryTieredBufferPool(defaultBufferPoolSizeExponents...)
+	if err != nil {
+		panic(fmt.Sprintf("Failed to create default buffer pool: %v", err))
+	}
 
 	internal.SetDefaultBufferPoolForTesting = func(pool BufferPool) {
 		defaultBufferPool = pool
@@ -127,37 +135,33 @@ type binaryTieredBufferPool struct {
 	maxPoolCap                       int // Optimization: Cache max capacity
 }
 
-// NewBinaryTieredBufferPool returns a BufferPool implementation that uses
-// multiple underlying pools of the given pool sizes. The pool sizes must be
-// powers of 2. This enables O(1) lookup when getting or putting a buffer.
+// NewBinaryTieredBufferPool returns a BufferPool backed by multiple sub-pools.
+// This structure enables O(1) lookup time for Get and Put operations.
 //
-// Note that the argument passed to this functions are the powers of 2 of the
-// capacity of the buffers in the pool, not the capacities of the buffers
-// themselves. For example, if you wanted a pool that had buffers with a capacity
-// of 16kb, you would pass 14 as the argument to this function.
-func NewBinaryTieredBufferPool(powerOfTwoExponents ...uint8) BufferPool {
+// The arguments provided are the exponents for the buffer capacities (powers
+// of 2), not the raw byte sizes. For example, to create a pool of 16KB buffers
+// (2^14 bytes), pass 14 as the argument.
+func NewBinaryTieredBufferPool(powerOfTwoExponents ...uint8) (BufferPool, error) {
 	slices.Sort(powerOfTwoExponents)
 
-	// Determine the maximum exponent we need to support.
-	// bits.Len64(math.MaxUint64) is 63.
-	const maxExponent = 63
+	// Determine the maximum exponent we need to support. This depends on the
+	// word size (32-bit vs 64-bit).
+	maxExponent := uintSize - 1
 	indexOfNextLargestBit := slices.Repeat([]int{-1}, maxExponent+1)
 	indexOfPreviousLargestBit := slices.Repeat([]int{-1}, maxExponent+1)
 
-	maxCap := 0
+	maxTier := 0
 	pools := make([]*sizedBufferPool, 0, len(powerOfTwoExponents))
 
 	for i, exp := range powerOfTwoExponents {
-		// Allocating slices of size > 2^maxExponent isn't possible on 64-bit
-		// machines.
-		if exp > maxExponent {
-			continue
-		}
-		capSize := 1 << exp
-		pools = append(pools, newSizedBufferPool(capSize))
-		if capSize > maxCap {
-			maxCap = capSize
+		// Allocating slices of size > 2^maxExponent isn't possible on
+		// maxExponent-bit machines.
+		if int(exp) > maxExponent {
+			return nil, fmt.Errorf("allocating slice of size 2^%d is not possible", exp)
 		}
+		tierSize := 1 << exp
+		pools = append(pools, newSizedBufferPool(tierSize))
+		maxTier = max(maxTier, tierSize)
 
 		// Map the exact power of 2 to this pool index.
 		indexOfNextLargestBit[exp] = i
@@ -184,8 +188,8 @@ func NewBinaryTieredBufferPool(powerOfTwoExponents ...uint8) BufferPool {
 		exponentToNextLargestPoolMap:     indexOfNextLargestBit,
 		exponentToPreviousLargestPoolMap: indexOfPreviousLargestBit,
 		sizedPools:                       pools,
-		maxPoolCap:                       maxCap,
-	}
+		maxPoolCap:                       maxTier,
+	}, nil
 }
 
 func (b *binaryTieredBufferPool) Get(size int) *[]byte {

mem/buffer_pool_internal_test.go

@@ -0,0 +1,114 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package mem
+
+import (
+	"testing"
+)
+
+func TestNewBinaryTieredBufferPool_WordSize(t *testing.T) {
+	origUintSize := uintSize
+	defer func() { uintSize = origUintSize }()
+
+	tests := []struct {
+		name      string
+		wordSize  int
+		exponents []uint8
+		wantErr   bool
+	}{
+		{
+			name:      "32-bit valid exponent",
+			wordSize:  32,
+			exponents: []uint8{31},
+			wantErr:   false,
+		},
+		{
+			name:      "32-bit invalid exponent",
+			wordSize:  32,
+			exponents: []uint8{32},
+			wantErr:   true,
+		},
+		{
+			name:      "64-bit valid exponent",
+			wordSize:  64,
+			exponents: []uint8{63},
+			wantErr:   false,
+		},
+		{
+			name:      "64-bit invalid exponent",
+			wordSize:  64,
+			exponents: []uint8{64},
+			wantErr:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			uintSize = tt.wordSize
+			pool, err := NewBinaryTieredBufferPool(tt.exponents...)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("NewBinaryTieredBufferPool() error = %t, wantErr %t", err, tt.wantErr)
+				return
+			}
+			if err == nil {
+				bp := pool.(*binaryTieredBufferPool)
+				if len(bp.exponentToNextLargestPoolMap) != tt.wordSize {
+					t.Errorf("exponentToNextLargestPoolMap length = %d, want %d", len(bp.exponentToNextLargestPoolMap), tt.wordSize)
+				}
+				if len(bp.exponentToPreviousLargestPoolMap) != tt.wordSize {
+					t.Errorf("exponentToPreviousLargestPoolMap length = %d, want %d", len(bp.exponentToPreviousLargestPoolMap), tt.wordSize)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkTieredPool benchmarks the performance of the tiered buffer pool
+// implementations, specifically focusing on the overhead of selecting the
+// correct bucket for a given size.
+func BenchmarkTieredPool(b *testing.B) {
+	defaultBufferPoolSizes := make([]int, len(defaultBufferPoolSizeExponents))
+	for i, exp := range defaultBufferPoolSizeExponents {
+		defaultBufferPoolSizes[i] = 1 << exp
+	}
+	b.Run("pool=Tiered", func(b *testing.B) {
+		p := NewTieredBufferPool(defaultBufferPoolSizes...).(*tieredBufferPool)
+		for b.Loop() {
+			for size := range 1 << 19 {
+				// One for get, one for put.
+				_ = p.getPool(size)
+				_ = p.getPool(size)
+			}
+		}
+	})
+
+	b.Run("pool=BinaryTiered", func(b *testing.B) {
+		pool, err := NewBinaryTieredBufferPool(defaultBufferPoolSizeExponents...)
+		p := pool.(*binaryTieredBufferPool)
+		if err != nil {
+			b.Fatalf("Faield to create buffer pool: %v", err)
+		}
+		for b.Loop() {
+			for size := range 1 << 19 {
+				_ = p.poolForGet(size)
+				_ = p.poolForPut(size)
+			}
+		}
+	})
+}

mem/buffer_pool_test.go

@@ -131,7 +131,10 @@ func TestBinaryBufferPool(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(fmt.Sprintf("requestSize=%d", tc.requestSize), func(t *testing.T) {
-			pool := mem.NewBinaryTieredBufferPool(poolSizes...)
+			pool, err := mem.NewBinaryTieredBufferPool(poolSizes...)
+			if err != nil {
+				t.Fatalf("Failed to create buffer pool: %v", err)
+			}
 			buf := pool.Get(tc.requestSize)
 			if cap(*buf) != tc.wantCapacity {
 				t.Errorf("Get(%d) returned buffer with capacity: %d, want %d", tc.requestSize, cap(*buf), tc.wantCapacity)