pool.end() resolves before the last pool.query()

packages/pg-pool/index.js

@@ -156,16 +156,22 @@ class Pool extends EventEmitter {
     throw new Error('unexpected condition')
   }
 
-  _remove(client) {
+  _remove(client, callback) {
     const removed = removeWhere(this._idle, (item) => item.client === client)
 
     if (removed !== undefined) {
       clearTimeout(removed.timeoutId)
     }
 
     this._clients = this._clients.filter((c) => c !== client)
-    client.end()
-    this.emit('remove', client)
+    const context = this
+    client.end(() => {
+      context.emit('remove', client)
+
+      if (typeof callback === 'function') {
+        callback()
+      }
+    })
   }
 
   connect(cb) {
@@ -342,26 +348,23 @@ class Pool extends EventEmitter {
       if (client._poolUseCount >= this.options.maxUses) {
         this.log('remove expended client')
       }
-      this._remove(client)
-      this._pulseQueue()
-      return
+
+      return this._remove(client, this._pulseQueue.bind(this))
     }
 
     const isExpired = this._expired.has(client)
     if (isExpired) {
       this.log('remove expired client')
       this._expired.delete(client)
-      this._remove(client)
-      this._pulseQueue()
-      return
+      return this._remove(client, this._pulseQueue.bind(this))
     }
 
     // idle timeout
     let tid
     if (this.options.idleTimeoutMillis) {
       tid = setTimeout(() => {
         this.log('remove idle client')
-        this._remove(client)
+        this._remove(client, this._pulseQueue.bind(this))
       }, this.options.idleTimeoutMillis)
 
       if (this.options.allowExitOnIdle) {

packages/pg-pool/test/ending.js

@@ -37,4 +37,14 @@ describe('pool ending', () => {
       expect(res.rows[0].name).to.equal('brianc')
     })
   )
+
+  it('pool.end() - finish pending queries', async () => {
+    const pool = new Pool({ max: 20 })
+    let completed = 0
+    for (let x = 1; x <= 20; x++) {
+      pool.query('SELECT $1::text as name', ['brianc']).then(() => completed++)
+    }
+    await pool.end()
+    expect(completed).to.equal(20)
+  })
 })

packages/pg-pool/test/idle-timeout.js

@@ -25,6 +25,7 @@ describe('idle timeout', () => {
   it(
     'times out and removes clients when others are also removed',
     co.wrap(function* () {
+      let currentClient = 1
       const pool = new Pool({ idleTimeoutMillis: 10 })
       const clientA = yield pool.connect()
       const clientB = yield pool.connect()
@@ -35,7 +36,12 @@ describe('idle timeout', () => {
         pool.on('remove', () => {
           expect(pool.idleCount).to.equal(0)
           expect(pool.totalCount).to.equal(0)
-          resolve()
+
+          if (currentClient >= 2) {
+            resolve()
+          } else {
+            currentClient++
+          }
         })
       })
 
@@ -44,7 +50,7 @@ describe('idle timeout', () => {
       try {
         yield Promise.race([removal, timeout])
       } finally {
-        pool.end()
+        yield pool.end()
       }
     })
   )