From efdc289a96e5bae0b5a28d705b86a711f8458d04 Mon Sep 17 00:00:00 2001 From: Lucas Manning Date: Fri, 24 Jun 2022 13:46:42 -0700 Subject: [PATCH] Add nil receiver and io methods to bufferv2. Allowing for nil receivers in read-like operations lets views be treated more like slices, which is easier for the programmer generally. io methods are useful for copy operations to avoid unnecessary allocations. PiperOrigin-RevId: 457086370 --- pkg/bufferv2/buffer.go | 30 +++++----- pkg/bufferv2/buffer_test.go | 21 ++++++- pkg/bufferv2/view.go | 115 +++++++++++++++++++++++++++++++++--- pkg/bufferv2/view_test.go | 91 ++++++++++++++++++++++------ 4 files changed, 216 insertions(+), 41 deletions(-) diff --git a/pkg/bufferv2/buffer.go b/pkg/bufferv2/buffer.go index ea961ab18..881e2f2d7 100644 --- a/pkg/bufferv2/buffer.go +++ b/pkg/bufferv2/buffer.go @@ -294,13 +294,13 @@ func (b *Buffer) prependOwned(v *View) { } // PullUp makes the specified range contiguous and returns the backing memory. -func (b *Buffer) PullUp(offset, length int) (*View, bool) { +func (b *Buffer) PullUp(offset, length int) (View, bool) { if length == 0 { - return nil, true + return View{}, true } tgt := Range{begin: offset, end: offset + length} if tgt.Intersect(Range{end: int(b.size)}).Len() != length { - return nil, false + return View{}, false } curr := Range{} @@ -312,12 +312,13 @@ func (b *Buffer) PullUp(offset, length int) (*View, bool) { if x := curr.Intersect(tgt); x.Len() == tgt.Len() { // buf covers the whole requested target range. sub := x.Offset(-curr.begin) - new := viewPool.Get().(*View) - new.read = sub.begin - new.write = sub.end // Don't increment the reference count of the underlying chunk. Views // returned by PullUp are explicitly unowned and read only - new.chunk = v.chunk + new := View{ + read: v.read + sub.begin, + write: v.read + sub.end, + chunk: v.chunk, + } return new, true } else if x.Len() > 0 { // buf is pointing at the starting buffer we want to merge. @@ -357,10 +358,11 @@ func (b *Buffer) PullUp(offset, length int) (*View, bool) { b.removeView(v) r := tgt.Offset(-curr.begin) - pulled := viewPool.Get().(*View) - pulled.read = r.begin - pulled.write = r.end - pulled.chunk = merged.chunk + pulled := View{ + read: r.begin, + write: r.end, + chunk: merged.chunk, + } return pulled, true } @@ -418,11 +420,11 @@ func (b *Buffer) Apply(fn func(*View)) { // outside of b is ignored. func (b *Buffer) SubApply(offset, length int, fn func(*View)) { for v := b.data.Front(); length > 0 && v != nil; v = v.Next() { - d := v.Clone() - if offset >= d.Size() { - offset -= d.Size() + if offset >= v.Size() { + offset -= v.Size() continue } + d := v.Clone() if offset > 0 { d.TrimFront(offset) offset = 0 diff --git a/pkg/bufferv2/buffer_test.go b/pkg/bufferv2/buffer_test.go index a5a0a2ea1..c362a16f5 100644 --- a/pkg/bufferv2/buffer_test.go +++ b/pkg/bufferv2/buffer_test.go @@ -612,7 +612,7 @@ func TestBufferPullUp(t *testing.T) { got, gotOk := b.PullUp(tc.offset, tc.length) want, wantOk := []byte(tc.output), !tc.failed - if gotOk == wantOk && got == nil && len(want) == 0 { + if gotOk == wantOk && got.Size() == 0 && len(want) == 0 { return } if gotOk != wantOk || !bytes.Equal(got.AsSlice(), want) { @@ -630,6 +630,25 @@ func TestBufferPullUp(t *testing.T) { } } +func TestPullUpModifiedViews(t *testing.T) { + var b Buffer + defer b.Release() + for _, s := range []string{"abcdef", "123456", "ghijkl"} { + v := NewViewWithData([]byte(s)) + v.TrimFront(3) + b.appendOwned(v) + } + + v, ok := b.PullUp(3, 3) + if !ok { + t.Errorf("PullUp failed: want ok=true, got ok=false") + } + want := []byte("456") + if !bytes.Equal(v.AsSlice(), want) { + t.Errorf("PullUp failed: want %v, got %v", want, v.AsSlice()) + } +} + func TestBufferClone(t *testing.T) { const ( originalSize = 90 diff --git a/pkg/bufferv2/view.go b/pkg/bufferv2/view.go index 659493bcc..6e7d34928 100644 --- a/pkg/bufferv2/view.go +++ b/pkg/bufferv2/view.go @@ -21,6 +21,10 @@ import ( "gvisor.dev/gvisor/pkg/sync" ) +// ReadSize is the default amount that a View's size is increased by when an +// io.Reader has more data than a View can hold during calls to ReadFrom. +const ReadSize = 512 + var viewPool = sync.Pool{ New: func() interface{} { return &View{} @@ -42,8 +46,6 @@ var viewPool = sync.Pool{ // 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 @@ -84,6 +86,9 @@ func NewViewWithData(data []byte) *View { // 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 { + if v == nil { + panic("cannot clone a nil view") + } v.chunk.IncRef() newV := viewPool.Get().(*View) newV.chunk = v.chunk @@ -94,6 +99,9 @@ func (v *View) Clone() *View { // Release releases the chunk held by v and returns v to the pool. func (v *View) Release() { + if v == nil { + panic("cannot release a nil view") + } v.chunk.DecRef() *v = View{} viewPool.Put(v) @@ -112,16 +120,25 @@ func (v *View) Full() bool { // 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) } // Size returns the size of data written to the view. func (v *View) Size() int { + if v == nil { + return 0 + } return v.write - v.read } // TrimFront advances the read index by the given amount. func (v *View) TrimFront(n int) { + if v.read+n > v.write { + panic("cannot trim past the end of a view") + } v.read += n } @@ -135,6 +152,9 @@ func (v *View) AsSlice() []byte { // AvailableSize returns the number of bytes available for writing. func (v *View) AvailableSize() int { + if v == nil { + return 0 + } return len(v.chunk.data) - v.write } @@ -153,6 +173,26 @@ func (v *View) Read(p []byte) (int, error) { return n, nil } +// WriteTo writes data to w until the view is empty or an error occurs. The +// return value n is the number of bytes written. +// +// WriteTo implements the io.WriterTo interface. +func (v *View) WriteTo(w io.Writer) (n int64, err error) { + if v.Size() > 0 { + sz := v.Size() + m, e := w.Write(v.AsSlice()) + v.TrimFront(m) + n = int64(m) + if e != nil { + return n, e + } + if m != sz { + return n, io.ErrShortWrite + } + } + return n, nil +} + // ReadAt reads data to the p starting at offset. // // Implements the io.ReaderAt interface. @@ -170,24 +210,62 @@ func (v *View) ReadAt(p []byte, off int) (int, error) { // // Implements the io.Writer interface. func (v *View) Write(p []byte) (int, error) { - if v.sharesChunk() { + if v == nil { + panic("cannot write to a nil view") + } + if v.AvailableSize() < len(p) { + v.growCap(len(p) - v.AvailableSize()) + } else 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, io.ErrShortWrite } return n, nil } +// ReadFrom reads data from r until EOF and appends it to the buffer, growing +// the buffer as needed. The return value n is the number of bytes read. Any +// error except io.EOF encountered during the read is also returned. +// +// ReadFrom implements the io.ReaderFrom interface. +func (v *View) ReadFrom(r io.Reader) (n int64, err error) { + if v == nil { + panic("cannot write to a nil view") + } + if v.sharesChunk() { + defer v.chunk.DecRef() + v.chunk = v.chunk.Clone() + } + for { + if v.AvailableSize() == 0 { + v.growCap(ReadSize) + } + m, e := r.Read(v.availableSlice()) + v.write += m + n += int64(m) + + if e == io.EOF { + return n, nil + } + if e != nil { + return n, e + } + } +} + // 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 v == nil { + panic("cannot write to a nil view") + } 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) } @@ -197,22 +275,43 @@ func (v *View) WriteAt(p []byte, off int) (int, error) { } 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, io.ErrShortWrite } return n, nil } -// Grow advances the write index by the given amount. +// Grow increases the size of the view. If the new size is greater than the +// view's current capacity, Grow will reallocate the view with an increased +// capacity. func (v *View) Grow(n int) { - if n+v.write > v.Capacity() { - panic("cannot grow view past capacity") + if v == nil { + panic("cannot grow a nil view") + } + if v.write+n > v.Capacity() { + v.growCap(n) } v.write += n } +// growCap increases the capacity of the view by at least n. +func (v *View) growCap(n int) { + if v == nil { + panic("cannot grow a nil view") + } + defer v.chunk.DecRef() + old := v.AsSlice() + v.chunk = newChunk(v.Capacity() + n) + copy(v.chunk.data, old) + v.read = 0 + v.write = len(old) +} + // 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 v == nil { + panic("cannot resize a nil view") + } if n < 0 { panic("n must be >= 0") } diff --git a/pkg/bufferv2/view_test.go b/pkg/bufferv2/view_test.go index 25eda9e22..aae67178c 100644 --- a/pkg/bufferv2/view_test.go +++ b/pkg/bufferv2/view_test.go @@ -15,6 +15,7 @@ package buffer import ( + "bytes" "math/rand" "testing" @@ -80,7 +81,7 @@ func TestWrite(t *testing.T) { for _, tc := range []struct { name string view *View - initData []byte + initSize int writeSize int }{ { @@ -91,45 +92,41 @@ func TestWrite(t *testing.T) { { name: "full view", view: NewView(100), - initData: make([]byte, 100), + initSize: 100, writeSize: 50, }, { name: "full write to partially full view", view: NewView(100), - initData: make([]byte, 20), + initSize: 20, writeSize: 50, }, { name: "partial write to partially full view", view: NewView(100), - initData: make([]byte, 80), + initSize: 80, writeSize: 50, }, } { t.Run(tc.name, func(t *testing.T) { - tc.view.Write(tc.initData) + tc.view.Grow(tc.initSize) defer tc.view.Release() - origWriteSize := tc.view.AvailableSize() - - var orig []byte - orig = append(orig, tc.view.AsSlice()...) + orig := append([]byte(nil), 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) + n, err := tc.view.Write(toWrite) + if err != nil { + t.Errorf("Write failed: %s", err) } - if tc.writeSize > origWriteSize { - toWrite = toWrite[:origWriteSize] + if n != tc.writeSize { + t.Errorf("got n=%d, want %d", n, tc.writeSize) } - 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 tc.view.Size() != len(orig)+tc.writeSize { + t.Errorf("got Size()=%d, want %d", tc.view.Size(), len(orig)+tc.writeSize) } if !cmp.Equal(tc.view.AsSlice(), append(orig, toWrite...)) { - t.Errorf("got tc.view.ReadSlice() = %d, want %d", tc.view.AsSlice(), toWrite) + t.Errorf("got tc.view.AsSlice() = %d, want %d", tc.view.AsSlice(), toWrite) } }) } @@ -172,3 +169,61 @@ func TestWriteAt(t *testing.T) { t.Errorf("got v.AsSlice()[:off] = %v, want %v", v.AsSlice()[:off], orig.AsSlice()[:off]) } } + +func TestWriteTo(t *testing.T) { + writeToSize := 100 + v := NewViewSize(writeToSize) + defer v.Release() + + w := bytes.NewBuffer(make([]byte, 100)) + + n, err := v.WriteTo(w) + if err != nil { + t.Errorf("WriteTo failed: %s", err) + } + if n != int64(writeToSize) { + t.Errorf("got n=%d, want 100", n) + } + if v.Size() != 0 { + t.Errorf("got v.Size()=%d, want 0", v.Size()) + } +} + +func TestReadFrom(t *testing.T) { + for _, tc := range []struct { + name string + data []byte + view *View + }{ + { + name: "basic", + data: []byte{1, 2, 3}, + view: NewView(10), + }, + { + name: "requires grow", + data: []byte{4, 5, 6}, + view: NewViewSize(63), + }, + } { + defer tc.view.Release() + clone := tc.view.Clone() + defer clone.Release() + r := bytes.NewReader(tc.data) + + n, err := tc.view.ReadFrom(r) + + if err != nil { + t.Errorf("v.ReadFrom failed: %s", err) + } + if int(n) != len(tc.data) { + t.Errorf("v.ReadFrom failed: want n=%d, got %d", len(tc.data), n) + } + if tc.view.Size() == clone.Size() { + t.Errorf("expected clone.Size() != v.Size(), got match") + } + if !bytes.Equal(tc.view.AsSlice(), append(clone.AsSlice(), tc.data...)) { + t.Errorf("v.ReadFrom failed: want %v, got %v", tc.data, tc.view.AsSlice()) + } + } +}