From c0dfce3bb4a0d64523ce6bbb2c9e197a3d3ea861 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Tue, 14 Mar 2017 16:08:57 -0700 Subject: [PATCH] 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 | 18 ++++--- buffer/buffer_test.go | 4 +- buffer/pool.go | 49 +++++++++++++++++++ .../bufferpool_test.go => buffer/pool_test.go | 7 +-- internal/bufferpool/bufferpool.go | 30 +++--------- internal/multierror/multierror.go | 2 +- stacktrace.go | 2 +- zapcore/console_encoder.go | 2 +- zapcore/console_encoder_bench_test.go | 3 +- zapcore/core.go | 4 +- zapcore/entry.go | 4 +- zapcore/json_encoder_bench_test.go | 3 +- 12 files changed, 81 insertions(+), 47 deletions(-) create mode 100644 buffer/pool.go rename internal/bufferpool/bufferpool_test.go => buffer/pool_test.go (96%) diff --git a/buffer/buffer.go b/buffer/buffer.go index e5015841d..ea6fdc8d9 100644 --- a/buffer/buffer.go +++ b/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) +} diff --git a/buffer/buffer_test.go b/buffer/buffer_test.go index c22b2bd32..59bc08a6a 100644 --- a/buffer/buffer_test.go +++ 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...) diff --git a/buffer/pool.go b/buffer/pool.go new file mode 100644 index 000000000..8fb3e202c --- /dev/null +++ b/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) +} diff --git a/internal/bufferpool/bufferpool_test.go b/buffer/pool_test.go similarity index 96% rename from internal/bufferpool/bufferpool_test.go rename to buffer/pool_test.go index 9f847b437..a219815b5 100644 --- a/internal/bufferpool/bufferpool_test.go +++ b/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() }() diff --git a/internal/bufferpool/bufferpool.go b/internal/bufferpool/bufferpool.go index 7a5500024..dad583aaa 100644 --- a/internal/bufferpool/bufferpool.go +++ b/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) -} diff --git a/internal/multierror/multierror.go b/internal/multierror/multierror.go index b87467f5a..e5639877f 100644 --- a/internal/multierror/multierror.go +++ b/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 } diff --git a/stacktrace.go b/stacktrace.go index a214c162d..3d2a22c29 100644 --- a/stacktrace.go +++ b/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) diff --git a/zapcore/console_encoder.go b/zapcore/console_encoder.go index 724fa15c1..05a359a30 100644 --- a/zapcore/console_encoder.go +++ b/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() diff --git a/zapcore/console_encoder_bench_test.go b/zapcore/console_encoder_bench_test.go index 71dec428c..62feaea71 100644 --- a/zapcore/console_encoder_bench_test.go +++ b/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() } }) } diff --git a/zapcore/core.go b/zapcore/core.go index 5daf6fdf7..a1ef8b034 100644 --- a/zapcore/core.go +++ b/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 } diff --git a/zapcore/entry.go b/zapcore/entry.go index 9bab935b2..b09bbbea6 100644 --- a/zapcore/entry.go +++ b/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 } diff --git a/zapcore/json_encoder_bench_test.go b/zapcore/json_encoder_bench_test.go index 5d1e4e160..4bd5033b4 100644 --- a/zapcore/json_encoder_bench_test.go +++ b/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() } }) }