Handle invalid connection hash entries (#4362)

If MemoryContextAlloc errors out -e.g. during an OOM-, ConnectionHashEntry->connections
stays as NULL.

With this commit, we add isValid flag to ConnectionHashEntry that should be set to true
right after we allocate & initialize ConnectionHashEntry->connections list properly, and we
check it before accesing to ConnectionHashEntry->connections.

src/backend/distributed/connection/connection_management.c

@@ -161,6 +161,12 @@ AfterXactConnectionHandling(bool isCommit)
 	hash_seq_init(&status, ConnectionHash);
 	while ((entry = (ConnectionHashEntry *) hash_seq_search(&status)) != 0)
 	{
+		if (!entry->isValid)
+		{
+			/* skip invalid connection hash entries */
+			continue;
+		}
+
 		AfterXactHostConnectionHandling(entry, isCommit);
 
 		/*
@@ -289,11 +295,24 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 	 */
 
 	ConnectionHashEntry *entry = hash_search(ConnectionHash, &key, HASH_ENTER, &found);
-	if (!found)
+	if (!found || !entry->isValid)
 	{
+		/*
+		 * We are just building hash entry or previously it was left in an
+		 * invalid state as we couldn't allocate memory for it.
+		 * So initialize entry->connections list here.
+		 */
+		entry->isValid = false;
 		entry->connections = MemoryContextAlloc(ConnectionContext,
 												sizeof(dlist_head));
 		dlist_init(entry->connections);
+
+		/*
+		 * If MemoryContextAlloc errors out -e.g. during an OOM-, entry->connections
+		 * stays as NULL. So entry->isValid should be set to true right after we
+		 * initialize entry->connections properly.
+		 */
+		entry->isValid = true;
 	}
 
 	/* if desired, check whether there's a usable connection */
@@ -449,6 +468,12 @@ CloseAllConnectionsAfterTransaction(void)
 	hash_seq_init(&status, ConnectionHash);
 	while ((entry = (ConnectionHashEntry *) hash_seq_search(&status)) != 0)
 	{
+		if (!entry->isValid)
+		{
+			/* skip invalid connection hash entries */
+			continue;
+		}
+
 		dlist_iter iter;
 
 		dlist_head *connections = entry->connections;
@@ -483,7 +508,7 @@ ConnectionAvailableToNode(char *hostName, int nodePort, const char *userName,
 	ConnectionHashEntry *entry =
 		(ConnectionHashEntry *) hash_search(ConnectionHash, &key, HASH_FIND, &found);
 
-	if (!found)
+	if (!found || !entry->isValid)
 	{
 		return false;
 	}
@@ -509,6 +534,12 @@ CloseNodeConnectionsAfterTransaction(char *nodeName, int nodePort)
 	hash_seq_init(&status, ConnectionHash);
 	while ((entry = (ConnectionHashEntry *) hash_seq_search(&status)) != 0)
 	{
+		if (!entry->isValid)
+		{
+			/* skip invalid connection hash entries */
+			continue;
+		}
+
 		dlist_iter iter;
 
 		if (strcmp(entry->key.hostname, nodeName) != 0 || entry->key.port != nodePort)
@@ -584,6 +615,12 @@ ShutdownAllConnections(void)
 	hash_seq_init(&status, ConnectionHash);
 	while ((entry = (ConnectionHashEntry *) hash_seq_search(&status)) != NULL)
 	{
+		if (!entry->isValid)
+		{
+			/* skip invalid connection hash entries */
+			continue;
+		}
+
 		dlist_iter iter;
 
 		dlist_foreach(iter, entry->connections)
@@ -1194,6 +1231,12 @@ FreeConnParamsHashEntryFields(ConnParamsHashEntry *entry)
 static void
 AfterXactHostConnectionHandling(ConnectionHashEntry *entry, bool isCommit)
 {
+	if (!entry || !entry->isValid)
+	{
+		/* callers only pass valid hash entries but let's be on the safe side */
+		ereport(ERROR, (errmsg("connection hash entry is NULL or invalid")));
+	}
+
 	dlist_mutable_iter iter;
 	int cachedConnectionCount = 0;
 

src/include/distributed/connection_management.h

@@ -178,6 +178,9 @@ typedef struct ConnectionHashEntry
 {
 	ConnectionHashKey key;
 	dlist_head *connections;
+
+	/* connections list is valid or not */
+	bool isValid;
 } ConnectionHashEntry;
 
 /* hash entry for cached connection parameters */