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

* Pass callback to client.end

* Add test for pool.end method

* fix: remove excessive _pulseQueue call

* fix: context problem

* fix: test resolve should be called when the last client is removed

* fix: wait for pool.end()

Because when you don't pass a callback to .end() it always returns a promise

* fix: handle idle timeout test data race

---------

Co-authored-by: Asadbek Raimov <asadbekraimov642@gmail.com>

Author: hjr3

packages/pg-pool/index.js

@@ -161,16 +161,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) {
@@ -351,26 +357,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 && this._isAboveMin()) {
       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

@@ -28,11 +28,19 @@ describe('idle timeout', () => {
       const pool = new Pool({ idleTimeoutMillis: 10 })
       const clientA = yield pool.connect()
       const clientB = yield pool.connect()
-      clientA.release()
-      clientB.release(new Error())
+      clientA.release() // this will put clientA in the idle pool
+      clientB.release(new Error()) // an error will cause clientB to be removed immediately
 
       const removal = new Promise((resolve) => {
-        pool.on('remove', () => {
+        pool.on('remove', (client) => {
+          // clientB's stream may take a while to close, so we may get a remove
+          // event for it
+          // we only want to handle the remove event for clientA when it times out
+          // due to being idle
+          if (client !== clientA) {
+            return
+          }
+
           expect(pool.idleCount).to.equal(0)
           expect(pool.totalCount).to.equal(0)
           resolve()