7244/multi db connection improvements

src/backend/distributed/commands/multi_copy.c

@@ -2160,6 +2160,7 @@ CitusCopyDestReceiverStartup(DestReceiver *dest, int operation,
 	copyDest->connectionStateHash = CreateConnectionStateHash(TopTransactionContext);
 
 	RecordRelationAccessIfNonDistTable(tableId, PLACEMENT_ACCESS_DML);
+	uint32 connectionFlags = 0;
 
 	/*
 	 * Colocated intermediate results do not honor citus.max_shared_pool_size,
@@ -2181,7 +2182,7 @@ CitusCopyDestReceiverStartup(DestReceiver *dest, int operation,
 		 * and cannot switch to local execution (e.g., disabled by user),
 		 * COPY would fail hinting the user to change the relevant settiing.
 		 */
-		EnsureConnectionPossibilityForRemotePrimaryNodes();
+		EnsureConnectionPossibilityForRemotePrimaryNodes(connectionFlags);
 	}
 
 	LocalCopyStatus localCopyStatus = GetLocalCopyStatus();
@@ -2211,7 +2212,8 @@ CitusCopyDestReceiverStartup(DestReceiver *dest, int operation,
 		 */
 		if (ShardIntervalListHasLocalPlacements(shardIntervalList))
 		{
-			bool reservedConnection = TryConnectionPossibilityForLocalPrimaryNode();
+			bool reservedConnection = TryConnectionPossibilityForLocalPrimaryNode(
+				connectionFlags);
 			copyDest->shouldUseLocalCopy = !reservedConnection;
 		}
 	}
@@ -3634,7 +3636,7 @@ CopyGetPlacementConnection(HTAB *connectionStateHash, ShardPlacement *placement,
 		return connection;
 	}
 
-	if (IsReservationPossible())
+	if (IsReservationPossible(connectionFlags))
 	{
 		/*
 		 * Enforce the requirements for adaptive connection management

src/backend/distributed/connection/connection_management.c

@@ -61,8 +61,8 @@ static MultiConnection * FindAvailableConnection(dlist_head *connections, uint32
 static void ErrorIfMultipleMetadataConnectionExists(dlist_head *connections);
 static void FreeConnParamsHashEntryFields(ConnParamsHashEntry *entry);
 static void AfterXactHostConnectionHandling(ConnectionHashEntry *entry, bool isCommit);
-static bool ShouldShutdownConnection(MultiConnection *connection, const int
-									 cachedConnectionCount);
+static bool ShouldShutdownConnection(MultiConnection *connection,
+									 const int cachedConnectionCount);
 static bool RemoteTransactionIdle(MultiConnection *connection);
 static int EventSetSizeForConnectionList(List *connections);
 
@@ -354,6 +354,11 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 		MultiConnection *connection = FindAvailableConnection(entry->connections, flags);
 		if (connection)
 		{
+			if (flags & REQUIRE_MAINTENANCE_CONNECTION)
+			{
+				/* Maintenance database may have changed, so cached connection should be closed */
+				connection->forceCloseAtTransactionEnd = true;
+			}
 			return connection;
 		}
 	}
@@ -377,9 +382,12 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 	/* these two flags are by nature cannot happen at the same time */
 	Assert(!((flags & WAIT_FOR_CONNECTION) && (flags & OPTIONAL_CONNECTION)));
 
+	int sharedCounterFlags = (flags & REQUIRE_MAINTENANCE_CONNECTION)
+							 ? MAINTENANCE_CONNECTION
+							 : 0;
 	if (flags & WAIT_FOR_CONNECTION)
 	{
-		WaitLoopForSharedConnection(hostname, port);
+		WaitLoopForSharedConnection(sharedCounterFlags, hostname, port);
 	}
 	else if (flags & OPTIONAL_CONNECTION)
 	{
@@ -389,7 +397,7 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 		 * cannot reserve the right to establish a connection, we prefer to
 		 * error out.
 		 */
-		if (!TryToIncrementSharedConnectionCounter(hostname, port))
+		if (!TryToIncrementSharedConnectionCounter(sharedCounterFlags, hostname, port))
 		{
 			/* do not track the connection anymore */
 			dlist_delete(&connection->connectionNode);
@@ -409,7 +417,7 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 		 *
 		 * Still, we keep track of the connection counter.
 		 */
-		IncrementSharedConnectionCounter(hostname, port);
+		IncrementSharedConnectionCounter(sharedCounterFlags, hostname, port);
 	}
 
 
@@ -423,11 +431,15 @@ StartNodeUserDatabaseConnection(uint32 flags, const char *hostname, int32 port,
 
 	ResetShardPlacementAssociation(connection);
 
-
-	if ((flags & REQUIRE_METADATA_CONNECTION))
+	if (flags & REQUIRE_METADATA_CONNECTION)
 	{
 		connection->useForMetadataOperations = true;
 	}
+	else if (flags & REQUIRE_MAINTENANCE_CONNECTION)
+	{
+		connection->useForMaintenanceOperations = true;
+		connection->forceCloseAtTransactionEnd = true;
+	}
 
 	/* fully initialized the connection, record it */
 	connection->initializationState = POOL_STATE_INITIALIZED;
@@ -495,6 +507,12 @@ FindAvailableConnection(dlist_head *connections, uint32 flags)
 			continue;
 		}
 
+		if ((flags & REQUIRE_MAINTENANCE_CONNECTION) &&
+			!connection->useForMaintenanceOperations)
+		{
+			continue;
+		}
+
 		if ((flags & REQUIRE_METADATA_CONNECTION) &&
 			!connection->useForMetadataOperations)
 		{
@@ -1191,7 +1209,11 @@ CitusPQFinish(MultiConnection *connection)
 	/* behave idempotently, there is no gurantee that CitusPQFinish() is called once */
 	if (connection->initializationState >= POOL_STATE_COUNTER_INCREMENTED)
 	{
-		DecrementSharedConnectionCounter(connection->hostname, connection->port);
+		int sharedCounterFlags = (connection->useForMaintenanceOperations)
+								 ? MAINTENANCE_CONNECTION
+								 : 0;
+		DecrementSharedConnectionCounter(sharedCounterFlags, connection->hostname,
+										 connection->port);
 		connection->initializationState = POOL_STATE_NOT_INITIALIZED;
 	}
 }

src/backend/distributed/connection/locally_reserved_shared_connections.c

@@ -92,9 +92,11 @@ static ReservedConnectionHashEntry * AllocateOrGetReservedConnectionEntry(char *
 																		  userId, Oid
 																		  databaseOid,
 																		  bool *found);
-static void EnsureConnectionPossibilityForNodeList(List *nodeList);
+static void EnsureConnectionPossibilityForNodeList(List *nodeList, uint32
+												   connectionFlags);
 static bool EnsureConnectionPossibilityForNode(WorkerNode *workerNode,
-											   bool waitForConnection);
+											   bool waitForConnection,
+											   uint32 connectionFlags);
 static uint32 LocalConnectionReserveHashHash(const void *key, Size keysize);
 static int LocalConnectionReserveHashCompare(const void *a, const void *b, Size keysize);
 
@@ -240,7 +242,9 @@ DeallocateReservedConnections(void)
 			 * We have not used this reservation, make sure to clean-up from
 			 * the shared memory as well.
 			 */
-			DecrementSharedConnectionCounter(entry->key.hostname, entry->key.port);
+			int sharedCounterFlags = 0;
+			DecrementSharedConnectionCounter(sharedCounterFlags, entry->key.hostname,
+											 entry->key.port);
 
 			/* for completeness, set it to true */
 			entry->usedReservation = true;
@@ -295,7 +299,7 @@ MarkReservedConnectionUsed(const char *hostName, int nodePort, Oid userId,
  * EnsureConnectionPossibilityForNodeList.
  */
 void
-EnsureConnectionPossibilityForRemotePrimaryNodes(void)
+EnsureConnectionPossibilityForRemotePrimaryNodes(uint32 connectionFlags)
 {
 	/*
 	 * By using NoLock there is a tiny risk of that we miss to reserve a
@@ -304,7 +308,7 @@ EnsureConnectionPossibilityForRemotePrimaryNodes(void)
 	 * going to access would be on the new node.
 	 */
 	List *remoteNodeList = ActivePrimaryRemoteNodeList(NoLock);
-	EnsureConnectionPossibilityForNodeList(remoteNodeList);
+	EnsureConnectionPossibilityForNodeList(remoteNodeList, connectionFlags);
 }
 
 
@@ -314,7 +318,7 @@ EnsureConnectionPossibilityForRemotePrimaryNodes(void)
  * If not, the function returns false.
  */
 bool
-TryConnectionPossibilityForLocalPrimaryNode(void)
+TryConnectionPossibilityForLocalPrimaryNode(uint32 connectionFlags)
 {
 	bool nodeIsInMetadata = false;
 	WorkerNode *localNode =
@@ -330,7 +334,8 @@ TryConnectionPossibilityForLocalPrimaryNode(void)
 	}
 
 	bool waitForConnection = false;
-	return EnsureConnectionPossibilityForNode(localNode, waitForConnection);
+	return EnsureConnectionPossibilityForNode(localNode, waitForConnection,
+											  connectionFlags);
 }
 
 
@@ -344,7 +349,7 @@ TryConnectionPossibilityForLocalPrimaryNode(void)
  *    single reservation per backend)
  */
 static void
-EnsureConnectionPossibilityForNodeList(List *nodeList)
+EnsureConnectionPossibilityForNodeList(List *nodeList, uint32 connectionFlags)
 {
 	/*
 	 * We sort the workerList because adaptive connection management
@@ -363,7 +368,8 @@ EnsureConnectionPossibilityForNodeList(List *nodeList)
 	foreach_ptr(workerNode, nodeList)
 	{
 		bool waitForConnection = true;
-		EnsureConnectionPossibilityForNode(workerNode, waitForConnection);
+		EnsureConnectionPossibilityForNode(workerNode, waitForConnection,
+										   connectionFlags);
 	}
 }
 
@@ -382,9 +388,10 @@ EnsureConnectionPossibilityForNodeList(List *nodeList)
  *   return false.
  */
 static bool
-EnsureConnectionPossibilityForNode(WorkerNode *workerNode, bool waitForConnection)
+EnsureConnectionPossibilityForNode(WorkerNode *workerNode, bool waitForConnection, uint32
+								   connectionFlags)
 {
-	if (!IsReservationPossible())
+	if (!IsReservationPossible(connectionFlags))
 	{
 		return false;
 	}
@@ -439,13 +446,17 @@ EnsureConnectionPossibilityForNode(WorkerNode *workerNode, bool waitForConnectio
 		 * Increment the shared counter, we may need to wait if there are
 		 * no space left.
 		 */
-		WaitLoopForSharedConnection(workerNode->workerName, workerNode->workerPort);
+		int sharedCounterFlags = 0;
+		WaitLoopForSharedConnection(sharedCounterFlags, workerNode->workerName,
+									workerNode->workerPort);
 	}
 	else
 	{
-		bool incremented =
-			TryToIncrementSharedConnectionCounter(workerNode->workerName,
-												  workerNode->workerPort);
+		int sharedCounterFlags = 0;
+		bool incremented = TryToIncrementSharedConnectionCounter(
+			sharedCounterFlags,
+			workerNode->workerName,
+			workerNode->workerPort);
 		if (!incremented)
 		{
 			/*
@@ -475,9 +486,13 @@ EnsureConnectionPossibilityForNode(WorkerNode *workerNode, bool waitForConnectio
  * session is eligible for shared connection reservation.
  */
 bool
-IsReservationPossible(void)
+IsReservationPossible(uint32 connectionFlags)
 {
-	if (GetMaxSharedPoolSize() == DISABLE_CONNECTION_THROTTLING)
+	bool connectionThrottlingDisabled =
+		connectionFlags & REQUIRE_MAINTENANCE_CONNECTION
+		? GetMaxMaintenanceSharedPoolSize() <= 0
+		: GetMaxSharedPoolSize() == DISABLE_CONNECTION_THROTTLING;
+	if (connectionThrottlingDisabled)
 	{
 		/* connection throttling disabled */
 		return false;