From d217b5c69fba5f213d1f19b894f1950c16e492d6 Mon Sep 17 00:00:00 2001 From: Lucas Manning Date: Mon, 11 Jul 2022 10:11:32 -0700 Subject: [PATCH] Refine Buffer io methods to utilize pooled views. PiperOrigin-RevId: 460241285 --- pkg/bufferv2/buffer.go | 187 +++++++++++++++++------------------- pkg/bufferv2/buffer_test.go | 67 +++++++++++++ pkg/bufferv2/chunk.go | 10 +- pkg/bufferv2/view.go | 4 + pkg/bufferv2/view_test.go | 4 +- pkg/bufferv2/view_unsafe.go | 6 -- 6 files changed, 168 insertions(+), 110 deletions(-) diff --git a/pkg/bufferv2/buffer.go b/pkg/bufferv2/buffer.go index 0b05abacd..a4105dbd4 100644 --- a/pkg/bufferv2/buffer.go +++ b/pkg/bufferv2/buffer.go @@ -452,124 +452,117 @@ func (b *Buffer) Merge(other *Buffer) { other.size = 0 } -// WriteFromReader writes to the buffer from an io.Reader. -// -// A minimum read size equal to unsafe.Sizeof(unintptr) is enforced, -// provided that count is greater than or equal to unsafe.Sizeof(uintptr). +// WriteFromReader writes to the buffer from an io.Reader. A maximum read size +// of MaxChunkSize is enforced to prevent allocating views from the heap. func (b *Buffer) WriteFromReader(r io.Reader, count int64) (int64, error) { - var ( - done int64 - n int - err error - ) + var done int64 for done < count { - view := b.data.Back() - - // Ensure we have an empty buffer. - if view.Full() { - view = NewView(int(count - done)) - b.data.PushBack(view) + vsize := count - done + if vsize > MaxChunkSize { + vsize = MaxChunkSize } - - // Is this less than the minimum batch? - if view.AvailableSize() < minBatch && (count-done) >= int64(minBatch) { - tmp := NewView(minBatch) - n, err = r.Read(tmp.availableSlice()) - tmp.Grow(n) - b.Append(tmp) - done += int64(n) - if err != nil { - break - } - continue - } - - // Limit the read, if necessary. - sz := view.AvailableSize() - if left := count - done; int64(sz) > left { - sz = int(left) - } - - // Pass the relevant portion of the buffer. - n, err = r.Read(view.availableSlice()[:sz]) - view.Grow(n) - done += int64(n) - b.size += int64(n) + v := NewView(int(vsize)) + lr := io.LimitedReader{R: r, N: vsize} + n, err := io.Copy(v, &lr) + b.Append(v) + done += n if err == io.EOF { - err = nil // Short write allowed. - break - } else if err != nil { break } + if err != nil { + return done, err + } } - return done, err + return done, nil } // ReadToWriter reads from the buffer into an io.Writer. // // N.B. This does not consume the bytes read. TrimFront should // be called appropriately after this call in order to do so. -// -// A minimum write size equal to unsafe.Sizeof(unintptr) is enforced, -// provided that count is greater than or equal to unsafe.Sizeof(uintptr). func (b *Buffer) ReadToWriter(w io.Writer, count int64) (int64, error) { - var ( - done int64 - n int - err error - ) - offset := 0 // Spill-over for batching. - for view := b.data.Front(); view != nil && done < count; view = view.Next() { - // Has this been consumed? Skip it. - sz := view.Size() - if sz <= offset { - offset -= sz - continue + bytesLeft := int(count) + for v := b.data.Front(); v != nil && bytesLeft > 0; v = v.Next() { + view := v.Clone() + if view.Size() > bytesLeft { + view.CapLength(bytesLeft) } - sz -= offset - - // Is this less than the minimum batch? - left := count - done - if sz < minBatch && left >= int64(minBatch) && (b.size-done) >= int64(minBatch) { - tmp := NewView(minBatch) - n, err = b.ReadAt(tmp.availableSlice()[:minBatch], done) - tmp.Grow(n) - w.Write(tmp.AsSlice()) - tmp.Release() - done += int64(n) - offset = n - sz // Reset below. - if err != nil { - break - } - continue - } - - // Limit the write if necessary. - if int64(sz) >= left { - sz = int(left) - } - - // Perform the actual write. - n, err = w.Write(view.AsSlice()[offset : offset+sz]) - done += int64(n) + n, err := io.Copy(w, view) + bytesLeft -= int(n) + view.Release() if err != nil { - break + return count - int64(bytesLeft), err } - - // Reset spill-over. - offset = 0 } - return done, err + return count - int64(bytesLeft), nil } -// AsSlices returns a list of each of Buffer's underlying Views as Slices. -// The slices returned should not be modifed. -func (b *Buffer) AsSlices() [][]byte { - slices := make([][]byte, 0, b.data.Len()) - for v := b.data.Front(); v != nil; v = v.Next() { - slices = append(slices, v.AsSlice()) +// read implements the io.Reader interface. This method is used by BufferReader +// to consume its underlying buffer. To perform io operations on buffers +// directly, use ReadToWriter or WriteToReader. +func (b *Buffer) read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil } - return slices + if b.Size() == 0 { + return 0, io.EOF + } + done := 0 + v := b.data.Front() + for v != nil && done < len(p) { + n, err := v.Read(p[done:]) + done += n + next := v.Next() + if v.Size() == 0 { + b.removeView(v) + } + b.size -= int64(n) + if err != nil && err != io.EOF { + return done, err + } + v = next + } + return done, nil +} + +// readByte implements the io.ByteReader interface. This method is used by +// BufferReader to consume its underlying buffer. To perform io operations on +// buffers directly, use ReadToWriter or WriteToReader. +func (b *Buffer) readByte() (byte, error) { + if b.Size() == 0 { + return 0, io.EOF + } + v := b.data.Front() + bt := v.AsSlice()[0] + b.TrimFront(1) + return bt, nil +} + +// AsBufferReader returns the Buffer as a BufferReader capabable of io methods. +// The new BufferReader takes ownership of b. +func (b *Buffer) AsBufferReader() BufferReader { + return BufferReader{b} +} + +// BufferReader implements io methods on Buffer. Users must call Close() +// when finished with the buffer to free the underlying memory. +type BufferReader struct { + b *Buffer +} + +// Read implements the io.Reader interface. +func (br *BufferReader) Read(p []byte) (int, error) { + return br.b.read(p) +} + +// ReadByte implements the io.ByteReader interface. +func (br *BufferReader) ReadByte() (byte, error) { + return br.b.readByte() +} + +// Close implements the io.Closer interface. +func (br *BufferReader) Close() { + br.b.Release() } // Range specifies a range of buffer. diff --git a/pkg/bufferv2/buffer_test.go b/pkg/bufferv2/buffer_test.go index a3c44489e..5448d953a 100644 --- a/pkg/bufferv2/buffer_test.go +++ b/pkg/bufferv2/buffer_test.go @@ -630,6 +630,73 @@ func TestBufferPullUp(t *testing.T) { } } +func TestReadFromLargeWriter(t *testing.T) { + writeSize := int64(1 << 20) + largeWriter := bytes.NewBuffer(make([]byte, writeSize)) + b := Buffer{} + // Expect this write to be buffered into several MaxChunkSize sized views. + n, err := b.WriteFromReader(largeWriter, writeSize) + if err != nil { + t.Fatalf("b.WriteFromReader() failed: want err=nil, got %v", err) + } + if n != writeSize { + t.Errorf("got b.WriteFromReader()=%d, want %d", n, writeSize) + } + nChunks := int(writeSize / MaxChunkSize) + if b.data.Len() != nChunks { + t.Errorf("b.WriteFromReader() failed, got b.data.Len()=%d, want %d", b.data.Len(), nChunks) + } +} + +func TestRead(t *testing.T) { + readStrings := []string{"abcdef", "123456", "ghijkl"} + totalSize := len(readStrings) * len(readStrings[0]) + for readSz := 0; readSz < totalSize+1; readSz++ { + b := Buffer{} + for _, s := range readStrings { + v := NewViewWithData([]byte(s)) + b.appendOwned(v) + } + orig := b.Clone() + orig.Truncate(int64(readSz)) + p := make([]byte, readSz) + _, err := b.read(p) + if err != nil { + t.Fatalf("Read([]byte(%d)) failed: %v", readSz, err) + } + if !bytes.Equal(p, orig.Flatten()) { + t.Errorf("Read([]byte(%d)) failed, want p=%v, got %v", readSz, orig.Flatten(), p) + } + if int(b.Size()) != totalSize-readSz { + t.Errorf("Read([]byte(%d)) failed, want b.Size()=%v, got %v", readSz, totalSize-readSz, b.Size()) + } + } +} + +func TestReadByte(t *testing.T) { + readString := "abcdef123456ghijkl" + b := Buffer{} + nViews := 3 + for i := 0; i < nViews; i++ { + vLen := len(readString) / nViews + v := NewViewWithData([]byte(readString[i*vLen : (i+1)*vLen])) + b.appendOwned(v) + } + for i := 0; i < len(readString); i++ { + orig := readString[i] + bt, err := b.readByte() + if err != nil { + t.Fatalf("readByte() failed: %v", err) + } + if bt != orig { + t.Errorf("readByte() failed, want %v, got %v", orig, bt) + } + if int(b.Size()) != len(readString[i+1:]) { + t.Errorf("readByte() failed, want b.Size()=%v, got %v", len(readString[i+1:]), b.Size()) + } + } +} + func TestPullUpModifiedViews(t *testing.T) { var b Buffer defer b.Release() diff --git a/pkg/bufferv2/chunk.go b/pkg/bufferv2/chunk.go index 20758827a..b88eceb2b 100644 --- a/pkg/bufferv2/chunk.go +++ b/pkg/bufferv2/chunk.go @@ -38,9 +38,9 @@ const ( // 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 + // MaxChunkSize is largest payload size that we pool. Payloads larger than + // this will be allocated from the heap and garbage collected as normal. + MaxChunkSize = baseChunkSize << (numPools - 1) // 64k // The number of chunk pools we have for use. numPools = 11 @@ -84,7 +84,7 @@ type chunk struct { func newChunk(size int) *chunk { var c *chunk - if !PoolingEnabled || size > maxChunkSize { + if !PoolingEnabled || size > MaxChunkSize { c = &chunk{ data: make([]byte, size), } @@ -100,7 +100,7 @@ func newChunk(size int) *chunk { } func (c *chunk) destroy() { - if !PoolingEnabled || len(c.data) > maxChunkSize { + if !PoolingEnabled || len(c.data) > MaxChunkSize { c.data = nil return } diff --git a/pkg/bufferv2/view.go b/pkg/bufferv2/view.go index d64820201..7fcd5dbff 100644 --- a/pkg/bufferv2/view.go +++ b/pkg/bufferv2/view.go @@ -241,6 +241,10 @@ func (v *View) ReadFrom(r io.Reader) (n int64, err error) { v.chunk = v.chunk.Clone() } for { + // Check for EOF to avoid an unnnecesary allocation. + if _, e := r.Read(nil); e == io.EOF { + return n, nil + } if v.AvailableSize() == 0 { v.growCap(ReadSize) } diff --git a/pkg/bufferv2/view_test.go b/pkg/bufferv2/view_test.go index fd4a6b405..7a1c6def7 100644 --- a/pkg/bufferv2/view_test.go +++ b/pkg/bufferv2/view_test.go @@ -23,7 +23,7 @@ import ( ) func TestNewView(t *testing.T) { - for sz := baseChunkSize; sz < maxChunkSize; sz <<= 1 { + for sz := baseChunkSize; sz < MaxChunkSize; sz <<= 1 { v := NewView(sz - 1) defer v.Release() @@ -53,7 +53,7 @@ func TestNewView(t *testing.T) { // 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 + viewSize := MaxChunkSize + 1 v := NewView(viewSize) defer v.Release() if v.Capacity() != viewSize { diff --git a/pkg/bufferv2/view_unsafe.go b/pkg/bufferv2/view_unsafe.go index 1be859b4f..d2b6d618e 100644 --- a/pkg/bufferv2/view_unsafe.go +++ b/pkg/bufferv2/view_unsafe.go @@ -19,12 +19,6 @@ import ( "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))