From 17080453a1a94d317247197b639472f87e2aaa03 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Thu, 21 Mar 2019 14:05:35 -0700
Subject: [PATCH] Support more coupled DBs in
 AtomicSectionUpdate/AutoCommitUpdate

This is useful for updating a secondary database based on a
commit to the main one, especially for using auto-commit.

Change-Id: Iad9084d2fb0490ecdfd5fcc97db33bfbcbfa5d21
---
 includes/deferred/AtomicSectionUpdate.php | 18 +++++++++++++-----
 includes/deferred/AutoCommitUpdate.php    | 18 +++++++++++++-----
 includes/deferred/MWCallableUpdate.php    |  4 ++++
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/includes/deferred/AtomicSectionUpdate.php b/includes/deferred/AtomicSectionUpdate.php
index 8b62989b53..69f09e3928 100644
--- a/includes/deferred/AtomicSectionUpdate.php
+++ b/includes/deferred/AtomicSectionUpdate.php
@@ -15,18 +15,22 @@ class AtomicSectionUpdate implements DeferrableUpdate, DeferrableCallback {
 	private $callback;
 
 	/**
-	 * @param IDatabase $dbw
+	 * @param IDatabase $dbw DB handle; update aborts if a transaction now this rolls back
 	 * @param string $fname Caller name (usually __METHOD__)
 	 * @param callable $callback
+	 * @param IDatabase[] $conns Abort if a transaction now on one of these rolls back [optional]
 	 * @see IDatabase::doAtomicSection()
 	 */
-	public function __construct( IDatabase $dbw, $fname, callable $callback ) {
+	public function __construct( IDatabase $dbw, $fname, callable $callback, array $conns = [] ) {
 		$this->dbw = $dbw;
 		$this->fname = $fname;
 		$this->callback = $callback;
-
-		if ( $this->dbw->trxLevel() ) {
-			$this->dbw->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+		// Register DB connections for which uncommitted changes are related to this update
+		$conns[] = $dbw;
+		foreach ( $conns as $conn ) {
+			if ( $conn->trxLevel() ) {
+				$conn->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+			}
 		}
 	}
 
@@ -36,6 +40,10 @@ class AtomicSectionUpdate implements DeferrableUpdate, DeferrableCallback {
 		}
 	}
 
+	/**
+	 * @private This method is public so that it works with onTransactionResolution()
+	 * @param int $trigger
+	 */
 	public function cancelOnRollback( $trigger ) {
 		if ( $trigger === IDatabase::TRIGGER_ROLLBACK ) {
 			$this->callback = null;
diff --git a/includes/deferred/AutoCommitUpdate.php b/includes/deferred/AutoCommitUpdate.php
index 85071576b0..ddfd98716d 100644
--- a/includes/deferred/AutoCommitUpdate.php
+++ b/includes/deferred/AutoCommitUpdate.php
@@ -15,17 +15,21 @@ class AutoCommitUpdate implements DeferrableUpdate, DeferrableCallback {
 	private $callback;
 
 	/**
-	 * @param IDatabase $dbw
+	 * @param IDatabase $dbw DB handle; update aborts if a transaction now this rolls back
 	 * @param string $fname Caller name (usually __METHOD__)
 	 * @param callable $callback Callback that takes (IDatabase, method name string)
+	 * @param IDatabase[] $conns Abort if a transaction now on one of these rolls back [optional]
 	 */
-	public function __construct( IDatabase $dbw, $fname, callable $callback ) {
+	public function __construct( IDatabase $dbw, $fname, callable $callback, array $conns = [] ) {
 		$this->dbw = $dbw;
 		$this->fname = $fname;
 		$this->callback = $callback;
-
-		if ( $this->dbw->trxLevel() ) {
-			$this->dbw->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+		// Register DB connections for which uncommitted changes are related to this update
+		$conns[] = $dbw;
+		foreach ( $conns as $conn ) {
+			if ( $conn->trxLevel() ) {
+				$conn->onTransactionResolution( [ $this, 'cancelOnRollback' ], $fname );
+			}
 		}
 	}
 
@@ -50,6 +54,10 @@ class AutoCommitUpdate implements DeferrableUpdate, DeferrableCallback {
 		}
 	}
 
+	/**
+	 * @private This method is public so that it works with onTransactionResolution()
+	 * @param int $trigger
+	 */
 	public function cancelOnRollback( $trigger ) {
 		if ( $trigger === IDatabase::TRIGGER_ROLLBACK ) {
 			$this->callback = null;
diff --git a/includes/deferred/MWCallableUpdate.php b/includes/deferred/MWCallableUpdate.php
index 9803b7a491..efca406881 100644
--- a/includes/deferred/MWCallableUpdate.php
+++ b/includes/deferred/MWCallableUpdate.php
@@ -35,6 +35,10 @@ class MWCallableUpdate implements DeferrableUpdate, DeferrableCallback {
 		}
 	}
 
+	/**
+	 * @private This method is public so that it works with onTransactionResolution()
+	 * @param int $trigger
+	 */
 	public function cancelOnRollback( $trigger ) {
 		if ( $trigger === IDatabase::TRIGGER_ROLLBACK ) {
 			$this->callback = null;
-- 
2.20.1