From c387e5cd84f1b7dfac68327714815197fd32e773 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 24 Aug 2016 14:02:15 -0700
Subject: [PATCH] Cleanups to SqlBagOStuff

* Keep track of the custom LoadBalancer when it makes one.
* Use the custom LoadBalancer to wait for slaves if one
  was used, rather than the main singleton.
* Only wait on the slaves in the LoadBalancer if the
  main DBs are being used.

Change-Id: I11de814306c44f27e0c33b08b5921c0fd4cdc24f
---
 includes/objectcache/SqlBagOStuff.php | 44 +++++++++++++++++++--------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/includes/objectcache/SqlBagOStuff.php b/includes/objectcache/SqlBagOStuff.php
index 9ab28aa44e..7a89991679 100644
--- a/includes/objectcache/SqlBagOStuff.php
+++ b/includes/objectcache/SqlBagOStuff.php
@@ -48,6 +48,8 @@ class SqlBagOStuff extends BagOStuff {
 	/** @var int */
 	protected $syncTimeout = 3;
 
+	/** @var LoadBalancer|null */
+	protected $separateMainLB;
 	/** @var array */
 	protected $conns;
 	/** @var array UNIX timestamps */
@@ -114,6 +116,7 @@ class SqlBagOStuff extends BagOStuff {
 			$this->serverInfos = [ $params['server'] ];
 			$this->numServers = count( $this->serverInfos );
 		} else {
+			// Default to using the main wiki's database servers
 			$this->serverInfos = false;
 			$this->numServers = 1;
 		}
@@ -132,6 +135,23 @@ class SqlBagOStuff extends BagOStuff {
 		$this->slaveOnly = !empty( $params['slaveOnly'] );
 	}
 
+	protected function getSeparateMainLB() {
+		global $wgDBtype;
+
+		if ( $wgDBtype === 'mysql' && $this->usesMainDB() ) {
+			if ( !$this->separateMainLB ) {
+				// We must keep a separate connection to MySQL in order to avoid deadlocks
+				$lbFactory = MediaWikiServices::getInstance()->getDBLoadBalancerFactory();
+				$this->separateMainLB = $lbFactory->newMainLB();
+			}
+			return $this->separateMainLB;
+		} else {
+			// However, SQLite has an opposite behavior. And PostgreSQL needs to know
+			// if we are in transaction or not (@TODO: find some PostgreSQL work-around).
+			return null;
+		}
+	}
+
 	/**
 	 * Get a connection to the specified database
 	 *
@@ -163,16 +183,13 @@ class SqlBagOStuff extends BagOStuff {
 				$db = DatabaseBase::factory( $type, $info );
 				$db->clearFlag( DBO_TRX );
 			} else {
-				// We must keep a separate connection to MySQL in order to avoid deadlocks
-				// However, SQLite has an opposite behavior. And PostgreSQL needs to know
-				// if we are in transaction or not (@TODO: find some work-around).
 				$index = $this->slaveOnly ? DB_SLAVE : DB_MASTER;
-				if ( wfGetDB( $index )->getType() == 'mysql' ) {
-					$lb = wfGetLBFactory()->newMainLB();
-					$db = $lb->getConnection( $index );
+				if ( $this->getSeparateMainLB() ) {
+					$db = $this->getSeparateMainLB()->getConnection( $index );
 					$db->clearFlag( DBO_TRX ); // auto-commit mode
 				} else {
 					$db = wfGetDB( $index );
+					// Can't mess with transaction rounds (DBO_TRX) :(
 				}
 			}
 			$this->logger->debug( sprintf( "Connection %s will be used for SqlBagOStuff", $db ) );
@@ -782,17 +799,20 @@ class SqlBagOStuff extends BagOStuff {
 
 	protected function waitForSlaves() {
 		if ( $this->usesMainDB() ) {
+			$lb = $this->getSeparateMainLB()
+				?: MediaWikiServices::getInstance()->getDBLoadBalancer();
 			// Main LB is used; wait for any slaves to catch up
 			try {
-				$lbFactory = MediaWikiServices::getInstance()->getDBLoadBalancerFactory();
-				$lbFactory->waitForReplication( [ 'wiki' => wfWikiID() ] );
-				return true;
+				$pos = $lb->getMasterPos();
+				if ( $pos ) {
+					return $lb->waitForAll( $pos, 3 );
+				}
 			} catch ( DBReplicationWaitError $e ) {
 				return false;
 			}
-		} else {
-			// Custom DB server list; probably doesn't use replication
-			return true;
 		}
+
+		// Custom DB server list; probably doesn't use replication
+		return true;
 	}
 }
-- 
2.20.1