From 7f1a2fcd6ba2cd8e7ba18441fe45e7311a58c220 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 24 Oct 2018 04:35:14 -0700
Subject: [PATCH] rdbms: fix connection reuse logic in LoadBalancer for
 postgres

Change-Id: Ic4b7af156fbed54aae0a34a0bbba76490cd505e0
---
 .../libs/rdbms/loadbalancer/LoadBalancer.php  | 46 +++++++++++++------
 1 file changed, 31 insertions(+), 15 deletions(-)

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 1f3fe4cf54..d157a231b9 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -1001,6 +1001,8 @@ class LoadBalancer implements ILoadBalancer {
 		}
 
 		/** @var Database $conn */
+		$conn = null;
+
 		if ( isset( $this->conns[$connInUseKey][$i][$domain] ) ) {
 			// Reuse an in-use connection for the same domain
 			$conn = $this->conns[$connInUseKey][$i][$domain];
@@ -1012,22 +1014,36 @@ class LoadBalancer implements ILoadBalancer {
 			$this->conns[$connInUseKey][$i][$domain] = $conn;
 			$this->connLogger->debug( __METHOD__ . ": reusing free connection $i/$domain" );
 		} elseif ( !empty( $this->conns[$connFreeKey][$i] ) ) {
-			// Reuse a free connection from another domain
-			$conn = reset( $this->conns[$connFreeKey][$i] );
-			$oldDomain = key( $this->conns[$connFreeKey][$i] );
-			if ( $domainInstance->getDatabase() !== null ) {
-				$conn->selectDomain( $domainInstance );
-			} else {
-				// Stay on the current database, but update the schema/prefix
-				$conn->dbSchema( $domainInstance->getSchema() );
-				$conn->tablePrefix( $domainInstance->getTablePrefix() );
+			// Reuse a free connection from another domain if possible
+			foreach ( $this->conns[$connFreeKey][$i] as $oldDomain => $conn ) {
+				if ( $domainInstance->getDatabase() !== null ) {
+					// Check if changing the database will require a new connection.
+					// In that case, leave the connection handle alone and keep looking.
+					// This prevents connections from being closed mid-transaction and can
+					// also avoid overhead if the same database will later be requested.
+					if (
+						$conn->databasesAreIndependent() &&
+						$conn->getDBname() !== $domainInstance->getDatabase()
+					) {
+						continue;
+					}
+					// Select the new database, schema, and prefix
+					$conn->selectDomain( $domainInstance );
+				} else {
+					// Stay on the current database, but update the schema/prefix
+					$conn->dbSchema( $domainInstance->getSchema() );
+					$conn->tablePrefix( $domainInstance->getTablePrefix() );
+				}
+				unset( $this->conns[$connFreeKey][$i][$oldDomain] );
+				// Note that if $domain is an empty string, getDomainID() might not match it
+				$this->conns[$connInUseKey][$i][$conn->getDomainId()] = $conn;
+				$this->connLogger->debug( __METHOD__ .
+					": reusing free connection from $oldDomain for $domain" );
+				break;
 			}
-			unset( $this->conns[$connFreeKey][$i][$oldDomain] );
-			// Note that if $domain is an empty string, getDomainID() might not match it
-			$this->conns[$connInUseKey][$i][$conn->getDomainId()] = $conn;
-			$this->connLogger->debug( __METHOD__ .
-				": reusing free connection from $oldDomain for $domain" );
-		} else {
+		}
+
+		if ( !$conn ) {
 			if ( !isset( $this->servers[$i] ) || !is_array( $this->servers[$i] ) ) {
 				throw new InvalidArgumentException( "No server with index '$i'." );
 			}
-- 
2.20.1