Improve performance/reduce allocs of bytespipe

Creates a `fixedBuffer` type that is used to encapsulate functionality
for reading/writing from the underlying byte slices.

Uses lazily-loaded set of sync.Pools for storing buffers that are no
longer needed so they can be re-used.

```
benchmark                     old ns/op     new ns/op     delta
BenchmarkBytesPipeWrite-8     138469        48985         -64.62%
BenchmarkBytesPipeRead-8      130922        56601         -56.77%

benchmark                     old allocs     new allocs     delta
BenchmarkBytesPipeWrite-8     18             8              -55.56%
BenchmarkBytesPipeRead-8      0              0              +0.00%

benchmark                     old bytes     new bytes     delta
BenchmarkBytesPipeWrite-8     66903         1649          -97.54%
BenchmarkBytesPipeRead-8      0             1             +Inf%
```

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

ioutils/bytespipe.go

@@ -9,12 +9,19 @@ import (
 // maxCap is the highest capacity to use in byte slices that buffer data.
 const maxCap = 1e6
 
+// minCap is the lowest capacity to use in byte slices that buffer data
+const minCap = 64
+
 // blockThreshold is the minimum number of bytes in the buffer which will cause
 // a write to BytesPipe to block when allocating a new slice.
 const blockThreshold = 1e6
 
-// ErrClosed is returned when Write is called on a closed BytesPipe.
-var ErrClosed = errors.New("write to closed BytesPipe")
+var (
+	// ErrClosed is returned when Write is called on a closed BytesPipe.
+	ErrClosed = errors.New("write to closed BytesPipe")
+
+	bufPools = make(map[int]*sync.Pool)
+)
 
 // BytesPipe is io.ReadWriteCloser which works similarly to pipe(queue).
 // All written data may be read at most once. Also, BytesPipe allocates
@@ -23,22 +30,17 @@ var ErrClosed = errors.New("write to closed BytesPipe")
 type BytesPipe struct {
 	mu       sync.Mutex
 	wait     *sync.Cond
-	buf      [][]byte // slice of byte-slices of buffered data
-	lastRead int      // index in the first slice to a read point
-	bufLen   int      // length of data buffered over the slices
-	closeErr error    // error to return from next Read. set to nil if not closed.
+	buf      []*fixedBuffer
+	bufLen   int
+	closeErr error // error to return from next Read. set to nil if not closed.
 }
 
 // NewBytesPipe creates new BytesPipe, initialized by specified slice.
 // If buf is nil, then it will be initialized with slice which cap is 64.
 // buf will be adjusted in a way that len(buf) == 0, cap(buf) == cap(buf).
-func NewBytesPipe(buf []byte) *BytesPipe {
-	if cap(buf) == 0 {
-		buf = make([]byte, 0, 64)
-	}
-	bp := &BytesPipe{
-		buf: [][]byte{buf[:0]},
-	}
+func NewBytesPipe() *BytesPipe {
+	bp := &BytesPipe{}
+	bp.buf = append(bp.buf, getBuffer(minCap))
 	bp.wait = sync.NewCond(&bp.mu)
 	return bp
 }
@@ -47,22 +49,30 @@ func NewBytesPipe(buf []byte) *BytesPipe {
 // It can allocate new []byte slices in a process of writing.
 func (bp *BytesPipe) Write(p []byte) (int, error) {
 	bp.mu.Lock()
-	defer bp.mu.Unlock()
+
 	written := 0
 	for {
 		if bp.closeErr != nil {
+			bp.mu.Unlock()
 			return written, ErrClosed
 		}
-		// write data to the last buffer
-		b := bp.buf[len(bp.buf)-1]
-		// copy data to the current empty allocated area
-		n := copy(b[len(b):cap(b)], p)
-		// increment buffered data length
-		bp.bufLen += n
-		// include written data in last buffer
-		bp.buf[len(bp.buf)-1] = b[:len(b)+n]
 
+		if len(bp.buf) == 0 {
+			bp.buf = append(bp.buf, getBuffer(64))
+		}
+		// get the last buffer
+		b := bp.buf[len(bp.buf)-1]
+
+		n, err := b.Write(p)
 		written += n
+		bp.bufLen += n
+
+		// errBufferFull is an error we expect to get if the buffer is full
+		if err != nil && err != errBufferFull {
+			bp.wait.Broadcast()
+			bp.mu.Unlock()
+			return written, err
+		}
 
 		// if there was enough room to write all then break
 		if len(p) == n {
@@ -72,20 +82,20 @@ func (bp *BytesPipe) Write(p []byte) (int, error) {
 		// more data: write to the next slice
 		p = p[n:]
 
-		// block if too much data is still in the buffer
+		// make sure the buffer doesn't grow too big from this write
 		for bp.bufLen >= blockThreshold {
 			bp.wait.Wait()
 		}
 
-		// allocate slice that has twice the size of the last unless maximum reached
-		nextCap := 2 * cap(bp.buf[len(bp.buf)-1])
+		// add new byte slice to the buffers slice and continue writing
+		nextCap := b.Cap() * 2
 		if nextCap > maxCap {
 			nextCap = maxCap
 		}
-		// add new byte slice to the buffers slice and continue writing
-		bp.buf = append(bp.buf, make([]byte, 0, nextCap))
+		bp.buf = append(bp.buf, getBuffer(nextCap))
 	}
 	bp.wait.Broadcast()
+	bp.mu.Unlock()
 	return written, nil
 }
 
@@ -107,46 +117,60 @@ func (bp *BytesPipe) Close() error {
 	return bp.CloseWithError(nil)
 }
 
-func (bp *BytesPipe) len() int {
-	return bp.bufLen - bp.lastRead
-}
-
 // Read reads bytes from BytesPipe.
 // Data could be read only once.
 func (bp *BytesPipe) Read(p []byte) (n int, err error) {
 	bp.mu.Lock()
-	defer bp.mu.Unlock()
-	if bp.len() == 0 {
+	if bp.bufLen == 0 {
 		if bp.closeErr != nil {
+			bp.mu.Unlock()
 			return 0, bp.closeErr
 		}
 		bp.wait.Wait()
-		if bp.len() == 0 && bp.closeErr != nil {
+		if bp.bufLen == 0 && bp.closeErr != nil {
+			bp.mu.Unlock()
 			return 0, bp.closeErr
 		}
 	}
-	for {
-		read := copy(p, bp.buf[0][bp.lastRead:])
+
+	for bp.bufLen > 0 {
+		b := bp.buf[0]
+		read, _ := b.Read(p) // ignore error since fixedBuffer doesn't really return an error
 		n += read
-		bp.lastRead += read
-		if bp.len() == 0 {
-			// we have read everything. reset to the beginning.
-			bp.lastRead = 0
-			bp.bufLen -= len(bp.buf[0])
-			bp.buf[0] = bp.buf[0][:0]
-			break
+		bp.bufLen -= read
+
+		if b.Len() == 0 {
+			// it's empty so return it to the pool and move to the next one
+			returnBuffer(b)
+			bp.buf[0] = nil
+			bp.buf = bp.buf[1:]
 		}
-		// break if everything was read
+
 		if len(p) == read {
 			break
 		}
-		// more buffered data and more asked. read from next slice.
+
 		p = p[read:]
-		bp.lastRead = 0
-		bp.bufLen -= len(bp.buf[0])
-		bp.buf[0] = nil     // throw away old slice
-		bp.buf = bp.buf[1:] // switch to next
 	}
+
 	bp.wait.Broadcast()
+	bp.mu.Unlock()
 	return
 }
+
+func returnBuffer(b *fixedBuffer) {
+	b.Reset()
+	pool := bufPools[b.Cap()]
+	if pool != nil {
+		pool.Put(b)
+	}
+}
+
+func getBuffer(size int) *fixedBuffer {
+	pool, ok := bufPools[size]
+	if !ok {
+		pool = &sync.Pool{New: func() interface{} { return &fixedBuffer{buf: make([]byte, 0, size)} }}
+		bufPools[size] = pool
+	}
+	return pool.Get().(*fixedBuffer)
+}