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

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)
+	}
+}