From da510a705526ba84a64a6435096936e16ac4960e Mon Sep 17 00:00:00 2001
From: ChrisG0x20 <position0x45@hotmail.com>
Date: Tue, 14 Dec 2021 09:33:26 -0800
Subject: [PATCH] Add connection lifetime limit option and tests

---
 packages/pg-pool/index.js                 | 27 +++++++++++++
 packages/pg-pool/test/lifetime-timeout.js | 46 +++++++++++++++++++++++
 2 files changed, 73 insertions(+)
 create mode 100644 packages/pg-pool/test/lifetime-timeout.js

diff --git a/packages/pg-pool/index.js b/packages/pg-pool/index.js
index 48bf5c788..46d2aab0c 100644
--- a/packages/pg-pool/index.js
+++ b/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 = new WeakSet()
     this._pendingQueue = []
     this._endCallback = undefined
     this.ending = false
@@ -123,6 +125,7 @@ class Pool extends EventEmitter {
       }
       return
     }
+
     // if we don't have any waiting, do nothing
     if (!this._pendingQueue.length) {
       this.log('no queued requests')
@@ -248,6 +251,17 @@ class Pool extends EventEmitter {
       } else {
         this.log('new client connected')
 
+        if (this.options.maxLifetimeSeconds !== 0) {
+          setTimeout(() => {
+            this.log('ending client due to expired lifetime')
+            this._expired.add(client)
+            const idleIndex = this._idle.findIndex(idleItem => idleItem.client === client)
+            if (idleIndex !== -1) {
+              this._acquireClient(client, new PendingItem((err, client, clientRelease) => clientRelease()), idleListener, false)
+            }
+          }, this.options.maxLifetimeSeconds * 1000)
+        }
+
         return this._acquireClient(client, pendingItem, idleListener, true)
       }
     })
@@ -318,6 +332,15 @@ class Pool extends EventEmitter {
       return
     }
 
+    const isExpired = this._expired.has(client)
+    if (isExpired) {
+      this.log('remove expired client')
+      this._expired.delete(client)
+      this._remove(client)
+      this._pulseQueue()
+      return
+    }
+
     // idle timeout
     let tid
     if (this.options.idleTimeoutMillis) {
@@ -414,6 +437,10 @@ class Pool extends EventEmitter {
     return this._idle.length
   }
 
+  get expiredCount() {
+    return this._clients.reduce((acc, client) => acc + (this._expired.has(client) ? 1 : 0), 0)
+  }
+
   get totalCount() {
     return this._clients.length
   }
diff --git a/packages/pg-pool/test/lifetime-timeout.js b/packages/pg-pool/test/lifetime-timeout.js
new file mode 100644
index 000000000..986161625
--- /dev/null
+++ b/packages/pg-pool/test/lifetime-timeout.js
@@ -0,0 +1,46 @@
+'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('../')
+
+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)
+    })
+  )
+})