Merge pull request #3631 from citusdata/improve_at_exit

Properly terminate connections at the end session

src/backend/distributed/connection/connection_management.c

@@ -494,6 +494,37 @@ CloseConnection(MultiConnection *connection)
 }
 
 
+/*
+ * ShutdownAllConnections shutdowns all the MultiConnections in the
+ * ConnectionHash.
+ *
+ * This function is intended to be called atexit() of the backend, so
+ * that the cached connections are closed properly. Calling this function
+ * at another point in the code could be dangerous, so think twice if you
+ * need to call this function.
+ */
+void
+ShutdownAllConnections(void)
+{
+	ConnectionHashEntry *entry = NULL;
+	HASH_SEQ_STATUS status;
+
+	hash_seq_init(&status, ConnectionHash);
+	while ((entry = (ConnectionHashEntry *) hash_seq_search(&status)) != NULL)
+	{
+		dlist_iter iter;
+
+		dlist_foreach(iter, entry->connections)
+		{
+			MultiConnection *connection =
+				dlist_container(MultiConnection, connectionNode, iter.cur);
+
+			ShutdownConnection(connection);
+		}
+	}
+}
+
+
 /*
  * ShutdownConnection, if necessary cancels the currently running statement,
  * and then closes the underlying libpq connection.  The MultiConnection

src/backend/distributed/shared_library_init.c

@@ -87,6 +87,7 @@ static char *CitusVersion = CITUS_VERSION;
 
 void _PG_init(void);
 
+static void CitusBackendAtExit(void);
 static void ResizeStackToMaximumDepth(void);
 static void multi_log_hook(ErrorData *edata);
 static void CreateRequiredDirectories(void);
@@ -270,6 +271,8 @@ _PG_init(void)
 	InitPlacementConnectionManagement();
 	InitializeCitusQueryStats();
 
+	atexit(CitusBackendAtExit);
+
 	/* enable modification of pg_catalog tables during pg_upgrade */
 	if (IsBinaryUpgrade)
 	{
@@ -279,6 +282,18 @@ _PG_init(void)
 }
 
 
+/*
+ * CitusBackendAtExit is called atexit of the backend for the purposes of
+ * any clean-up needed.
+ */
+static void
+CitusBackendAtExit(void)
+{
+	/* properly close all the cached connections */
+	ShutdownAllConnections();
+}
+
+
 /*
  * Stack size increase during high memory load may cause unexpected crashes.
  * With this alloca call, we are increasing stack size explicitly, so that if

src/include/distributed/connection_management.h

@@ -221,6 +221,7 @@ extern MultiConnection * StartNodeUserDatabaseConnection(uint32 flags,
 														 const char *database);
 extern void CloseNodeConnectionsAfterTransaction(char *nodeName, int nodePort);
 extern void CloseConnection(MultiConnection *connection);
+extern void ShutdownAllConnections(void);
 extern void ShutdownConnection(MultiConnection *connection);
 
 /* dealing with a connection */