Add TryAcquire and ErrNotAvailable
Fixes https://github.com/jackc/puddle/issues/10
This commit is contained in:
@@ -18,6 +18,10 @@ const (
|
||||
// or a pool that is closed while the acquire is waiting.
|
||||
var ErrClosedPool = errors.New("closed pool")
|
||||
|
||||
// ErrNotAvailable occurs on an attempt to acquire a resource from a pool
|
||||
// that is at maximum capacity and has no available resources.
|
||||
var ErrNotAvailable = errors.New("resource not available")
|
||||
|
||||
// Constructor is a function called by the pool to construct a resource.
|
||||
type Constructor func(ctx context.Context) (res interface{}, err error)
|
||||
|
||||
@@ -258,6 +262,19 @@ func (p *Pool) Stat() *Stat {
|
||||
// maximum capacity it will block until a resource is available. ctx can be used
|
||||
// to cancel the Acquire.
|
||||
func (p *Pool) Acquire(ctx context.Context) (*Resource, error) {
|
||||
return p.doAcquire(ctx, true)
|
||||
}
|
||||
|
||||
// TryAcquire gets a resource from the pool. TryAcquire is the same as Acquire except
|
||||
// it returns ErrNotAvailable if the pool is at maximum capacity and no resources are available.
|
||||
func (p *Pool) TryAcquire(ctx context.Context) (*Resource, error) {
|
||||
return p.doAcquire(ctx, false)
|
||||
}
|
||||
|
||||
// doAcquire implements shared logic behind Acquire and TryAcquire. If block is true
|
||||
// doAcquire will block until a resource becomes available. If block is false, doAcquire
|
||||
// will return ErrNotAvailable if no resource is available.
|
||||
func (p *Pool) doAcquire(ctx context.Context, block bool) (*Resource, error) {
|
||||
startNano := nanotime()
|
||||
p.cond.L.Lock()
|
||||
if doneChan := ctx.Done(); doneChan != nil {
|
||||
@@ -327,6 +344,10 @@ func (p *Pool) Acquire(ctx context.Context) (*Resource, error) {
|
||||
p.acquireDuration += time.Duration(nanotime() - startNano)
|
||||
p.cond.L.Unlock()
|
||||
return res, nil
|
||||
} else if !block {
|
||||
// If the pool is at maximum capacity and we're not blocking
|
||||
p.cond.L.Unlock()
|
||||
return nil, ErrNotAvailable
|
||||
}
|
||||
|
||||
if ctx.Done() == nil {
|
||||
|
||||
@@ -175,6 +175,23 @@ func TestPoolAcquireReusesResources(t *testing.T) {
|
||||
assert.Equal(t, 1, createCounter.Value())
|
||||
}
|
||||
|
||||
func TestPoolTryAcquireDoesNotBlock(t *testing.T) {
|
||||
constructor, createCounter := createConstructor()
|
||||
pool := puddle.NewPool(constructor, stubDestructor, 1)
|
||||
|
||||
res, err := pool.TryAcquire(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, res.Value())
|
||||
|
||||
defer res.Release()
|
||||
|
||||
res, err = pool.TryAcquire(context.Background())
|
||||
require.EqualError(t, err, puddle.ErrNotAvailable.Error())
|
||||
assert.Nil(t, res)
|
||||
|
||||
assert.Equal(t, 1, createCounter.Value())
|
||||
}
|
||||
|
||||
func TestPoolAcquireContextAlreadyCanceled(t *testing.T) {
|
||||
constructor := func(ctx context.Context) (interface{}, error) {
|
||||
panic("should never be called")
|
||||
|
||||
Reference in New Issue
Block a user