diff --git a/pool/pool.go b/pool/pool.go index 4fdc5efe..afbaf98a 100644 --- a/pool/pool.go +++ b/pool/pool.go @@ -88,6 +88,17 @@ func (p *Pool) Acquire(ctx context.Context) (*Conn, error) { return &Conn{res: res}, nil } +// AcquireAllIdle atomically acquires all currently idle connections. Its intended use is for health check and +// keep-alive functionality. It does not update pool statistics. +func (p *Pool) AcquireAllIdle() []*Conn { + resources := p.p.AcquireAllIdle() + conns := make([]*Conn, len(resources)) + for i := range conns { + conns[i] = &Conn{res: resources[i]} + } + return conns +} + func (p *Pool) Stat() *Stat { return &Stat{s: p.p.Stat()} } diff --git a/pool/pool_test.go b/pool/pool_test.go index 34ba424d..b026b2ee 100644 --- a/pool/pool_test.go +++ b/pool/pool_test.go @@ -51,6 +51,42 @@ func TestPoolAcquireAndConnRelease(t *testing.T) { c.Release() } +func TestPoolAcquireAllIdle(t *testing.T) { + t.Parallel() + + db, err := pool.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE")) + require.NoError(t, err) + defer db.Close() + + conns := db.AcquireAllIdle() + assert.Len(t, conns, 1) + + for _, c := range conns { + c.Release() + } + waitForReleaseToComplete() + + conns = make([]*pool.Conn, 3) + for i := range conns { + conns[i], err = db.Acquire(context.Background()) + assert.NoError(t, err) + } + + for _, c := range conns { + if c != nil { + c.Release() + } + } + waitForReleaseToComplete() + + conns = db.AcquireAllIdle() + assert.Len(t, conns, 3) + + for _, c := range conns { + c.Release() + } +} + func TestPoolExec(t *testing.T) { t.Parallel()