diff --git a/pool.go b/pool.go index b128adb..d3e76c2 100644 --- a/pool.go +++ b/pool.go @@ -302,6 +302,23 @@ func (p *Pool) Acquire(ctx context.Context) (*Resource, error) { } } +// AcquireAllIdle atomically acquires all currently idle resources. Its intended +// use is for health check and keep-alive functionality. It does not update pool +// statistics. +func (p *Pool) AcquireAllIdle() []*Resource { + p.cond.L.Lock() + + for _, res := range p.idleResources { + res.status = resourceStatusAcquired + } + resources := make([]*Resource, len(p.idleResources)) + copy(resources, p.idleResources) + p.idleResources = p.idleResources[0:0] + + p.cond.L.Unlock() + return resources +} + // releaseAcquiredResource returns res to the the pool. func (p *Pool) releaseAcquiredResource(res *Resource) { p.cond.L.Lock() diff --git a/pool_test.go b/pool_test.go index 542dc6d..226c018 100644 --- a/pool_test.go +++ b/pool_test.go @@ -197,6 +197,43 @@ func TestPoolAcquireContextCanceledDuringCreate(t *testing.T) { assert.Nil(t, res) } +func TestPoolAcquireAllIdle(t *testing.T) { + constructor, _ := createConstructor() + pool := puddle.NewPool(constructor, stubDestructor, 10) + defer pool.Close() + + resources := make([]*puddle.Resource, 4) + var err error + + resources[0], err = pool.Acquire(context.Background()) + require.NoError(t, err) + resources[1], err = pool.Acquire(context.Background()) + require.NoError(t, err) + resources[2], err = pool.Acquire(context.Background()) + require.NoError(t, err) + resources[3], err = pool.Acquire(context.Background()) + require.NoError(t, err) + + assert.Len(t, pool.AcquireAllIdle(), 0) + + resources[0].Release() + resources[3].Release() + + assert.ElementsMatch(t, []*puddle.Resource{resources[0], resources[3]}, pool.AcquireAllIdle()) + + resources[0].Release() + resources[3].Release() + resources[1].Release() + resources[2].Release() + + assert.ElementsMatch(t, resources, pool.AcquireAllIdle()) + + resources[0].Release() + resources[1].Release() + resources[2].Release() + resources[3].Release() +} + func TestPoolCloseClosesAllIdleResources(t *testing.T) { constructor, _ := createConstructor()