Documentation
¶
Overview ¶
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.
Index ¶
Examples ¶
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func IsBelowBufferPoolingThreshold ¶
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 NewWriter ¶
func NewWriter(buffers *BufferSlice, pool BufferPool) io.Writer
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.
Example ¶
/*
*
* 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 main
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 main() {
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()))
}
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)
}
Output: Wrote 4 bytes, err: <nil> Wrote 4 bytes, err: <nil> Wrote 4 bytes, err: <nil> abcdabcdabcd
Types ¶
type Buffer ¶
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
// contains filtered or unexported methods
}
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.
func Copy ¶
func Copy(data []byte, pool BufferPool) Buffer
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 NewBuffer ¶
func NewBuffer(data *[]byte, pool BufferPool) 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 ReadUnsafe ¶
ReadUnsafe reads bytes from the given Buffer into the provided slice. It does not perform safety checks.
func SplitUnsafe ¶
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().
type BufferPool ¶
type BufferPool interface {
// Get returns a buffer with specified length from the pool.
Get(length int) *[]byte
// Put returns a buffer to the pool.
Put(*[]byte)
}
BufferPool is a pool of buffers that can be shared and reused, resulting in decreased memory allocation.
func DefaultBufferPool ¶
func DefaultBufferPool() BufferPool
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 NewTieredBufferPool ¶
func NewTieredBufferPool(poolSizes ...int) BufferPool
NewTieredBufferPool returns a BufferPool implementation that uses multiple underlying pools of the given pool sizes.
type BufferSlice ¶
type BufferSlice []Buffer
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.
func ReadAll ¶ added in v1.69.0
func ReadAll(r io.Reader, pool BufferPool) (BufferSlice, error)
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 (BufferSlice) CopyTo ¶
func (s BufferSlice) CopyTo(dst []byte) int
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 (BufferSlice) Free ¶
func (s BufferSlice) Free()
Free invokes Buffer.Free() on each Buffer in the slice.
func (BufferSlice) Len ¶
func (s BufferSlice) Len() int
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 (BufferSlice) Materialize ¶
func (s BufferSlice) Materialize() []byte
Materialize concatenates all the underlying Buffer's data into a single contiguous buffer using CopyTo.
func (BufferSlice) MaterializeToBuffer ¶
func (s BufferSlice) MaterializeToBuffer(pool BufferPool) Buffer
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 (BufferSlice) Reader ¶
func (s BufferSlice) Reader() Reader
Reader returns a new Reader for the input slice after taking references to each underlying buffer.
type NopBufferPool ¶
type NopBufferPool struct{}
NopBufferPool is a buffer pool that returns new buffers without pooling.
func (NopBufferPool) Get ¶
func (NopBufferPool) Get(length int) *[]byte
Get returns a buffer with specified length from the pool.
type Reader ¶
type Reader interface {
io.Reader
io.ByteReader
// Close frees the underlying BufferSlice and never returns an error. Subsequent
// calls to Read will return (0, io.EOF).
Close() error
// Remaining returns the number of unread bytes remaining in the slice.
Remaining() int
}
Reader exposes a BufferSlice's data as an io.Reader, allowing it to interface with other parts systems. It also provides an additional convenience method Remaining(), which returns the number of unread bytes remaining in the slice. Buffers will be freed as they are read.
type SliceBuffer ¶
type SliceBuffer []byte
SliceBuffer is a Buffer implementation that wraps a byte slice. It provides methods for reading, splitting, and managing the byte slice.
func (SliceBuffer) ReadOnlyData ¶
func (s SliceBuffer) ReadOnlyData() []byte
ReadOnlyData returns the byte slice.
Source Files