From 56464531eb62b66d02a71e0efc64711d3ab45a91 Mon Sep 17 00:00:00 2001
From: Brian Crowell <brian@fluggo.com>
Date: Fri, 2 Jul 2021 15:06:35 -0500
Subject: [PATCH] Allow Node to exit if the pool is idle

Based on the suggestion from #2078. This adds ref/unref methods to the
Connection and Client classes and then uses them to allow the process to
exit if all of the connections in the pool are idle. This behavior is
controlled by the allowExitOnIdle flag to the Pool constructor; it defaults
to the old behavior.
---
 packages/pg-pool/index.js                  | 11 ++++++++
 packages/pg-pool/test/idle-timeout-exit.js | 16 +++++++++++
 packages/pg-pool/test/idle-timeout.js      | 31 ++++++++++++++++++++++
 packages/pg/lib/client.js                  |  8 ++++++
 packages/pg/lib/connection.js              |  8 ++++++
 5 files changed, 74 insertions(+)
 create mode 100644 packages/pg-pool/test/idle-timeout-exit.js

diff --git a/packages/pg-pool/index.js b/packages/pg-pool/index.js
index 403d05a19..5557de5c0 100644
--- a/packages/pg-pool/index.js
+++ b/packages/pg-pool/index.js
@@ -83,6 +83,7 @@ class Pool extends EventEmitter {
 
     this.options.max = this.options.max || this.options.poolSize || 10
     this.options.maxUses = this.options.maxUses || Infinity
+    this.options.allowExitOnIdle = this.options.allowExitOnIdle || false
     this.log = this.options.log || function () {}
     this.Client = this.options.Client || Client || require('pg').Client
     this.Promise = this.options.Promise || global.Promise
@@ -136,6 +137,7 @@ class Pool extends EventEmitter {
       const idleItem = this._idle.pop()
       clearTimeout(idleItem.timeoutId)
       const client = idleItem.client
+      client.ref()
       const idleListener = idleItem.idleListener
 
       return this._acquireClient(client, pendingItem, idleListener, false)
@@ -323,6 +325,15 @@ class Pool extends EventEmitter {
         this.log('remove idle client')
         this._remove(client)
       }, this.options.idleTimeoutMillis)
+
+      if (this.options.allowExitOnIdle) {
+        // allow Node to exit if this is all that's left
+        tid.unref()
+      }
+    }
+
+    if (this.options.allowExitOnIdle) {
+      client.unref()
     }
 
     this._idle.push(new IdleItem(client, idleListener, tid))
diff --git a/packages/pg-pool/test/idle-timeout-exit.js b/packages/pg-pool/test/idle-timeout-exit.js
new file mode 100644
index 000000000..1292634a8
--- /dev/null
+++ b/packages/pg-pool/test/idle-timeout-exit.js
@@ -0,0 +1,16 @@
+// This test is meant to be spawned from idle-timeout.js
+if (module === require.main) {
+  const allowExitOnIdle = process.env.ALLOW_EXIT_ON_IDLE === '1'
+  const Pool = require('../index')
+
+  const pool = new Pool({ idleTimeoutMillis: 200, ...(allowExitOnIdle ? { allowExitOnIdle: true } : {}) })
+  pool.query('SELECT NOW()', (err, res) => console.log('completed first'))
+  pool.on('remove', () => {
+    console.log('removed')
+    done()
+  })
+
+  setTimeout(() => {
+    pool.query('SELECT * from generate_series(0, 1000)', (err, res) => console.log('completed second'))
+  }, 50)
+}
diff --git a/packages/pg-pool/test/idle-timeout.js b/packages/pg-pool/test/idle-timeout.js
index fd9fba4a4..0bb097565 100644
--- a/packages/pg-pool/test/idle-timeout.js
+++ b/packages/pg-pool/test/idle-timeout.js
@@ -4,6 +4,8 @@ const expect = require('expect.js')
 
 const describe = require('mocha').describe
 const it = require('mocha').it
+const { fork } = require('child_process')
+const path = require('path')
 
 const Pool = require('../')
 
@@ -84,4 +86,33 @@ describe('idle timeout', () => {
       return pool.end()
     })
   )
+
+  it('unrefs the connections and timeouts so the program can exit when idle when the allowExitOnIdle option is set', function (done) {
+    const child = fork(path.join(__dirname, 'idle-timeout-exit.js'), [], {
+      silent: true,
+      env: { ...process.env, ALLOW_EXIT_ON_IDLE: '1' },
+    })
+    let result = ''
+    child.stdout.setEncoding('utf8')
+    child.stdout.on('data', (chunk) => (result += chunk))
+    child.on('error', (err) => done(err))
+    child.on('close', () => {
+      expect(result).to.equal('completed first\ncompleted second\n')
+      done()
+    })
+  })
+
+  it('keeps old behavior when allowExitOnIdle option is not set', function (done) {
+    const child = fork(path.join(__dirname, 'idle-timeout-exit.js'), [], {
+      silent: true,
+    })
+    let result = ''
+    child.stdout.setEncoding('utf8')
+    child.stdout.on('data', (chunk) => (result += chunk))
+    child.on('error', (err) => done(err))
+    child.on('close', () => {
+      expect(result).to.equal('completed first\ncompleted second\nremoved\n')
+      done()
+    })
+  })
 })
diff --git a/packages/pg/lib/client.js b/packages/pg/lib/client.js
index 1e1e83374..589aa9f84 100644
--- a/packages/pg/lib/client.js
+++ b/packages/pg/lib/client.js
@@ -577,6 +577,14 @@ class Client extends EventEmitter {
     return result
   }
 
+  ref() {
+    this.connection.ref()
+  }
+
+  unref() {
+    this.connection.unref()
+  }
+
   end(cb) {
     this._ending = true
 
diff --git a/packages/pg/lib/connection.js b/packages/pg/lib/connection.js
index 7d45de2b7..ebb2f099d 100644
--- a/packages/pg/lib/connection.js
+++ b/packages/pg/lib/connection.js
@@ -177,6 +177,14 @@ class Connection extends EventEmitter {
     this._send(syncBuffer)
   }
 
+  ref() {
+    this.stream.ref()
+  }
+
+  unref() {
+    this.stream.unref()
+  }
+
   end() {
     // 0x58 = 'X'
     this._ending = true