snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

mem/buffer_pool.go

@@ -0,0 +1,92 @@
+/*
+ *
+ * Copyright 2024 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 (
+	"fmt"
+
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/mem"
+)
+
+// BufferPool is a pool of buffers that can be shared and reused, resulting in
+// decreased memory allocation.
+type BufferPool interface {
+	// Get returns a buffer with specified length from the pool.
+	Get(length int) *[]byte
+
+	// Put returns a buffer to the pool.
+	//
+	// The provided pointer must hold a prefix of the buffer obtained via
+	// BufferPool.Get to ensure the buffer's entire capacity can be re-used.
+	Put(*[]byte)
+}
+
+var (
+	defaultBufferPoolSizeExponents = []uint8{
+		8,
+		12, // Go page size, 4KB
+		14, // 16KB (max HTTP/2 frame size used by gRPC)
+		15, // 32KB (default buffer size for io.Copy)
+		20, // 1MB
+	}
+	defaultBufferPool BufferPool
+)
+
+func init() {
+	var err error
+	defaultBufferPool, err = NewBinaryTieredBufferPool(defaultBufferPoolSizeExponents...)
+	if err != nil {
+		panic(fmt.Sprintf("Failed to create default buffer pool: %v", err))
+	}
+
+	internal.SetDefaultBufferPool = func(pool BufferPool) {
+		defaultBufferPool = pool
+	}
+}
+
+// DefaultBufferPool returns the current default buffer pool. It is a BufferPool
+// created with NewBufferPool that uses a set of default sizes optimized for
+// expected workflows.
+func DefaultBufferPool() BufferPool {
+	return defaultBufferPool
+}
+
+// NewTieredBufferPool returns a BufferPool implementation that uses multiple
+// underlying pools of the given pool sizes.
+func NewTieredBufferPool(poolSizes ...int) BufferPool {
+	return mem.NewTieredBufferPool(poolSizes...)
+}
+
+// NewBinaryTieredBufferPool returns a BufferPool backed by multiple sub-pools.
+// This structure enables O(1) lookup time for Get and Put operations.
+//
+// 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) {
+	return mem.NewBinaryTieredBufferPool(powerOfTwoExponents...)
+}
+
+// NopBufferPool is a buffer pool that returns new buffers without pooling.
+type NopBufferPool struct {
+	mem.NopBufferPool
+}
+
+var _ BufferPool = NopBufferPool{}

mem/buffer_pool_test.go

@@ -0,0 +1,148 @@
+/*
+ *
+ * Copyright 2023 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_test
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+	"unsafe"
+
+	"google.golang.org/grpc/mem"
+)
+
+func (s) TestBufferPool(t *testing.T) {
+	var poolSizes = []int{4, 8, 16, 32}
+	pools := []mem.BufferPool{
+		mem.NopBufferPool{},
+		mem.NewTieredBufferPool(poolSizes...),
+	}
+
+	testSizes := append([]int{1}, poolSizes...)
+	testSizes = append(testSizes, 64)
+
+	for _, p := range pools {
+		for _, l := range testSizes {
+			bs := p.Get(l)
+			if len(*bs) != l {
+				t.Fatalf("Get(%d) returned buffer of length %d, want %d", l, len(*bs), l)
+			}
+
+			p.Put(bs)
+		}
+	}
+}
+
+func (s) TestBufferPoolClears(t *testing.T) {
+	const poolSize = 4
+	pool := mem.NewTieredBufferPool(poolSize)
+	tests := []struct {
+		name       string
+		bufferSize int
+	}{
+		{
+			name:       "sized_buffer_pool",
+			bufferSize: poolSize,
+		},
+		{
+			name:       "simple_buffer_pool",
+			bufferSize: poolSize + 1,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			for {
+				buf1 := pool.Get(tc.bufferSize)
+				copy(*buf1, "1234")
+				pool.Put(buf1)
+
+				buf2 := pool.Get(tc.bufferSize)
+				if unsafe.SliceData(*buf1) != unsafe.SliceData(*buf2) {
+					pool.Put(buf2)
+					// This test is only relevant if a buffer is reused, otherwise try again. This
+					// can happen if a GC pause happens between putting the buffer back in the pool
+					// and getting a new one.
+					continue
+				}
+
+				if !bytes.Equal(*buf1, make([]byte, tc.bufferSize)) {
+					t.Fatalf("buffer not cleared")
+				}
+				break
+			}
+		})
+	}
+}
+
+func (s) TestBufferPoolIgnoresShortBuffers(t *testing.T) {
+	pool := mem.NewTieredBufferPool(10, 20)
+	buf := pool.Get(1)
+	if cap(*buf) != 10 {
+		t.Fatalf("Get(1) returned buffer with capacity: %d, want 10", cap(*buf))
+	}
+
+	// Insert a short buffer into the pool, which is currently empty.
+	short := make([]byte, 1)
+	pool.Put(&short)
+	// Then immediately request a buffer that would be pulled from the pool where the
+	// short buffer would have been returned. If the short buffer is pulled from the
+	// pool, it could cause a panic.
+	pool.Get(10)
+}
+
+func TestBinaryBufferPool(t *testing.T) {
+	poolSizes := []uint8{0, 2, 3, 4, 2, 3, 4} // duplicates will be ignored.
+
+	testCases := []struct {
+		requestSize  int
+		wantCapacity int
+	}{
+		{requestSize: 0, wantCapacity: 0},
+		{requestSize: 1, wantCapacity: 1},
+		{requestSize: 2, wantCapacity: 4},
+		{requestSize: 3, wantCapacity: 4},
+		{requestSize: 4, wantCapacity: 4},
+		{requestSize: 5, wantCapacity: 8},
+		{requestSize: 6, wantCapacity: 8},
+		{requestSize: 7, wantCapacity: 8},
+		{requestSize: 8, wantCapacity: 8},
+		{requestSize: 9, wantCapacity: 16},
+		{requestSize: 15, wantCapacity: 16},
+		{requestSize: 16, wantCapacity: 16},
+		{requestSize: 17, wantCapacity: 4096}, // fallback pool returns sizes in multiples of 4096.
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("requestSize=%d", tc.requestSize), func(t *testing.T) {
+			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)
+			}
+			if len(*buf) != tc.requestSize {
+				t.Errorf("Get(%d) returned buffer with length: %d, want %d", tc.requestSize, len(*buf), tc.requestSize)
+			}
+			pool.Put(buf)
+		})
+	}
+}

mem/buffer_slice.go

@@ -0,0 +1,345 @@
+/*
+ *
+ * Copyright 2024 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 (
+	"fmt"
+	"io"
+)
+
+const (
+	// 32 KiB is what io.Copy uses.
+	readAllBufSize = 32 * 1024
+)
+
+// BufferSlice offers a means to represent data that spans one or more Buffer
+// instances. A BufferSlice is meant to be immutable after creation, and methods
+// like Ref create and return copies of the slice. This is why all methods have
+// value receivers rather than pointer receivers.
+//
+// Note that any of the methods that read the underlying buffers such as Ref,
+// Len or CopyTo etc., will panic if any underlying buffers have already been
+// freed. It is recommended to not directly interact with any of the underlying
+// buffers directly, rather such interactions should be mediated through the
+// various methods on this type.
+//
+// By convention, any APIs that return (mem.BufferSlice, error) should reduce
+// the burden on the caller by never returning a mem.BufferSlice that needs to
+// be freed if the error is non-nil, unless explicitly stated.
+type BufferSlice []Buffer
+
+// Len returns the sum of the length of all the Buffers in this slice.
+//
+// # Warning
+//
+// Invoking the built-in len on a BufferSlice will return the number of buffers
+// in the slice, and *not* the value returned by this function.
+func (s BufferSlice) Len() int {
+	var length int
+	for _, b := range s {
+		length += b.Len()
+	}
+	return length
+}
+
+// Ref invokes Ref on each buffer in the slice.
+func (s BufferSlice) Ref() {
+	for _, b := range s {
+		b.Ref()
+	}
+}
+
+// Free invokes Buffer.Free() on each Buffer in the slice.
+func (s BufferSlice) Free() {
+	for _, b := range s {
+		b.Free()
+	}
+}
+
+// CopyTo copies each of the underlying Buffer's data into the given buffer,
+// returning the number of bytes copied. Has the same semantics as the copy
+// builtin in that it will copy as many bytes as it can, stopping when either dst
+// is full or s runs out of data, returning the minimum of s.Len() and len(dst).
+func (s BufferSlice) CopyTo(dst []byte) int {
+	off := 0
+	for _, b := range s {
+		off += copy(dst[off:], b.ReadOnlyData())
+	}
+	return off
+}
+
+// Materialize concatenates all the underlying Buffer's data into a single
+// contiguous buffer using CopyTo.
+func (s BufferSlice) Materialize() []byte {
+	l := s.Len()
+	if l == 0 {
+		return nil
+	}
+	out := make([]byte, l)
+	s.CopyTo(out)
+	return out
+}
+
+// MaterializeToBuffer functions like Materialize except that it writes the data
+// to a single Buffer pulled from the given BufferPool.
+//
+// As a special case, if the input BufferSlice only actually has one Buffer, this
+// function simply increases the refcount before returning said Buffer. Freeing this
+// buffer won't release it until the BufferSlice is itself released.
+func (s BufferSlice) MaterializeToBuffer(pool BufferPool) Buffer {
+	if len(s) == 1 {
+		s[0].Ref()
+		return s[0]
+	}
+	sLen := s.Len()
+	if sLen == 0 {
+		return emptyBuffer{}
+	}
+	buf := pool.Get(sLen)
+	s.CopyTo(*buf)
+	return NewBuffer(buf, pool)
+}
+
+// Reader returns a new Reader for the input slice after taking references to
+// each underlying buffer.
+func (s BufferSlice) Reader() *Reader {
+	s.Ref()
+	return &Reader{
+		data: s,
+		len:  s.Len(),
+	}
+}
+
+// Reader exposes a BufferSlice's data as an io.Reader, allowing it to interface
+// with other systems.
+//
+// Buffers will be freed as they are read.
+//
+// A Reader can be constructed from a BufferSlice; alternatively the zero value
+// of a Reader may be used after calling Reset on it.
+type Reader struct {
+	data BufferSlice
+	len  int
+	// The index into data[0].ReadOnlyData().
+	bufferIdx int
+}
+
+// Remaining returns the number of unread bytes remaining in the slice.
+func (r *Reader) Remaining() int {
+	return r.len
+}
+
+// Reset frees the currently held buffer slice and starts reading from the
+// provided slice. This allows reusing the reader object.
+func (r *Reader) Reset(s BufferSlice) {
+	r.data.Free()
+	s.Ref()
+	r.data = s
+	r.len = s.Len()
+	r.bufferIdx = 0
+}
+
+// Close frees the underlying BufferSlice and never returns an error. Subsequent
+// calls to Read will return (0, io.EOF).
+func (r *Reader) Close() error {
+	r.data.Free()
+	r.data = nil
+	r.len = 0
+	return nil
+}
+
+func (r *Reader) freeFirstBufferIfEmpty() bool {
+	if len(r.data) == 0 || r.bufferIdx != r.data[0].Len() {
+		return false
+	}
+
+	r.data[0].Free()
+	r.data = r.data[1:]
+	r.bufferIdx = 0
+	return true
+}
+
+func (r *Reader) Read(buf []byte) (n int, _ error) {
+	if r.len == 0 {
+		return 0, io.EOF
+	}
+
+	for len(buf) != 0 && r.len != 0 {
+		// Copy as much as possible from the first Buffer in the slice into the
+		// given byte slice.
+		data := r.data[0].ReadOnlyData()
+		copied := copy(buf, data[r.bufferIdx:])
+		r.len -= copied       // Reduce len by the number of bytes copied.
+		r.bufferIdx += copied // Increment the buffer index.
+		n += copied           // Increment the total number of bytes read.
+		buf = buf[copied:]    // Shrink the given byte slice.
+
+		// If we have copied all the data from the first Buffer, free it and advance to
+		// the next in the slice.
+		r.freeFirstBufferIfEmpty()
+	}
+
+	return n, nil
+}
+
+// ReadByte reads a single byte.
+func (r *Reader) ReadByte() (byte, error) {
+	if r.len == 0 {
+		return 0, io.EOF
+	}
+
+	// There may be any number of empty buffers in the slice, clear them all until a
+	// non-empty buffer is reached. This is guaranteed to exit since r.len is not 0.
+	for r.freeFirstBufferIfEmpty() {
+	}
+
+	b := r.data[0].ReadOnlyData()[r.bufferIdx]
+	r.len--
+	r.bufferIdx++
+	// Free the first buffer in the slice if the last byte was read
+	r.freeFirstBufferIfEmpty()
+	return b, nil
+}
+
+var _ io.Writer = (*writer)(nil)
+
+type writer struct {
+	buffers *BufferSlice
+	pool    BufferPool
+}
+
+func (w *writer) Write(p []byte) (n int, err error) {
+	b := Copy(p, w.pool)
+	*w.buffers = append(*w.buffers, b)
+	return b.Len(), nil
+}
+
+// NewWriter wraps the given BufferSlice and BufferPool to implement the
+// io.Writer interface. Every call to Write copies the contents of the given
+// buffer into a new Buffer pulled from the given pool and the Buffer is
+// added to the given BufferSlice.
+func NewWriter(buffers *BufferSlice, pool BufferPool) io.Writer {
+	return &writer{buffers: buffers, pool: pool}
+}
+
+// ReadAll reads from r until an error or EOF and returns the data it read.
+// A successful call returns err == nil, not err == EOF. Because ReadAll is
+// defined to read from src until EOF, it does not treat an EOF from Read
+// as an error to be reported.
+//
+// Important: A failed call returns a non-nil error and may also return
+// partially read buffers. It is the responsibility of the caller to free the
+// BufferSlice returned, or its memory will not be reused.
+func ReadAll(r io.Reader, pool BufferPool) (BufferSlice, error) {
+	var result BufferSlice
+	if wt, ok := r.(io.WriterTo); ok {
+		// This is more optimal since wt knows the size of chunks it wants to
+		// write and, hence, we can allocate buffers of an optimal size to fit
+		// them. E.g. might be a single big chunk, and we wouldn't chop it
+		// into pieces.
+		w := NewWriter(&result, pool)
+		_, err := wt.WriteTo(w)
+		return result, err
+	}
+nextBuffer:
+	for {
+		buf := pool.Get(readAllBufSize)
+		// We asked for 32KiB but may have been given a bigger buffer.
+		// Use all of it if that's the case.
+		*buf = (*buf)[:cap(*buf)]
+		usedCap := 0
+		for {
+			n, err := r.Read((*buf)[usedCap:])
+			usedCap += n
+			if err != nil {
+				if usedCap == 0 {
+					// Nothing in this buf, put it back
+					pool.Put(buf)
+				} else {
+					*buf = (*buf)[:usedCap]
+					result = append(result, NewBuffer(buf, pool))
+				}
+				if err == io.EOF {
+					err = nil
+				}
+				return result, err
+			}
+			if len(*buf) == usedCap {
+				result = append(result, NewBuffer(buf, pool))
+				continue nextBuffer
+			}
+		}
+	}
+}
+
+// Discard skips the next n bytes, returning the number of bytes discarded.
+//
+// It frees buffers as they are fully consumed.
+//
+// If Discard skips fewer than n bytes, it also returns an error.
+func (r *Reader) Discard(n int) (discarded int, err error) {
+	total := n
+	for n > 0 && r.len > 0 {
+		curData := r.data[0].ReadOnlyData()
+		curSize := min(n, len(curData)-r.bufferIdx)
+		n -= curSize
+		r.len -= curSize
+		r.bufferIdx += curSize
+		if r.bufferIdx >= len(curData) {
+			r.data[0].Free()
+			r.data = r.data[1:]
+			r.bufferIdx = 0
+		}
+	}
+	discarded = total - n
+	if n > 0 {
+		return discarded, fmt.Errorf("insufficient bytes in reader")
+	}
+	return discarded, nil
+}
+
+// Peek returns the next n bytes without advancing the reader.
+//
+// Peek appends results to the provided res slice and returns the updated slice.
+// This pattern allows re-using the storage of res if it has sufficient
+// capacity.
+//
+// The returned subslices are views into the underlying buffers and are only
+// valid until the reader is advanced past the corresponding buffer.
+//
+// If Peek returns fewer than n bytes, it also returns an error.
+func (r *Reader) Peek(n int, res [][]byte) ([][]byte, error) {
+	for i := 0; n > 0 && i < len(r.data); i++ {
+		curData := r.data[i].ReadOnlyData()
+		start := 0
+		if i == 0 {
+			start = r.bufferIdx
+		}
+		curSize := min(n, len(curData)-start)
+		if curSize == 0 {
+			continue
+		}
+		res = append(res, curData[start:start+curSize])
+		n -= curSize
+	}
+	if n > 0 {
+		return nil, fmt.Errorf("insufficient bytes in reader")
+	}
+	return res, nil
+}

mem/buffer_slice_test.go

@@ -0,0 +1,676 @@
+/*
+ *
+ * Copyright 2024 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_test
+
+import (
+	"bytes"
+	"crypto/rand"
+	"errors"
+	"fmt"
+	"io"
+	"testing"
+
+	"google.golang.org/grpc/mem"
+)
+
+const (
+	minReadSize = 1
+	// Should match the constant in buffer_slice.go (another package)
+	readAllBufSize = 32 * 1024 // 32 KiB
+)
+
+func newBuffer(data []byte, pool mem.BufferPool) mem.Buffer {
+	return mem.NewBuffer(&data, pool)
+}
+
+func (s) TestBufferSlice_Len(t *testing.T) {
+	tests := []struct {
+		name string
+		in   mem.BufferSlice
+		want int
+	}{
+		{
+			name: "empty",
+			in:   nil,
+			want: 0,
+		},
+		{
+			name: "single",
+			in:   mem.BufferSlice{newBuffer([]byte("abcd"), nil)},
+			want: 4,
+		},
+		{
+			name: "multiple",
+			in: mem.BufferSlice{
+				newBuffer([]byte("abcd"), nil),
+				newBuffer([]byte("abcd"), nil),
+				newBuffer([]byte("abcd"), nil),
+			},
+			want: 12,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.in.Len(); got != tt.want {
+				t.Errorf("BufferSlice.Len() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func (s) TestBufferSlice_Ref(t *testing.T) {
+	// Create a new buffer slice and a reference to it.
+	bs := mem.BufferSlice{
+		newBuffer([]byte("abcd"), nil),
+		newBuffer([]byte("abcd"), nil),
+	}
+	bs.Ref()
+
+	// Free the original buffer slice and verify that the reference can still
+	// read data from it.
+	bs.Free()
+	got := bs.Materialize()
+	want := []byte("abcdabcd")
+	if !bytes.Equal(got, want) {
+		t.Errorf("BufferSlice.Materialize() = %s, want %s", string(got), string(want))
+	}
+}
+
+func (s) TestBufferSlice_MaterializeToBuffer(t *testing.T) {
+	tests := []struct {
+		name     string
+		in       mem.BufferSlice
+		pool     mem.BufferPool
+		wantData []byte
+	}{
+		{
+			name:     "single",
+			in:       mem.BufferSlice{newBuffer([]byte("abcd"), nil)},
+			pool:     nil, // MaterializeToBuffer should not use the pool in this case.
+			wantData: []byte("abcd"),
+		},
+		{
+			name: "multiple",
+			in: mem.BufferSlice{
+				newBuffer([]byte("abcd"), nil),
+				newBuffer([]byte("abcd"), nil),
+				newBuffer([]byte("abcd"), nil),
+			},
+			pool:     mem.DefaultBufferPool(),
+			wantData: []byte("abcdabcdabcd"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer tt.in.Free()
+			got := tt.in.MaterializeToBuffer(tt.pool)
+			defer got.Free()
+			if !bytes.Equal(got.ReadOnlyData(), tt.wantData) {
+				t.Errorf("BufferSlice.MaterializeToBuffer() = %s, want %s", string(got.ReadOnlyData()), string(tt.wantData))
+			}
+		})
+	}
+}
+
+func (s) TestBufferSlice_Reader(t *testing.T) {
+	bs := mem.BufferSlice{
+		newBuffer([]byte("abcd"), nil),
+		newBuffer([]byte("abcd"), nil),
+		newBuffer([]byte("abcd"), nil),
+	}
+	wantData := []byte("abcdabcdabcd")
+
+	reader := bs.Reader()
+	var gotData []byte
+	// Read into a buffer of size 1 until EOF, and verify that the data matches.
+	for {
+		buf := make([]byte, 1)
+		n, err := reader.Read(buf)
+		if n > 0 {
+			gotData = append(gotData, buf[:n]...)
+		}
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Fatalf("BufferSlice.Reader() failed unexpectedly: %v", err)
+		}
+	}
+	if !bytes.Equal(gotData, wantData) {
+		t.Errorf("BufferSlice.Reader() returned data %v, want %v", string(gotData), string(wantData))
+	}
+
+	// Reader should have released its references to the underlying buffers, but
+	// bs still holds its reference and it should be able to read data from it.
+	gotData = bs.Materialize()
+	if !bytes.Equal(gotData, wantData) {
+		t.Errorf("BufferSlice.Materialize() = %s, want %s", string(gotData), string(wantData))
+	}
+}
+
+// TestBufferSlice_ReadAll_Reads exercises ReadAll by allowing it to read
+// various combinations of data, empty data, EOF.
+func (s) TestBufferSlice_ReadAll_Reads(t *testing.T) {
+	testcases := []struct {
+		name     string
+		reads    []readStep
+		wantErr  string
+		wantBufs int
+	}{
+		{
+			name: "EOF",
+			reads: []readStep{
+				{
+					err: io.EOF,
+				},
+			},
+		},
+		{
+			name: "data,EOF",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data+EOF",
+			reads: []readStep{
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "0,data+EOF",
+			reads: []readStep{
+				{},
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "0,data,EOF",
+			reads: []readStep{
+				{},
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+EOF",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "error",
+			reads: []readStep{
+				{
+					err: errors.New("boom"),
+				},
+			},
+			wantErr: "boom",
+		},
+		{
+			name: "data+error",
+			reads: []readStep{
+				{
+					n:   minReadSize,
+					err: errors.New("boom"),
+				},
+			},
+			wantErr:  "boom",
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+error",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   minReadSize,
+					err: errors.New("boom"),
+				},
+			},
+			wantErr:  "boom",
+			wantBufs: 1,
+		},
+		{
+			name: "data,data+EOF - whole buf",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n:   readAllBufSize - minReadSize,
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data,EOF - whole buf",
+			reads: []readStep{
+				{
+					n: minReadSize,
+				},
+				{
+					n: readAllBufSize - minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 1,
+		},
+		{
+			name: "data,data,EOF - 2 bufs",
+			reads: []readStep{
+				{
+					n: readAllBufSize,
+				},
+				{
+					n: minReadSize,
+				},
+				{
+					n: readAllBufSize - minReadSize,
+				},
+				{
+					n: minReadSize,
+				},
+				{
+					err: io.EOF,
+				},
+			},
+			wantBufs: 3,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			pool := &testPool{
+				allocated: make(map[*[]byte]struct{}),
+			}
+			r := &stepReader{
+				reads: tc.reads,
+			}
+			data, err := mem.ReadAll(r, pool)
+			if tc.wantErr != "" {
+				if err == nil || err.Error() != tc.wantErr {
+					t.Fatalf("ReadAll() returned err %v, wanted %q", err, tc.wantErr)
+				}
+			} else {
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			gotData := data.Materialize()
+			if !bytes.Equal(r.read, gotData) {
+				t.Fatalf("ReadAll() returned data %q, wanted %q", gotData, r.read)
+			}
+			if len(data) != tc.wantBufs {
+				t.Fatalf("ReadAll() returned %d bufs, wanted %d bufs", len(data), tc.wantBufs)
+			}
+			// all but last should be full buffers
+			for i := 0; i < len(data)-1; i++ {
+				if data[i].Len() != readAllBufSize {
+					t.Fatalf("ReadAll() returned data length %d, wanted %d", data[i].Len(), readAllBufSize)
+				}
+			}
+			data.Free()
+			if len(pool.allocated) > 0 {
+				t.Fatalf("got %d allocated buffers, wanted none", len(pool.allocated))
+			}
+		})
+	}
+}
+
+func (s) TestBufferSlice_ReadAll_WriteTo(t *testing.T) {
+	testcases := []struct {
+		name string
+		size int
+	}{
+		{
+			name: "small",
+			size: minReadSize,
+		},
+		{
+			name: "exact size",
+			size: readAllBufSize,
+		},
+		{
+			name: "big",
+			size: readAllBufSize * 3,
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			pool := &testPool{
+				allocated: make(map[*[]byte]struct{}),
+			}
+			buf := make([]byte, tc.size)
+			_, err := rand.Read(buf)
+			if err != nil {
+				t.Fatal(err)
+			}
+			r := bytes.NewBuffer(buf)
+			data, err := mem.ReadAll(r, pool)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			gotData := data.Materialize()
+			if !bytes.Equal(buf, gotData) {
+				t.Fatalf("ReadAll() = %q, wanted %q", gotData, buf)
+			}
+			data.Free()
+			if len(pool.allocated) > 0 {
+				t.Fatalf("wanted no allocated buffers, got %d", len(pool.allocated))
+			}
+		})
+	}
+}
+
+func ExampleNewWriter() {
+	var bs mem.BufferSlice
+	pool := mem.DefaultBufferPool()
+	writer := mem.NewWriter(&bs, pool)
+
+	for _, data := range [][]byte{
+		[]byte("abcd"),
+		[]byte("abcd"),
+		[]byte("abcd"),
+	} {
+		n, err := writer.Write(data)
+		fmt.Printf("Wrote %d bytes, err: %v\n", n, err)
+	}
+	fmt.Println(string(bs.Materialize()))
+	// Output:
+	// Wrote 4 bytes, err: <nil>
+	// Wrote 4 bytes, err: <nil>
+	// Wrote 4 bytes, err: <nil>
+	// abcdabcdabcd
+}
+
+var (
+	_ io.Reader      = (*stepReader)(nil)
+	_ mem.BufferPool = (*testPool)(nil)
+)
+
+// readStep describes what a single stepReader.Read should do - how much data
+// to return and what error to return.
+type readStep struct {
+	n   int
+	err error
+}
+
+// stepReader implements io.Reader that reads specified amount of data and/or
+// returns the specified error in specified steps.
+// The read data is accumulated in the read field.
+type stepReader struct {
+	reads []readStep
+	read  []byte
+}
+
+func (s *stepReader) Read(buf []byte) (int, error) {
+	if len(s.reads) == 0 {
+		panic("unexpected Read() call")
+	}
+	read := s.reads[0]
+	s.reads = s.reads[1:]
+	_, err := rand.Read(buf[:read.n])
+	if err != nil {
+		panic(err)
+	}
+	s.read = append(s.read, buf[:read.n]...)
+	return read.n, read.err
+}
+
+// testPool is an implementation of BufferPool that allows to ensure that:
+// - there are matching Put calls for all Get calls.
+// - there are no unexpected Put calls.
+type testPool struct {
+	allocated map[*[]byte]struct{}
+}
+
+func (t *testPool) Get(length int) *[]byte {
+	buf := make([]byte, length)
+	t.allocated[&buf] = struct{}{}
+	return &buf
+}
+
+func (t *testPool) Put(buf *[]byte) {
+	if _, ok := t.allocated[buf]; !ok {
+		panic("unexpected put")
+	}
+	delete(t.allocated, buf)
+}
+
+func (s) TestBufferSlice_Iteration(t *testing.T) {
+	tests := []struct {
+		name       string
+		buffers    [][]byte
+		operations func(t *testing.T, c *mem.Reader)
+	}{
+		{
+			name: "empty",
+			operations: func(t *testing.T, r *mem.Reader) {
+				if r.Remaining() != 0 {
+					t.Fatalf("Remaining() = %v, want 0", r.Remaining())
+				}
+				_, err := r.Peek(1, nil)
+				if err == nil {
+					t.Fatalf("Peek(1) returned error <nil>, want non-nil")
+				}
+				discarded, err := r.Discard(1)
+				if got, want := discarded, 0; got != want {
+					t.Fatalf("Discard(1) = %d, want %d", got, want)
+				}
+				if err == nil {
+					t.Fatalf("Discard(1) returned error <nil>, want non-nil")
+				}
+				if r.Remaining() != 0 {
+					t.Fatalf("Remaining() after Discard = %v, want 0", r.Remaining())
+				}
+			},
+		},
+		{
+			name:    "single_buffer",
+			buffers: [][]byte{[]byte("0123456789")},
+			operations: func(t *testing.T, r *mem.Reader) {
+				if r.Remaining() != 10 {
+					t.Fatalf("Remaining() = %v, want 10", r.Remaining())
+				}
+
+				res := make([][]byte, 0, 10)
+				res, err := r.Peek(5, res)
+				if err != nil {
+					t.Fatalf("Peek(5) return error %v, want <nil>", err)
+				}
+				if len(res) != 1 || !bytes.Equal(res[0], []byte("01234")) {
+					t.Fatalf("Peek(5) = %v, want [[01234]]", res)
+				}
+				if cap(res) != 10 {
+					t.Fatalf("Peek(5) did not use the provided slice.")
+				}
+
+				discarded, err := r.Discard(5)
+				if got, want := discarded, 5; got != want {
+					t.Fatalf("Discard(5) = %d, want %d", got, want)
+				}
+				if err != nil {
+					t.Fatalf("Discard(5) return error %v, want <nil>", err)
+				}
+				if r.Remaining() != 5 {
+					t.Fatalf("Remaining() after Discard(5) = %v, want 5", r.Remaining())
+				}
+				res, err = r.Peek(5, res[:0])
+				if err != nil {
+					t.Fatalf("Peek(5) return error %v, want <nil>", err)
+				}
+				if len(res) != 1 || !bytes.Equal(res[0], []byte("56789")) {
+					t.Fatalf("Peek(5) after Discard(5) = %v, want [[56789]]", res)
+				}
+
+				discarded, err = r.Discard(100)
+				if got, want := discarded, 5; got != want {
+					t.Fatalf("Discard(100) = %d, want %d", got, want)
+				}
+				if err == nil {
+					t.Fatalf("Discard(100) returned error <nil>, want non-nil")
+				}
+				if r.Remaining() != 0 {
+					t.Fatalf("Remaining() after Discard(100) = %v, want 0", r.Remaining())
+				}
+			},
+		},
+		{
+			name:    "multiple_buffers",
+			buffers: [][]byte{[]byte("012"), []byte("345"), []byte("6789")},
+			operations: func(t *testing.T, r *mem.Reader) {
+				if r.Remaining() != 10 {
+					t.Fatalf("Remaining() = %v, want 10", r.Remaining())
+				}
+
+				res, err := r.Peek(5, nil)
+				if err != nil {
+					t.Fatalf("Peek(5) return error %v, want <nil>", err)
+				}
+				if len(res) != 2 || !bytes.Equal(res[0], []byte("012")) || !bytes.Equal(res[1], []byte("34")) {
+					t.Fatalf("Peek(5) = %v, want [[012] [34]]", res)
+				}
+
+				discarded, err := r.Discard(5)
+				if got, want := discarded, 5; got != want {
+					t.Fatalf("Discard(5) = %d, want %d", got, want)
+				}
+				if err != nil {
+					t.Fatalf("Discard(5) return error %v, want <nil>", err)
+				}
+				if r.Remaining() != 5 {
+					t.Fatalf("Remaining() after Discard(5) = %v, want 5", r.Remaining())
+				}
+
+				res, err = r.Peek(5, res[:0])
+				if err != nil {
+					t.Fatalf("Peek(5) return error %v, want <nil>", err)
+				}
+				if len(res) != 2 || !bytes.Equal(res[0], []byte("5")) || !bytes.Equal(res[1], []byte("6789")) {
+					t.Fatalf("Peek(5) after advance = %v, want [[5] [6789]]", res)
+				}
+			},
+		},
+		{
+			name:    "close",
+			buffers: [][]byte{[]byte("0123456789")},
+			operations: func(t *testing.T, r *mem.Reader) {
+				r.Close()
+				if r.Remaining() != 0 {
+					t.Fatalf("Remaining() after Close = %v, want 0", r.Remaining())
+				}
+			},
+		},
+		{
+			name:    "reset",
+			buffers: [][]byte{[]byte("0123")},
+			operations: func(t *testing.T, r *mem.Reader) {
+				newSlice := mem.BufferSlice{mem.SliceBuffer([]byte("56789"))}
+				r.Reset(newSlice)
+				if r.Remaining() != 5 {
+					t.Fatalf("Remaining() after Reset = %v, want 5", r.Remaining())
+				}
+				res, err := r.Peek(5, nil)
+				if err != nil {
+					t.Fatalf("Peek(5) return error %v, want <nil>", err)
+				}
+				if len(res) != 1 || !bytes.Equal(res[0], []byte("56789")) {
+					t.Fatalf("Peek(5) after Reset = %v, want [[56789]]", res)
+				}
+			},
+		},
+		{
+			name:    "zero_ops",
+			buffers: [][]byte{[]byte("01234")},
+			operations: func(t *testing.T, c *mem.Reader) {
+				if c.Remaining() != 5 {
+					t.Fatalf("Remaining() = %v, want 5", c.Remaining())
+				}
+				res, err := c.Peek(0, nil)
+				if err != nil {
+					t.Fatalf("Peek(0) return error %v, want <nil>", err)
+				}
+				if len(res) != 0 {
+					t.Fatalf("Peek(0) got slices: %v, want empty", res)
+				}
+				discarded, err := c.Discard(0)
+				if err != nil {
+					t.Fatalf("Discard(0) return error %v, want <nil>", err)
+				}
+				if got, want := discarded, 0; got != want {
+					t.Fatalf("Discard(0) = %d, want %d", got, want)
+				}
+				if c.Remaining() != 5 {
+					t.Fatalf("Remaining() after Discard(0) = %v, want 5", c.Remaining())
+				}
+				res, err = c.Peek(2, res[:0])
+				if err != nil {
+					t.Fatalf("Peek(2) return error %v, want <nil>", err)
+				}
+				if len(res) != 1 || !bytes.Equal(res[0], []byte("01")) {
+					t.Fatalf("Peek(2) after zero ops = %v, want [[01]]", res)
+				}
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var slice mem.BufferSlice
+			for _, b := range tt.buffers {
+				slice = append(slice, mem.SliceBuffer(b))
+			}
+			c := slice.Reader()
+			slice.Free()
+			defer c.Close()
+			tt.operations(t, c)
+		})
+	}
+}

mem/buffers.go

@@ -0,0 +1,317 @@
+/*
+ *
+ * Copyright 2024 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 provides utilities that facilitate memory reuse in byte slices
+// that are used as buffers.
+//
+// # Experimental
+//
+// Notice: All APIs in this package are EXPERIMENTAL and may be changed or
+// removed in a later release.
+package mem
+
+import (
+	"fmt"
+	"sync"
+	"sync/atomic"
+
+	"google.golang.org/grpc/internal/mem"
+)
+
+// A Buffer represents a reference counted piece of data (in bytes) that can be
+// acquired by a call to NewBuffer() or Copy(). A reference to a Buffer may be
+// released by calling Free(), which invokes the free function given at creation
+// only after all references are released.
+//
+// Note that a Buffer is not safe for concurrent access and instead each
+// goroutine should use its own reference to the data, which can be acquired via
+// a call to Ref().
+//
+// Attempts to access the underlying data after releasing the reference to the
+// Buffer will panic.
+type Buffer interface {
+	// ReadOnlyData returns the underlying byte slice. Note that it is undefined
+	// behavior to modify the contents of this slice in any way.
+	ReadOnlyData() []byte
+	// Ref increases the reference counter for this Buffer.
+	Ref()
+	// Free decrements this Buffer's reference counter and frees the underlying
+	// byte slice if the counter reaches 0 as a result of this call.
+	Free()
+	// Len returns the Buffer's size.
+	Len() int
+	// Slice returns a new Buffer that is a view into this buffer's data
+	// from [start:end). The buffer is not modified. Panics if the buffer
+	// has been freed or if start/end are out of bounds.
+	Slice(start, end int) Buffer
+
+	split(n int) (left, right Buffer)
+	read(buf []byte) (int, Buffer)
+}
+
+var (
+	bufferObjectPool = sync.Pool{New: func() any { return new(buffer) }}
+)
+
+// IsBelowBufferPoolingThreshold returns true if the given size is less than or
+// equal to the threshold for buffer pooling. This is used to determine whether
+// to pool buffers or allocate them directly.
+func IsBelowBufferPoolingThreshold(size int) bool {
+	return size <= mem.BufferPoolingThreshold
+}
+
+type buffer struct {
+	refs atomic.Int32
+	data []byte
+
+	// rootBuf is the buffer responsible for returning origData to the pool
+	// once the reference count drops to 0.
+	//
+	// When a buffer is split, the new buffer inherits the rootBuf of the
+	// original and increments the root's reference count. For the
+	// initial buffer (the root), this field points to itself.
+	rootBuf *buffer
+
+	// The following fields are only set for root buffers.
+	origData *[]byte
+	pool     BufferPool
+}
+
+func newBuffer() *buffer {
+	return bufferObjectPool.Get().(*buffer)
+}
+
+// NewBuffer creates a new Buffer from the given data, initializing the reference
+// counter to 1. The data will then be returned to the given pool when all
+// references to the returned Buffer are released. As a special case to avoid
+// additional allocations, if the given buffer pool is nil, the returned buffer
+// will be a "no-op" Buffer where invoking Buffer.Free() does nothing and the
+// underlying data is never freed.
+//
+// Note that the backing array of the given data is not copied.
+func NewBuffer(data *[]byte, pool BufferPool) Buffer {
+	// Use the buffer's capacity instead of the length, otherwise buffers may
+	// not be reused under certain conditions. For example, if a large buffer
+	// is acquired from the pool, but fewer bytes than the buffering threshold
+	// are written to it, the buffer will not be returned to the pool.
+	if pool == nil || IsBelowBufferPoolingThreshold(cap(*data)) {
+		return (SliceBuffer)(*data)
+	}
+	b := newBuffer()
+	b.origData = data
+	b.data = *data
+	b.pool = pool
+	b.rootBuf = b
+	b.refs.Store(1)
+	return b
+}
+
+// Copy creates a new Buffer from the given data, initializing the reference
+// counter to 1.
+//
+// It acquires a []byte from the given pool and copies over the backing array
+// of the given data. The []byte acquired from the pool is returned to the
+// pool when all references to the returned Buffer are released.
+func Copy(data []byte, pool BufferPool) Buffer {
+	if IsBelowBufferPoolingThreshold(len(data)) {
+		buf := make(SliceBuffer, len(data))
+		copy(buf, data)
+		return buf
+	}
+
+	buf := pool.Get(len(data))
+	copy(*buf, data)
+	return NewBuffer(buf, pool)
+}
+
+func (b *buffer) ReadOnlyData() []byte {
+	if b.rootBuf == nil {
+		panic("Cannot read freed buffer")
+	}
+	return b.data
+}
+
+func (b *buffer) Ref() {
+	if b.refs.Add(1) <= 1 {
+		panic("Cannot ref freed buffer")
+	}
+}
+
+func (b *buffer) Free() {
+	refs := b.refs.Add(-1)
+	if refs < 0 {
+		panic("Cannot free freed buffer")
+	}
+	if refs > 0 {
+		return
+	}
+
+	b.data = nil
+	if b.rootBuf == b {
+		// This buffer is the owner of the data slice and its ref count reached
+		// 0, free the slice.
+		if b.pool != nil {
+			b.pool.Put(b.origData)
+			b.pool = nil
+		}
+		b.origData = nil
+	} else {
+		// This buffer doesn't own the data slice, decrement a ref on the root
+		// buffer.
+		b.rootBuf.Free()
+	}
+
+	b.rootBuf = nil
+	bufferObjectPool.Put(b)
+}
+
+func (b *buffer) Len() int {
+	return len(b.ReadOnlyData())
+}
+
+func (b *buffer) Slice(start, end int) Buffer {
+	if b.rootBuf == nil {
+		panic("Cannot slice freed buffer")
+	}
+
+	data := b.data[start:end] // access the data to check slice bounds
+
+	if len(data) == 0 {
+		return emptyBuffer{}
+	}
+	if len(data) == len(b.data) {
+		b.Ref()
+		return b
+	}
+	// We are creating a new reference (view) to a portion of the root buffer's
+	// data. Therefore, we must increment the reference count of the root buffer
+	// to ensure the underlying data is not freed while this view is still in
+	// use.
+	b.rootBuf.Ref()
+	s := newBuffer()
+	s.data = data
+	s.rootBuf = b.rootBuf
+	s.refs.Store(1)
+	return s
+}
+
+func (b *buffer) split(n int) (Buffer, Buffer) {
+	if b.rootBuf == nil || b.rootBuf.refs.Add(1) <= 1 {
+		panic("Cannot split freed buffer")
+	}
+
+	split := newBuffer()
+	split.data = b.data[n:]
+	split.rootBuf = b.rootBuf
+	split.refs.Store(1)
+
+	b.data = b.data[:n]
+
+	return b, split
+}
+
+func (b *buffer) read(buf []byte) (int, Buffer) {
+	if b.rootBuf == nil {
+		panic("Cannot read freed buffer")
+	}
+
+	n := copy(buf, b.data)
+	if n == len(b.data) {
+		b.Free()
+		return n, nil
+	}
+
+	b.data = b.data[n:]
+	return n, b
+}
+
+func (b *buffer) String() string {
+	return fmt.Sprintf("mem.Buffer(%p, data: %p, length: %d)", b, b.ReadOnlyData(), len(b.ReadOnlyData()))
+}
+
+// ReadUnsafe reads bytes from the given Buffer into the provided slice.
+// It does not perform safety checks.
+func ReadUnsafe(dst []byte, buf Buffer) (int, Buffer) {
+	return buf.read(dst)
+}
+
+// SplitUnsafe modifies the receiver to point to the first n bytes while it
+// returns a new reference to the remaining bytes. The returned Buffer
+// functions just like a normal reference acquired using Ref().
+func SplitUnsafe(buf Buffer, n int) (left, right Buffer) {
+	return buf.split(n)
+}
+
+type emptyBuffer struct{}
+
+func (e emptyBuffer) ReadOnlyData() []byte {
+	return nil
+}
+
+func (e emptyBuffer) Ref()  {}
+func (e emptyBuffer) Free() {}
+
+func (e emptyBuffer) Len() int {
+	return 0
+}
+
+func (e emptyBuffer) Slice(start, end int) Buffer {
+	if start != 0 || end != 0 {
+		panic(fmt.Sprintf("slice bounds out of range [%d:%d] with length 0", start, end))
+	}
+	return e
+}
+
+func (e emptyBuffer) split(int) (left, right Buffer) {
+	return e, e
+}
+
+func (e emptyBuffer) read([]byte) (int, Buffer) {
+	return 0, e
+}
+
+// SliceBuffer is a Buffer implementation that wraps a byte slice. It provides
+// methods for reading, splitting, and managing the byte slice.
+type SliceBuffer []byte
+
+// ReadOnlyData returns the byte slice.
+func (s SliceBuffer) ReadOnlyData() []byte { return s }
+
+// Ref is a noop implementation of Ref.
+func (s SliceBuffer) Ref() {}
+
+// Free is a noop implementation of Free.
+func (s SliceBuffer) Free() {}
+
+// Len is a noop implementation of Len.
+func (s SliceBuffer) Len() int { return len(s) }
+
+// Slice returns a new SliceBuffer that is a view into the receiver from [start:end).
+func (s SliceBuffer) Slice(start, end int) Buffer { return s[start:end] }
+
+func (s SliceBuffer) split(n int) (left, right Buffer) {
+	return s[:n], s[n:]
+}
+
+func (s SliceBuffer) read(buf []byte) (int, Buffer) {
+	n := copy(buf, s)
+	if n == len(s) {
+		return n, nil
+	}
+	return n, s[n:]
+}

mem/buffers_test.go

@@ -0,0 +1,493 @@
+/*
+ *
+ * Copyright 2024 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_test
+
+import (
+	"bytes"
+	"fmt"
+	"slices"
+	"testing"
+
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/mem"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	internal.SetBufferPoolingThresholdForTesting.(func(int))(0)
+
+	grpctest.RunSubTests(t, s{})
+}
+
+// Tests that a buffer created with NewBuffer, which when later freed, invokes
+// the free function with the correct data.
+func (s) TestBuffer_NewBufferAndFree(t *testing.T) {
+	data := "abcd"
+	freed := false
+	freeF := poolFunc(func(got *[]byte) {
+		if !bytes.Equal(*got, []byte(data)) {
+			t.Fatalf("Free function called with bytes %s, want %s", string(*got), data)
+		}
+		freed = true
+	})
+
+	buf := newBuffer([]byte(data), freeF)
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
+		t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	// Verify that the free function is invoked when all references are freed.
+	buf.Free()
+	if !freed {
+		t.Fatalf("Buffer not freed")
+	}
+}
+
+// Tests that a buffer created with NewBuffer, on which an additional reference
+// is acquired, which when later freed, invokes the free function with the
+// correct data, but only after all references are released.
+func (s) TestBuffer_NewBufferRefAndFree(t *testing.T) {
+	data := "abcd"
+	freed := false
+	freeF := poolFunc(func(got *[]byte) {
+		if !bytes.Equal(*got, []byte(data)) {
+			t.Fatalf("Free function called with bytes %s, want %s", string(*got), string(data))
+		}
+		freed = true
+	})
+
+	buf := newBuffer([]byte(data), freeF)
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
+		t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	buf.Ref()
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
+		t.Fatalf("New reference to the Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	// Verify that the free function is not invoked when all references are yet
+	// to be freed.
+	buf.Free()
+	if freed {
+		t.Fatalf("Free function called before all references freed")
+	}
+
+	// Verify that the free function is invoked when all references are freed.
+	buf.Free()
+	if !freed {
+		t.Fatalf("Buffer not freed")
+	}
+}
+
+func (s) TestBuffer_NewBufferHandlesShortBuffers(t *testing.T) {
+	const threshold = 100
+
+	// Update the pooling threshold, since that's what's being tested.
+	internal.SetBufferPoolingThresholdForTesting.(func(int))(threshold)
+	t.Cleanup(func() {
+		internal.SetBufferPoolingThresholdForTesting.(func(int))(0)
+	})
+
+	// Make a pool with a buffer whose capacity is larger than the pooling
+	// threshold, but whose length is less than the threshold.
+	b := make([]byte, threshold/2, threshold*2)
+	pool := &singleBufferPool{
+		t:    t,
+		data: &b,
+	}
+
+	// Get a Buffer, then free it. If NewBuffer decided that the Buffer
+	// shouldn't get pooled, Free will be a noop and singleBufferPool will not
+	// have been updated.
+	mem.NewBuffer(&b, pool).Free()
+
+	if pool.data != nil {
+		t.Fatalf("Buffer not returned to pool")
+	}
+}
+
+func (s) TestBuffer_FreeAfterFree(t *testing.T) {
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	if buf.Len() != 4 {
+		t.Fatalf("Buffer length is %d, want 4", buf.Len())
+	}
+
+	// Ensure that a double free does panic.
+	buf.Free()
+	defer checkForPanic(t, "Cannot free freed buffer")
+	buf.Free()
+}
+
+type singleBufferPool struct {
+	t    *testing.T
+	data *[]byte
+}
+
+func (s *singleBufferPool) Get(length int) *[]byte {
+	if len(*s.data) != length {
+		s.t.Fatalf("Invalid requested length, got %d want %d", length, len(*s.data))
+	}
+	return s.data
+}
+
+func (s *singleBufferPool) Put(b *[]byte) {
+	if s.data != b {
+		s.t.Fatalf("Wrong buffer returned to pool, got %p want %p", b, s.data)
+	}
+	s.data = nil
+}
+
+// Tests that a buffer created with Copy, which when later freed, returns the underlying
+// byte slice to the buffer pool.
+func (s) TestBuffer_CopyAndFree(t *testing.T) {
+	data := []byte("abcd")
+	testPool := &singleBufferPool{
+		t:    t,
+		data: &data,
+	}
+
+	buf := mem.Copy(data, testPool)
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, data) {
+		t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	// Verify that the free function is invoked when all references are freed.
+	buf.Free()
+	if testPool.data != nil {
+		t.Fatalf("Buffer not freed")
+	}
+}
+
+// Tests that a buffer created with Copy, on which an additional reference is
+// acquired, which when later freed, returns the underlying byte slice to the
+// buffer pool.
+func (s) TestBuffer_CopyRefAndFree(t *testing.T) {
+	data := []byte("abcd")
+	testPool := &singleBufferPool{
+		t:    t,
+		data: &data,
+	}
+
+	buf := mem.Copy(data, testPool)
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, data) {
+		t.Fatalf("Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	buf.Ref()
+	if got := buf.ReadOnlyData(); !bytes.Equal(got, []byte(data)) {
+		t.Fatalf("New reference to the Buffer contains data %s, want %s", string(got), string(data))
+	}
+
+	// Verify that the free function is not invoked when all references are yet
+	// to be freed.
+	buf.Free()
+	if testPool.data == nil {
+		t.Fatalf("Free function called before all references freed")
+	}
+
+	// Verify that the free function is invoked when all references are freed.
+	buf.Free()
+	if testPool.data != nil {
+		t.Fatalf("Free never called")
+	}
+}
+
+func (s) TestBuffer_ReadOnlyDataAfterFree(t *testing.T) {
+	// Verify that reading before freeing does not panic.
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	buf.ReadOnlyData()
+
+	buf.Free()
+	defer checkForPanic(t, "Cannot read freed buffer")
+	buf.ReadOnlyData()
+}
+
+func (s) TestBuffer_RefAfterFree(t *testing.T) {
+	// Verify that acquiring a ref before freeing does not panic.
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	buf.Ref()
+
+	// This first call should not panic and bring the ref counter down to 1
+	buf.Free()
+	// This second call actually frees the buffer
+	buf.Free()
+	defer checkForPanic(t, "Cannot ref freed buffer")
+	buf.Ref()
+}
+
+func (s) TestBuffer_SplitAfterFree(t *testing.T) {
+	// Verify that splitting before freeing does not panic.
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	buf, bufSplit := mem.SplitUnsafe(buf, 2)
+
+	bufSplit.Free()
+	buf.Free()
+	defer checkForPanic(t, "Cannot split freed buffer")
+	mem.SplitUnsafe(buf, 2)
+}
+
+type poolFunc func(*[]byte)
+
+func (p poolFunc) Get(int) *[]byte {
+	panic("Get should never be called")
+}
+
+func (p poolFunc) Put(i *[]byte) {
+	p(i)
+}
+
+func (s) TestBuffer_Split(t *testing.T) {
+	freed := make(chan struct{})
+	data := []byte{1, 2, 3, 4}
+	buf := mem.NewBuffer(&data, poolFunc(func(*[]byte) {
+		close(freed)
+	}))
+
+	buf, split1 := mem.SplitUnsafe(buf, 2)
+	if !bytes.Equal(buf.ReadOnlyData(), data[:2]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", buf.ReadOnlyData(), data[:2])
+	}
+	if !bytes.Equal(split1.ReadOnlyData(), data[2:]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", split1.ReadOnlyData(), data[2:])
+	}
+
+	// Check that splitting the buffer more than once works as intended.
+	split1, split2 := mem.SplitUnsafe(split1, 1)
+	if !bytes.Equal(split1.ReadOnlyData(), data[2:3]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", split1.ReadOnlyData(), data[2:3])
+	}
+	if !bytes.Equal(split2.ReadOnlyData(), data[3:]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", split2.ReadOnlyData(), data[3:])
+	}
+
+	// If any of the following frees actually free the buffer, the test will fail.
+	buf.Free()
+	split2.Free()
+	select {
+	case <-freed:
+		t.Fatalf("Freed too early")
+	default:
+	}
+
+	split1.Free()
+
+	select {
+	case <-freed:
+	default:
+		t.Fatalf("Still not freed")
+	}
+}
+
+func (s) TestBuffer_Slice(t *testing.T) {
+	freed := make(chan struct{})
+	data := []byte{1, 2, 3, 4}
+	buf := mem.NewBuffer(&data, poolFunc(func(*[]byte) {
+		close(freed)
+	}))
+
+	// Slice the buffer and verify the data.
+	slice1 := buf.Slice(1, 3)
+	if !bytes.Equal(slice1.ReadOnlyData(), data[1:3]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", slice1.ReadOnlyData(), data[1:3])
+	}
+
+	// Verify the original buffer is not modified.
+	if !bytes.Equal(buf.ReadOnlyData(), data) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", buf.ReadOnlyData(), data)
+	}
+
+	// Slice the slice.
+	slice2 := slice1.Slice(0, 1)
+	if !bytes.Equal(slice2.ReadOnlyData(), data[1:2]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", slice2.ReadOnlyData(), data[1:2])
+	}
+
+	// Free original and first slice — root should not be freed yet.
+	buf.Free()
+	slice1.Free()
+	select {
+	case <-freed:
+		t.Fatalf("Freed too early")
+	default:
+	}
+
+	// The last slice keeps the root alive.
+	if !bytes.Equal(slice2.ReadOnlyData(), data[1:2]) {
+		t.Fatalf("Buffer did not contain expected data. got %v, want %v", slice2.ReadOnlyData(), data[1:2])
+	}
+
+	slice2.Free()
+
+	select {
+	case <-freed:
+	default:
+		t.Fatalf("Still not freed")
+	}
+}
+
+func (s) TestBuffer_SliceAfterFree(t *testing.T) {
+	buf := newBuffer([]byte("abcd"), mem.NopBufferPool{})
+	buf.Free()
+	defer checkForPanic(t, "Cannot slice freed buffer")
+	buf.Slice(0, 2)
+}
+
+type namedCtor struct {
+	name   string
+	newBuf func([]byte) mem.Buffer
+}
+
+var ctors = []namedCtor{
+	{name: "pooled", newBuf: newPooledBuffer},
+	{name: "slice", newBuf: newSliceBuf},
+}
+
+func (s) TestBuffer_SliceBasic(t *testing.T) {
+	type sliceCase struct {
+		start, end int
+		want       []byte
+	}
+	cases := []sliceCase{
+		{1, 3, []byte{2, 3}},
+		{0, 4, []byte{1, 2, 3, 4}},
+		{4, 4, []byte{}},
+	}
+	for _, c := range ctors {
+		for _, tc := range cases {
+			t.Run(fmt.Sprintf("%s[%d:%d]", c.name, tc.start, tc.end), func(t *testing.T) {
+				buf := c.newBuf([]byte{1, 2, 3, 4})
+				got := buf.Slice(tc.start, tc.end)
+				if !bytes.Equal(got.ReadOnlyData(), tc.want) {
+					t.Fatalf("Buffer did not contain expected data. got %v, want %v", got.ReadOnlyData(), tc.want)
+				}
+			})
+		}
+	}
+}
+
+func (s) TestBuffer_SliceSubslice(t *testing.T) {
+	for _, c := range ctors {
+		t.Run(c.name+" subslice", func(t *testing.T) {
+			buf := c.newBuf([]byte{1, 2, 3, 4})
+			slice := buf.Slice(1, 3)
+			slice2 := slice.Slice(0, 1)
+			if !bytes.Equal(slice2.ReadOnlyData(), []byte{2}) {
+				t.Fatalf("Buffer did not contain expected data. got %v, want %v", buf.ReadOnlyData(), []byte{2})
+			}
+		})
+	}
+}
+
+func (s) TestBuffer_SliceEmpty(t *testing.T) {
+	allCtors := append(slices.Clip(ctors), namedCtor{name: "empty", newBuf: newEmptyBuf})
+	for _, c := range allCtors {
+		t.Run(c.name, func(t *testing.T) {
+			buf := c.newBuf([]byte{1, 2, 3, 4})
+			got := buf.Slice(0, 0)
+			if !bytes.Equal(got.ReadOnlyData(), nil) {
+				t.Fatalf("Buffer did not contain expected data. got %v, want %v", got.ReadOnlyData(), nil)
+			}
+		})
+	}
+}
+
+func (s) TestBuffer_SliceBoundsCheck(t *testing.T) {
+	type panicCase struct {
+		name       string
+		start, end int
+	}
+	nonEmptyCases := []panicCase{
+		{"end_out_of_bounds", 0, 5},
+		{"start_negative", -1, 3},
+		{"start_greater_than_end", 3, 0},
+	}
+	tests := []struct {
+		namedCtor
+		panicCases []panicCase
+	}{
+		{
+			namedCtor: namedCtor{
+				name:   "buffer",
+				newBuf: newPooledBuffer,
+			},
+			panicCases: nonEmptyCases,
+		},
+		{
+			namedCtor: namedCtor{
+				name:   "SliceBuffer",
+				newBuf: newSliceBuf,
+			},
+			panicCases: nonEmptyCases,
+		},
+		{
+			namedCtor: namedCtor{
+				name:   "emptyBuffer",
+				newBuf: newEmptyBuf,
+			},
+			panicCases: []panicCase{
+				{"end_out_of_bounds", 0, 1},
+				{"start_negative", -1, 0},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for _, pc := range tt.panicCases {
+				t.Run(pc.name, func(t *testing.T) {
+					buf := tt.newBuf([]byte{1, 2, 3, 4})
+					defer func() {
+						if recover() == nil {
+							t.Fatalf("Slice(%d, %d) did not panic", pc.start, pc.end)
+						}
+					}()
+					buf.Slice(pc.start, pc.end)
+				})
+			}
+		})
+	}
+}
+
+func newPooledBuffer(data []byte) mem.Buffer {
+	return newBuffer(data, mem.NopBufferPool{})
+}
+
+func newSliceBuf(data []byte) mem.Buffer {
+	return mem.SliceBuffer(data)
+}
+
+func newEmptyBuf([]byte) mem.Buffer {
+	var bs mem.BufferSlice
+	return bs.MaterializeToBuffer(mem.NopBufferPool{})
+}
+
+func checkForPanic(t *testing.T, wantErr string) {
+	t.Helper()
+	r := recover()
+	if r == nil {
+		t.Fatalf("Use after free did not panic")
+	}
+	if msg, ok := r.(string); !ok || msg != wantErr {
+		t.Fatalf("panic called with %v, want %s", r, wantErr)
+	}
+}