Add connection lifetime limit option and tests

Author: ChrisG0x20

packages/pg-pool/index.js

@@ -84,6 +84,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.options.maxLifetimeSeconds = this.options.maxLifetimeSeconds || 0
     this.log = this.options.log || function () {}
     this.Client = this.options.Client || Client || require('pg').Client
     this.Promise = this.options.Promise || global.Promise
@@ -94,6 +95,7 @@ class Pool extends EventEmitter {
 
     this._clients = []
     this._idle = []
+    this._expired = []
     this._pendingQueue = []
     this._endCallback = undefined
     this.ending = false
@@ -123,6 +125,24 @@ class Pool extends EventEmitter {
       }
       return
     }
+
+    // if there are idle clients; see if any lifetimes have expired
+    if (this._idle.length && this._expired.length) {
+      let newlyExpired = []
+      this._idle.forEach((idleItem) => {
+        const client = idleItem.client
+        const expiredIndex = this._expired.findIndex(c => c === client)
+        if (-1 !== expiredIndex) {
+          newlyExpired.push(expiredIndex)
+        }
+      })
+      newlyExpired.forEach((expiredIndex) => {
+        this.log('remove expired client')
+        const expiredClient = this._expired.splice(expiredIndex, 1)[0]
+        this._remove(expiredClient)
+      })
+    }
+
     // if we don't have any waiting, do nothing
     if (!this._pendingQueue.length) {
       this.log('no queued requests')
@@ -248,6 +268,17 @@ class Pool extends EventEmitter {
       } else {
         this.log('new client connected')
 
+        if (0 !== this.options.maxLifetimeSeconds) {
+          setTimeout(() => {
+            const expiredClient = this._clients.findIndex(c => c === client)
+            if (-1 !== expiredClient) {
+              this.log('ending client due to expired lifetime')
+              this._expired.push(client)
+              this._pulseQueue()
+            }
+          }, this.options.maxLifetimeSeconds * 1000)
+        }
+
         return this._acquireClient(client, pendingItem, idleListener, true)
       }
     })
@@ -318,6 +349,15 @@ class Pool extends EventEmitter {
       return
     }
 
+    const expiredClient = this._expired.findIndex(c => c === client)
+    if (-1 !== expiredClient) {
+      this.log('remove expired client')
+      this._expired.splice(expiredClient, 1)
+      this._remove(client)
+      this._pulseQueue()
+      return
+    }
+
     // idle timeout
     let tid
     if (this.options.idleTimeoutMillis) {
@@ -414,6 +454,10 @@ class Pool extends EventEmitter {
     return this._idle.length
   }
 
+  get expiredCount() {
+    return this._expired.length
+  }
+
   get totalCount() {
     return this._clients.length
   }

packages/pg-pool/test/lifetime-timeout.js

@@ -0,0 +1,48 @@
+'use strict'
+const co = require('co')
+const expect = require('expect.js')
+
+const describe = require('mocha').describe
+const it = require('mocha').it
+const path = require('path')
+
+const Pool = require('../')
+
+const wait = (time) => new Promise((resolve) => setTimeout(resolve, time))
+
+describe('lifetime timeout', () => {
+  it('connection lifetime should expire and remove the client', (done) => {
+    const pool = new Pool({ maxLifetimeSeconds: 1 })
+    pool.query('SELECT NOW()')
+    pool.on('remove', () => {
+      console.log('expired while idle - on-remove event')
+      expect(pool.expiredCount).to.equal(0)
+      expect(pool.totalCount).to.equal(0)
+      done()
+    })
+  })
+  it('connection lifetime should expire and remove the client after the client is done working', (done) => {
+    const pool = new Pool({ maxLifetimeSeconds: 1 })
+    pool.query('SELECT pg_sleep(1.01)')
+    pool.on('remove', () => {
+      console.log('expired while busy - on-remove event')
+      expect(pool.expiredCount).to.equal(0)
+      expect(pool.totalCount).to.equal(0)
+      done()
+    })
+  })
+  it('can remove expired clients and recreate them',
+    co.wrap(function* () {
+      const pool = new Pool({ maxLifetimeSeconds: 1 })
+      let query = pool.query('SELECT pg_sleep(1)')
+      expect(pool.expiredCount).to.equal(0)
+      expect(pool.totalCount).to.equal(1)
+      yield query
+      expect(pool.expiredCount).to.equal(0)
+      expect(pool.totalCount).to.equal(0)
+      yield pool.query('SELECT NOW()')
+      expect(pool.expiredCount).to.equal(0)
+      expect(pool.totalCount).to.equal(1)
+    })
+  )
+})