Fix allocation when returning []byte to pool (uber-go#289)

In order to pool byte slices without incurring an allocation (by casting into
interface{}), we need to switch to using a pointer type. The natural candidate
is the standard library's bytes.Buffer; it's a bit slower than straight
slices, but that's probably okay. However, we use strconv.AppendFloat and
friends to avoid allocations in the encoder, and those functions won't work with
byte.Buffer.

Instead, this PR adds a small Buffer type that implements just the bits of
bytes.Buffer and strconv that we need. It renames the private buffers
package to bufferpool to clarify the difference between the two.

Makefile

@@ -3,7 +3,7 @@ export GO15VENDOREXPERIMENT=1
 BENCH_FLAGS ?= -cpuprofile=cpu.pprof -memprofile=mem.pprof -benchmem
 PKGS ?= $(shell glide novendor)
 # Many Go tools take file globs or directories as arguments instead of packages.
-PKG_FILES ?= *.go zapcore benchmarks testutils internal/buffers internal/exit internal/multierror
+PKG_FILES ?= *.go zapcore benchmarks buffer testutils internal/bufferpool internal/exit internal/multierror
 
 # The linting tools evolve with each Go version, so run them only on the latest
 # stable release.

buffer/buffer.go

@@ -0,0 +1,102 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// Package buffer provides a thin wrapper around a byte slice. Unlike the
+// standard library's bytes.Buffer, it supports a portion of the strconv
+// package's zero-allocation formatters.
+package buffer
+
+import "strconv"
+
+const _size = 1024 // by default, create 1 KiB buffers
+
+// Buffer is a thin wrapper around a byte slice.
+type Buffer struct {
+	bs []byte
+}
+
+// New creates a new Buffer of the default size.
+func New() *Buffer {
+	return &Buffer{make([]byte, 0, _size)}
+}
+
+// AppendByte writes a single byte to the Buffer.
+func (b *Buffer) AppendByte(v byte) {
+	b.bs = append(b.bs, v)
+}
+
+// AppendString writes a string to the Buffer.
+func (b *Buffer) AppendString(s string) {
+	b.bs = append(b.bs, s...)
+}
+
+// AppendInt appends an integer to the underlying buffer (assuming base 10).
+func (b *Buffer) AppendInt(i int64) {
+	b.bs = strconv.AppendInt(b.bs, i, 10)
+}
+
+// AppendUint appends an unsigned integer to the underlying buffer (assuming
+// base 10).
+func (b *Buffer) AppendUint(i uint64) {
+	b.bs = strconv.AppendUint(b.bs, i, 10)
+}
+
+// AppendBool appends a bool to the underlying buffer.
+func (b *Buffer) AppendBool(v bool) {
+	b.bs = strconv.AppendBool(b.bs, v)
+}
+
+// AppendFloat appends a float to the underlying buffer. It doesn't quote NaN
+// or +/- Inf.
+func (b *Buffer) AppendFloat(f float64, bitSize int) {
+	b.bs = strconv.AppendFloat(b.bs, f, 'f', -1, bitSize)
+}
+
+// Len returns the length of the underlying byte slice.
+func (b *Buffer) Len() int {
+	return len(b.bs)
+}
+
+// Cap returns the capacity of the underlying byte slice.
+func (b *Buffer) Cap() int {
+	return cap(b.bs)
+}
+
+// Bytes returns a mutable reference to the underlying byte slice.
+func (b *Buffer) Bytes() []byte {
+	return b.bs
+}
+
+// String returns a string copy of the underlying byte slice.
+func (b *Buffer) String() string {
+	return string(b.bs)
+}
+
+// Reset resets the underlying byte slice. Subsequent writes re-use the slice's
+// backing array.
+func (b *Buffer) Reset() {
+	b.bs = b.bs[:0]
+}
+
+// Write implements io.Writer.
+func (b *Buffer) Write(bs []byte) (int, error) {
+	b.bs = append(b.bs, bs...)
+	return len(bs), nil
+}

buffer/buffer_test.go

@@ -0,0 +1,91 @@
+// Copyright (c) 2016 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package buffer
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBufferWrites(t *testing.T) {
+	buf := New()
+
+	tests := []struct {
+		desc string
+		f    func()
+		want string
+	}{
+		{"AppendByte", func() { buf.AppendByte('v') }, "v"},
+		{"AppendString", func() { buf.AppendString("foo") }, "foo"},
+		{"AppendIntPositive", func() { buf.AppendInt(42) }, "42"},
+		{"AppendIntNegative", func() { buf.AppendInt(-42) }, "-42"},
+		{"AppendUint", func() { buf.AppendUint(42) }, "42"},
+		{"AppendBool", func() { buf.AppendBool(true) }, "true"},
+		{"AppendFloat64", func() { buf.AppendFloat(3.14, 64) }, "3.14"},
+		// Intenationally introduce some floating-point error.
+		{"AppendFloat32", func() { buf.AppendFloat(float64(float32(3.14)), 32) }, "3.14"},
+		{"AppendWrite", func() { buf.Write([]byte("foo")) }, "foo"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			buf.Reset()
+			tt.f()
+			assert.Equal(t, tt.want, buf.String(), "Unexpected buffer.String().")
+			assert.Equal(t, tt.want, string(buf.Bytes()), "Unexpected string(buffer.Bytes()).")
+			assert.Equal(t, len(tt.want), buf.Len(), "Unexpected buffer length.")
+			// We're not writing more than a kibibyte in tests.
+			assert.Equal(t, _size, buf.Cap(), "Expected buffer capacity to remain constant.")
+		})
+	}
+}
+
+func BenchmarkBuffers(b *testing.B) {
+	// Because we use the strconv.AppendFoo functions so liberally, we can't
+	// use the standard library's bytes.Buffer anyways (without incurring a
+	// bunch of extra allocations). Nevertheless, let's make sure that we're
+	// not losing any precious nanoseconds.
+	str := strings.Repeat("a", 1024)
+	slice := make([]byte, 1024)
+	buf := bytes.NewBuffer(slice)
+	custom := New()
+	b.Run("ByteSlice", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			slice = append(slice, str...)
+			slice = slice[:0]
+		}
+	})
+	b.Run("BytesBuffer", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			buf.WriteString(str)
+			buf.Reset()
+		}
+	})
+	b.Run("CustomBuffer", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			custom.AppendString(str)
+			custom.Reset()
+		}
+	})
+}

field.go

@@ -25,7 +25,7 @@ import (
 	"math"
 	"time"
 
-	"go.uber.org/zap/internal/buffers"
+	"go.uber.org/zap/internal/bufferpool"
 	"go.uber.org/zap/zapcore"
 )
 
@@ -188,13 +188,14 @@ func Error(err error) zapcore.Field {
 // allocation and takes ~10 microseconds.
 func Stack() zapcore.Field {
 	// Try to avoid allocating a buffer.
-	buf := buffers.Get()
+	buf := bufferpool.Get()
+	bs := buf.Bytes()
 	// Returning the stacktrace as a string costs an allocation, but saves us
 	// from expanding the zapcore.Field union struct to include a byte slice. Since
 	// taking a stacktrace is already so expensive (~10us), the extra allocation
 	// is okay.
-	field := String("stacktrace", takeStacktrace(buf[:cap(buf)], false))
-	buffers.Put(buf)
+	field := String("stacktrace", takeStacktrace(bs[:cap(bs)], false))
+	bufferpool.Put(buf)
 	return field
 }
 

internal/buffers/buffers.go → internal/bufferpool/bufferpool.go

@@ -18,32 +18,30 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-// Package buffers provides strongly-typed functions to interact with a shared
-// pool of byte slices.
-//
-// It's used heavily inside zap, but callers may also take advantage of it;
-// it's particularly useful when implementing json.Marshaler, text.Marshaler,
-// and similar interfaces.
-package buffers
+// Package bufferpool provides strongly-typed functions to interact with a shared
+// pool of byte buffers.
+package bufferpool
 
-import "sync"
+import (
+	"sync"
 
-const _size = 1024 // create 1 KiB buffers
+	"go.uber.org/zap/buffer"
+)
 
 var _pool = sync.Pool{
 	New: func() interface{} {
-		return make([]byte, 0, _size)
+		return buffer.New()
 	},
 }
 
-// Get retrieves a slice from the pool, creating one if necessary.
-// Newly-created slices have a 1 KiB capacity.
-func Get() []byte {
-	buf := _pool.Get().([]byte)
-	return buf[:0]
+// Get retrieves a buffer from the pool, creating one if necessary.
+func Get() *buffer.Buffer {
+	buf := _pool.Get().(*buffer.Buffer)
+	buf.Reset()
+	return buf
 }
 
 // Put returns a slice to the pool.
-func Put(buf []byte) {
+func Put(buf *buffer.Buffer) {
 	_pool.Put(buf)
 }

internal/buffers/buffers_test.go → internal/bufferpool/bufferpool_test.go

@@ -18,7 +18,7 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-package buffers
+package bufferpool
 
 import (
 	"sync"
@@ -34,8 +34,8 @@ func TestBuffers(t *testing.T) {
 		go func() {
 			for i := 0; i < 100; i++ {
 				buf := Get()
-				assert.Zero(t, len(buf), "Expected truncated buffer")
-				assert.NotZero(t, cap(buf), "Expected non-zero capacity")
+				assert.Zero(t, buf.Len(), "Expected truncated buffer")
+				assert.NotZero(t, buf.Cap(), "Expected non-zero capacity")
 				Put(buf)
 			}
 			wg.Done()

internal/multierror/multierror.go

@@ -22,22 +22,23 @@
 // a single error.
 package multierror
 
-import "go.uber.org/zap/internal/buffers"
+import "go.uber.org/zap/internal/bufferpool"
 
 // implement the standard lib's error interface on a private type so that we
 // can't forget to call Error.AsError().
 type errSlice []error
 
 func (es errSlice) Error() string {
-	b := buffers.Get()
+	b := bufferpool.Get()
 	for i, err := range es {
 		if i > 0 {
-			b = append(b, ';', ' ')
+			b.AppendByte(';')
+			b.AppendByte(' ')
 		}
-		b = append(b, err.Error()...)
+		b.AppendString(err.Error())
 	}
-	ret := string(b)
-	buffers.Put(b)
+	ret := b.String()
+	bufferpool.Put(b)
 	return ret
 }