mem/buffer_slice_test.go - external/github.com/grpc/grpc-go - Git at Google

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