From 641750bfaa586f61dfd471c7652ff138b4cc3f57 Mon Sep 17 00:00:00 2001
From: Tugberk Ugurlu <tugberk.ugurlu@deliveroo.co.uk>
Date: Wed, 14 Oct 2020 19:03:46 +0100
Subject: [PATCH] Add MaxConnLifetime config setting to pool

in light of the issue we are experiencing with https://github.com/jackc/pgx/issues/845,
and the fact that we want to have the ability to recycle our connections,
this patch adds the MaxConnLifetime config value and its behaviour to v3 releases.

The work done in this PR is trying to mimic what has done for v4 track: https://github.com/jackc/pgx/commit/c604afba82679df5ca91fbd6da922fdd205d4310

The only difference is that this one doesn't have a default maxConnLifetime value defined,
which means that the current v3 usage out there shouldn't be impacted by this change.
---
 conn.go           |  7 ++++++-
 conn_pool.go      | 11 +++++++----
 conn_pool_test.go | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/conn.go b/conn.go
index b67eb98ac..23d358c3c 100644
--- a/conn.go
+++ b/conn.go
@@ -224,6 +224,7 @@ func (cc *ConnConfig) networkAddresses() ([]net.Addr, error) {
 // goroutines.
 type Conn struct {
 	conn               net.Conn  // the underlying TCP or unix domain socket connection
+	creationTime       time.Time // the time indicating when the connection is established
 	lastActivityTime   time.Time // the last time the connection was used
 	wbuf               []byte
 	pid                uint32            // backend pid
@@ -457,7 +458,7 @@ func connect(config ConnConfig, connInfo *pgtype.ConnInfo) (c *Conn, err error)
 		}
 
 		c.addr = addr
-
+		c.creationTime = time.Now()
 		return c, nil
 	}
 
@@ -602,6 +603,10 @@ func (c *Conn) connect(config ConnConfig, network, address string, tlsConfig *tl
 	}
 }
 
+func (c *Conn) CreationTime() time.Time {
+	return c.creationTime
+}
+
 func initPostgresql(c *Conn) (*pgtype.ConnInfo, error) {
 	const (
 		namedOIDQuery = `select t.oid,
diff --git a/conn_pool.go b/conn_pool.go
index 95e1b015e..cd09cb31e 100644
--- a/conn_pool.go
+++ b/conn_pool.go
@@ -13,9 +13,10 @@ import (
 
 type ConnPoolConfig struct {
 	ConnConfig
-	MaxConnections int               // max simultaneous connections to use, default 5, must be at least 2
-	AfterConnect   func(*Conn) error // function to call on every new connection
-	AcquireTimeout time.Duration     // max wait time when all connections are busy (0 means no timeout)
+	MaxConnections  int               // max simultaneous connections to use, default 5, must be at least 2
+	MaxConnLifetime time.Duration     // the duration since creation after which a connection will be automatically closed.
+	AfterConnect    func(*Conn) error // function to call on every new connection
+	AcquireTimeout  time.Duration     // max wait time when all connections are busy (0 means no timeout)
 }
 
 type ConnPool struct {
@@ -25,6 +26,7 @@ type ConnPool struct {
 	config               ConnConfig // config used when establishing connection
 	inProgressConnects   int
 	maxConnections       int
+	maxConnLifetime      time.Duration
 	resetCount           int
 	afterConnect         func(*Conn) error
 	logger               Logger
@@ -60,6 +62,7 @@ func NewConnPool(config ConnPoolConfig) (p *ConnPool, err error) {
 	p.config = config.ConnConfig
 	p.connInfo = minimalConnInfo
 	p.maxConnections = config.MaxConnections
+	p.maxConnLifetime = config.MaxConnLifetime
 	if p.maxConnections == 0 {
 		p.maxConnections = 5
 	}
@@ -226,7 +229,7 @@ func (p *ConnPool) Release(conn *Conn) {
 
 	p.cond.L.Lock()
 
-	if conn.poolResetCount != p.resetCount {
+	if conn.poolResetCount != p.resetCount || (p.maxConnLifetime != 0 && time.Now().Sub(conn.CreationTime()) > p.maxConnLifetime) {
 		conn.Close()
 		p.cond.L.Unlock()
 		p.cond.Signal()
diff --git a/conn_pool_test.go b/conn_pool_test.go
index db645e637..27924ee53 100644
--- a/conn_pool_test.go
+++ b/conn_pool_test.go
@@ -15,6 +15,7 @@ import (
 
 func createConnPool(t *testing.T, maxConnections int) *pgx.ConnPool {
 	config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: maxConnections}
+	config.MaxConnLifetime = 250 * time.Millisecond
 	pool, err := pgx.NewConnPool(config)
 	if err != nil {
 		t.Fatalf("Unable to create connection pool: %v", err)
@@ -150,6 +151,32 @@ func TestPoolAcquireAndReleaseCycle(t *testing.T) {
 	releaseAllConnections(pool, allConnections)
 }
 
+func TestPoolReleaseChecksMaxConnLifetime(t *testing.T) {
+	t.Parallel()
+
+	config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig}
+	pool, err := pgx.NewConnPool(config)
+	if err != nil {
+		t.Fatal("Unable to establish connection pool")
+	}
+	defer pool.Close()
+
+	c, err := pool.Acquire()
+	if err != nil {
+		t.Fatal("Unable to acquire connection from the pool")
+	}
+
+	time.Sleep(config.MaxConnLifetime)
+
+	pool.Release(c)
+	waitForReleaseToComplete()
+
+	stats := pool.Stat()
+	if stats.CurrentConnections != 0 && stats.AvailableConnections != 0 {
+		t.Fatal("Unable to recycle connection from the pool")
+	}
+}
+
 func TestPoolNonBlockingConnections(t *testing.T) {
 	t.Parallel()
 
@@ -1227,3 +1254,10 @@ func TestConnPoolBeginEx(t *testing.T) {
 		t.Fatal("Should not be able to create a tx")
 	}
 }
+
+// Conn.Release is an asynchronous process that returns immediately. There is no signal when the actual work is
+// completed. To test something that relies on the actual work for Conn.Release being completed we must simply wait.
+// This function wraps the sleep so there is more meaning for the callers.
+func waitForReleaseToComplete() {
+	time.Sleep(5 * time.Millisecond)
+}
\ No newline at end of file