Quick fix for possible segfault in PurgeConnection

Now that connections can be acquired without going through the cache,
we have to handle cases where functions assume the cache has been ini-
tialized.

src/backend/distributed/test/connection_cache.c

@@ -13,7 +13,6 @@
 #include "postgres.h"
 #include "c.h"
 #include "fmgr.h"
-
 #include "libpq-fe.h"
 
 #include <stddef.h>
@@ -22,7 +21,9 @@
 
 #include "catalog/pg_type.h"
 #include "distributed/connection_cache.h"
+#include "distributed/metadata_cache.h"
 #include "distributed/test_helper_functions.h" /* IWYU pragma: keep */
+#include "utils/elog.h"
 #include "utils/lsyscache.h"
 
 
@@ -34,6 +35,7 @@ static Datum ExtractIntegerDatum(char *input);
 PG_FUNCTION_INFO_V1(initialize_remote_temp_table);
 PG_FUNCTION_INFO_V1(count_remote_temp_table_rows);
 PG_FUNCTION_INFO_V1(get_and_purge_connection);
+PG_FUNCTION_INFO_V1(connect_and_purge_connection);
 PG_FUNCTION_INFO_V1(set_connection_status_bad);
 
 
@@ -128,6 +130,30 @@ get_and_purge_connection(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * get_and_purge_connection first gets a connection using the provided hostname
+ * and port before immediately passing that connection to PurgeConnection. This
+ * is to test PurgeConnection behvaior when circumventing the cache.
+ */
+Datum
+connect_and_purge_connection(PG_FUNCTION_ARGS)
+{
+	char *nodeName = PG_GETARG_CSTRING(0);
+	int32 nodePort = PG_GETARG_INT32(1);
+	char *nodeUser = CurrentUserName();
+
+	PGconn *connection = ConnectToNode(nodeName, nodePort, nodeUser);
+	if (connection == NULL)
+	{
+		PG_RETURN_BOOL(false);
+	}
+
+	PurgeConnection(connection);
+
+	PG_RETURN_BOOL(true);
+}
+
+
 /*
  * set_connection_status_bad does not remove the given connection from the connection hash.
  * It simply shuts down the underlying socket. On success, it returns true.

src/backend/distributed/utils/connection_cache.c

@@ -14,16 +14,14 @@
 #include "libpq-fe.h"
 #include "miscadmin.h"
 
-
 #include <stddef.h>
+#include <stdio.h>
 #include <string.h>
 
 #include "commands/dbcommands.h"
 #include "distributed/connection_cache.h"
 #include "distributed/metadata_cache.h"
-#include "lib/stringinfo.h"
 #include "mb/pg_wchar.h"
-#include "nodes/pg_list.h"
 #include "utils/builtins.h"
 #include "utils/elog.h"
 #include "utils/errcodes.h"
@@ -172,10 +170,6 @@ PurgeConnection(PGconn *connection)
 	 */
 	if (purgedConnection != connection)
 	{
-		ereport(WARNING, (errmsg("hash entry for \"%s:%d\" contained different "
-								 "connection than that provided by caller",
-								 nodeConnectionKey.nodeName,
-								 nodeConnectionKey.nodePort)));
 		PQfinish(connection);
 	}
 }
@@ -186,22 +180,21 @@ PurgeConnectionByKey(NodeConnectionKey *nodeConnectionKey)
 {
 	bool entryFound = false;
 	NodeConnectionEntry *nodeConnectionEntry = NULL;
+	PGconn *connection = NULL;
+
+	if (NodeConnectionHash != NULL)
+	{
+		nodeConnectionEntry = hash_search(NodeConnectionHash, nodeConnectionKey,
+										  HASH_REMOVE, &entryFound);
+	}
 
-	nodeConnectionEntry = hash_search(NodeConnectionHash, nodeConnectionKey, HASH_REMOVE,
-									  &entryFound);
 	if (entryFound)
 	{
+		connection = nodeConnectionEntry->connection;
 		PQfinish(nodeConnectionEntry->connection);
 	}
-	else
-	{
-		ereport(WARNING, (errcode(ERRCODE_NO_DATA),
-						  errmsg("could not find hash entry for connection to \"%s:%d\"",
-								 nodeConnectionKey->nodeName,
-								 nodeConnectionKey->nodePort)));
-	}
 
-	return nodeConnectionEntry->connection;
+	return connection;
 }
 
 

src/include/distributed/test_helper_functions.h

@@ -39,6 +39,7 @@ extern Datum fake_fdw_handler(PG_FUNCTION_ARGS);
 extern Datum initialize_remote_temp_table(PG_FUNCTION_ARGS);
 extern Datum count_remote_temp_table_rows(PG_FUNCTION_ARGS);
 extern Datum get_and_purge_connection(PG_FUNCTION_ARGS);
+extern Datum connect_and_purge_connection(PG_FUNCTION_ARGS);
 extern Datum set_connection_status_bad(PG_FUNCTION_ARGS);
 
 /* function declarations for exercising metadata functions */

src/test/regress/expected/multi_connection_cache.out

@@ -15,6 +15,10 @@ CREATE FUNCTION get_and_purge_connection(cstring, integer)
 	RETURNS bool
 	AS 'citus'
 	LANGUAGE C STRICT;
+CREATE FUNCTION connect_and_purge_connection(cstring, integer)
+	RETURNS bool
+	AS 'citus'
+	LANGUAGE C STRICT;
 CREATE FUNCTION set_connection_status_bad(cstring, integer)
 	RETURNS bool
 	AS 'citus'
@@ -104,3 +108,11 @@ SELECT get_and_purge_connection('localhost', :worker_port);
 (1 row)
 
 SET client_min_messages TO DEFAULT;
+\c
+-- purge existing connection to localhost
+SELECT connect_and_purge_connection('localhost', :worker_port);
+ connect_and_purge_connection 
+------------------------------
+ t
+(1 row)
+

src/test/regress/sql/multi_connection_cache.sql

@@ -22,6 +22,11 @@ CREATE FUNCTION get_and_purge_connection(cstring, integer)
 	AS 'citus'
 	LANGUAGE C STRICT;
 
+CREATE FUNCTION connect_and_purge_connection(cstring, integer)
+	RETURNS bool
+	AS 'citus'
+	LANGUAGE C STRICT;
+
 CREATE FUNCTION set_connection_status_bad(cstring, integer)
 	RETURNS bool
 	AS 'citus'
@@ -76,3 +81,8 @@ SELECT count_remote_temp_table_rows('localhost', :worker_port);
 SELECT get_and_purge_connection('localhost', :worker_port);
 
 SET client_min_messages TO DEFAULT;
+
+\c
+
+-- purge existing connection to localhost
+SELECT connect_and_purge_connection('localhost', :worker_port);