Allow Node to exit if the pool is idle

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