From f1b861dec81ac636f0183b085ad79993962ffa5e Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Wed, 12 Apr 2017 16:48:40 -0700
Subject: [PATCH] Make LoadBalancer::waitFor() and friends  more robust via
 try/finally

Change-Id: I50742752bb94ca17997138b440b911849bcbadf3
---
 .../libs/rdbms/loadbalancer/LoadBalancer.php  | 75 ++++++++++---------
 1 file changed, 40 insertions(+), 35 deletions(-)

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index d268f83b5e..ffa4f78a76 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -430,60 +430,65 @@ class LoadBalancer implements ILoadBalancer {
 
 	public function waitFor( $pos ) {
 		$oldPos = $this->mWaitForPos;
-		$this->mWaitForPos = $pos;
-
-		// If a generic reader connection was already established, then wait now
-		$i = $this->mReadIndex;
-		if ( $i > 0 ) {
-			if ( !$this->doWait( $i ) ) {
-				$this->laggedReplicaMode = true;
+		try {
+			$this->mWaitForPos = $pos;
+			// If a generic reader connection was already established, then wait now
+			$i = $this->mReadIndex;
+			if ( $i > 0 ) {
+				if ( !$this->doWait( $i ) ) {
+					$this->laggedReplicaMode = true;
+				}
 			}
+		} finally {
+			// Restore the older position if it was higher
+			$this->setWaitForPositionIfHigher( $oldPos );
 		}
-
-		// Restore the older position if it was higher
-		$this->setWaitForPositionIfHigher( $oldPos );
 	}
 
 	public function waitForOne( $pos, $timeout = null ) {
 		$oldPos = $this->mWaitForPos;
-		$this->mWaitForPos = $pos;
+		try {
+			$this->mWaitForPos = $pos;
 
-		$i = $this->mReadIndex;
-		if ( $i <= 0 ) {
-			// Pick a generic replica DB if there isn't one yet
-			$readLoads = $this->mLoads;
-			unset( $readLoads[$this->getWriterIndex()] ); // replica DBs only
-			$readLoads = array_filter( $readLoads ); // with non-zero load
-			$i = ArrayUtils::pickRandom( $readLoads );
-		}
+			$i = $this->mReadIndex;
+			if ( $i <= 0 ) {
+				// Pick a generic replica DB if there isn't one yet
+				$readLoads = $this->mLoads;
+				unset( $readLoads[$this->getWriterIndex()] ); // replica DBs only
+				$readLoads = array_filter( $readLoads ); // with non-zero load
+				$i = ArrayUtils::pickRandom( $readLoads );
+			}
 
-		if ( $i > 0 ) {
-			$ok = $this->doWait( $i, true, $timeout );
-		} else {
-			$ok = true; // no applicable loads
+			if ( $i > 0 ) {
+				$ok = $this->doWait( $i, true, $timeout );
+			} else {
+				$ok = true; // no applicable loads
+			}
+		} finally {
+			// Restore the older position if it was higher
+			$this->setWaitForPositionIfHigher( $oldPos );
 		}
 
-		// Restore the older position if it was higher
-		$this->setWaitForPositionIfHigher( $oldPos );
-
 		return $ok;
 	}
 
 	public function waitForAll( $pos, $timeout = null ) {
 		$oldPos = $this->mWaitForPos;
-		$this->mWaitForPos = $pos;
-		$serverCount = count( $this->mServers );
+		try {
+			$this->mWaitForPos = $pos;
+			$serverCount = count( $this->mServers );
 
-		$ok = true;
-		for ( $i = 1; $i < $serverCount; $i++ ) {
-			if ( $this->mLoads[$i] > 0 ) {
-				$ok = $this->doWait( $i, true, $timeout ) && $ok;
+			$ok = true;
+			for ( $i = 1; $i < $serverCount; $i++ ) {
+				if ( $this->mLoads[$i] > 0 ) {
+					$ok = $this->doWait( $i, true, $timeout ) && $ok;
+				}
 			}
+		} finally {
+			// Restore the older position if it was higher
+			$this->setWaitForPositionIfHigher( $oldPos );
 		}
 
-		// Restore the older position if it was higher
-		$this->setWaitForPositionIfHigher( $oldPos );
-
 		return $ok;
 	}
 
-- 
2.20.1