Add use before initialization checks to the ObjectPool

pool/object.go

@@ -21,13 +21,20 @@
 package pool
 
 import (
+	"errors"
 	"math"
 	"sync/atomic"
 	"time"
 
 	"github.com/uber-go/tally"
 )
 
+var (
+	errPoolAlreadyInitialized   = errors.New("object pool already initialized")
+	errPoolGetBeforeInitialized = errors.New("object pool get before initialized")
+	errPoolPutBeforeInitialized = errors.New("object pool put before initialized")
+)
+
 const (
 	// TODO(r): Use tally sampling when available
 	sampleObjectPoolLengthEvery = 100
@@ -40,7 +47,8 @@ type objectPool struct {
 	size                int
 	refillLowWatermark  int
 	refillHighWatermark int
-	filling             int64
+	filling             int32
+	initialized         int32
 	metrics             objectPoolMetrics
 }
 
@@ -81,6 +89,12 @@ func NewObjectPool(opts ObjectPoolOptions) ObjectPool {
 }
 
 func (p *objectPool) Init(alloc Allocator) {
+	if !atomic.CompareAndSwapInt32(&p.initialized, 0, 1) {
+		fn := p.opts.OnPoolAccessError()
+		fn(errPoolAlreadyInitialized)
+		return
+	}
+
 	p.alloc = alloc
 
 	for i := 0; i < cap(p.values); i++ {
@@ -89,8 +103,13 @@ func (p *objectPool) Init(alloc Allocator) {
 }
 
 func (p *objectPool) Get() interface{} {
-	var v interface{}
+	if atomic.LoadInt32(&p.initialized) != 1 {
+		fn := p.opts.OnPoolAccessError()
+		fn(errPoolGetBeforeInitialized)
+		return p.alloc()
+	}
 
+	var v interface{}
 	select {
 	case v = <-p.values:
 	default:
@@ -108,6 +127,12 @@ func (p *objectPool) Get() interface{} {
 }
 
 func (p *objectPool) Put(obj interface{}) {
+	if atomic.LoadInt32(&p.initialized) != 1 {
+		fn := p.opts.OnPoolAccessError()
+		fn(errPoolPutBeforeInitialized)
+		return
+	}
+
 	select {
 	case p.values <- obj:
 	default:
@@ -129,12 +154,12 @@ func (p *objectPool) setGauges() {
 }
 
 func (p *objectPool) tryFill() {
-	if !atomic.CompareAndSwapInt64(&p.filling, 0, 1) {
+	if !atomic.CompareAndSwapInt32(&p.filling, 0, 1) {
 		return
 	}
 
 	go func() {
-		defer atomic.StoreInt64(&p.filling, 0)
+		defer atomic.StoreInt32(&p.filling, 0)
 
 		for len(p.values) < p.refillHighWatermark {
 			select {

pool/object_test.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestObjectPoolRefillOnLowWaterMark(t *testing.T) {
@@ -60,3 +61,58 @@ func TestObjectPoolRefillOnLowWaterMark(t *testing.T) {
 	// Assert refilled
 	assert.Equal(t, 75, len(pool.values))
 }
+
+func TestObjectPoolInitTwiceError(t *testing.T) {
+	var accessErr error
+	opts := NewObjectPoolOptions().SetOnPoolAccessError(func(err error) {
+		accessErr = err
+	})
+
+	pool := NewObjectPool(opts)
+	pool.Init(func() interface{} {
+		return 1
+	})
+
+	require.NoError(t, accessErr)
+
+	pool.Init(func() interface{} {
+		return 1
+	})
+
+	assert.Error(t, accessErr)
+	assert.Equal(t, errPoolAlreadyInitialized, accessErr)
+}
+
+func TestObjectPoolGetBeforeInitError(t *testing.T) {
+	var accessErr error
+	opts := NewObjectPoolOptions().SetOnPoolAccessError(func(err error) {
+		accessErr = err
+	})
+
+	pool := NewObjectPool(opts)
+
+	require.NoError(t, accessErr)
+
+	assert.Panics(t, func() {
+		pool.Get()
+	})
+
+	assert.Error(t, accessErr)
+	assert.Equal(t, errPoolGetBeforeInitialized, accessErr)
+}
+
+func TestObjectPoolPutBeforeInitError(t *testing.T) {
+	var accessErr error
+	opts := NewObjectPoolOptions().SetOnPoolAccessError(func(err error) {
+		accessErr = err
+	})
+
+	pool := NewObjectPool(opts)
+
+	require.NoError(t, accessErr)
+
+	pool.Put(1)
+
+	assert.Error(t, accessErr)
+	assert.Equal(t, errPoolPutBeforeInitialized, accessErr)
+}

pool/options.go

@@ -32,6 +32,7 @@ type objectPoolOptions struct {
 	refillLowWatermark  float64
 	refillHighWatermark float64
 	instrumentOpts      instrument.Options
+	onPoolAccessError   OnPoolAccessError
 }
 
 // NewObjectPoolOptions creates a new set of object pool options
@@ -40,6 +41,7 @@ func NewObjectPoolOptions() ObjectPoolOptions {
 		size:               defaultSize,
 		refillLowWatermark: defaultRefillLowWatermark,
 		instrumentOpts:     instrument.NewOptions(),
+		onPoolAccessError:  func(err error) { panic(err) },
 	}
 }
 
@@ -82,3 +84,13 @@ func (o *objectPoolOptions) SetInstrumentOptions(value instrument.Options) Objec
 func (o *objectPoolOptions) InstrumentOptions() instrument.Options {
 	return o.instrumentOpts
 }
+
+func (o *objectPoolOptions) SetOnPoolAccessError(value OnPoolAccessError) ObjectPoolOptions {
+	opts := *o
+	opts.onPoolAccessError = value
+	return &opts
+}
+
+func (o *objectPoolOptions) OnPoolAccessError() OnPoolAccessError {
+	return o.onPoolAccessError
+}

pool/types.go

@@ -56,6 +56,10 @@ type CheckedObjectPool interface {
 	Get() checked.ReadWriteRef
 }
 
+// OnPoolAccessError is a function to call when a pool access error occurs,
+// such as get or put before the pool is initialized.
+type OnPoolAccessError func(err error)
+
 // ObjectPoolOptions provides options for an object pool.
 type ObjectPoolOptions interface {
 	// SetSize sets the size of the object pool.
@@ -85,6 +89,14 @@ type ObjectPoolOptions interface {
 
 	// InstrumentOptions returns the instrument options.
 	InstrumentOptions() instrument.Options
+
+	// SetOnPoolAccessError sets the on pool access error callback, by
+	// default this is a panic.
+	SetOnPoolAccessError(value OnPoolAccessError) ObjectPoolOptions
+
+	// OnPoolAccessError returns the on pool access error callback, by
+	// default this is a panic.
+	OnPoolAccessError() OnPoolAccessError
 }
 
 // Bucket specifies a pool bucket.