From c9f1f64cad140571728b395b95cbc8fb43d36670 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Sun, 28 Aug 2016 08:56:17 -0700
Subject: [PATCH] Add LoadBalancer::setTransactionListener()

This is useful for setting callbacks without making a new
DB connection. It also catches extra connections due to
foriegn wiki usage.

Change-Id: Ief009a0c1c2b4245fe3594fdf0d306dcc4956151
---
 includes/db/loadbalancer/LoadBalancer.php | 29 +++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/includes/db/loadbalancer/LoadBalancer.php b/includes/db/loadbalancer/LoadBalancer.php
index 30dbfc53ee..c554f17109 100644
--- a/includes/db/loadbalancer/LoadBalancer.php
+++ b/includes/db/loadbalancer/LoadBalancer.php
@@ -72,6 +72,8 @@ class LoadBalancer {
 	private $connsOpened = 0;
 	/** @var string|bool String if a requested DBO_TRX transaction round is active */
 	private $trxRoundId = false;
+	/** @var array[] Map of (name => callable) */
+	private $trxRecurringCallbacks = [];
 
 	/** @var integer Warn when this many connection are held */
 	const CONN_HELD_WARN_THRESHOLD = 10;
@@ -869,6 +871,12 @@ class LoadBalancer {
 			$this->applyTransactionRoundFlags( $db );
 		}
 
+		if ( $server['serverIndex'] === $this->getWriterIndex() ) {
+			foreach ( $this->trxRecurringCallbacks as $name => $callback ) {
+				$db->setTransactionListener( $name, $callback );
+			}
+		}
+
 		return $db;
 	}
 
@@ -1666,4 +1674,25 @@ class LoadBalancer {
 	public function clearLagTimeCache() {
 		$this->getLoadMonitor()->clearCaches();
 	}
+
+	/**
+	 * Set a callback via DatabaseBase::setTransactionListener() on
+	 * all current and future master connections of this load balancer
+	 *
+	 * @param string $name Callback name
+	 * @param callable|null $callback
+	 * @since 1.28
+	 */
+	public function setTransactionListener( $name, callable $callback = null ) {
+		if ( $callback ) {
+			$this->trxRecurringCallbacks[$name] = $callback;
+		} else {
+			unset( $this->trxRecurringCallbacks[$name] );
+		}
+		$this->forEachOpenMasterConnection(
+			function ( DatabaseBase $conn ) use ( $name, $callback ) {
+				$conn->setTransactionListener( $name, $callback );
+			}
+		);
+	}
 }
-- 
2.20.1