internal: Add and use a generic pool instead of using sync.Pool directly

buffer/pool.go

@@ -20,25 +20,29 @@
 
 package buffer
 
-import "sync"
+import (
+	"go.uber.org/zap/internal/pool"
+)
 
 // A Pool is a type-safe wrapper around a sync.Pool.
 type Pool struct {
-	p *sync.Pool
+	p *pool.Pool[*Buffer]
 }
 
 // NewPool constructs a new Pool.
 func NewPool() Pool {
-	return Pool{p: &sync.Pool{
-		New: func() interface{} {
-			return &Buffer{bs: make([]byte, 0, _size)}
-		},
-	}}
+	return Pool{
+		p: pool.New(func() *Buffer {
+			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 := p.p.Get()
 	buf.Reset()
 	buf.pool = p
 	return buf

error.go

@@ -21,14 +21,13 @@
 package zap
 
 import (
-	"sync"
-
+	"go.uber.org/zap/internal/pool"
 	"go.uber.org/zap/zapcore"
 )
 
-var _errArrayElemPool = sync.Pool{New: func() interface{} {
+var _errArrayElemPool = pool.New(func() *errArrayElem {
 	return &errArrayElem{}
-}}
+})
 
 // Error is shorthand for the common idiom NamedError("error", err).
 func Error(err error) Field {
@@ -60,7 +59,7 @@ func (errs errArray) MarshalLogArray(arr zapcore.ArrayEncoder) error {
 		// potentially an "errorVerbose" attribute, we need to wrap it in a
 		// type that implements LogObjectMarshaler. To prevent this from
 		// allocating, pool the wrapper type.
-		elem := _errArrayElemPool.Get().(*errArrayElem)
+		elem := _errArrayElemPool.Get()
 		elem.error = errs[i]
 		arr.AppendObject(elem)
 		elem.error = nil

internal/pool/pool.go

@@ -0,0 +1,64 @@
+// Copyright (c) 2023 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 pool provides internal pool utilities.
+package pool
+
+import (
+	"sync"
+)
+
+// A Constructor is a function that creates a T when a [Pool] is empty.
+type Constructor[T any] func() T
+
+// A Pool is a generic wrapper around [sync.Pool] to provide strongly-typed
+// object pooling.
+type Pool[T any] struct {
+	pool sync.Pool
+	ctor Constructor[T]
+}
+
+// New returns a new [Pool] for T, and will use ctor to construct new Ts when
+// the pool is empty.
+func New[T any](ctor Constructor[T]) *Pool[T] {
+	return &Pool[T]{
+		pool: sync.Pool{
+			New: func() any {
+				return ctor()
+			},
+		},
+		ctor: ctor,
+	}
+}
+
+// Get gets a new T from the pool, or creates a new one if the pool is empty.
+func (p *Pool[T]) Get() T {
+	if x, ok := p.pool.Get().(T); ok {
+		return x
+	}
+
+	// n.b. This branch is effectively unreachable.
+	return p.ctor()
+}
+
+// Put puts x into the pool.
+func (p *Pool[T]) Put(x T) {
+	p.pool.Put(x)
+}

internal/pool/pool_internals_test.go

@@ -0,0 +1,81 @@
+// Copyright (c) 2023 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 pool
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPool_ConstructorCalledIfWrongType(t *testing.T) {
+	cases := []struct {
+		input      any
+		expectCall bool
+	}{
+		{
+			input:      int64(123),
+			expectCall: false,
+		},
+		{
+			input:      uint64(123),
+			expectCall: true,
+		},
+		{
+			input:      int(123),
+			expectCall: true,
+		},
+		{
+			input:      uint(123),
+			expectCall: true,
+		},
+		{
+			input:      struct{}{},
+			expectCall: true,
+		},
+		{
+			input:      nil,
+			expectCall: true,
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(fmt.Sprintf("%T", tt.input), func(t *testing.T) {
+			var (
+				called bool
+				pool   = New(func() int64 {
+					called = true
+					return 0
+				})
+			)
+
+			// Override the internal pool to provide unexpected types.
+			pool.pool.New = func() any {
+				return tt.input
+			}
+
+			pool.pool.Put(tt.input)
+			pool.Get()
+			require.Equal(t, tt.expectCall, called)
+		})
+	}
+}

internal/pool/pool_norace_test.go

@@ -0,0 +1,56 @@
+// Copyright (c) 2023 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.
+
+//go:build !race
+
+package pool_test
+
+import (
+	"bytes"
+	"runtime/debug"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/internal/pool"
+)
+
+func TestNew(t *testing.T) {
+	// n.b. Disable GC to avoid the victim cache during the test.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+
+	p := pool.New(func() *bytes.Buffer {
+		return bytes.NewBuffer([]byte("new"))
+	})
+	p.Put(bytes.NewBuffer([]byte(t.Name())))
+
+	// Ensure that we always get the expected value.
+	for i := 0; i < 1_000; i++ {
+		func() {
+			buf := p.Get()
+			defer p.Put(buf)
+			require.Equal(t, t.Name(), buf.String())
+		}()
+	}
+
+	// Depool an extra object to ensure that the constructor is called and
+	// produces an expected value.
+	require.Equal(t, t.Name(), p.Get().String())
+	require.Equal(t, "new", p.Get().String())
+}

internal/pool/pool_race_test.go

@@ -0,0 +1,69 @@
+// Copyright (c) 2023 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.
+
+//go:build race
+
+package pool_test
+
+import (
+	"sync"
+	"testing"
+
+	"go.uber.org/zap/internal/pool"
+)
+
+type pooledValue struct {
+	n int64
+}
+
+func TestNew(t *testing.T) {
+	// n.b. [sync.Pool] will randomly drop re-pooled objects when race is
+	//      enabled, so rather than testing nondeterminsitic behavior, we use
+	//      this test solely to prove that there are no races. See pool_test.go
+	//      for correctness testing.
+
+	var (
+		p = pool.New(func() *pooledValue {
+			return &pooledValue{
+				n: -1,
+			}
+		})
+		wg sync.WaitGroup
+	)
+
+	defer wg.Wait()
+
+	for i := int64(0); i < 1_000; i++ {
+		i := i
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			x := p.Get()
+			defer p.Put(x)
+
+			// n.b. Must read and write the field.
+			if n := x.n; n >= -1 {
+				x.n = int64(i)
+			}
+		}()
+	}
+}

stacktrace.go

@@ -22,19 +22,17 @@ package zap
 
 import (
 	"runtime"
-	"sync"
 
 	"go.uber.org/zap/buffer"
 	"go.uber.org/zap/internal/bufferpool"
+	"go.uber.org/zap/internal/pool"
 )
 
-var _stacktracePool = sync.Pool{
-	New: func() interface{} {
-		return &stacktrace{
-			storage: make([]uintptr, 64),
-		}
-	},
-}
+var _stacktracePool = pool.New(func() *stacktrace {
+	return &stacktrace{
+		storage: make([]uintptr, 64),
+	}
+})
 
 type stacktrace struct {
 	pcs    []uintptr // program counters; always a subslice of storage
@@ -68,7 +66,7 @@ const (
 //
 // The caller must call Free on the returned stacktrace after using it.
 func captureStacktrace(skip int, depth stacktraceDepth) *stacktrace {
-	stack := _stacktracePool.Get().(*stacktrace)
+	stack := _stacktracePool.Get()
 
 	switch depth {
 	case stacktraceFirst: