From f0433cc2d6dae7fd01a2f6e65a756c855600edcd Mon Sep 17 00:00:00 2001
From: Aaron Schulz <aschulz@wikimedia.org>
Date: Sun, 21 Jul 2019 16:50:50 -0700
Subject: [PATCH] rdbms: make LoadBalancer::waitForAll() include servers with
 load in any group

This avoids needing the hack of a server needing a generic load of 1
in configuration just to force lag checks to include it.

Change-Id: I3de41efae427acb05beddb85dd5b5943b39f1b22
---
 .../libs/rdbms/loadbalancer/LoadBalancer.php  | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 1ef1d09b5f..69bd51f744 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -656,7 +656,7 @@ class LoadBalancer implements ILoadBalancer {
 				$ok = true; // no applicable loads
 			}
 		} finally {
-			# Restore the old position, as this is not used for lag-protection but for throttling
+			// Restore the old position; this is used for throttling, not lag-protection
 			$this->waitForPos = $oldPos;
 		}
 
@@ -673,7 +673,7 @@ class LoadBalancer implements ILoadBalancer {
 
 			$ok = true;
 			for ( $i = 1; $i < $serverCount; $i++ ) {
-				if ( $this->groupLoads[self::GROUP_GENERIC][$i] > 0 ) {
+				if ( $this->serverHasLoadInAnyGroup( $i ) ) {
 					$start = microtime( true );
 					$ok = $this->doWait( $i, true, $timeout ) && $ok;
 					$timeout -= intval( microtime( true ) - $start );
@@ -683,13 +683,27 @@ class LoadBalancer implements ILoadBalancer {
 				}
 			}
 		} finally {
-			# Restore the old position, as this is not used for lag-protection but for throttling
+			// Restore the old position; this is used for throttling, not lag-protection
 			$this->waitForPos = $oldPos;
 		}
 
 		return $ok;
 	}
 
+	/**
+	 * @param int $i Specific server index
+	 * @return bool
+	 */
+	private function serverHasLoadInAnyGroup( $i ) {
+		foreach ( $this->groupLoads as $loadsByIndex ) {
+			if ( ( $loadsByIndex[$i] ?? 0 ) > 0 ) {
+				return true;
+			}
+		}
+
+		return false;
+	}
+
 	/**
 	 * @param DBMasterPos|bool $pos
 	 */
-- 
2.20.1