fix: allow reusing a pool after end

docs/pages/apis/pool.mdx

@@ -204,7 +204,7 @@ assert(pool.totalCount === 0)
 
 ## pool.end
 
-Calling `pool.end` will drain the pool of all active clients, disconnect them, and shut down any internal timers in the pool. It is common to call this at the end of a script using the pool or when your process is attempting to shut down cleanly.
+Calling `pool.end` will drain the pool of all active clients, disconnect them, and shut down any internal timers in the pool. It is common to call this at the end of a script using the pool or when your process is attempting to shut down cleanly. After the returned promise resolves, the same pool can be used again and will create new connections as needed.
 
 ```js
 // again both promises and callbacks are supported:
@@ -240,7 +240,8 @@ The number of queued requests waiting on a client when all clients are checked o
 Whenever the pool establishes a new client connection to the PostgreSQL backend it will emit the `connect` event with the newly connected client.
 
 <Alert>
-  The event listener does not wait for promises or async functions. If you want to run setup commands on each new client, use the `onConnect` option. (See documentation above.)
+  The event listener does not wait for promises or async functions. If you want to run setup commands on each new
+  client, use the `onConnect` option. (See documentation above.)
 </Alert>
 
 ### acquire

docs/pages/features/pooling.mdx

@@ -105,5 +105,6 @@ pool has drained
 ```
 
 <Info>
-  The pool will return errors when attempting to check out a client after you've called pool.end() on the pool.
+  The pool will return errors when attempting to check out a client while `pool.end()` is still draining. After the
+  returned promise resolves, the pool can create connections again.
 </Info>

packages/pg-pool/index.js

@@ -103,6 +103,7 @@ class Pool extends EventEmitter {
     this._idle = []
     this._expired = new WeakSet()
     this._pendingQueue = []
+    this._pendingEnds = 0
     this._endCallback = undefined
     this.ending = false
     this.ended = false
@@ -132,12 +133,25 @@ class Pool extends EventEmitter {
     }
     if (this.ending) {
       this.log('pulse queue on ending')
+      if (this._pendingQueue.length) {
+        const pendingItem = this._pendingQueue.shift()
+        if (this._idle.length) {
+          const idleItem = this._idle.pop()
+          clearTimeout(idleItem.timeoutId)
+          return this._acquireClient(idleItem.client, pendingItem, idleItem.idleListener, false)
+        }
+        if (!this._isFull()) {
+          return this.newClient(pendingItem)
+        }
+        this._pendingQueue.unshift(pendingItem)
+        return
+      }
       if (this._idle.length) {
         this._idle.slice().map((item) => {
-          this._remove(item.client)
+          this._remove(item.client, this._pulseQueue.bind(this))
         })
       }
-      if (!this._clients.length) {
+      if (!this._clients.length && !this._pendingEnds) {
         this.ended = true
         this._endCallback()
       }
@@ -178,17 +192,25 @@ class Pool extends EventEmitter {
 
     this._clients = this._clients.filter((c) => c !== client)
     const context = this
+    this._pendingEnds++
     client.end(() => {
+      context._pendingEnds--
       context.emit('remove', client)
 
       if (typeof callback === 'function') {
         callback()
+      } else if (context.ending) {
+        context._pulseQueue()
       }
     })
   }
 
   connect(cb) {
-    if (this.ending) {
+    if (this.ended) {
+      this.ending = false
+      this.ended = false
+      this._endCallback = undefined
+    } else if (this.ending) {
       const err = new Error('Cannot use a pool after calling end on the pool')
       return cb ? cb(err) : this.Promise.reject(err)
     }
@@ -389,7 +411,13 @@ class Pool extends EventEmitter {
     this.emit('release', err, client)
 
     // TODO(bmc): expose a proper, public interface _queryable and _ending
-    if (err || this.ending || !client._queryable || client._ending || client._poolUseCount >= this.options.maxUses) {
+    if (
+      err ||
+      (this.ending && !this._pendingQueue.length) ||
+      !client._queryable ||
+      client._ending ||
+      client._poolUseCount >= this.options.maxUses
+    ) {
       if (client._poolUseCount >= this.options.maxUses) {
         this.log('remove expended client')
       }
@@ -404,6 +432,11 @@ class Pool extends EventEmitter {
       return this._remove(client, this._pulseQueue.bind(this))
     }
 
+    if (this.ending) {
+      this._idle.push(new IdleItem(client, idleListener, undefined))
+      return this._pulseQueue()
+    }
+
     // idle timeout
     let tid
     if (this.options.idleTimeoutMillis && this._isAboveMin()) {

packages/pg-pool/test/ending.js

@@ -1,12 +1,60 @@
 'use strict'
 const co = require('co')
 const expect = require('expect.js')
+const EventEmitter = require('events')
 
 const describe = require('mocha').describe
 const it = require('mocha').it
 
 const Pool = require('../')
 
+class MockClient extends EventEmitter {
+  constructor() {
+    super()
+    this._queryable = true
+    this._ending = false
+  }
+
+  connect(callback) {
+    process.nextTick(callback)
+  }
+
+  query(_text, _values, callback) {
+    process.nextTick(() => callback(undefined, { rows: [{ value: 1 }] }))
+  }
+
+  end(callback) {
+    this._ending = true
+    process.nextTick(() => {
+      this.emit('end')
+      callback?.()
+    })
+  }
+}
+
+class DeferredEndClient extends MockClient {
+  end(callback) {
+    this._ending = true
+    this.endCallback = callback
+  }
+
+  finishEnd() {
+    this.emit('end')
+    this.endCallback?.()
+  }
+}
+
+class DeferredQueryClient extends MockClient {
+  query(_text, _values, callback) {
+    this.queryCallbacks ??= []
+    this.queryCallbacks.push(callback)
+  }
+
+  finishQuery(value) {
+    this.queryCallbacks.shift()(undefined, { rows: [{ value }] })
+  }
+}
+
 describe('pool ending', () => {
   it('ends without being used', (done) => {
     const pool = new Pool()
@@ -47,4 +95,102 @@ describe('pool ending', () => {
     await pool.end()
     expect(completed).to.equal(20)
   })
+
+  it('can be used again after end completes', async () => {
+    const pool = new Pool({ Client: MockClient })
+
+    expect((await pool.query('SELECT 1')).rows[0].value).to.equal(1)
+    await pool.end()
+    expect(pool.totalCount).to.equal(0)
+
+    expect((await pool.query('SELECT 1')).rows[0].value).to.equal(1)
+    await pool.end()
+  })
+
+  it('can be used again with callbacks while duplicate end still errors', (done) => {
+    const pool = new Pool({ Client: MockClient })
+
+    pool.end(() => {
+      pool.end((err) => {
+        expect(err).to.be.an(Error)
+        pool.query('SELECT 1', (err, result) => {
+          if (err) return done(err)
+          expect(result.rows[0].value).to.equal(1)
+          pool.connect((err, client) => {
+            if (err) return done(err)
+            client.release()
+            pool.end(done)
+          })
+        })
+      })
+    })
+  })
+
+  it('cannot be used while end is still draining', async () => {
+    const pool = new Pool({ Client: DeferredEndClient })
+    const client = await pool.connect()
+    client.release()
+
+    let endCompleted = false
+    const endPromise = pool.end().then(() => {
+      endCompleted = true
+    })
+    await new Promise((resolve) => setImmediate(resolve))
+    expect(endCompleted).to.equal(false)
+    let error
+    try {
+      await pool.connect()
+    } catch (caughtError) {
+      error = caughtError
+    }
+    expect(error.message).to.contain('Cannot use a pool after calling end')
+
+    client.finishEnd()
+    await endPromise
+    expect(endCompleted).to.equal(true)
+    const reusedClient = await pool.connect()
+    reusedClient.release()
+    const secondEndPromise = pool.end()
+    reusedClient.finishEnd()
+    await secondEndPromise
+  })
+
+  it('finishes queued queries before end resolves', async () => {
+    const pool = new Pool({ Client: DeferredQueryClient, max: 1 })
+    const firstQuery = pool.query('SELECT 1')
+    const secondQuery = pool.query('SELECT 2')
+    const client = pool._clients[0]
+    const endPromise = pool.end()
+
+    await new Promise((resolve) => setImmediate(resolve))
+    client.finishQuery(1)
+    expect((await firstQuery).rows[0].value).to.equal(1)
+    await new Promise((resolve) => setImmediate(resolve))
+    client.finishQuery(2)
+    expect((await secondQuery).rows[0].value).to.equal(2)
+    await endPromise
+
+    expect(pool.waitingCount).to.equal(0)
+    expect(pool.totalCount).to.equal(0)
+  })
+
+  it('waits for every idle client and completes end once', async () => {
+    const pool = new Pool({ Client: DeferredEndClient, max: 2 })
+    const firstClient = await pool.connect()
+    const secondClient = await pool.connect()
+    firstClient.release()
+    secondClient.release()
+
+    let endCount = 0
+    const endPromise = pool.end(() => {
+      endCount++
+    })
+    firstClient.finishEnd()
+    await new Promise((resolve) => setImmediate(resolve))
+    expect(endCount).to.equal(0)
+    secondClient.finishEnd()
+    await new Promise((resolve) => setImmediate(resolve))
+    expect(endCount).to.equal(1)
+    expect(endPromise).to.equal(undefined)
+  })
 })

packages/pg-pool/test/error-handling.js

@@ -77,33 +77,28 @@ describe('pool error handling', function () {
     })
   })
 
-  describe('using an ended pool', () => {
-    it('rejects all additional promises', (done) => {
+  describe('using a pool after end', () => {
+    it('can be reused with promises', async () => {
       const pool = new Pool()
-      const promises = []
-      pool.end().then(() => {
-        const squash = (promise) => promise.catch((e) => 'okay!')
-        promises.push(squash(pool.connect()))
-        promises.push(squash(pool.query('SELECT NOW()')))
-        promises.push(squash(pool.end()))
-        Promise.all(promises).then((res) => {
-          expect(res).to.eql(['okay!', 'okay!', 'okay!'])
-          done()
-        })
-      })
+      await pool.end()
+
+      const result = await pool.query('SELECT NOW()')
+      expect(result.rows).to.have.length(1)
+      const client = await pool.connect()
+      client.release()
+      await pool.end()
     })
 
-    it('returns an error on all additional callbacks', (done) => {
+    it('can be reused with callbacks', (done) => {
       const pool = new Pool()
       pool.end(() => {
-        pool.query('SELECT *', (err) => {
-          expect(err).to.be.an(Error)
-          pool.connect((err) => {
-            expect(err).to.be.an(Error)
-            pool.end((err) => {
-              expect(err).to.be.an(Error)
-              done()
-            })
+        pool.query('SELECT NOW()', (err, result) => {
+          if (err) return done(err)
+          expect(result.rows).to.have.length(1)
+          pool.connect((err, client) => {
+            if (err) return done(err)
+            client.release()
+            pool.end(done)
           })
         })
       })