From c0eb67b24faf2a3b69f9f8b088eb4f2f24a809eb Mon Sep 17 00:00:00 2001
From: Halil Ozan Akgul <hozanakgul@gmail.com>
Date: Mon, 1 Nov 2021 13:47:19 +0300
Subject: [PATCH] Skip forceCloseAtTransactionEnd connections only if BEGIN was
 not sent on them

---
 .../connection/connection_management.c        | 14 +++--
 .../regress/expected/node_conninfo_reload.out | 51 +++++++++++++++++++
 src/test/regress/sql/node_conninfo_reload.sql | 24 +++++++++
 3 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/src/backend/distributed/connection/connection_management.c b/src/backend/distributed/connection/connection_management.c
index 4ce015361..1768575c7 100644
--- a/src/backend/distributed/connection/connection_management.c
+++ b/src/backend/distributed/connection/connection_management.c
@@ -429,12 +429,18 @@ FindAvailableConnection(dlist_head *connections, uint32 flags)
 			continue;
 		}
 
-		if (connection->forceCloseAtTransactionEnd)
+		if (connection->forceCloseAtTransactionEnd &&
+			!connection->remoteTransaction.beginSent)
 		{
 			/*
-			 * This is a connection that should be closed, probabably because
-			 * of old connection options. So we ignore it. It will
-			 * automatically be closed at the end of the transaction.
+			 * This is a connection that should be closed, probably because
+			 * of old connection options or removing a node. This will
+			 * automatically be closed at the end of the transaction. But, if we are still
+			 * inside a transaction, we should keep using this connection as long as a remote
+			 * transaction is in progress over the connection. The main use for this case
+			 * is having some commands inside a transaction block after removing nodes. And, we
+			 * currently allow very limited operations after removing a node inside a
+			 * transaction block (e.g., no placement access can happen).
 			 */
 			continue;
 		}
diff --git a/src/test/regress/expected/node_conninfo_reload.out b/src/test/regress/expected/node_conninfo_reload.out
index 8a82eebee..f3c0f561d 100644
--- a/src/test/regress/expected/node_conninfo_reload.out
+++ b/src/test/regress/expected/node_conninfo_reload.out
@@ -458,5 +458,56 @@ select count(*) from test where a = 0;
      0
 (1 row)
 
+-- Test connecting all the shards
+ALTER SYSTEM SET citus.node_conninfo = 'sslmode=doesnotexist';
+BEGIN;
+ALTER TABLE test ADD COLUMN b INT;
+select pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
+select pg_sleep(0.1); -- wait for config reload to apply
+ pg_sleep
+---------------------------------------------------------------------
+
+(1 row)
+
+show citus.node_conninfo;
+ citus.node_conninfo
+---------------------------------------------------------------------
+ sslmode=doesnotexist
+(1 row)
+
+-- Should work since connections to the same shards that BEGIN is sent
+-- are reused.
+ALTER TABLE test ADD COLUMN c INT;
+COMMIT;
+-- Should fail now, when transaction is finished
+ALTER TABLE test ADD COLUMN d INT;
+ERROR:  connection to the remote node localhost:xxxxx failed with the following error: invalid sslmode value: "doesnotexist"
+-- Reset it again
+ALTER SYSTEM RESET citus.node_conninfo;
+select pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
+select pg_sleep(0.1); -- wait for config reload to apply
+ pg_sleep
+---------------------------------------------------------------------
+
+(1 row)
+
+show citus.node_conninfo;
+ citus.node_conninfo
+---------------------------------------------------------------------
+ sslmode=require
+(1 row)
+
+-- Should work again
+ALTER TABLE test ADD COLUMN e INT;
 DROP SCHEMA node_conninfo_reload CASCADE;
 NOTICE:  drop cascades to table test
diff --git a/src/test/regress/sql/node_conninfo_reload.sql b/src/test/regress/sql/node_conninfo_reload.sql
index 3754130a7..cd76399b5 100644
--- a/src/test/regress/sql/node_conninfo_reload.sql
+++ b/src/test/regress/sql/node_conninfo_reload.sql
@@ -171,4 +171,28 @@ show citus.node_conninfo;
 -- Should work
 select count(*) from test where a = 0;
 
+-- Test connecting all the shards
+ALTER SYSTEM SET citus.node_conninfo = 'sslmode=doesnotexist';
+BEGIN;
+ALTER TABLE test ADD COLUMN b INT;
+select pg_reload_conf();
+select pg_sleep(0.1); -- wait for config reload to apply
+show citus.node_conninfo;
+-- Should work since connections to the same shards that BEGIN is sent
+-- are reused.
+ALTER TABLE test ADD COLUMN c INT;
+COMMIT;
+
+-- Should fail now, when transaction is finished
+ALTER TABLE test ADD COLUMN d INT;
+
+-- Reset it again
+ALTER SYSTEM RESET citus.node_conninfo;
+select pg_reload_conf();
+select pg_sleep(0.1); -- wait for config reload to apply
+show citus.node_conninfo;
+
+-- Should work again
+ALTER TABLE test ADD COLUMN e INT;
+
 DROP SCHEMA node_conninfo_reload CASCADE;