Add new buffer implementation.

PiperOrigin-RevId: 455003647

pkg/bufferv2/BUILD

@@ -0,0 +1,65 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools/go_generics:defs.bzl", "go_template_instance")
+
+package(licenses = ["notice"])
+
+go_template_instance(
+    name = "chunk_refs",
+    out = "chunk_refs.go",
+    package = "buffer",
+    prefix = "chunk",
+    template = "//pkg/refsvfs2:refs_template",
+    types = {
+        "T": "chunk",
+    },
+)
+
+go_template_instance(
+    name = "view_list",
+    out = "view_list.go",
+    package = "buffer",
+    prefix = "view",
+    template = "//pkg/ilist:generic_list",
+    types = {
+        "Element": "*View",
+        "Linker": "*View",
+    },
+)
+
+go_library(
+    name = "buffer",
+    srcs = [
+        "buffer.go",
+        "buffer_state.go",
+        "chunk.go",
+        "chunk_refs.go",
+        "view.go",
+        "view_list.go",
+        "view_unsafe.go",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/atomicbitops",
+        "//pkg/bits",
+        "//pkg/context",
+        "//pkg/ilist",
+        "//pkg/log",
+        "//pkg/pool",
+        "//pkg/refsvfs2",
+        "//pkg/sync",
+    ],
+)
+
+go_test(
+    name = "buffer_test",
+    size = "small",
+    srcs = [
+        "buffer_test.go",
+        "view_test.go",
+    ],
+    library = ":buffer",
+    deps = [
+        "//pkg/state",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+    ],
+)

pkg/bufferv2/buffer_state.go

@@ -0,0 +1,26 @@
+// Copyright 2022 The gVisor 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 buffer
+
+// saveBuf is invoked by stateify.
+func (b *Buffer) saveData() []byte {
+	return b.Flatten()
+}
+
+// loadBuf is invoked by stateify.
+func (b *Buffer) loadData(data []byte) {
+	v := NewViewWithData(data)
+	b.Append(v)
+}

pkg/bufferv2/chunk.go

@@ -0,0 +1,113 @@
+// Copyright 2022 The gVisor 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 buffer
+
+import (
+	"fmt"
+
+	"gvisor.dev/gvisor/pkg/bits"
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+const (
+	// This is log2(baseChunkSize). This number is used to calculate which pool
+	// to use for a payload size by right shifting the payload size by this
+	// number and passing the result to MostSignificantOne64.
+	baseChunkSizeLog2 = 6
+
+	// This is the size of the buffers in the first pool. Each subsquent pool
+	// creates payloads 2^(pool index) times larger than the first pool's
+	// payloads.
+	baseChunkSize = 1 << baseChunkSizeLog2 // 64
+
+	// The largest payload size that we pool. Payloads larger than this will
+	// allocated from the heap and garbage collected as normal.
+	maxChunkSize = baseChunkSize << (numPools - 1) // 65536
+
+	// The number of chunk pools we have for use.
+	numPools = 11
+)
+
+// chunkPools is a collection of pools for payloads of different sizes. The
+// size of the payloads doubles in each successive pool.
+var chunkPools [numPools]sync.Pool
+
+func init() {
+	for i := 0; i < numPools; i++ {
+		chunkSize := baseChunkSize * (1 << i)
+		chunkPools[i].New = func() interface{} {
+			return &chunk{
+				data: make([]byte, chunkSize),
+			}
+		}
+	}
+}
+
+// Precondition: 0 <= size <= maxChunkSize
+func getChunkPool(size int) *sync.Pool {
+	idx := 0
+	if size > baseChunkSize {
+		idx = bits.MostSignificantOne64(uint64(size) >> baseChunkSizeLog2)
+		if size > 1<<(idx+baseChunkSizeLog2) {
+			idx++
+		}
+	}
+	if idx >= numPools {
+		panic(fmt.Sprintf("pool for chunk size %d does not exist", size))
+	}
+	return &chunkPools[idx]
+}
+
+// Chunk represents a slice of pooled memory.
+type chunk struct {
+	chunkRefs
+	data []byte
+}
+
+func newChunk(size int) *chunk {
+	var c *chunk
+	if size > maxChunkSize {
+		c = &chunk{
+			data: make([]byte, size),
+		}
+	} else {
+		pool := getChunkPool(size)
+		c = pool.Get().(*chunk)
+		for i := range c.data {
+			c.data[i] = 0
+		}
+	}
+	c.InitRefs()
+	return c
+}
+
+func (c *chunk) destroy() {
+	if len(c.data) > maxChunkSize {
+		c.data = nil
+		return
+	}
+	pool := getChunkPool(len(c.data))
+	pool.Put(c)
+}
+
+func (c *chunk) DecRef() {
+	c.chunkRefs.DecRef(c.destroy)
+}
+
+func (c *chunk) Clone() *chunk {
+	cpy := newChunk(len(c.data))
+	copy(cpy.data, c.data)
+	return cpy
+}

pkg/bufferv2/view.go

@@ -0,0 +1,232 @@
+// Copyright 2022 The gVisor 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 buffer
+
+import (
+	"fmt"
+	"io"
+
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+var viewPool = sync.Pool{
+	New: func() interface{} {
+		return &View{}
+	},
+}
+
+// View is a window into a shared chunk. Views are held by Buffers in
+// viewLists to represent contiguous memory.
+//
+// A View must be created with NewView, NewViewWithData, or Clone. Owners are
+// responsible for maintaining ownership over their views. When Views need to be
+// shared or copied, the owner should create a new View with Clone. Clone must
+// only ever be called on a owned View, not a borrowed one.
+//
+// Users are responsible for calling Release when finished with their View so
+// that its resources can be returned to the pool.
+//
+// Users must not write directly to slices returned by AsSlice. Instead, they
+// must use Write/WriteAt/CopyIn to modify the underlying View. This preserves
+// the safety guarantees of copy-on-write.
+type View struct {
+	sync.NoCopy
+
+	viewEntry
+	read  int
+	write int
+	chunk *chunk
+}
+
+// NewView creates a new view with capacity at least as big as cap. It is
+// analogous to make([]byte, 0, cap).
+func NewView(cap int) *View {
+	c := newChunk(cap)
+	v := viewPool.Get().(*View)
+	*v = View{chunk: c}
+	return v
+}
+
+// NewViewSize creates a new view with capacity at least as big as size and
+// length that is exactly size. It is analogous to make([]byte, size).
+func NewViewSize(size int) *View {
+	v := NewView(size)
+	v.Grow(size)
+	return v
+}
+
+// NewViewWithData creates a new view and initializes it with data. This
+// function should be used with caution to avoid unnecessary []byte allocations.
+// When in doubt use NewWithView to maximize chunk reuse in production
+// environments.
+func NewViewWithData(data []byte) *View {
+	c := newChunk(len(data))
+	v := viewPool.Get().(*View)
+	*v = View{chunk: c}
+	v.Write(data)
+	return v
+}
+
+// Clone creates a shallow clone of v where the underlying chunk is shared.
+//
+// The caller must own the View to call Clone. It is not safe to call Clone
+// on a borrowed or shared View because it can race with other View methods.
+func (v *View) Clone() *View {
+	v.chunk.IncRef()
+	newV := viewPool.Get().(*View)
+	newV.chunk = v.chunk
+	newV.read = v.read
+	newV.write = v.write
+	return newV
+}
+
+// Release releases the chunk held by v and returns v to the pool.
+func (v *View) Release() {
+	v.chunk.DecRef()
+	*v = View{}
+	viewPool.Put(v)
+}
+
+func (v *View) sharesChunk() bool {
+	return v.chunk.refCount.Load() > 1
+}
+
+// Full indicates the chunk is full.
+//
+// This indicates there is no capacity left to write.
+func (v *View) Full() bool {
+	return v == nil || v.write == len(v.chunk.data)
+}
+
+// Capacity returns the total size of this view's chunk.
+func (v *View) Capacity() int {
+	return len(v.chunk.data)
+}
+
+// Size returns the size of data written to the view.
+func (v *View) Size() int {
+	return v.write - v.read
+}
+
+// TrimFront advances the read index by the given amount.
+func (v *View) TrimFront(n int) {
+	v.read += n
+}
+
+// AsSlice returns a slice of the data written to this view.
+func (v *View) AsSlice() []byte {
+	if v == nil {
+		return nil
+	}
+	return v.chunk.data[v.read:v.write]
+}
+
+// AvailableSize returns the number of bytes available for writing.
+func (v *View) AvailableSize() int {
+	return len(v.chunk.data) - v.write
+}
+
+// Read reads v's data into p.
+//
+// Implements the io.Reader interface.
+func (v *View) Read(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if v.Size() == 0 {
+		return 0, io.EOF
+	}
+	n := copy(p, v.AsSlice())
+	v.TrimFront(n)
+	return n, nil
+}
+
+// ReadAt reads data to the p starting at offset.
+//
+// Implements the io.ReaderAt interface.
+func (v *View) ReadAt(p []byte, off int) (int, error) {
+	if off < 0 || off > v.Size() {
+		return 0, fmt.Errorf("ReadAt(): offset out of bounds: want 0 < off < %d, got off=%d", v.Size(), off)
+	}
+	n := copy(p, v.AsSlice()[off:])
+	return n, nil
+}
+
+// Write writes data to the view's chunk starting at the v.write index. If the
+// view's chunk has a reference count greater than 1, the chunk is copied first
+// and then written to.
+//
+// Implements the io.Writer interface.
+func (v *View) Write(p []byte) (int, error) {
+	if v.sharesChunk() {
+		defer v.chunk.DecRef()
+		v.chunk = v.chunk.Clone()
+	}
+	n := copy(v.chunk.data[v.write:], p)
+	v.write += n
+	if n < len(p) {
+		return n, fmt.Errorf("could not finish write: want len(p) <= v.AvailableSize(), got len(p)=%d, v.AvailableSize()=%d", len(p), v.AvailableSize())
+	}
+	return n, nil
+}
+
+// WriteAt writes data to the views's chunk starting at start. If the
+// view's chunk has a reference count greater than 1, the chunk is copied first
+// and then written to.
+//
+// Implements the io.WriterAt interface.
+func (v *View) WriteAt(p []byte, off int) (int, error) {
+	if off < 0 || off > v.Size() {
+		return 0, fmt.Errorf("write offset out of bounds: want 0 < off < %d, got off=%d", v.Size(), off)
+	}
+	if v.sharesChunk() {
+		defer v.chunk.DecRef()
+		v.chunk = v.chunk.Clone()
+	}
+	n := copy(v.AsSlice()[off:], p)
+	if n < len(p) {
+		return n, fmt.Errorf("could not finish write: want off + len(p) < v.Capacity(), got off=%d, len(p)=%d ,v.Size() = %d", off, len(p), v.Size())
+	}
+	return n, nil
+}
+
+// Grow advances the write index by the given amount.
+func (v *View) Grow(n int) {
+	if n+v.write > v.Capacity() {
+		panic("cannot grow view past capacity")
+	}
+	v.write += n
+}
+
+// CapLength caps the length of the view's read slice to n. If n > v.Size(),
+// the function is a no-op.
+func (v *View) CapLength(n int) {
+	if n < 0 {
+		panic("n must be >= 0")
+	}
+	if n > v.Size() {
+		n = v.Size()
+	}
+	v.write = v.read + n
+}
+
+func (v *View) availableSlice() []byte {
+	if v.sharesChunk() {
+		defer v.chunk.DecRef()
+		c := v.chunk.Clone()
+		v.chunk = c
+	}
+	return v.chunk.data[v.write:]
+}

pkg/bufferv2/view_test.go

@@ -0,0 +1,174 @@
+// Copyright 2022 The gVisor 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 buffer
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestNewView(t *testing.T) {
+	for sz := baseChunkSize; sz < maxChunkSize; sz <<= 1 {
+		v := NewView(sz - 1)
+		defer v.Release()
+
+		if v.Capacity() != sz {
+			t.Errorf("v.Capacity() = %d, want %d", v.Capacity(), sz)
+		}
+		if v.AvailableSize() != sz {
+			t.Errorf("v.WriteSize() = %d, want %d", v.AvailableSize(), sz)
+		}
+		if v.Size() != 0 {
+			t.Errorf("v.ReadSize() = %d, want %d", v.Size(), 0)
+		}
+
+		v1 := NewView(sz)
+		defer v1.Release()
+
+		if v1.Capacity() != sz {
+			t.Errorf("v.Capacity() = %d, want %d", v.Capacity(), sz)
+		}
+		if v1.AvailableSize() != sz {
+			t.Errorf("v.WriteSize() = %d, want %d", v.AvailableSize(), sz)
+		}
+		if v1.Size() != 0 {
+			t.Errorf("v.ReadSize() = %d, want %d", v.Size(), 0)
+		}
+	}
+
+	// Allocating from heap should produce a chunk with the exact size requested
+	// instead of a chunk where the size is contingent on the pool it came from.
+	viewSize := maxChunkSize + 1
+	v := NewView(viewSize)
+	defer v.Release()
+	if v.Capacity() != viewSize {
+		t.Errorf("v.Capacity() = %d, want %d", v.Capacity(), viewSize)
+	}
+}
+
+func TestClone(t *testing.T) {
+	orig := NewView(100)
+	clone := orig.Clone()
+	if orig.chunk != clone.chunk {
+		t.Errorf("orig.Clone().chunk = %p, want %p", clone.chunk, orig.chunk)
+	}
+	if orig.chunk.refCount.Load() != 2 {
+		t.Errorf("got orig.chunk.chunkRefs.Load() = %d, want 2", orig.chunk.refCount.Load())
+	}
+	orig.Release()
+	if clone.chunk.refCount.Load() != 1 {
+		t.Errorf("got clone.chunk.chunkRefs.Load() = %d, want 1", clone.chunk.refCount.Load())
+	}
+	clone.Release()
+}
+
+func TestWrite(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		view      *View
+		initData  []byte
+		writeSize int
+	}{
+		{
+			name:      "empty view",
+			view:      NewView(100),
+			writeSize: 50,
+		},
+		{
+			name:      "full view",
+			view:      NewView(100),
+			initData:  make([]byte, 100),
+			writeSize: 50,
+		},
+		{
+			name:      "full write to partially full view",
+			view:      NewView(100),
+			initData:  make([]byte, 20),
+			writeSize: 50,
+		},
+		{
+			name:      "partial write to partially full view",
+			view:      NewView(100),
+			initData:  make([]byte, 80),
+			writeSize: 50,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.view.Write(tc.initData)
+			defer tc.view.Release()
+			origWriteSize := tc.view.AvailableSize()
+
+			var orig []byte
+			orig = append(orig, tc.view.AsSlice()...)
+			toWrite := make([]byte, tc.writeSize)
+			rand.Read(toWrite)
+
+			n, _ := tc.view.Write(toWrite)
+
+			if n > origWriteSize {
+				t.Errorf("got tc.view.Write() = %d, want <=%d", n, origWriteSize)
+			}
+			if tc.writeSize > origWriteSize {
+				toWrite = toWrite[:origWriteSize]
+			}
+			if tc.view.AvailableSize() != tc.view.Capacity()-(len(toWrite)+len(orig)) {
+				t.Errorf("got tc.view.WriteSize() = %d, want %d", tc.view.AvailableSize(), tc.view.Capacity()-(len(toWrite)+len(orig)))
+			}
+			if !cmp.Equal(tc.view.AsSlice(), append(orig, toWrite...)) {
+				t.Errorf("got tc.view.ReadSlice() = %d, want %d", tc.view.AsSlice(), toWrite)
+			}
+		})
+	}
+}
+
+func TestWriteToCloned(t *testing.T) {
+	orig := NewView(100)
+	toWrite := make([]byte, 20)
+	rand.Read(toWrite)
+	orig.Write(toWrite)
+
+	clone := orig.Clone()
+	clone.Write(toWrite)
+
+	if !cmp.Equal(orig.AsSlice(), toWrite) {
+		t.Errorf("got orig.ReadSlice() = %v, want %v", orig.AsSlice(), toWrite)
+	}
+
+	toWrite = append(toWrite, toWrite...)
+	if !cmp.Equal(clone.AsSlice(), toWrite) {
+		t.Errorf("got clone.ReadSlice() = %v, want %v", clone.AsSlice(), toWrite)
+	}
+}
+
+func TestWriteAt(t *testing.T) {
+	size := 10
+	off := 5
+	v := NewViewSize(size)
+	p := make([]byte, 20)
+	rand.Read(p)
+	orig := v.Clone()
+
+	if n, _ := v.WriteAt(p, off); n != size-off {
+		t.Errorf("got v.CopyIn()= %v, want %v", n, size-off)
+	}
+	if !cmp.Equal(v.AsSlice()[off:], p[:size-off]) {
+		t.Errorf("got v.AsSlice()[off:] = %v, want %v", v.AsSlice()[off:], p[off:size])
+	}
+	if !cmp.Equal(v.AsSlice()[:off], orig.AsSlice()[:off]) {
+		t.Errorf("got v.AsSlice()[:off] = %v, want %v", v.AsSlice()[:off], orig.AsSlice()[:off])
+	}
+}

pkg/bufferv2/view_unsafe.go

@@ -0,0 +1,32 @@
+// Copyright 2020 The gVisor 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 buffer
+
+import (
+	"reflect"
+	"unsafe"
+)
+
+// minBatch is the smallest Read or Write operation that the
+// WriteFromReader and ReadToWriter functions will use.
+//
+// This is defined as the size of a native pointer.
+const minBatch = int(unsafe.Sizeof(uintptr(0)))
+
+// BasePtr returns a pointer to the view's chunk.
+func (v *View) BasePtr() *byte {
+	hdr := (*reflect.SliceHeader)(unsafe.Pointer(&v.chunk.data))
+	return (*byte)(unsafe.Pointer(hdr.Data))
+}