Allow Node to exit if the pool is idle (brianc#2568)

Based on the suggestion from brianc#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.

Author: Brian Crowell

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))

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)
+}

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()
+    })
+  })
 })

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
 

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