From dc17c8ddbf2a51d0061698d1cde9f17136ff0444 Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Thu, 13 Apr 2017 19:15:42 -0700
Subject: [PATCH] Fix some LoadBalancer::waitFor*() inconsistencies

* Make sure waitFor() always handles laggedReplicaMode updates,
  even if the actual waiting was deferred until a connection was
  needed.
* Restore the old mWaitForPos in waitForOne()/waitForAll() since
  this do not care about the generic reader index or the sort of
  ChronologyProtector logic related to it.

Change-Id: I0767e9831b8fd7fd115a472354977e3c1e12114a
---
 .../libs/rdbms/loadbalancer/ILoadBalancer.php    |  4 ++++
 .../libs/rdbms/loadbalancer/LoadBalancer.php     | 16 ++++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
index 8b5a98db06..cbd0ff3401 100644
--- a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
@@ -135,6 +135,10 @@ interface ILoadBalancer {
 	 * If a DB_REPLICA connection has been opened already, then wait immediately.
 	 * Otherwise sets a variable telling it to wait if such a connection is opened.
 	 *
+	 * This only applies to connections to the generic replica DB for this request.
+	 * If a timeout happens when waiting, then getLaggedReplicaMode()/laggedReplicaUsed()
+	 * will return true.
+	 *
 	 * @param DBMasterPos|bool $pos Master position or false
 	 */
 	public function waitFor( $pos );
diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index ffa4f78a76..e8069c0110 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -411,7 +411,11 @@ class LoadBalancer implements ILoadBalancer {
 			# Replica DB connection successful.
 			# Wait for the session master pos for a short time.
 			if ( $this->mWaitForPos && $i > 0 ) {
-				$this->doWait( $i );
+				# When LoadBalancer::waitFor() set mWaitForPos, the wait will happen here.
+				# Be sure to update laggedReplicaMode accordingly for consistency.
+				if ( !$this->doWait( $i ) ) {
+					$laggedReplicaMode = true;
+				}
 			}
 			if ( $this->mReadIndex <= 0 && $this->mLoads[$i] > 0 && $group === false ) {
 				$this->mReadIndex = $i;
@@ -440,7 +444,7 @@ class LoadBalancer implements ILoadBalancer {
 				}
 			}
 		} finally {
-			// Restore the older position if it was higher
+			// Restore the older position if it was higher since this is used for lag-protection
 			$this->setWaitForPositionIfHigher( $oldPos );
 		}
 	}
@@ -465,8 +469,8 @@ class LoadBalancer implements ILoadBalancer {
 				$ok = true; // no applicable loads
 			}
 		} finally {
-			// Restore the older position if it was higher
-			$this->setWaitForPositionIfHigher( $oldPos );
+			# Restore the old position, as this is not used for lag-protection but for throttling
+			$this->mWaitForPos = $oldPos;
 		}
 
 		return $ok;
@@ -485,8 +489,8 @@ class LoadBalancer implements ILoadBalancer {
 				}
 			}
 		} finally {
-			// Restore the older position if it was higher
-			$this->setWaitForPositionIfHigher( $oldPos );
+			# Restore the old position, as this is not used for lag-protection but for throttling
+			$this->mWaitForPos = $oldPos;
 		}
 
 		return $ok;
-- 
2.20.1