From 34b17b1020860fe151c8c5318b39e1d261bee00c Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onderkalaci@gmail.com>
Date: Wed, 16 Feb 2022 20:40:47 +0100
Subject: [PATCH] Fix isolation tester

---
 .../test/run_from_same_connection.c            | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/backend/distributed/test/run_from_same_connection.c b/src/backend/distributed/test/run_from_same_connection.c
index 59a30a850..2eeca472e 100644
--- a/src/backend/distributed/test/run_from_same_connection.c
+++ b/src/backend/distributed/test/run_from_same_connection.c
@@ -131,11 +131,6 @@ start_session_level_connection_to_node(PG_FUNCTION_ARGS)
 static void
 EstablishSingleConnnection(char *nodeNameString, int nodePort)
 {
-	if (LocalConnectionContext != NULL)
-	{
-		MemoryContextDelete(LocalConnectionContext);
-		LocalConnectionContext = NULL;
-	}
 
 	LocalConnectionContext =
 		AllocSetContextCreateExtended(ConnectionContext,
@@ -253,12 +248,19 @@ stop_session_level_connection_to_node(PG_FUNCTION_ARGS)
 {
 	allowNonIdleRemoteTransactionOnXactHandling = false;
 
-	if (singleConnection != NULL)
+	if (LocalConnectionContext != NULL)
 	{
-		PQfinish(singleConnection->pgConn);
-		singleConnection = NULL;
+		if (singleConnection != NULL)
+		{
+			PQfinish(singleConnection->pgConn);
+			singleConnection = NULL;
+		}
+
+		MemoryContextDelete(LocalConnectionContext);
+		LocalConnectionContext = NULL;
 	}
 
+
 	PG_RETURN_VOID();
 }