mem: Allow overriding the default buffer pool.

benchmark/benchmain/main.go

@@ -177,7 +177,7 @@ func (p swappableBufferPool) Put(i *[]byte) {
 }
 
 func init() {
-	internal.SetDefaultBufferPoolForTesting.(func(mem.BufferPool))(swappableBufferPool{mem.DefaultBufferPool()})
+	internal.SetDefaultBufferPool.(func(mem.BufferPool))(swappableBufferPool{mem.DefaultBufferPool()})
 }
 
 var (

experimental/experimental.go

@@ -31,6 +31,18 @@ import (
 	"google.golang.org/grpc/mem"
 )
 
+// SetDefaultBufferPool sets the default buffer pool used by all grpc clients
+// and servers that do not have a buffer pool configured via WithBufferPool or
+// BufferPool. It also changes the buffer pool used by the proto codec, which
+// can't be changed otherwise. The provided buffer pool must be non-nil. The
+// default value is mem.DefaultBufferPool.
+//
+// NOTE: this function must only be called during initialization time (i.e. in
+// an init() function), and is not thread-safe. The last caller wins.
+func SetDefaultBufferPool(bufferPool mem.BufferPool) {
+	internal.SetDefaultBufferPool.(func(mem.BufferPool))(bufferPool)
+}
+
 // WithBufferPool returns a grpc.DialOption that configures the use of bufferPool
 // for parsing incoming messages on a grpc.ClientConn, and for temporary buffers
 // when marshaling outgoing messages. By default, mem.DefaultBufferPool is used,

internal/experimental.go

@@ -26,6 +26,9 @@ var (
 	// option to configure a shared buffer pool for a grpc.Server.
 	BufferPool any // func (grpc.SharedBufferPool) grpc.ServerOption
 
+	// SetDefaultBufferPool updates the default buffer pool.
+	SetDefaultBufferPool any // func(mem.BufferPool)
+
 	// AcceptCompressors is implemented by the grpc package and returns
 	// a call option that restricts the grpc-accept-encoding header for a call.
 	AcceptCompressors any // func(...string) grpc.CallOption

internal/internal.go

@@ -223,10 +223,6 @@ var (
 	// original state. Only called in testing functions.
 	SnapshotMetricRegistryForTesting func() func()
 
-	// SetDefaultBufferPoolForTesting updates the default buffer pool, for
-	// testing purposes.
-	SetDefaultBufferPoolForTesting any // func(mem.BufferPool)
-
 	// SetBufferPoolingThresholdForTesting updates the buffer pooling threshold, for
 	// testing purposes.
 	SetBufferPoolingThresholdForTesting any // func(int)

internal/leakcheck/leakcheck.go

@@ -51,7 +51,7 @@ var failTestsOnLeakedBuffers = false
 func init() {
 	defaultPool := mem.DefaultBufferPool()
 	globalPool.Store(&defaultPool)
-	(internal.SetDefaultBufferPoolForTesting.(func(mem.BufferPool)))(&globalPool)
+	internal.SetDefaultBufferPool.(func(mem.BufferPool))(&globalPool)
 }
 
 var globalPool swappableBufferPool

mem/buffer_pool.go

@@ -53,7 +53,7 @@ var defaultBufferPool BufferPool
 func init() {
 	defaultBufferPool = NewTieredBufferPool(defaultBufferPoolSizes...)
 
-	internal.SetDefaultBufferPoolForTesting = func(pool BufferPool) {
+	internal.SetDefaultBufferPool = func(pool BufferPool) {
 		defaultBufferPool = pool
 	}