From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 21 Sep 2016 08:33:29 +0000 (-0700)
Subject: Add sanity check to LoadBalancer::setDomainPrefix()
X-Git-Tag: 1.31.0-rc.0~5420^2
X-Git-Url: http://git.cyclocoop.org/%24href?a=commitdiff_plain;h=c36ef27c53f00408430c637cfcdfa892b7f34f2c;p=lhc%2Fweb%2Fwiklou.git

Add sanity check to LoadBalancer::setDomainPrefix()

Fixed some errors that popped up in CI:
* Also cleanup $domain handling in reuseConnection().
* Fix empty string handling in openForeignConnection() where
  the empty string check against $dbName failed since an empty
  string $domain results in $dbName being null.

Change-Id: Ie78fefa1acb401fe4e8bdc96b75053692aa0a925
---

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 9c07043c0c..c030cb20ae 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -136,9 +136,10 @@ class LoadBalancer implements ILoadBalancer {
 
 		$this->mReadIndex = -1;
 		$this->mConns = [
-			'local' => [],
+			'local'       => [],
 			'foreignUsed' => [],
-			'foreignFree' => [] ];
+			'foreignFree' => []
+		];
 		$this->mLoads = [];
 		$this->mWaitForPos = false;
 		$this->mErrorConnection = false;
@@ -598,21 +599,18 @@ class LoadBalancer implements ILoadBalancer {
 			return;
 		}
 
-		$dbName = $conn->getDBname();
-		$prefix = $conn->tablePrefix();
-		if ( strval( $prefix ) !== '' ) {
-			$domain = "$dbName-$prefix";
-		} else {
-			$domain = $dbName;
-		}
+		$domain = $conn->getDomainID();
 		if ( $this->mConns['foreignUsed'][$serverIndex][$domain] !== $conn ) {
-			throw new InvalidArgumentException( __METHOD__ . ": connection not found, has " .
-				"the connection been freed already?" );
+			throw new InvalidArgumentException( __METHOD__ .
+				": connection not found, has the connection been freed already?" );
 		}
 		$conn->setLBInfo( 'foreignPoolRefCount', --$refCount );
 		if ( $refCount <= 0 ) {
 			$this->mConns['foreignFree'][$serverIndex][$domain] = $conn;
 			unset( $this->mConns['foreignUsed'][$serverIndex][$domain] );
+			if ( !$this->mConns['foreignUsed'][$serverIndex] ) {
+				unset( $this->mConns[ 'foreignUsed' ][$serverIndex] ); // clean up
+			}
 			$this->connLogger->debug( __METHOD__ . ": freed connection $serverIndex/$domain" );
 		} else {
 			$this->connLogger->debug( __METHOD__ .
@@ -707,11 +705,10 @@ class LoadBalancer implements ILoadBalancer {
 			// Reuse a connection from another domain
 			$conn = reset( $this->mConns['foreignFree'][$i] );
 			$oldDomain = key( $this->mConns['foreignFree'][$i] );
-
 			// The empty string as a DB name means "don't care".
 			// DatabaseMysqlBase::open() already handle this on connection.
-			if ( $dbName !== '' && !$conn->selectDB( $dbName ) ) {
-				$this->mLastError = "Error selecting database $dbName on server " .
+			if ( strlen( $dbName ) && !$conn->selectDB( $dbName ) ) {
+				$this->mLastError = "Error selecting database '$dbName' on server " .
 					$conn->getServer() . " from client host {$this->host}";
 				$this->mErrorConnection = $conn;
 				$conn = false;
@@ -770,7 +767,7 @@ class LoadBalancer implements ILoadBalancer {
 	 * @access private
 	 *
 	 * @param array $server
-	 * @param bool $dbNameOverride
+	 * @param string|bool $dbNameOverride Use "" to not select any database
 	 * @return IDatabase
 	 * @throws DBAccessError
 	 * @throws InvalidArgumentException
@@ -1473,6 +1470,17 @@ class LoadBalancer implements ILoadBalancer {
 	}
 
 	public function setDomainPrefix( $prefix ) {
+		if ( $this->mConns['foreignUsed'] ) {
+			// Do not switch connections to explicit foreign domains unless marked as free
+			$domains = [];
+			foreach ( $this->mConns['foreignUsed'] as $i => $connsByDomain ) {
+				$domains = array_merge( $domains, array_keys( $connsByDomain ) );
+			}
+			$domains = implode( ', ', $domains );
+			throw new DBUnexpectedError( null,
+				"Foreign domain connections are still in use ($domains)." );
+		}
+
 		$this->localDomain = new DatabaseDomain(
 			$this->localDomain->getDatabase(),
 			null,
diff --git a/tests/phpunit/includes/db/LBFactoryTest.php b/tests/phpunit/includes/db/LBFactoryTest.php
index adf8a40108..d72768d434 100644
--- a/tests/phpunit/includes/db/LBFactoryTest.php
+++ b/tests/phpunit/includes/db/LBFactoryTest.php
@@ -272,6 +272,7 @@ class LBFactoryTest extends MediaWikiTestCase {
 
 		/** @var DatabaseBase $db */
 		$db = $lb->getConnection( DB_MASTER, [], '' );
+		$lb->reuseConnection( $db ); // don't care
 
 		$this->assertEquals(
 			'',
@@ -323,6 +324,7 @@ class LBFactoryTest extends MediaWikiTestCase {
 		$lb = $factory->getMainLB();
 		/** @var DatabaseBase $db */
 		$db = $lb->getConnection( DB_MASTER, [], '' );
+		$lb->reuseConnection( $db ); // don't care
 
 		$this->assertEquals(
 			'',