mem/buffer_slice_test.go

/*
 *
 * 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)
		})
	}
}