Refine Buffer io methods to utilize pooled views.

PiperOrigin-RevId: 460241285

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.

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()

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
 	}

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)
 		}

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 {

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))