From f4e0c720a83946b6b8716731f220783731be2804 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 26 Jul 2017 02:04:34 -0700
Subject: [PATCH] rdbms: Ensure onTransactionPreCommitOrIdle() callbacks don't
 lead transactions

If no writes started a transaction yet, the callback would run
but not commit (by design, joining the request round). Later
writes will then pile on top of it.

The point of this method is to avoid such cases, so this edge
case has been fixed.

Change-Id: I9b44b19261d679de4aff6e44a9cfeb4f684ce02e
---
 includes/libs/rdbms/database/Database.php     |  7 +-
 .../libs/rdbms/database/DatabaseTest.php      | 66 +++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/includes/libs/rdbms/database/Database.php b/includes/libs/rdbms/database/Database.php
index 723a4a607b..b8b44e698e 100644
--- a/includes/libs/rdbms/database/Database.php
+++ b/includes/libs/rdbms/database/Database.php
@@ -2664,10 +2664,13 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 	}
 
 	final public function onTransactionPreCommitOrIdle( callable $callback, $fname = __METHOD__ ) {
-		if ( $this->mTrxLevel ) {
+		if ( $this->mTrxLevel || $this->getFlag( self::DBO_TRX ) ) {
+			// As long as DBO_TRX is set, writes will accumulate until the load balancer issues
+			// an implicit commit of all peer databases. This is true even if a transaction has
+			// not yet been triggered by writes; make sure $callback runs *after* any such writes.
 			$this->mTrxPreCommitCallbacks[] = [ $callback, $fname ];
 		} else {
-			// If no transaction is active, then make one for this callback
+			// No transaction is active nor will start implicitly, so make one for this callback
 			$this->startAtomic( __METHOD__ );
 			try {
 				call_user_func( $callback );
diff --git a/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php b/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php
index 9bea7ff0ea..70b6c36032 100644
--- a/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php
+++ b/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php
@@ -1,6 +1,7 @@
 <?php
 
 use Wikimedia\Rdbms\IDatabase;
+use Wikimedia\Rdbms\LBFactorySingle;
 use Wikimedia\Rdbms\TransactionProfiler;
 use Wikimedia\TestingAccessWrapper;
 
@@ -135,6 +136,71 @@ class DatabaseTest extends PHPUnit_Framework_TestCase {
 		$this->assertFalse( $db->getFlag( DBO_TRX ), 'DBO_TRX restored to default' );
 	}
 
+	/**
+	 * @covers Wikimedia\Rdbms\Database::onTransactionPreCommitOrIdle
+	 * @covers Wikimedia\Rdbms\Database::runOnTransactionPreCommitCallbacks
+	 */
+	public function testTransactionPreCommitOrIdle() {
+		$db = $this->getMockDB( [ 'isOpen' ] );
+		$db->method( 'isOpen' )->willReturn( true );
+		$db->clearFlag( DBO_TRX );
+
+		$this->assertFalse( $db->getFlag( DBO_TRX ), 'DBO_TRX is not set' );
+
+		$called = false;
+		$db->onTransactionPreCommitOrIdle(
+			function () use ( &$called ) {
+				$called = true;
+			},
+			__METHOD__
+		);
+		$this->assertTrue( $called, 'Called when idle' );
+
+		$db->begin( __METHOD__ );
+		$called = false;
+		$db->onTransactionPreCommitOrIdle(
+			function () use ( &$called ) {
+				$called = true;
+			},
+			__METHOD__
+		);
+		$this->assertFalse( $called, 'Not called when transaction is active' );
+		$db->commit( __METHOD__ );
+		$this->assertTrue( $called, 'Called when transaction is committed' );
+	}
+
+	/**
+	 * @covers Wikimedia\Rdbms\Database::onTransactionPreCommitOrIdle
+	 * @covers Wikimedia\Rdbms\Database::runOnTransactionPreCommitCallbacks
+	 */
+	public function testTransactionPreCommitOrIdle_TRX() {
+		$db = $this->getMockDB( [ 'isOpen' ] );
+		$db->method( 'isOpen' )->willReturn( true );
+		$db->setFlag( DBO_TRX );
+
+		$lbFactory = LBFactorySingle::newFromConnection( $db );
+		// Ask for the connectin so that LB sets internal state
+		// about this connection being the master connection
+		$lb = $lbFactory->getMainLB();
+		$conn = $lb->openConnection( $lb->getWriterIndex() );
+		$this->assertSame( $db, $conn, 'Same DB instance' );
+		$this->assertTrue( $db->getFlag( DBO_TRX ), 'DBO_TRX is set' );
+
+		$called = false;
+		$db->onTransactionPreCommitOrIdle(
+			function () use ( &$called ) {
+				$called = true;
+			}
+		);
+		$this->assertFalse( $called, 'Not called when idle if DBO_TRX is set' );
+
+		$lbFactory->beginMasterChanges( __METHOD__ );
+		$this->assertFalse( $called, 'Not called when lb-transaction is active' );
+
+		$lbFactory->commitMasterChanges( __METHOD__ );
+		$this->assertTrue( $called, 'Called when lb-transaction is committed' );
+	}
+
 	/**
 	 * @covers Wikimedia\Rdbms\Database::onTransactionResolution
 	 * @covers Wikimedia\Rdbms\Database::runOnTransactionIdleCallbacks
-- 
2.20.1