Retry Cached Remote Connections for postgres_fdw in case remote backend gets killed/goes away

From b3725acc916e0daa2797b9d159fc6cc5cdcf43c6 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Sat, 11 Jul 2020 14:57:58 +0530
Subject: [PATCH v2] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local backend.
There are high chances that the remote backend can no logner be
avaiable i.e. it can get killed, the subsequent foreign queries
from local backend/session that uses cached connection fails as
the remote backend woule have become unavailable/killed. So,
this patch, solves this problem,
1. local backend/session uses cached connection, but recieves error
2. upon receiving the first error, it deletes the cached entry
3. try to get new connection at the start of begin remote xact
---
 contrib/postgres_fdw/connection.c             | 130 +++++++++++++++++-
 .../postgres_fdw/expected/postgres_fdw.out    |  30 ++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     |  19 +++
 3 files changed, 175 insertions(+), 4 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 52d1fe3563..e3658dd704 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -45,6 +45,19 @@
  */
 typedef Oid ConnCacheKey;
 
+typedef enum ConnectionRetryType
+{
+	/* initial value for all cache entries */
+	CONN_RETRY_NONE,
+	/* indicates that the caller is ready to retry connection */
+	CONN_RETRY_READY,
+	/*
+	 * indicates to the caller that the connection may have
+	 * broken, it's okay to retry to get a new connection.
+	*/
+	CONN_RETRY_DO
+}ConnectionRetryType;
+
 typedef struct ConnCacheEntry
 {
 	ConnCacheKey key;			/* hash key (must be first) */
@@ -58,8 +71,15 @@ typedef struct ConnCacheEntry
 	bool		invalidated;	/* true if reconnect is pending */
 	uint32		server_hashvalue;	/* hash value of foreign server OID */
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
+	/* conn retry status, default is CONN_RETRY_NONE for all entries */
+	uint8       conn_retry;
 } ConnCacheEntry;
 
+#define IsRemoteXactBegin(sql) \
+ ((strcmp(sql, "START TRANSACTION ISOLATION LEVEL SERIALIZABLE") == 0 || \
+  strcmp(sql, "START TRANSACTION ISOLATION LEVEL REPEATABLE READ") == 0) ? \
+  1 : 0)
+
 /*
  * Connection cache (initialized on first use)
  */
@@ -109,6 +129,8 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	bool		found;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
+	bool	dobeginxact = true;
+	bool	doretry		= false;
 
 	/* First time through, initialize connection cache hashtable */
 	if (ConnectionHash == NULL)
@@ -170,10 +192,54 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
+	 * We check the health of cached connection here, while the remote xact is
+	 * goting to begin.  Broken connection will be detected and if so, clear the
+	 * existing conn, get a new connection and resubmit the remote xact begin.
+	 * We retry the broken connections only once during each begin remote xact
+	 * call, still the broken connection exist after that then, the query fails.
 	 */
+	if (entry->conn != NULL &&
+		entry->xact_depth <= 0 &&
+		entry->conn_retry == CONN_RETRY_NONE)
+	{
+		/*
+		 * This is before begin remote xact, so safe to do retry, hence indicate
+		 * that in the cached entry.
+		 */
+		entry->conn_retry = CONN_RETRY_READY;
+		begin_remote_xact(entry);
+
+		/*
+		 * The previously set RETRY_READY status gets changed to RETRY_DO, if there
+		 * exists a broken connection, if so, then clear the entry.
+		 */
+		if (entry->conn_retry == CONN_RETRY_DO)
+		{
+			elog(DEBUG3, "postgres_fdw cached connection %p is broken, hence closing it", entry->conn);
+			/*
+			 * It is okay to call disconnect here thought the connection is broken.
+			 * No termninate message will actually be sent to remote backend by
+			 * sendTerminateConn() as the connection status will be CONNECTION_BAD.
+			 * So disconnect will just do the clearing and resetting of parameters
+			 * in entry data structure.
+			 */
+			disconnect_pg_server(entry);
+			doretry = true;
+		}
+		else if (entry->conn_retry == CONN_RETRY_READY)
+		{
+			elog(DEBUG3, "postgres_fdw cached connection %p is healthy, hence using it", entry->conn);
+			/*
+			 * Connection is not broken, so remote xact being command is executed
+			 * successfully. Mark that here, to avoid further begin remote xact call
+			 * in this funciton.
+			 */
+			dobeginxact = false;
+		}
+
+		/* Reset the conn_retry to default value. */
+		entry->conn_retry = CONN_RETRY_NONE;
+	}
 
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
@@ -199,15 +265,20 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 
 		/* Now try to make the connection */
 		entry->conn = connect_pg_server(server, user);
+		entry->conn_retry = CONN_RETRY_NONE;
 
 		elog(DEBUG3, "new postgres_fdw connection %p for server \"%s\" (user mapping oid %u, userid %u)",
 			 entry->conn, server->servername, user->umid, user->userid);
+
+		if (doretry)
+			elog(DEBUG3, "postgres_fdw connection retry is successful");
 	}
 
 	/*
 	 * Start a new transaction or subtransaction if needed.
 	 */
-	begin_remote_xact(entry);
+	if (dobeginxact)
+		begin_remote_xact(entry);
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
@@ -457,6 +528,17 @@ do_sql_command(PGconn *conn, const char *sql)
 	if (!PQsendQuery(conn, sql))
 		pgfdw_report_error(ERROR, NULL, conn, false, sql);
 	res = pgfdw_get_result(conn, sql);
+
+	/*
+	 * If this function is called from begin_remote_xact and returned res is null
+	 * from pgfdw_get_result, which means that the remote connection is broken, so
+	 * just return from here, the cached entry's retry status would have been set
+	 * appropriately to indicate the broken connection.
+	 */
+	if (res == NULL &&
+		IsRemoteXactBegin(sql))
+		return;
+
 	if (PQresultStatus(res) != PGRES_COMMAND_OK)
 		pgfdw_report_error(ERROR, res, conn, true, sql);
 	PQclear(res);
@@ -591,6 +673,7 @@ PGresult *
 pgfdw_get_result(PGconn *conn, const char *query)
 {
 	PGresult   *volatile last_res = NULL;
+	bool 		retry  = false;
 
 	/* In what follows, do not leak any PGresults on an error. */
 	PG_TRY();
@@ -617,10 +700,49 @@ pgfdw_get_result(PGconn *conn, const char *query)
 				if (wc & WL_SOCKET_READABLE)
 				{
 					if (!PQconsumeInput(conn))
+					{
+						/*
+						 * Do this on error and for only begin remote xact
+						 * commands.
+						 */
+						if (IsRemoteXactBegin(query) &&
+							ConnectionHash != NULL)
+						{
+							HASH_SEQ_STATUS scan;
+							ConnCacheEntry *entry;
+
+							hash_seq_init(&scan, ConnectionHash);
+							while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+							{
+								if (entry->conn == conn &&
+									entry->conn_retry == CONN_RETRY_READY)
+								{
+									/* got the required connection from the cache. */
+									entry->conn_retry = CONN_RETRY_DO;
+									hash_seq_term(&scan);
+									retry = true;
+									break;
+								}
+							}
+						}
+
+						/* if retry to be done, don't report error. */
+						if (retry == true)
+							break;
+
 						pgfdw_report_error(ERROR, NULL, conn, false, query);
+					}
 				}
 			}
 
+			if (retry == true)
+			{
+				/* clear if there is any previous res. */
+				PQclear(last_res);
+				last_res = NULL;
+				break;
+			}
+
 			res = PQgetResult(conn);
 			if (res == NULL)
 				break;			/* query is complete */
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 82fc1290ef..a813c4918f 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8974,3 +8974,33 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 83971665e3..9552625d2e 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2634,3 +2634,22 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
-- 
2.25.1

From 5be13fbfd31ab5af9265ac9cb3adcf0589bf2a8e Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Tue, 14 Jul 2020 07:23:14 +0530
Subject: [PATCH v3] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local backend.
There are high chances that the remote backend can no logner be
avaiable i.e. it can get killed, the subsequent foreign queries
from local backend/session that uses cached connection fails as
the remote backend woule have become unavailable/killed. So,
this patch, solves this problem,
1. local backend/session uses cached connection, but recieves error
2. upon receiving the first error, it deletes the cached entry
3. try to get new connection at the start of begin remote xact
---
 contrib/postgres_fdw/connection.c             | 65 ++++++++++++++++---
 .../postgres_fdw/expected/postgres_fdw.out    | 55 ++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 31 +++++++++
 3 files changed, 143 insertions(+), 8 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 52d1fe3563..b7ae8dbbac 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -169,12 +169,6 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -205,9 +199,64 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of cached connection here, while the remote xact is
+	 * going to begin.  Broken connection will be detected and if so, get a
+	 * new connection and resubmit the remote xact begin. We retry the broken
+	 * connections only once during each begin remote xact call, still the broken
+	 * connection exist after that then, the query fails.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+	}
+	PG_CATCH();
+	{
+		if (entry->conn &&
+			PQstatus(entry->conn) == CONNECTION_BAD &&
+			entry->xact_depth <= 0)
+		{
+			/* server object will be used for log messages. */
+			ForeignServer *server = GetForeignServer(user->serverid);
+
+			/*
+			 * xact_depth <=0, which means we are retrying only at the beginning
+			 * of remote xact. This is sensible, since we can't retry in the
+			 * middle of a remote xact.
+			 */
+			elog(DEBUG3, "postgres_fdw connection %p to the server \"%s\" is broken. retrying to connect",
+				entry->conn, server->servername);
+
+			PQreset(entry->conn);
+
+			if (entry->conn && PQstatus(entry->conn) == CONNECTION_OK)
+				elog(DEBUG3, "postgres_fdw connection retry to the server  \"%s\" is successful",
+					server->servername);
+			else
+				ereport(ERROR,
+					(errcode(ERRCODE_SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION),
+					 errmsg("could not connect to server \"%s\"",
+							server->servername),
+					 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn)))));
+
+			/* Retry starting a new transaction. */
+			begin_remote_xact(entry);
+		}
+		else
+		{
+			/*
+			 * We are here, due to either some error other than CONNECTION_BAD
+			 * occurred or connection may have broken during start of a
+			 * subtransacation. Just, clear off any result, try rethrowing the
+			 * error, so that it will be caught appropriately.
+			 */
+			PGresult *res = NULL;
+			res = PQgetResult(entry->conn);
+			PQclear(res);
+			PG_RE_THROW();
+		}
+	}
+	PG_END_TRY();
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 90db550b92..46511a340e 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8974,3 +8974,58 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 83971665e3..6f647450e7 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2634,3 +2634,34 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
-- 
2.25.1

From cc911b79ed1bbcf1a330977b69cc4fe291c2a3a0 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 25 Sep 2020 10:11:59 +0530
Subject: [PATCH] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 62 ++++++++++++++++---
 .../postgres_fdw/expected/postgres_fdw.out    | 55 ++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 31 ++++++++++
 3 files changed, 140 insertions(+), 8 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..ab8d505d64 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -170,12 +170,6 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +200,61 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote
+	 * xact is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+	}
+	PG_CATCH();
+	{
+		if (entry->conn &&
+			PQstatus(entry->conn) == CONNECTION_BAD &&
+			entry->xact_depth <= 0)
+		{
+			/* server object will be used for log messages. */
+			ForeignServer *server = GetForeignServer(user->serverid);
+
+			/*
+			 * xact_depth <=0, which means we are retrying only at
+			 * the beginning of remote xact. This is sensible,
+			 * since we can't retry in the middle of a remote xact.
+			 */
+			elog(DEBUG3, "postgres_fdw connection %p to the server \"%s\" is broken, retrying to connect",
+				entry->conn, server->servername);
+
+			disconnect_pg_server(entry);
+
+			/*
+			 * Reset the transient state fields that may have been
+			 * modified in the first attempt.
+			 */
+			entry->changing_xact_state = false;
+
+			entry->conn = connect_pg_server(server, user);
+
+			if (entry->conn && PQstatus(entry->conn) == CONNECTION_OK)
+				elog(DEBUG3, "postgres_fdw connection retry to the server  \"%s\" is successful",
+					server->servername);
+			else
+				ereport(ERROR,
+					(errcode(ERRCODE_SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION),
+					 errmsg("could not connect to server \"%s\"",
+							server->servername),
+					 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn)))));
+
+			/* Retry starting a new transaction. */
+			begin_remote_xact(entry);
+		}
+		else
+			PG_RE_THROW();
+	}
+	PG_END_TRY();
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..d72b11a3ac 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,58 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..7a1b4c78bd 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,34 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
-- 
2.25.1

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..9f76261d99 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	bool		need_reset_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+reset:
 	/*
 	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * soon as we're out of all transactions. Also if previous attempt to start
+	 * new remote transaction failed on the cached connection, disconnect
+	 * it to reestablish new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated && entry->xact_depth == 0) ||
+		need_reset_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (need_reset_conn)
+			elog(DEBUG3, "closing connection %p to reestablish connection",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,31 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote
+	 * xact is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		need_reset_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (!entry->conn ||
+			PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			need_reset_conn)
+			PG_RE_THROW();
+		need_reset_conn = true;
+	}
+	PG_END_TRY();
+
+	if (need_reset_conn)
+		goto reset;
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..d72b11a3ac 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,58 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..7a1b4c78bd 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,34 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+COMMIT;

From 0df20a1cf6d084f8bf3d8768d29893416fd3c83e Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 25 Sep 2020 17:30:35 +0530
Subject: [PATCH v5] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 54 +++++++++++++-----
 .../postgres_fdw/expected/postgres_fdw.out    | 55 +++++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 31 +++++++++++
 3 files changed, 127 insertions(+), 13 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..435cf656f7 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
-	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * If the connection needs to be remade due to invalidation, disconnect
+	 * as soon as we're out of all transactions. Also, if previous attempt
+	 * to start new remote transaction failed on the cached connection,
+	 * disconnect it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,34 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote
+	 * xact is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (!entry->conn ||
+			PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		elog(DEBUG3, "could not connect to postgres_fdw connection %p", entry->conn);
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..d72b11a3ac 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,58 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+  c1  | c2 | c3 | c4 | c5 | c6 |     c7     | c8 
+------+----+----+----+----+----+------------+----
+ 1111 |  2 |    |    |    |    | ft1        | 
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..7a1b4c78bd 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,34 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT * FROM ft1 LIMIT 1;
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT * FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+SELECT * FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
-- 
2.25.1

From 6c5545c4deb1783a09ecb869a9b189add325099e Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Tue, 29 Sep 2020 21:16:30 +0530
Subject: [PATCH v6] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 57 ++++++++++++++-----
 .../postgres_fdw/expected/postgres_fdw.out    | 55 ++++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 31 ++++++++++
 3 files changed, 130 insertions(+), 13 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..16eca788ac 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
-	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * If the connection needs to be remade due to invalidation, disconnect
+	 * as soon as we're out of all transactions. Also, if previous attempt
+	 * to start new remote transaction failed on the cached connection,
+	 * disconnect it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,37 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote
+	 * xact is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (!entry->conn ||
+			PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		ereport(DEBUG3,
+				(errmsg("could not start remote transaction on connection %p",
+				 entry->conn)),
+				 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn))));
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..e16a5b7e04 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,58 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+ ?column? 
+----------
+        1
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..db18d88323 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,34 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Retry cached connections at the beginning of the remote xact
+-- in case remote backend is killed.
+-- Let's use a different application name for remote connection,
+-- so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check
+-- and kill it intentionally here. Note that, local backend still has
+-- the remote connection/backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Next query using the same foreign server should succeed if connection
+-- retry works.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Subtransaction - remote backend is killed in the middle of a remote
+-- xact, and we don't do retry connection, hence the subsequent query
+-- using the same foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
-- 
2.25.1

From e2a85116116e79c76a7b0fcfc1a25eff57627ed4 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Thu, 1 Oct 2020 15:16:26 +0530
Subject: [PATCH v7] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 55 +++++++++---
 .../postgres_fdw/expected/postgres_fdw.out    | 85 +++++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 64 ++++++++++++++
 3 files changed, 192 insertions(+), 12 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..441124963f 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
 	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * soon as we're out of all transactions. Also, if previous attempt to
+	 * start new remote transaction failed on the cached connection, disconnect
+	 * it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,37 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote xact
+	 * is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (!entry->conn ||
+			PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		ereport(DEBUG3,
+				(errmsg("could not start remote transaction on connection %p",
+				 entry->conn)),
+				 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn))));
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..09d0578d05 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,88 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Connection retry feature
+-- Retry cached connections at the beginning of the remote xact in case remote
+-- backend is killed. Let's use a different application name for remote
+-- connection, so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- This procedure waits until the terminated backend has really gone away. We
+-- also had to check for wait event type and wait event of the terminated
+-- backend's row from pg_stat_activity. Otherwise, the usage of this procedure
+-- in an explicit txn block(BEGIN - END) after terminating the backend would
+-- still have an entry with wait event type wait event null, as we have not yet
+-- committed the txn. The entry of the terminated backend from pg_stat_activity
+-- would be removed only after the txn commit.
+CREATE OR REPLACE PROCEDURE wait_for_backend_termination()
+LANGUAGE plpgsql
+AS $$
+DECLARE proccnt int;
+BEGIN
+    LOOP
+		SELECT count(*) INTO proccnt FROM pg_stat_activity
+				WHERE backend_type = 'client backend' AND
+				  application_name = 'fdw_retry_check' AND
+				  wait_event_type IS NOT NULL AND
+				  wait_event IS NOT NULL;
+		EXIT WHEN proccnt = 0;
+    END LOOP;
+END;
+$$;
+-- Connection retry feature test case 001:
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check and
+-- kill intentionally here. Note that, local backend still has the remote
+-- backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Despite issuing terminate, there are chances that the backend is still not
+-- actually terminated. Hence, wait for the terminated backend to go away.
+CALL wait_for_backend_termination();
+-- Next query using the same foreign server should succeed if connection retry
+-- works.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Connection retry feature test case 002:
+-- Subtransaction - remote backend is killed in the middle of a remote xact and
+-- we don't do retry connection, hence the subsequent query using the same
+-- foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+ ?column? 
+----------
+        1
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+CALL wait_for_backend_termination();
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
+-- Clean up
+DROP PROCEDURE wait_for_backend_termination();
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..17246f4176 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,67 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Connection retry feature
+-- Retry cached connections at the beginning of the remote xact in case remote
+-- backend is killed. Let's use a different application name for remote
+-- connection, so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- This procedure waits until the terminated backend has really gone away. We
+-- also had to check for wait event type and wait event of the terminated
+-- backend's row from pg_stat_activity. Otherwise, the usage of this procedure
+-- in an explicit txn block(BEGIN - END) after terminating the backend would
+-- still have an entry with wait event type wait event null, as we have not yet
+-- committed the txn. The entry of the terminated backend from pg_stat_activity
+-- would be removed only after the txn commit.
+CREATE OR REPLACE PROCEDURE wait_for_backend_termination()
+LANGUAGE plpgsql
+AS $$
+DECLARE proccnt int;
+BEGIN
+    LOOP
+		SELECT count(*) INTO proccnt FROM pg_stat_activity
+				WHERE backend_type = 'client backend' AND
+				  application_name = 'fdw_retry_check' AND
+				  wait_event_type IS NOT NULL AND
+				  wait_event IS NOT NULL;
+		EXIT WHEN proccnt = 0;
+    END LOOP;
+END;
+$$;
+
+-- Connection retry feature test case 001:
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check and
+-- kill intentionally here. Note that, local backend still has the remote
+-- backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Despite issuing terminate, there are chances that the backend is still not
+-- actually terminated. Hence, wait for the terminated backend to go away.
+CALL wait_for_backend_termination();
+
+-- Next query using the same foreign server should succeed if connection retry
+-- works.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Connection retry feature test case 002:
+-- Subtransaction - remote backend is killed in the middle of a remote xact and
+-- we don't do retry connection, hence the subsequent query using the same
+-- foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+CALL wait_for_backend_termination();
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
+
+-- Clean up
+DROP PROCEDURE wait_for_backend_termination();
\ No newline at end of file
-- 
2.25.1

From 7aa1b7e1b9820a92beded06a965275fd237040cd Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Thu, 1 Oct 2020 21:13:05 +0530
Subject: [PATCH v8] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 55 +++++++++++---
 .../postgres_fdw/expected/postgres_fdw.out    | 76 +++++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 55 ++++++++++++++
 3 files changed, 174 insertions(+), 12 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..441124963f 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
 	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * soon as we're out of all transactions. Also, if previous attempt to
+	 * start new remote transaction failed on the cached connection, disconnect
+	 * it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,37 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here, while the remote xact
+	 * is going to begin.  Broken connection will be detected and the
+	 * connection is retried. We do this only once during each begin remote
+	 * xact call, if the connection is still broken after the retry attempt,
+	 * then the query will fail.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (!entry->conn ||
+			PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		ereport(DEBUG3,
+				(errmsg("could not start remote transaction on connection %p",
+				 entry->conn)),
+				 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn))));
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..db039237e1 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,79 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- Connection retry feature
+-- Retry cached connections at the beginning of the remote xact in case remote
+-- backend is killed. Let's use a different application name for remote
+-- connection, so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+-- This procedure waits until the terminated backend has really gone away. We
+-- do pg_stat_clear_snapshot() to fetch the new snapshot of the stats.
+CREATE OR REPLACE PROCEDURE wait_for_backend_termination()
+LANGUAGE plpgsql
+AS $$
+BEGIN
+     LOOP
+         PERFORM * FROM pg_stat_activity WHERE application_name = 'fdw_retry_check';
+         EXIT WHEN NOT FOUND;
+         PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+     END LOOP;
+END;
+$$;
+-- Connection retry feature test case 001:
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Retrieve pid of remote backend with application name fdw_retry_check and
+-- kill intentionally here. Note that, local backend still has the remote
+-- backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+-- Despite issuing terminate, there are chances that the backend is still not
+-- actually terminated. Hence, wait for the terminated backend to go away.
+CALL wait_for_backend_termination();
+-- Next query using the same foreign server should succeed if connection retry
+-- works.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Connection retry feature test case 002:
+-- Subtransaction - remote backend is killed in the middle of a remote xact and
+-- we don't do retry connection, hence the subsequent query using the same
+-- foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+ ?column? 
+----------
+        1
+(1 row)
+
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+ pg_terminate_backend 
+----------------------
+ t
+(1 row)
+
+CALL wait_for_backend_termination();
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
+-- Clean up
+DROP PROCEDURE wait_for_backend_termination();
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..3d4eb8414d 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,58 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- Connection retry feature
+-- Retry cached connections at the beginning of the remote xact in case remote
+-- backend is killed. Let's use a different application name for remote
+-- connection, so that this test will not kill other backends wrongly.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+
+-- This procedure waits until the terminated backend has really gone away. We
+-- do pg_stat_clear_snapshot() to fetch the new snapshot of the stats.
+CREATE OR REPLACE PROCEDURE wait_for_backend_termination()
+LANGUAGE plpgsql
+AS $$
+BEGIN
+     LOOP
+         PERFORM * FROM pg_stat_activity WHERE application_name = 'fdw_retry_check';
+         EXIT WHEN NOT FOUND;
+         PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+     END LOOP;
+END;
+$$;
+
+-- Connection retry feature test case 001:
+-- Generate a connection to remote. Local backend will cache it.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Retrieve pid of remote backend with application name fdw_retry_check and
+-- kill intentionally here. Note that, local backend still has the remote
+-- backend info in it's cache.
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+
+-- Despite issuing terminate, there are chances that the backend is still not
+-- actually terminated. Hence, wait for the terminated backend to go away.
+CALL wait_for_backend_termination();
+
+-- Next query using the same foreign server should succeed if connection retry
+-- works.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Connection retry feature test case 002:
+-- Subtransaction - remote backend is killed in the middle of a remote xact and
+-- we don't do retry connection, hence the subsequent query using the same
+-- foreign server should fail.
+BEGIN;
+DECLARE c CURSOR FOR SELECT 1 FROM ft1 LIMIT 1;
+FETCH c;
+SAVEPOINT s;
+SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE
+backend_type = 'client backend' AND application_name = 'fdw_retry_check';
+CALL wait_for_backend_termination();
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
+
+-- Clean up
+DROP PROCEDURE wait_for_backend_termination();
-- 
2.25.1

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..901d3a4661 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
 	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * soon as we're out of all transactions. Also, if previous attempt to
+	 * start new remote transaction failed on the cached connection, disconnect
+	 * it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,35 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here when starting
+	 * a new remote transaction. If broken connection is detected,
+	 * we try to reestablish a new connection. If broken connection is
+	 * detected again here, we give up getting a connection.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		ereport(DEBUG3,
+				(errmsg_internal("could not start remote transaction on connection %p",
+				 entry->conn)),
+				 errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn))));
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..2c5614073f 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,51 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- ===================================================================
+-- reestablish new connection
+-- ===================================================================
+-- Terminate the backend having the specified application_name and wait for
+-- the termination to complete.
+CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+BEGIN
+    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
+    WHERE application_name = appname;
+    LOOP
+        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
+        EXIT WHEN NOT FOUND;
+        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+    END LOOP;
+END;
+$$ LANGUAGE plpgsql;
+-- Change application_name of remote connection to special one
+-- so that we can easily terminate the connection later.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Terminate the remote connection.
+CALL terminate_backend_and_wait('fdw_retry_check');
+-- This query should detect the broken connection when starting new remote
+-- transaction, reestablish new connection, and then succeed.
+BEGIN;
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- If the query detects the broken connection when starting new remote
+-- subtransaction, it doesn't reestablish new connection and should fail.
+CALL terminate_backend_and_wait('fdw_retry_check');
+SAVEPOINT s;
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
+-- Clean up
+DROP PROCEDURE terminate_backend_and_wait(text);
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..4da1f78956 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,44 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- ===================================================================
+-- reestablish new connection
+-- ===================================================================
+
+-- Terminate the backend having the specified application_name and wait for
+-- the termination to complete.
+CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+BEGIN
+    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
+    WHERE application_name = appname;
+    LOOP
+        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
+        EXIT WHEN NOT FOUND;
+        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+    END LOOP;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Change application_name of remote connection to special one
+-- so that we can easily terminate the connection later.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Terminate the remote connection.
+CALL terminate_backend_and_wait('fdw_retry_check');
+
+-- This query should detect the broken connection when starting new remote
+-- transaction, reestablish new connection, and then succeed.
+BEGIN;
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- If the query detects the broken connection when starting new remote
+-- subtransaction, it doesn't reestablish new connection and should fail.
+CALL terminate_backend_and_wait('fdw_retry_check');
+SAVEPOINT s;
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
+
+-- Clean up
+DROP PROCEDURE terminate_backend_and_wait(text);

From 68d764690fc0409ee15afbbbd00974e0d43116cf Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Sat, 3 Oct 2020 05:49:18 +0530
Subject: [PATCH v10] Retry Cached Remote Connections For postgres_fdw

Remote connections are cached for postgres_fdw in the local
backend. There are high chances that the remote backend may
not be available i.e. it can get killed, the subsequent foreign
queries from local backend/session that use cached connection
will fail as the remote backend would have become unavailable.

This patch proposes a retry mechanism. Local backend/session
uses cached connection and if it receives connection bad error,
it deletes the cached entry and retry to get a new connection.
This retry happens only at the start of remote xact.
---
 contrib/postgres_fdw/connection.c             | 53 ++++++++++++++-----
 .../postgres_fdw/expected/postgres_fdw.out    | 48 +++++++++++++++++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 41 ++++++++++++++
 3 files changed, 130 insertions(+), 12 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 08daf26fdf..5fa4b634c3 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -108,6 +108,7 @@ PGconn *
 GetConnection(UserMapping *user, bool will_prep_stmt)
 {
 	bool		found;
+	volatile bool retry_conn = false;
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 
@@ -159,23 +160,25 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	/* Reject further use of connections which failed abort cleanup. */
 	pgfdw_reject_incomplete_xact_state_change(entry);
 
+retry:
 	/*
 	 * If the connection needs to be remade due to invalidation, disconnect as
-	 * soon as we're out of all transactions.
+	 * soon as we're out of all transactions. Also, if previous attempt to
+	 * start new remote transaction failed on the cached connection,
+	 * disconnect it to retry a new connection.
 	 */
-	if (entry->conn != NULL && entry->invalidated && entry->xact_depth == 0)
+	if ((entry->conn != NULL && entry->invalidated &&
+		 entry->xact_depth == 0) || retry_conn)
 	{
-		elog(DEBUG3, "closing connection %p for option changes to take effect",
-			 entry->conn);
+		if (retry_conn)
+			elog(DEBUG3, "closing connection %p to reestablish a new one",
+				 entry->conn);
+		else
+			elog(DEBUG3, "closing connection %p for option changes to take effect",
+				 entry->conn);
 		disconnect_pg_server(entry);
 	}
 
-	/*
-	 * We don't check the health of cached connection here, because it would
-	 * require some overhead.  Broken connection will be detected when the
-	 * connection is actually used.
-	 */
-
 	/*
 	 * If cache entry doesn't have a connection, we have to establish a new
 	 * connection.  (If connect_pg_server throws an error, the cache entry
@@ -206,9 +209,35 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	}
 
 	/*
-	 * Start a new transaction or subtransaction if needed.
+	 * We check the health of the cached connection here when starting a new
+	 * remote transaction. If a broken connection is detected in the first
+	 * attempt, we try to reestablish a new connection. If broken connection
+	 * is detected again here, we give up getting a connection.
 	 */
-	begin_remote_xact(entry);
+	PG_TRY();
+	{
+		/* Start a new transaction or subtransaction if needed. */
+		begin_remote_xact(entry);
+		retry_conn = false;
+	}
+	PG_CATCH();
+	{
+		if (PQstatus(entry->conn) != CONNECTION_BAD ||
+			entry->xact_depth > 0 ||
+			retry_conn)
+			PG_RE_THROW();
+		retry_conn = true;
+	}
+	PG_END_TRY();
+
+	if (retry_conn)
+	{
+		ereport(DEBUG3,
+				(errmsg_internal("could not start remote transaction on connection %p",
+								 entry->conn)),
+				errdetail_internal("%s", pchomp(PQerrorMessage(entry->conn))));
+		goto retry;
+	}
 
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 10e23d02ed..2c5614073f 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8987,3 +8987,51 @@ PREPARE TRANSACTION 'fdw_tpc';
 ERROR:  cannot PREPARE a transaction that has operated on postgres_fdw foreign tables
 ROLLBACK;
 WARNING:  there is no transaction in progress
+-- ===================================================================
+-- reestablish new connection
+-- ===================================================================
+-- Terminate the backend having the specified application_name and wait for
+-- the termination to complete.
+CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+BEGIN
+    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
+    WHERE application_name = appname;
+    LOOP
+        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
+        EXIT WHEN NOT FOUND;
+        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+    END LOOP;
+END;
+$$ LANGUAGE plpgsql;
+-- Change application_name of remote connection to special one
+-- so that we can easily terminate the connection later.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Terminate the remote connection.
+CALL terminate_backend_and_wait('fdw_retry_check');
+-- This query should detect the broken connection when starting new remote
+-- transaction, reestablish new connection, and then succeed.
+BEGIN;
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- If the query detects the broken connection when starting new remote
+-- subtransaction, it doesn't reestablish new connection and should fail.
+CALL terminate_backend_and_wait('fdw_retry_check');
+SAVEPOINT s;
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+ERROR:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  remote SQL command: SAVEPOINT s2
+COMMIT;
+-- Clean up
+DROP PROCEDURE terminate_backend_and_wait(text);
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 78156d10b4..4da1f78956 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2653,3 +2653,44 @@ SELECT count(*) FROM ft1;
 -- error here
 PREPARE TRANSACTION 'fdw_tpc';
 ROLLBACK;
+
+-- ===================================================================
+-- reestablish new connection
+-- ===================================================================
+
+-- Terminate the backend having the specified application_name and wait for
+-- the termination to complete.
+CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+BEGIN
+    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
+    WHERE application_name = appname;
+    LOOP
+        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
+        EXIT WHEN NOT FOUND;
+        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+    END LOOP;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Change application_name of remote connection to special one
+-- so that we can easily terminate the connection later.
+ALTER SERVER loopback OPTIONS (application_name 'fdw_retry_check');
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Terminate the remote connection.
+CALL terminate_backend_and_wait('fdw_retry_check');
+
+-- This query should detect the broken connection when starting new remote
+-- transaction, reestablish new connection, and then succeed.
+BEGIN;
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- If the query detects the broken connection when starting new remote
+-- subtransaction, it doesn't reestablish new connection and should fail.
+CALL terminate_backend_and_wait('fdw_retry_check');
+SAVEPOINT s;
+SELECT 1 FROM ft1 LIMIT 1;    -- should fail
+COMMIT;
+
+-- Clean up
+DROP PROCEDURE terminate_backend_and_wait(text);
-- 
2.25.1