From: Aaron Schulz <aschulz@wikimedia.org>
Date: Thu, 5 Apr 2018 02:44:21 +0000 (-0700)
Subject: rdbms: add $flags argument to ILoadBalancer::getAnyOpenConnection
X-Git-Tag: 1.31.0-rc.0~57
X-Git-Url: https://git.heureux-cyclage.org/?p=lhc%2Fweb%2Fwiklou.git;a=commitdiff_plain;h=e9d42c0e7ea9390f650de8e350a4f64ff19787b0

rdbms: add $flags argument to ILoadBalancer::getAnyOpenConnection

Make LoadMonitor use this flag for getting a connection.

Change-Id: I32ea9aadc0223c86db5979d6579d781a6af0ff53
---

diff --git a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
index a699b2341f..2f493c7dd8 100644
--- a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
@@ -167,10 +167,13 @@ interface ILoadBalancer {
 	/**
 	 * Get any open connection to a given server index, local or foreign
 	 *
+	 * Use CONN_TRX_AUTOCOMMIT to only look for connections opened with that flag
+	 *
 	 * @param int $i Server index or DB_MASTER/DB_REPLICA
+	 * @param int $flags Bitfield of CONN_* class constants
 	 * @return Database|bool False if no such connection is open
 	 */
-	public function getAnyOpenConnection( $i );
+	public function getAnyOpenConnection( $i, $flags = 0 );
 
 	/**
 	 * Get a connection handle by server index
diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 96ea949d8c..be7d0212e4 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -566,16 +566,17 @@ class LoadBalancer implements ILoadBalancer {
 		}
 	}
 
-	/**
-	 * @param int $i
-	 * @return IDatabase|bool
-	 */
-	public function getAnyOpenConnection( $i ) {
+	public function getAnyOpenConnection( $i, $flags = 0 ) {
+		$autocommit = ( ( $flags & self::CONN_TRX_AUTOCOMMIT ) == self::CONN_TRX_AUTOCOMMIT );
 		foreach ( $this->conns as $connsByServer ) {
-			if ( !empty( $connsByServer[$i] ) ) {
-				/** @var IDatabase[] $serverConns */
-				$serverConns = $connsByServer[$i];
-				return reset( $serverConns );
+			if ( !isset( $connsByServer[$i] ) ) {
+				continue;
+			}
+
+			foreach ( $connsByServer[$i] as $conn ) {
+				if ( !$autocommit || $conn->getLBInfo( 'autoCommitOnly' ) ) {
+					return $conn;
+				}
 			}
 		}
 
diff --git a/includes/libs/rdbms/loadmonitor/LoadMonitor.php b/includes/libs/rdbms/loadmonitor/LoadMonitor.php
index 3372839648..c7f807a0bc 100644
--- a/includes/libs/rdbms/loadmonitor/LoadMonitor.php
+++ b/includes/libs/rdbms/loadmonitor/LoadMonitor.php
@@ -156,13 +156,14 @@ class LoadMonitor implements ILoadMonitor {
 				continue;
 			}
 
-			$conn = $this->parent->getAnyOpenConnection( $i );
-			if ( $conn && !$conn->trxLevel() ) {
-				# Handles with open transactions are avoided since they might be subject
-				# to REPEATABLE-READ snapshots, which could affect the lag estimate query.
+			# Handles with open transactions are avoided since they might be subject
+			# to REPEATABLE-READ snapshots, which could affect the lag estimate query.
+			$flags = ILoadBalancer::CONN_TRX_AUTOCOMMIT;
+			$conn = $this->parent->getAnyOpenConnection( $i, $flags );
+			if ( $conn ) {
 				$close = false; // already open
 			} else {
-				$conn = $this->parent->openConnection( $i, '' );
+				$conn = $this->parent->openConnection( $i, ILoadBalancer::DOMAIN_ANY, $flags );
 				$close = true; // new connection
 			}
 
diff --git a/tests/phpunit/includes/db/LoadBalancerTest.php b/tests/phpunit/includes/db/LoadBalancerTest.php
index 88cf0e04c6..7e58555550 100644
--- a/tests/phpunit/includes/db/LoadBalancerTest.php
+++ b/tests/phpunit/includes/db/LoadBalancerTest.php
@@ -32,25 +32,27 @@ use Wikimedia\Rdbms\LoadMonitorNull;
  * @covers \Wikimedia\Rdbms\LoadBalancer
  */
 class LoadBalancerTest extends MediaWikiTestCase {
-	public function testWithoutReplica() {
+	private function makeServerConfig() {
 		global $wgDBserver, $wgDBname, $wgDBuser, $wgDBpassword, $wgDBtype, $wgSQLiteDataDir;
 
-		$servers = [
-			[
-				'host'        => $wgDBserver,
-				'dbname'      => $wgDBname,
-				'tablePrefix' => $this->dbPrefix(),
-				'user'        => $wgDBuser,
-				'password'    => $wgDBpassword,
-				'type'        => $wgDBtype,
-				'dbDirectory' => $wgSQLiteDataDir,
-				'load'        => 0,
-				'flags'       => DBO_TRX // REPEATABLE-READ for consistency
-			],
+		return [
+			'host' => $wgDBserver,
+			'dbname' => $wgDBname,
+			'tablePrefix' => $this->dbPrefix(),
+			'user' => $wgDBuser,
+			'password' => $wgDBpassword,
+			'type' => $wgDBtype,
+			'dbDirectory' => $wgSQLiteDataDir,
+			'load' => 0,
+			'flags' => DBO_TRX // REPEATABLE-READ for consistency
 		];
+	}
+
+	public function testWithoutReplica() {
+		global $wgDBname;
 
 		$lb = new LoadBalancer( [
-			'servers' => $servers,
+			'servers' => [ $this->makeServerConfig() ],
 			'queryLogger' => MediaWiki\Logger\LoggerFactory::getInstance( 'DBQuery' ),
 			'localDomain' => new DatabaseDomain( $wgDBname, null, $this->dbPrefix() )
 		] );
@@ -259,4 +261,39 @@ class LoadBalancerTest extends MediaWikiTestCase {
 
 		$this->assertFalse( $lb->getServerAttributes( 1 )[Database::ATTR_DB_LEVEL_LOCKING] );
 	}
+
+	/**
+	 * @covers LoadBalancer::openConnection()
+	 * @covers LoadBalancer::getAnyOpenConnection()
+	 */
+	function testOpenConnection() {
+		global $wgDBname;
+
+		$lb = new LoadBalancer( [
+			'servers' => [ $this->makeServerConfig() ],
+			'localDomain' => new DatabaseDomain( $wgDBname, null, $this->dbPrefix() )
+		] );
+
+		$i = $lb->getWriterIndex();
+		$this->assertEquals( null, $lb->getAnyOpenConnection( $i ) );
+		$conn1 = $lb->getConnection( $i );
+		$this->assertNotEquals( null, $conn1 );
+		$this->assertEquals( $conn1, $lb->getAnyOpenConnection( $i ) );
+		$conn2 = $lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
+		$this->assertNotEquals( null, $conn2 );
+		if ( $lb->getServerAttributes( $i )[Database::ATTR_DB_LEVEL_LOCKING] ) {
+			$this->assertEquals( null,
+				$lb->getAnyOpenConnection( $i, $lb::CONN_TRX_AUTOCOMMIT ) );
+			$this->assertEquals( $conn1,
+				$lb->getConnection(
+					$i, [], false, $lb::CONN_TRX_AUTOCOMMIT ), $lb::CONN_TRX_AUTOCOMMIT );
+		} else {
+			$this->assertEquals( $conn2,
+				$lb->getAnyOpenConnection( $i, $lb::CONN_TRX_AUTOCOMMIT ) );
+			$this->assertEquals( $conn2,
+				$lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT ) );
+		}
+
+		$lb->closeAll();
+	}
 }