From c23f326188d0d32e573e54ad28df729e61fbfcef Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 21 Sep 2016 00:34:23 -0700
Subject: [PATCH] Add some missing reuseConnection() to LoadBalancer

Change-Id: I99099f909e7759c061d44ad4e1ed049e038a4d7d
---
 .../libs/rdbms/loadbalancer/LoadBalancer.php  | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 7ba21ace6a..9c07043c0c 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -1303,6 +1303,9 @@ class LoadBalancer implements ILoadBalancer {
 				try {
 					$dbw = $conn ?: $this->getConnection( DB_MASTER, [], $domain );
 					$readOnly = (int)$dbw->serverIsReadOnly();
+					if ( !$conn ) {
+						$this->reuseConnection( $dbw );
+					}
 				} catch ( DBError $e ) {
 					$readOnly = 0;
 				}
@@ -1411,7 +1414,7 @@ class LoadBalancer implements ILoadBalancer {
 	}
 
 	public function safeGetLag( IDatabase $conn ) {
-		if ( $this->getServerCount() == 1 ) {
+		if ( $this->getServerCount() <= 1 ) {
 			return 0;
 		} else {
 			return $conn->getLag();
@@ -1419,23 +1422,30 @@ class LoadBalancer implements ILoadBalancer {
 	}
 
 	public function safeWaitForMasterPos( IDatabase $conn, $pos = false, $timeout = 10 ) {
-		if ( $this->getServerCount() == 1 || !$conn->getLBInfo( 'replica' ) ) {
+		if ( $this->getServerCount() <= 1 || !$conn->getLBInfo( 'replica' ) ) {
 			return true; // server is not a replica DB
 		}
 
-		$pos = $pos ?: $this->getConnection( DB_MASTER )->getMasterPos();
-		if ( !( $pos instanceof DBMasterPos ) ) {
-			return false; // something is misconfigured
+		if ( !$pos ) {
+			// Get the current master position
+			$dbw = $this->getConnection( DB_MASTER );
+			$pos = $dbw->getMasterPos();
+			$this->reuseConnection( $dbw );
 		}
 
-		$result = $conn->masterPosWait( $pos, $timeout );
-		if ( $result == -1 || is_null( $result ) ) {
-			$msg = __METHOD__ . ": Timed out waiting on {$conn->getServer()} pos {$pos}";
-			$this->replLogger->warning( "$msg" );
-			$ok = false;
+		if ( $pos instanceof DBMasterPos ) {
+			$result = $conn->masterPosWait( $pos, $timeout );
+			if ( $result == -1 || is_null( $result ) ) {
+				$msg = __METHOD__ . ": Timed out waiting on {$conn->getServer()} pos {$pos}";
+				$this->replLogger->warning( "$msg" );
+				$ok = false;
+			} else {
+				$this->replLogger->info( __METHOD__ . ": Done" );
+				$ok = true;
+			}
 		} else {
-			$this->replLogger->info( __METHOD__ . ": Done" );
-			$ok = true;
+			$ok = false; // something is misconfigured
+			$this->replLogger->error( "Could not get master pos for {$conn->getServer()}." );
 		}
 
 		return $ok;
-- 
2.20.1