From 749542e607005b70e8412113b3b3f70eed75de37 Mon Sep 17 00:00:00 2001
From: Blake Gentry <blakesgentry@gmail.com>
Date: Wed, 14 Jun 2023 10:02:35 -0500
Subject: [PATCH] Faster shutdown of pgxpool.Pool background goroutines

When a pool is closed, some background goroutines may be left open,
particularly for health checks as detailed in #1641. Two specific
examples have been refactored here to avoid a blocking sleep and instead
also select on the pool being closed to potentially return/continue
sooner.
---
 pgxpool/pool.go | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index f4c3a30ca..53af62918 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -93,6 +93,9 @@ type Pool struct {
 	maxConnIdleTime       time.Duration
 	healthCheckPeriod     time.Duration
 
+	healthCheckMu    sync.Mutex
+	healthCheckTimer *time.Timer
+
 	healthCheckChan chan struct{}
 
 	closeOnce sync.Once
@@ -381,15 +384,25 @@ func (p *Pool) isExpired(res *puddle.Resource[*connResource]) bool {
 }
 
 func (p *Pool) triggerHealthCheck() {
-	go func() {
+	const healthCheckDelay = 500 * time.Millisecond
+
+	p.healthCheckMu.Lock()
+	defer p.healthCheckMu.Unlock()
+
+	if p.healthCheckTimer == nil {
 		// Destroy is asynchronous so we give it time to actually remove itself from
 		// the pool otherwise we might try to check the pool size too soon
-		time.Sleep(500 * time.Millisecond)
-		select {
-		case p.healthCheckChan <- struct{}{}:
-		default:
-		}
-	}()
+		p.healthCheckTimer = time.AfterFunc(healthCheckDelay, func() {
+			select {
+			case <-p.closeChan:
+			case p.healthCheckChan <- struct{}{}:
+			default:
+			}
+		})
+		return
+	}
+
+	p.healthCheckTimer.Reset(healthCheckDelay)
 }
 
 func (p *Pool) backgroundHealthCheck() {
@@ -408,6 +421,9 @@ func (p *Pool) backgroundHealthCheck() {
 }
 
 func (p *Pool) checkHealth() {
+	ticker := time.NewTicker(500 * time.Millisecond)
+	defer ticker.Stop()
+
 	for {
 		// If checkMinConns failed we don't destroy any connections since we couldn't
 		// even get to minConns
@@ -424,7 +440,7 @@ func (p *Pool) checkHealth() {
 		select {
 		case <-p.closeChan:
 			return
-		case <-time.After(500 * time.Millisecond):
+		case <-ticker.C:
 		}
 	}
 }