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