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

diff --git a/includes/db/loadbalancer/LoadBalancer.php b/includes/db/loadbalancer/LoadBalancer.php
index 30dbfc5..c554f17 100644 (file)
--- 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 );
+                       }
+               );
+       }
 }