From a21924267cc0976f4fdd8b98fec0395c0ef64982 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Tue, 23 Jan 2018 16:06:32 +0100
Subject: [PATCH] deferred: Make MWCallableUpdate support a list of DB handles

Change-Id: I0371ecd20101c5e622497e5912676300b040865e
---
 includes/deferred/MWCallableUpdate.php        | 12 +++--
 .../deferred/MWCallableUpdateTest.php         | 46 +++++++++++++++++++
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/includes/deferred/MWCallableUpdate.php b/includes/deferred/MWCallableUpdate.php
index 5b822af492..9803b7a491 100644
--- a/includes/deferred/MWCallableUpdate.php
+++ b/includes/deferred/MWCallableUpdate.php
@@ -14,14 +14,18 @@ class MWCallableUpdate implements DeferrableUpdate, DeferrableCallback {
 	/**
 	 * @param callable $callback
 	 * @param string $fname Calling method
-	 * @param IDatabase|null $dbw Abort if this DB is rolled back [optional] (since 1.28)
+	 * @param IDatabase|IDatabase[]|null $dbws Abort if any of the specified DB handles have
+	 *   a currently pending transaction which later gets rolled back [optional] (since 1.28)
 	 */
-	public function __construct( callable $callback, $fname = 'unknown', IDatabase $dbw = null ) {
+	public function __construct( callable $callback, $fname = 'unknown', $dbws = [] ) {
 		$this->callback = $callback;
 		$this->fname = $fname;
 
-		if ( $dbw && $dbw->trxLevel() ) {
-			$dbw->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+		$dbws = is_array( $dbws ) ? $dbws : [ $dbws ];
+		foreach ( $dbws as $dbw ) {
+			if ( $dbw && $dbw->trxLevel() ) {
+				$dbw->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+			}
 		}
 	}
 
diff --git a/tests/phpunit/includes/deferred/MWCallableUpdateTest.php b/tests/phpunit/includes/deferred/MWCallableUpdateTest.php
index 088ab4f544..6977aef1a2 100644
--- a/tests/phpunit/includes/deferred/MWCallableUpdateTest.php
+++ b/tests/phpunit/includes/deferred/MWCallableUpdateTest.php
@@ -29,8 +29,54 @@ class MWCallableUpdateTest extends PHPUnit_Framework_TestCase {
 		// Emulate rollback
 		$db->rollback( __METHOD__ );
 
+		$update->doUpdate();
+
+		// Ensure it was cancelled
+		$this->assertSame( 0, $ran );
+	}
+
+	public function testCancelSome() {
+		// Prepare update and DB
+		$db1 = new DatabaseTestHelper( __METHOD__ );
+		$db1->begin( __METHOD__ );
+		$db2 = new DatabaseTestHelper( __METHOD__ );
+		$db2->begin( __METHOD__ );
+		$ran = 0;
+		$update = new MWCallableUpdate( function () use ( &$ran ) {
+			$ran++;
+		}, __METHOD__, [ $db1, $db2 ] );
+
+		// Emulate rollback
+		$db1->rollback( __METHOD__ );
+
+		$update->doUpdate();
+
+		// Prevents: "Notice: DB transaction writes or callbacks still pending"
+		$db2->rollback( __METHOD__ );
+
 		// Ensure it was cancelled
+		$this->assertSame( 0, $ran );
+	}
+
+	public function testCancelAll() {
+		// Prepare update and DB
+		$db1 = new DatabaseTestHelper( __METHOD__ );
+		$db1->begin( __METHOD__ );
+		$db2 = new DatabaseTestHelper( __METHOD__ );
+		$db2->begin( __METHOD__ );
+		$ran = 0;
+		$update = new MWCallableUpdate( function () use ( &$ran ) {
+			$ran++;
+		}, __METHOD__, [ $db1, $db2 ] );
+
+		// Emulate rollbacks
+		$db1->rollback( __METHOD__ );
+		$db2->rollback( __METHOD__ );
+
 		$update->doUpdate();
+
+		// Ensure it was cancelled
 		$this->assertSame( 0, $ran );
 	}
+
 }
-- 
2.20.1