PR #6728 / commit - 12

Force activated bare connections to close at transaction end.

src/backend/distributed/connection/connection_management.c

@@ -1202,6 +1202,17 @@ FinishConnectionEstablishment(MultiConnection *connection)
 }
 
 
+/*
+ * ForceConnectionCloseAtTransactionEnd marks connection to be closed at the end of the
+ * transaction.
+ */
+void
+ForceConnectionCloseAtTransactionEnd(MultiConnection *connection)
+{
+	connection->forceCloseAtTransactionEnd = true;
+}
+
+
 /*
  * ClaimConnectionExclusively signals that this connection is actively being
  * used. That means it'll not be, again, returned by

src/backend/distributed/metadata/metadata_sync.c

@@ -210,9 +210,6 @@ start_metadata_sync_to_node(PG_FUNCTION_ARGS)
 	ActivateNodeList(context);
 	TransactionModifiedNodeMetadata = true;
 
-	/* cleanup metadata memory context and connections */
-	DestroyMetadataSyncContext(context);
-
 	PG_RETURN_VOID();
 }
 
@@ -246,9 +243,6 @@ start_metadata_sync_to_all_nodes(PG_FUNCTION_ARGS)
 	ActivateNodeList(context);
 	TransactionModifiedNodeMetadata = true;
 
-	/* cleanup metadata memory context and connections */
-	DestroyMetadataSyncContext(context);
-
 	PG_RETURN_BOOL(true);
 }
 
@@ -4002,6 +3996,7 @@ EstablishAndSetMetadataSyncBareConnections(MetadataSyncContext *context)
 																	NULL);
 
 		Assert(connection != NULL);
+		ForceConnectionCloseAtTransactionEnd(connection);
 		bareConnectionList = lappend(bareConnectionList, connection);
 	}
 
@@ -4059,26 +4054,6 @@ CreateMetadataSyncContext(List *nodeList, bool collectCommands,
 }
 
 
-/*
- * DestroyMetadataSyncContext destroys the memory context inside metadataSyncContext
- * and also closes open connections if any.
- */
-void
-DestroyMetadataSyncContext(MetadataSyncContext *context)
-{
-	/* todo: make sure context is always cleanup by using resource release callback?? */
-	/* close connections */
-	MultiConnection *connection = NULL;
-	foreach_ptr(connection, context->activatedWorkerBareConnections)
-	{
-		CloseConnection(connection);
-	}
-
-	/* delete memory context */
-	MemoryContextDelete(context->context);
-}
-
-
 /*
  * ResetMetadataSyncMemoryContext resets memory context inside metadataSyncContext, if
  * we are not collecting commands.

src/backend/distributed/metadata/node_metadata.c

@@ -771,9 +771,6 @@ citus_activate_node(PG_FUNCTION_ARGS)
 	ActivateNodeList(context);
 	TransactionModifiedNodeMetadata = true;
 
-	/* cleanup metadata memory context and connections */
-	DestroyMetadataSyncContext(context);
-
 	PG_RETURN_INT32(workerNode->nodeId);
 }
 
@@ -2240,9 +2237,6 @@ AddNodeMetadataViaMetadataContext(char *nodeName, int32 nodePort,
 
 	ActivateNodeList(context);
 
-	/* cleanup metadata memory context and connections */
-	DestroyMetadataSyncContext(context);
-
 	return nodeId;
 }
 

src/backend/distributed/test/metadata_sync.c

@@ -86,9 +86,6 @@ activate_node_snapshot(PG_FUNCTION_ARGS)
 		activateNodeCommandCount,
 		ddlCommandTypeId);
 
-	/* cleanup metadata memory context and connections */
-	DestroyMetadataSyncContext(context);
-
 	PG_RETURN_ARRAYTYPE_P(activateNodeCommandArrayType);
 }
 

src/include/distributed/connection_management.h

@@ -323,6 +323,7 @@ extern void ShutdownConnection(MultiConnection *connection);
 /* dealing with a connection */
 extern void FinishConnectionListEstablishment(List *multiConnectionList);
 extern void FinishConnectionEstablishment(MultiConnection *connection);
+extern void ForceConnectionCloseAtTransactionEnd(MultiConnection *connection);
 extern void ClaimConnectionExclusively(MultiConnection *connection);
 extern void UnclaimConnection(MultiConnection *connection);
 extern void MarkConnectionConnected(MultiConnection *connection);

src/include/distributed/metadata_sync.h

@@ -141,7 +141,6 @@ extern void SyncDeleteColocationGroupToNodes(uint32 colocationId);
 extern MetadataSyncContext * CreateMetadataSyncContext(List *nodeList,
 													   bool collectCommands,
 													   bool nodesAddedInSameTransaction);
-extern void DestroyMetadataSyncContext(MetadataSyncContext *context);
 extern void EstablishAndSetMetadataSyncBareConnections(MetadataSyncContext *context);
 extern void SetMetadataSyncNodesFromNodeList(MetadataSyncContext *context,
 											 List *nodeList);