bytes: optimize Buffer's Write, WriteString, WriteByte, and WriteRune

In the common case, the grow method only needs to reslice the internal
buffer. Making another function call to grow can be expensive when Write
is called very often with small pieces of data (like a byte or rune).
Thus, we add a tryGrowByReslice method that is inlineable so that we can
avoid an extra call in most cases.

name                       old time/op    new time/op    delta
WriteByte-4                  35.5µs ± 0%    17.4µs ± 1%   -51.03%  (p=0.000 n=19+20)
WriteRune-4                  55.7µs ± 1%    38.7µs ± 1%   -30.56%  (p=0.000 n=18+19)
BufferNotEmptyWriteRead-4     304µs ± 5%     283µs ± 3%    -6.86%  (p=0.000 n=19+17)
BufferFullSmallReads-4       87.0µs ± 5%    66.8µs ± 2%   -23.26%  (p=0.000 n=17+17)

name                       old speed      new speed      delta
WriteByte-4                 115MB/s ± 0%   235MB/s ± 1%  +104.19%  (p=0.000 n=19+20)
WriteRune-4                 221MB/s ± 1%   318MB/s ± 1%   +44.01%  (p=0.000 n=18+19)

Fixes #17857

Change-Id: I08dfb10a1c7e001817729dbfcc951bda12fe8814
Reviewed-on: https://go-review.googlesource.com/42813
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

src/bytes/buffer.go

@@ -93,6 +93,17 @@ func (b *Buffer) Reset() {
 	b.lastRead = opInvalid
 }
 
+// tryGrowByReslice is a inlineable version of grow for the fast-case where the
+// internal buffer only needs to be resliced.
+// It returns the index where bytes should be written and whether it succeeded.
+func (b *Buffer) tryGrowByReslice(n int) (int, bool) {
+	if l := len(b.buf); l+n <= cap(b.buf) {
+		b.buf = b.buf[:l+n]
+		return l, true
+	}
+	return 0, false
+}
+
 // grow grows the buffer to guarantee space for n more bytes.
 // It returns the index where bytes should be written.
 // If the buffer can't grow it will panic with ErrTooLarge.
@@ -102,27 +113,31 @@ func (b *Buffer) grow(n int) int {
 	if m == 0 && b.off != 0 {
 		b.Reset()
 	}
-	if len(b.buf)+n > cap(b.buf) {
-		var buf []byte
+	// Try to grow by means of a reslice.
+	if i, ok := b.tryGrowByReslice(n); ok {
+		return i
+	}
+	// Check if we can make use of bootstrap array.
 	if b.buf == nil && n <= len(b.bootstrap) {
-			buf = b.bootstrap[0:]
-		} else if m+n <= cap(b.buf)/2 {
+		b.buf = b.bootstrap[:n]
+		return 0
+	}
+	if m+n <= cap(b.buf)/2 {
 		// We can slide things down instead of allocating a new
 		// slice. We only need m+n <= cap(b.buf) to slide, but
 		// we instead let capacity get twice as large so we
 		// don't spend all our time copying.
 		copy(b.buf[:], b.buf[b.off:])
-			buf = b.buf[:m]
 	} else {
-			// not enough space anywhere
-			buf = makeSlice(2*cap(b.buf) + n)
+		// Not enough space anywhere, we need to allocate.
+		buf := makeSlice(2*cap(b.buf) + n)
 		copy(buf, b.buf[b.off:])
-		}
 		b.buf = buf
-		b.off = 0
 	}
-	b.buf = b.buf[0 : b.off+m+n]
-	return b.off + m
+	// Restore b.off and len(b.buf).
+	b.off = 0
+	b.buf = b.buf[:m+n]
+	return m
 }
 
 // Grow grows the buffer's capacity, if necessary, to guarantee space for
@@ -143,7 +158,10 @@ func (b *Buffer) Grow(n int) {
 // buffer becomes too large, Write will panic with ErrTooLarge.
 func (b *Buffer) Write(p []byte) (n int, err error) {
 	b.lastRead = opInvalid
-	m := b.grow(len(p))
+	m, ok := b.tryGrowByReslice(len(p))
+	if !ok {
+		m = b.grow(len(p))
+	}
 	return copy(b.buf[m:], p), nil
 }
 
@@ -152,7 +170,10 @@ func (b *Buffer) Write(p []byte) (n int, err error) {
 // buffer becomes too large, WriteString will panic with ErrTooLarge.
 func (b *Buffer) WriteString(s string) (n int, err error) {
 	b.lastRead = opInvalid
-	m := b.grow(len(s))
+	m, ok := b.tryGrowByReslice(len(s))
+	if !ok {
+		m = b.grow(len(s))
+	}
 	return copy(b.buf[m:], s), nil
 }
 
@@ -244,7 +265,10 @@ func (b *Buffer) WriteTo(w io.Writer) (n int64, err error) {
 // ErrTooLarge.
 func (b *Buffer) WriteByte(c byte) error {
 	b.lastRead = opInvalid
-	m := b.grow(1)
+	m, ok := b.tryGrowByReslice(1)
+	if !ok {
+		m = b.grow(1)
+	}
 	b.buf[m] = c
 	return nil
 }
@@ -259,7 +283,10 @@ func (b *Buffer) WriteRune(r rune) (n int, err error) {
 		return 1, nil
 	}
 	b.lastRead = opInvalid
-	m := b.grow(utf8.UTFMax)
+	m, ok := b.tryGrowByReslice(utf8.UTFMax)
+	if !ok {
+		m = b.grow(utf8.UTFMax)
+	}
 	n = utf8.EncodeRune(b.buf[m:m+utf8.UTFMax], r)
 	b.buf = b.buf[:m+n]
 	return n, nil

src/bytes/buffer_test.go

@@ -6,8 +6,10 @@ package bytes_test
 
 import (
 	. "bytes"
+	"internal/testenv"
 	"io"
 	"math/rand"
+	"os/exec"
 	"runtime"
 	"testing"
 	"unicode/utf8"
@@ -546,6 +548,33 @@ func TestBufferGrowth(t *testing.T) {
 	}
 }
 
+// Test that tryGrowByReslice is inlined.
+func TestTryGrowByResliceInlined(t *testing.T) {
+	t.Parallel()
+	goBin := testenv.GoToolPath(t)
+	out, err := exec.Command(goBin, "tool", "nm", goBin).CombinedOutput()
+	if err != nil {
+		t.Fatalf("go tool nm: %v: %s", err, out)
+	}
+	// Verify this doesn't exist:
+	sym := "bytes.(*Buffer).tryGrowByReslice"
+	if Contains(out, []byte(sym)) {
+		t.Errorf("found symbol %q in cmd/go, but should be inlined", sym)
+	}
+}
+
+func BenchmarkWriteByte(b *testing.B) {
+	const n = 4 << 10
+	b.SetBytes(n)
+	buf := NewBuffer(make([]byte, n))
+	for i := 0; i < b.N; i++ {
+		buf.Reset()
+		for i := 0; i < n; i++ {
+			buf.WriteByte('x')
+		}
+	}
+}
+
 func BenchmarkWriteRune(b *testing.B) {
 	const n = 4 << 10
 	const r = '☺'