rdbms: fix finalization stage errors in LBFactory::commitMasterChanges

If a pre-commit callback caused a new LoadBalancer object to be created,
that object will be in the "cursory" state rather than the "finalized"
state. If any callbacks run on an LB instance, make LBFactory iterate
over them all again to finalize these new instances.

Make LoadBalancer::finializeMasterChanges allow calls to
already-finalized instances for simplicity.

Bug: T193668
Change-Id: I4493e9571625a350c0a102219081ce090967a4ac

includes/libs/rdbms/database/DBConnRef.php

@@ -113,6 +113,10 @@ class DBConnRef implements IDatabase {
 		return $this->__call( __FUNCTION__, func_get_args() );
 	}
 
+	public function preCommitCallbacksPending() {
+		return $this->__call( __FUNCTION__, func_get_args() );
+	}
+
 	public function writesOrCallbacksPending() {
 		return $this->__call( __FUNCTION__, func_get_args() );
 	}

includes/libs/rdbms/database/Database.php

@@ -677,6 +677,10 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 		);
 	}
 
+	public function preCommitCallbacksPending() {
+		return $this->trxLevel && $this->trxPreCommitCallbacks;
+	}
+
 	/**
 	 * @return string|null
 	 */

includes/libs/rdbms/database/IDatabase.php

@@ -254,8 +254,15 @@ interface IDatabase {
 	public function writesPending();
 
 	/**
-	 * Returns true if there is a transaction/round open with possible write
-	 * queries or transaction pre-commit/idle callbacks waiting on it to finish.
+	 * @return bool Whether there is a transaction open with pre-commit callbacks pending
+	 * @since 1.32
+	 */
+	public function preCommitCallbacksPending();
+
+	/**
+	 * Whether there is a transaction open with either possible write queries
+	 * or unresolved pre-commit/commit/resolution callbacks pending
+	 *
 	 * This does *not* count recurring callbacks, e.g. from setTransactionListener().
 	 *
 	 * @return bool

includes/libs/rdbms/lbfactory/LBFactory.php

@@ -246,7 +246,12 @@ abstract class LBFactory implements ILBFactory {
 		/** @noinspection PhpUnusedLocalVariableInspection */
 		$scope = $this->getScopedPHPBehaviorForCommit(); // try to ignore client aborts
 		// Run pre-commit callbacks and suppress post-commit callbacks, aborting on failure
-		$this->forEachLBCallMethod( 'finalizeMasterChanges' );
+		do {
+			$count = 0; // number of callbacks executed this iteration
+			$this->forEachLB( function ( ILoadBalancer $lb ) use ( &$count ) {
+				$count += $lb->finalizeMasterChanges();
+			} );
+		} while ( $count > 0 );
 		$this->trxRoundId = false;
 		// Perform pre-commit checks, aborting on failure
 		$this->forEachLBCallMethod( 'approveMasterChanges', [ $options ] );

includes/libs/rdbms/loadbalancer/ILoadBalancer.php

@@ -380,6 +380,8 @@ interface ILoadBalancer {
 	 * Run pre-commit callbacks and defer execution of post-commit callbacks
 	 *
 	 * Use this only for mutli-database commits
+	 *
+	 * @return int Number of pre-commit callbacks run (since 1.32)
 	 */
 	public function finalizeMasterChanges();
 
@@ -449,7 +451,7 @@ interface ILoadBalancer {
 	public function hasMasterConnection();
 
 	/**
-	 * Determine if there are pending changes in a transaction by this thread
+	 * Whether there are pending changes or callbacks in a transaction by this thread
 	 * @return bool
 	 */
 	public function hasMasterChanges();

includes/libs/rdbms/loadbalancer/LoadBalancer.php

@@ -1263,10 +1263,11 @@ class LoadBalancer implements ILoadBalancer {
 	}
 
 	public function finalizeMasterChanges() {
-		$this->assertTransactionRoundStage( self::ROUND_CURSORY );
+		$this->assertTransactionRoundStage( [ self::ROUND_CURSORY, self::ROUND_FINALIZED ] );
 
 		$this->trxRoundStage = self::ROUND_ERROR; // "failed" until proven otherwise
 		// Loop until callbacks stop adding callbacks on other connections
+		$total = 0;
 		do {
 			$count = 0; // callbacks execution attempts
 			$this->forEachOpenMasterConnection( function ( Database $conn ) use ( &$count ) {
@@ -1274,12 +1275,15 @@ class LoadBalancer implements ILoadBalancer {
 				// Any error should cause all (peer) transactions to be rolled back together.
 				$count += $conn->runOnTransactionPreCommitCallbacks();
 			} );
+			$total += $count;
 		} while ( $count > 0 );
 		// Defer post-commit callbacks until after COMMIT/ROLLBACK happens on all handles
 		$this->forEachOpenMasterConnection( function ( Database $conn ) {
 			$conn->setTrxEndCallbackSuppression( true );
 		} );
 		$this->trxRoundStage = self::ROUND_FINALIZED;
+
+		return $total;
 	}
 
 	public function approveMasterChanges( array $options ) {
@@ -1494,13 +1498,21 @@ class LoadBalancer implements ILoadBalancer {
 	}
 
 	/**
-	 * @param string $stage
+	 * @param string|string[] $stage
 	 */
 	private function assertTransactionRoundStage( $stage ) {
-		if ( $this->trxRoundStage !== $stage ) {
+		$stages = (array)$stage;
+
+		if ( !in_array( $this->trxRoundStage, $stages, true ) ) {
+			$stageList = implode(
+				'/',
+				array_map( function ( $v ) {
+					return "'$v'";
+				}, $stages )
+			);
 			throw new DBTransactionError(
 				null,
-				"Transaction round stage must be '$stage' (not '{$this->trxRoundStage}')"
+				"Transaction round stage must be $stageList (not '{$this->trxRoundStage}')"
 			);
 		}
 	}

tests/phpunit/includes/db/LBFactoryTest.php

@@ -157,12 +157,18 @@ class LBFactoryTest extends MediaWikiTestCase {
 		global $wgDBserver, $wgDBname, $wgDBuser, $wgDBpassword, $wgDBtype, $wgSQLiteDataDir;
 
 		$factory = new LBFactoryMulti( [
-			'sectionsByDB' => [],
+			'sectionsByDB' => [
+				's1wiki' => 's1',
+			],
 			'sectionLoads' => [
+				's1' => [
+					'test-db3' => 0,
+					'test-db4' => 100,
+				],
 				'DEFAULT' => [
 					'test-db1' => 0,
 					'test-db2' => 100,
-				],
+				]
 			],
 			'serverTemplate' => [
 				'dbname'      => $wgDBname,
@@ -174,7 +180,9 @@ class LBFactoryTest extends MediaWikiTestCase {
 			],
 			'hostsByName' => [
 				'test-db1'  => $wgDBserver,
-				'test-db2'  => $wgDBserver
+				'test-db2'  => $wgDBserver,
+				'test-db3'  => $wgDBserver,
+				'test-db4'  => $wgDBserver
 			],
 			'loadMonitorClass' => LoadMonitorNull::class
 		] );
@@ -186,8 +194,25 @@ class LBFactoryTest extends MediaWikiTestCase {
 		$dbr = $lb->getConnection( DB_REPLICA );
 		$this->assertTrue( $dbr->getLBInfo( 'replica' ), 'slave shows as slave' );
 
+		// Test that LoadBalancer instances made during commitMasterChanges() do not throw
+		// DBTransactionError due to transaction ROUND_* stages being mismatched.
+		$factory->beginMasterChanges( __METHOD__ );
+		$dbw->onTransactionPreCommitOrIdle( function () use ( $factory ) {
+			// Trigger s1 LoadBalancer instantiation during "finalize" stage.
+			// There is no s1wiki DB to select so it is not in getConnection(),
+			// but this fools getMainLB() at least.
+			$factory->getMainLB( 's1wiki' )->getConnection( DB_MASTER );
+		} );
+		$factory->commitMasterChanges( __METHOD__ );
+
+		$count = 0;
+		$factory->forEachLB( function () use ( &$count ) {
+			++$count;
+		} );
+		$this->assertEquals( 2, $count );
+
 		$factory->shutdown();
-		$lb->closeAll();
+		$factory->closeAll();
 	}
 
 	/**