diff --git a/benchmark/benchmain/main.go b/benchmark/benchmain/main.go index 443799f57871..297ba4d2722c 100644 --- a/benchmark/benchmain/main.go +++ b/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 ( diff --git a/experimental/experimental.go b/experimental/experimental.go index 3ba948bab316..a7c008a3022f 100644 --- a/experimental/experimental.go +++ b/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, diff --git a/internal/experimental.go b/internal/experimental.go index c90cc51bdd2b..8a999917d978 100644 --- a/internal/experimental.go +++ b/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 diff --git a/internal/internal.go b/internal/internal.go index 144b2803090e..47d69638540a 100644 --- a/internal/internal.go +++ b/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) diff --git a/internal/leakcheck/leakcheck.go b/internal/leakcheck/leakcheck.go index 3c5d905751be..fe4efd5d3dd2 100644 --- a/internal/leakcheck/leakcheck.go +++ b/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 diff --git a/mem/buffer_pool.go b/mem/buffer_pool.go index e37afdd1981d..2ea763a49a2f 100644 --- a/mem/buffer_pool.go +++ b/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 }