Support externally owned buffer chunks

pkg/buffer/buffer.go

@@ -28,7 +28,7 @@ import (
 //
 // +stateify savable
 type Buffer struct {
-	data ViewList `state:".([]byte)"`
+	data ViewList `state:".([]savedView)"`
 	size int64
 }
 

pkg/buffer/buffer_state.go

@@ -18,12 +18,44 @@ import (
 	"context"
 )
 
+// savedView keeps heap data compact while preserving external chunk ownership
+// and sharing. It stores either live heap bytes in data or an externally backed
+// View in external, never both.
+//
+// +stateify savable
+type savedView struct {
+	data     []byte
+	external *View
+}
+
 // saveData is invoked by stateify.
-func (b *Buffer) saveData() []byte {
-	return b.Flatten()
+func (b *Buffer) saveData() []savedView {
+	var views []savedView
+	for v := b.data.Front(); v != nil; v = v.Next() {
+		if v.chunk.external != nil {
+			// Save the owned View itself, without adding a chunk reference.
+			// Its storage implementation owns saving the external bytes.
+			views = append(views, savedView{external: v})
+		} else {
+			// A subslice would retain unused capacity and heap chunk aliases.
+			// Keep only the live bytes, as the flattened representation did.
+			views = append(views, savedView{data: v.ToSlice()})
+		}
+	}
+	return views
 }
 
 // loadData is invoked by stateify.
-func (b *Buffer) loadData(_ context.Context, data []byte) {
-	*b = MakeWithData(data)
+func (b *Buffer) loadData(_ context.Context, views []savedView) {
+	*b = Buffer{}
+	for _, saved := range views {
+		v := saved.external
+		if v == nil {
+			v = NewViewWithData(saved.data)
+		}
+		// Restore list ownership without acquiring another chunk reference or
+		// accessing external bytes that may not have been restored yet. Append
+		// could copy data or release an empty View, so only rebuild the links.
+		b.appendOwned(v)
+	}
 }

pkg/buffer/buffer_test.go

@@ -819,6 +819,170 @@ func TestGrowBufferCloned(t *testing.T) {
 	}
 }
 
+// externalTestStorage models an owner whose bytes are not available until its
+// restore lifecycle finishes, after state.Load has reconstructed the buffers.
+type externalTestStorage struct {
+	data     []byte
+	ready    bool
+	releases int
+}
+
+func (s *externalTestStorage) Bytes() []byte {
+	if !s.ready {
+		panic("external storage is not ready")
+	}
+	return s.data
+}
+
+func (s *externalTestStorage) Release() {
+	s.releases++
+	s.ready = false
+}
+
+func (*externalTestStorage) StateTypeName() string {
+	return "gvisor.dev/gvisor/pkg/buffer.externalTestStorage"
+}
+
+func (*externalTestStorage) StateFields() []string {
+	return []string{"data"}
+}
+
+func (s *externalTestStorage) StateSave(sink state.Sink) {
+	sink.Save(0, &s.data)
+}
+
+func (s *externalTestStorage) StateLoad(_ context.Context, source state.Source) {
+	source.Load(0, &s.data)
+}
+
+func init() {
+	state.Register((*externalTestStorage)(nil))
+}
+
+func TestSaveRestoreExternalBuffers(t *testing.T) {
+	storage := &externalTestStorage{data: bytes.Repeat([]byte("0123456789"), 1024), ready: true}
+	v := NewViewWithExternalStorage(storage)
+	v.TrimFront(3)
+	v.CapLength(8)
+	var original Buffer
+	original.appendOwned(NewViewWithData([]byte("head")))
+	original.appendOwned(v)
+	original.appendOwned(NewViewWithData([]byte("tail")))
+	defer original.Release()
+	shared := original.Clone()
+	defer shared.Release()
+	written := original.Clone()
+	defer written.Release()
+	if _, err := written.data.Front().Next().WriteAt([]byte("X"), 0); err != nil {
+		t.Fatal(err)
+	}
+
+	toSave := []*Buffer{&original, &shared, &written}
+	var encoded bytes.Buffer
+	ctx := t.Context()
+	if _, err := state.Save(ctx, &encoded, &toSave); err != nil {
+		t.Fatal(err)
+	}
+	var restored []*Buffer
+	if _, err := state.Load(ctx, &encoded, &restored); err != nil {
+		t.Fatal(err)
+	}
+	if len(restored) != len(toSave) {
+		t.Fatalf("restored %d buffers, want %d", len(restored), len(toSave))
+	}
+	for _, b := range restored {
+		defer b.Release()
+	}
+	first := restored[0].data.Front().Next()
+	second := restored[1].data.Front().Next()
+	owner := first.chunk.external.(*externalTestStorage)
+	if first.chunk != second.chunk {
+		t.Fatal("restored external views do not share their chunk")
+	}
+	if owner.ready || owner.releases != 0 {
+		t.Fatalf("storage touched during load: ready=%t, releases=%d", owner.ready, owner.releases)
+	}
+	owner.ready = true
+	for i, want := range []string{"head34567890tail", "head34567890tail", "headX4567890tail"} {
+		if got := string(restored[i].Flatten()); got != want {
+			t.Errorf("buffer %d = %q, want %q", i, got, want)
+		}
+	}
+	if _, err := first.WriteAt([]byte("Y"), 0); err != nil {
+		t.Fatal(err)
+	}
+	if got := string(second.AsSlice()); got != "34567890" {
+		t.Fatalf("restored COW changed shared bytes to %q", got)
+	}
+	restored[0].Release()
+	if owner.releases != 0 {
+		t.Fatalf("storage released with a surviving view: %d releases", owner.releases)
+	}
+	restored[1].Release()
+	restored[2].Release()
+	if owner.releases != 1 {
+		t.Fatalf("final storage releases = %d, want 1", owner.releases)
+	}
+}
+
+func TestSaveRestoreEmptyExternalView(t *testing.T) {
+	storage := &externalTestStorage{data: []byte("x"), ready: true}
+	v := NewViewWithExternalStorage(storage)
+	v.CapLength(0)
+	var original Buffer
+	original.appendOwned(v)
+	defer original.Release()
+	var restored Buffer
+	doSaveAndLoad(t, &original, &restored)
+	owner := restored.data.Front().chunk.external.(*externalTestStorage)
+	if owner.ready || owner.releases != 0 {
+		t.Fatal("loading released or accessed empty external storage")
+	}
+	restored.Release()
+	if owner.releases != 1 {
+		t.Fatalf("storage releases = %d, want 1", owner.releases)
+	}
+}
+
+func TestSaveRestoreTrimmedHeapBuffer(t *testing.T) {
+	b := MakeWithView(NewViewWithData(bytes.Repeat([]byte("x"), MaxChunkSize)))
+	defer b.Release()
+	b.TrimFront(MaxChunkSize - 1)
+	var encoded bytes.Buffer
+	ctx := t.Context()
+	if _, err := state.Save(ctx, &encoded, &b); err != nil {
+		t.Fatal(err)
+	}
+	// The snapshot should not retain the unused 64 KiB heap chunk capacity.
+	if got := encoded.Len(); got >= MaxChunkSize/2 {
+		t.Fatalf("one-byte buffer snapshot size = %d, retains unused chunk capacity", got)
+	}
+	var restored Buffer
+	if _, err := state.Load(ctx, &encoded, &restored); err != nil {
+		t.Fatal(err)
+	}
+	defer restored.Release()
+	if got := string(restored.Flatten()); got != "x" {
+		t.Fatalf("restored buffer = %q, want x", got)
+	}
+}
+
+func TestGrowExternalBuffer(t *testing.T) {
+	storage := &externalTestStorage{data: []byte("abcd"), ready: true}
+	original := MakeWithView(NewViewWithExternalStorage(storage))
+	defer original.Release()
+	clone := original.Clone()
+	defer clone.Release()
+	clone.Truncate(2)
+	clone.GrowTo(4, true)
+	if got := string(clone.Flatten()); got != "ab\x00\x00" {
+		t.Errorf("grown clone = %q, want ab followed by two zero bytes", got)
+	}
+	if got := string(original.Flatten()); got != "abcd" {
+		t.Errorf("growing clone changed original to %q", got)
+	}
+}
+
 func TestRangeIntersect(t *testing.T) {
 	for _, tc := range []struct {
 		desc       string

pkg/buffer/chunk.go

@@ -70,12 +70,41 @@ func getChunkPool(size int) *sync.Pool {
 	return &chunkPools[idx]
 }
 
-// Chunk represents a slice of pooled memory.
+// ExternalStorage owns memory supplied to a View by another package. The caller
+// must not otherwise access the memory after transferring ownership to a View.
+//
+// To support checkpointing, the concrete implementation must be savable and
+// own restoration of the backing memory.
+type ExternalStorage interface {
+	// Bytes returns the same writable, nonempty slice until Release. Its length
+	// must not exceed MaxChunkSize, which bounds the heap copy when a shared
+	// external chunk detaches on write. Access must not require a fallible
+	// operation.
+	//
+	// Bytes must be usable before restored buffers are accessed. Loading a
+	// Buffer does not call Bytes.
+	Bytes() []byte
+
+	// Release relinquishes the backing memory. It is called exactly once, when
+	// the last chunk reference is released, and must not require calling Bytes.
+	Release()
+}
+
+// chunk represents reference-counted heap or externally owned memory. It holds
+// either heap bytes in data or an owner in external, never both.
 //
 // +stateify savable
 type chunk struct {
 	chunkRefs
-	data []byte
+	data     []byte
+	external ExternalStorage
+}
+
+func (c *chunk) bytes() []byte {
+	if c.external != nil {
+		return c.external.Bytes()
+	}
+	return c.data
 }
 
 func newChunk(size int) *chunk {
@@ -94,6 +123,11 @@ func newChunk(size int) *chunk {
 }
 
 func (c *chunk) destroy() {
+	if c.external != nil {
+		c.external.Release()
+		c.external = nil
+		return
+	}
 	if len(c.data) > MaxChunkSize {
 		c.data = nil
 		return
@@ -107,7 +141,8 @@ func (c *chunk) DecRef() {
 }
 
 func (c *chunk) Clone() *chunk {
-	cpy := newChunk(len(c.data))
-	copy(cpy.data, c.data)
+	data := c.bytes()
+	cpy := newChunk(len(data))
+	copy(cpy.data, data)
 	return cpy
 }

pkg/buffer/view.go

@@ -34,7 +34,7 @@ var viewPool = sync.Pool{
 // 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
+// A View must be created with a NewView function 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.
@@ -83,6 +83,23 @@ func NewViewWithData(data []byte) *View {
 	return v
 }
 
+// NewViewWithExternalStorage creates a fully populated View without copying
+// storage's bytes. On success it takes ownership of storage, which must satisfy
+// the ExternalStorage contract. An invalid length panics without transferring
+// ownership. Each storage owner may be passed successfully to this function only
+// once; use Clone to share its chunk with another View.
+func NewViewWithExternalStorage(storage ExternalStorage) *View {
+	data := storage.Bytes()
+	if len(data) == 0 || len(data) > MaxChunkSize {
+		panic(fmt.Sprintf("external storage length %d is outside [1, %d]", len(data), MaxChunkSize))
+	}
+	c := &chunk{external: storage}
+	c.InitRefs()
+	v := viewPool.Get().(*View)
+	*v = View{write: len(data), chunk: c}
+	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
@@ -126,15 +143,15 @@ func (v *View) sharesChunk() bool {
 //
 // This indicates there is no capacity left to write.
 func (v *View) Full() bool {
-	return v == nil || v.write == len(v.chunk.data)
+	return v == nil || v.write == len(v.chunk.bytes())
 }
 
 // Capacity returns the total size of this view's chunk.
 func (v *View) Capacity() int {
 	if v == nil {
 		return 0
 	}
-	return len(v.chunk.data)
+	return len(v.chunk.bytes())
 }
 
 // Size returns the size of data written to the view.
@@ -158,7 +175,7 @@ func (v *View) AsSlice() []byte {
 	if v.Size() == 0 {
 		return nil
 	}
-	return v.chunk.data[v.read:v.write]
+	return v.chunk.bytes()[v.read:v.write]
 }
 
 // ToSlice returns an owned copy of the data in this view.
@@ -176,7 +193,7 @@ func (v *View) AvailableSize() int {
 	if v == nil {
 		return 0
 	}
-	return len(v.chunk.data) - v.write
+	return len(v.chunk.bytes()) - v.write
 }
 
 // Read reads v's data into p.
@@ -250,7 +267,7 @@ func (v *View) Write(p []byte) (int, error) {
 		defer v.chunk.DecRef()
 		v.chunk = v.chunk.Clone()
 	}
-	n := copy(v.chunk.data[v.write:], p)
+	n := copy(v.chunk.bytes()[v.write:], p)
 	v.write += n
 	if n < len(p) {
 		return n, io.ErrShortWrite
@@ -362,5 +379,5 @@ func (v *View) availableSlice() []byte {
 		c := v.chunk.Clone()
 		v.chunk = c
 	}
-	return v.chunk.data[v.write:]
+	return v.chunk.bytes()[v.write:]
 }