Expose Buffer pools to third-party encoders (#376)

Zap currently uses an internal package to house its buffer pool, and it returns
all buffers to that pool. This is problematic for two reasons:

1. First, third-party encoders can easily corrupt the pool by keeping a
   reference to a buffer.
2. Second, third-party encoders can't pool their buffers, since zap needs to
   know how to return them to their pools.

This PR introduces a `buffer.Pool` type to solve this problem. Individual
buffers can only be created via a pool, and each buffer keeps a reference to its
pool of origin. This allows zap to continue using a single internal pool, but it
allows third-party encoders to safely pool their buffers too.

buffer/buffer.go

@@ -27,14 +27,11 @@ import "strconv"
 
 const _size = 1024 // by default, create 1 KiB buffers
 
-// Buffer is a thin wrapper around a byte slice.
+// Buffer is a thin wrapper around a byte slice. It's intended to be pooled, so
+// the only way to construct one is via a Pool.
 type Buffer struct {
-	bs []byte
-}
-
-// New creates a new Buffer of the default size.
-func New() *Buffer {
-	return &Buffer{make([]byte, 0, _size)}
+	bs   []byte
+	pool Pool
 }
 
 // AppendByte writes a single byte to the Buffer.
@@ -100,3 +97,10 @@ func (b *Buffer) Write(bs []byte) (int, error) {
 	b.bs = append(b.bs, bs...)
 	return len(bs), nil
 }
+
+// Free returns the Buffer to its Pool.
+//
+// Callers must not retain references to the Buffer after calling Free.
+func (b *Buffer) Free() {
+	b.pool.put(b)
+}

buffer/buffer_test.go

@@ -29,7 +29,7 @@ import (
 )
 
 func TestBufferWrites(t *testing.T) {
-	buf := New()
+	buf := NewPool().Get()
 
 	tests := []struct {
 		desc string
@@ -69,7 +69,7 @@ func BenchmarkBuffers(b *testing.B) {
 	str := strings.Repeat("a", 1024)
 	slice := make([]byte, 1024)
 	buf := bytes.NewBuffer(slice)
-	custom := New()
+	custom := NewPool().Get()
 	b.Run("ByteSlice", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
 			slice = append(slice, str...)

buffer/pool.go

@@ -0,0 +1,49 @@
+// 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 "sync"
+
+// A Pool is a type-safe wrapper around a sync.Pool.
+type Pool struct {
+	p *sync.Pool
+}
+
+// NewPool constructs a new Pool.
+func NewPool() Pool {
+	return Pool{p: &sync.Pool{
+		New: func() interface{} {
+			return &Buffer{bs: make([]byte, 0, _size)}
+		},
+	}}
+}
+
+// Get retrieves a Buffer from the pool, creating one if necessary.
+func (p Pool) Get() *Buffer {
+	buf := p.p.Get().(*Buffer)
+	buf.Reset()
+	buf.pool = p
+	return buf
+}
+
+func (p Pool) put(buf *Buffer) {
+	p.p.Put(buf)
+}

internal/bufferpool/bufferpool_test.go → buffer/pool_test.go

@@ -18,7 +18,7 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-package bufferpool
+package buffer
 
 import (
 	"sync"
@@ -29,20 +29,21 @@ import (
 
 func TestBuffers(t *testing.T) {
 	const dummyData = "dummy data"
+	p := NewPool()
 
 	var wg sync.WaitGroup
 	for g := 0; g < 10; g++ {
 		wg.Add(1)
 		go func() {
 			for i := 0; i < 100; i++ {
-				buf := Get()
+				buf := p.Get()
 				assert.Zero(t, buf.Len(), "Expected truncated buffer")
 				assert.NotZero(t, buf.Cap(), "Expected non-zero capacity")
 
 				buf.AppendString(dummyData)
 				assert.Equal(t, buf.Len(), len(dummyData), "Expected buffer to contain dummy data")
 
-				Put(buf)
+				buf.Free()
 			}
 			wg.Done()
 		}()

internal/bufferpool/bufferpool.go

@@ -18,30 +18,14 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-// Package bufferpool provides strongly-typed functions to interact with a shared
-// pool of byte buffers.
+// Package bufferpool houses zap's shared internal buffer pool. Third-party
+// packages can recreate the same functionality with buffers.NewPool.
 package bufferpool
 
-import (
-	"sync"
+import "go.uber.org/zap/buffer"
 
-	"go.uber.org/zap/buffer"
+var (
+	_pool = buffer.NewPool()
+	// Get retrieves a buffer from the pool, creating one if necessary.
+	Get = _pool.Get
 )
-
-var _pool = sync.Pool{
-	New: func() interface{} {
-		return buffer.New()
-	},
-}
-
-// 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 *buffer.Buffer) {
-	_pool.Put(buf)
-}

internal/multierror/multierror.go

@@ -38,7 +38,7 @@ func (es errSlice) Error() string {
 		b.AppendString(err.Error())
 	}
 	ret := b.String()
-	bufferpool.Put(b)
+	b.Free()
 	return ret
 }
 

stacktrace.go

@@ -42,7 +42,7 @@ var (
 
 func takeStacktrace() string {
 	buffer := bufferpool.Get()
-	defer bufferpool.Put(buffer)
+	defer buffer.Free()
 	programCounters := _stacktracePool.Get().(*programCounters)
 	defer _stacktracePool.Put(programCounters)
 

zapcore/console_encoder.go

@@ -115,7 +115,7 @@ func (c consoleEncoder) EncodeEntry(ent Entry, fields []Field) (*buffer.Buffer,
 
 func (c consoleEncoder) writeContext(line *buffer.Buffer, extra []Field) {
 	context := c.jsonEncoder.Clone().(*jsonEncoder)
-	defer bufferpool.Put(context.buf)
+	defer context.buf.Free()
 
 	addFields(context, extra)
 	context.closeOpenNamespaces()

zapcore/console_encoder_bench_test.go

@@ -23,7 +23,6 @@ package zapcore_test
 import (
 	"testing"
 
-	"go.uber.org/zap/internal/bufferpool"
 	. "go.uber.org/zap/zapcore"
 )
 
@@ -44,7 +43,7 @@ func BenchmarkZapConsole(b *testing.B) {
 				Message: "fake",
 				Level:   DebugLevel,
 			}, nil)
-			bufferpool.Put(buf)
+			buf.Free()
 		}
 	})
 }

zapcore/core.go

@@ -20,8 +20,6 @@
 
 package zapcore
 
-import "go.uber.org/zap/internal/bufferpool"
-
 // Core is a minimal, fast logger interface. It's designed for library authors
 // to wrap in a more user-friendly API.
 type Core interface {
@@ -90,7 +88,7 @@ func (c *ioCore) Write(ent Entry, fields []Field) error {
 		return err
 	}
 	_, err = c.out.Write(buf.Bytes())
-	bufferpool.Put(buf)
+	buf.Free()
 	if err != nil {
 		return err
 	}

zapcore/entry.go

@@ -92,7 +92,7 @@ func (ec EntryCaller) FullPath() string {
 	buf.AppendByte(':')
 	buf.AppendInt(int64(ec.Line))
 	caller := buf.String()
-	bufferpool.Put(buf)
+	buf.Free()
 	return caller
 }
 
@@ -118,7 +118,7 @@ func (ec EntryCaller) TrimmedPath() string {
 	buf.AppendByte(':')
 	buf.AppendInt(int64(ec.Line))
 	caller := buf.String()
-	bufferpool.Put(buf)
+	buf.Free()
 	return caller
 }
 

zapcore/json_encoder_bench_test.go

@@ -25,7 +25,6 @@ import (
 	"testing"
 	"time"
 
-	"go.uber.org/zap/internal/bufferpool"
 	. "go.uber.org/zap/zapcore"
 )
 
@@ -56,7 +55,7 @@ func BenchmarkZapJSON(b *testing.B) {
 				Message: "fake",
 				Level:   DebugLevel,
 			}, nil)
-			bufferpool.Put(buf)
+			buf.Free()
 		}
 	})
 }