From faabb0696f92988c815fc675c61c1961c87a7047 Mon Sep 17 00:00:00 2001
From: Nathan Giardina <nathan.giardina@broadcom.com>
Date: Fri, 12 Aug 2022 20:40:44 +0000
Subject: [PATCH] Fix for timeout when a single node has timed out, created a
 new context to allow for each db node to timeout individually

---
 pgconn.go | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/pgconn.go b/pgconn.go
index 17f19e95..b23d7ded 100644
--- a/pgconn.go
+++ b/pgconn.go
@@ -128,19 +128,13 @@ func ConnectWithOptions(ctx context.Context, connString string, parseConfigOptio
 // authentication error will terminate the chain of attempts (like libpq:
 // https://www.postgresql.org/docs/11/libpq-connect.html#LIBPQ-MULTIPLE-HOSTS) and be returned as the error. Otherwise,
 // if all attempts fail the last error is returned.
-func ConnectConfig(ctx context.Context, config *Config) (pgConn *PgConn, err error) {
+func ConnectConfig(octx context.Context, config *Config) (pgConn *PgConn, err error) {
 	// Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from
 	// zero values.
 	if !config.createdByParseConfig {
 		panic("config must be created by ParseConfig")
 	}
 
-	// ConnectTimeout restricts the whole connection process.
-	if config.ConnectTimeout != 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, config.ConnectTimeout)
-		defer cancel()
-	}
 	// Simplify usage by treating primary config and fallbacks the same.
 	fallbackConfigs := []*FallbackConfig{
 		{
@@ -150,7 +144,7 @@ func ConnectConfig(ctx context.Context, config *Config) (pgConn *PgConn, err err
 		},
 	}
 	fallbackConfigs = append(fallbackConfigs, config.Fallbacks...)
-
+	ctx := octx
 	fallbackConfigs, err = expandWithIPs(ctx, config.LookupFunc, fallbackConfigs)
 	if err != nil {
 		return nil, &connectError{config: config, msg: "hostname resolving error", err: err}
@@ -163,6 +157,14 @@ func ConnectConfig(ctx context.Context, config *Config) (pgConn *PgConn, err err
 	foundBestServer := false
 	var fallbackConfig *FallbackConfig
 	for _, fc := range fallbackConfigs {
+		// ConnectTimeout restricts the whole connection process.
+		if config.ConnectTimeout != 0 {
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout)
+			defer cancel()
+		} else {
+			ctx = octx
+		}
 		pgConn, err = connect(ctx, config, fc, false)
 		if err == nil {
 			foundBestServer = true