From ed7703c26e871510a488d8e725dc4bf101cde5a2 Mon Sep 17 00:00:00 2001
From: Tommaso Tocci <tommaso.tocci@mongodb.com>
Date: Fri, 15 Jul 2022 18:44:22 +0000
Subject: SERVER-67733  ShardingTest::awaitBalancerRound() doesn't work in case
 of CSRS stepdowns

---
 src/mongo/shell/shardingtest.js | 35 -----------------------------------
 src/mongo/shell/utils_sh.js     | 26 ++++++++++++++++++--------
 2 files changed, 18 insertions(+), 43 deletions(-)

(limited to 'src/mongo/shell')

diff --git a/src/mongo/shell/shardingtest.js b/src/mongo/shell/shardingtest.js
index b9465970890..1b5af0bbc66 100644
--- a/src/mongo/shell/shardingtest.js
+++ b/src/mongo/shell/shardingtest.js
@@ -739,41 +739,6 @@ var ShardingTest = function(params) {
         return max - min;
     };
 
-    /**
-     * Waits up to the specified timeout (with a default of 60s) for the balancer to execute one
-     * round. If no round has been executed, throws an error.
-     *
-     * The mongosConnection parameter is optional and allows callers to specify a connection
-     * different than the first mongos instance in the list.
-     */
-    this.awaitBalancerRound = function(timeoutMs, mongosConnection) {
-        timeoutMs = timeoutMs || 60000;
-        mongosConnection = mongosConnection || self.s0;
-
-        // Get the balancer section from the server status of the config server primary
-        function getBalancerStatus() {
-            var balancerStatus =
-                assert.commandWorked(mongosConnection.adminCommand({balancerStatus: 1}));
-            if (balancerStatus.mode !== 'full') {
-                throw Error('Balancer is not enabled');
-            }
-
-            return balancerStatus;
-        }
-
-        var initialStatus = getBalancerStatus();
-        var currentStatus;
-        assert.soon(
-            function() {
-                currentStatus = getBalancerStatus();
-                return (currentStatus.numBalancerRounds - initialStatus.numBalancerRounds) != 0;
-            },
-            function() {
-                return 'Latest balancer status: ' + tojson(currentStatus);
-            },
-            timeoutMs);
-    };
-
     /**
      * Waits up to one minute for the difference in chunks between the most loaded shard and
      * least loaded shard to be 0 or 1, indicating that the collection is well balanced. This should
diff --git a/src/mongo/shell/utils_sh.js b/src/mongo/shell/utils_sh.js
index b4bd7175096..3bb6db08d49 100644
--- a/src/mongo/shell/utils_sh.js
+++ b/src/mongo/shell/utils_sh.js
@@ -255,18 +255,28 @@ sh.waitForPingChange = function(activePings, timeout, interval) {
     return remainingPings;
 };
 
-sh.waitForBalancer = function(wait, timeout, interval) {
-    if (typeof (wait) === 'undefined') {
-        wait = false;
-    }
+/**
+ * Waits up to the specified timeout (with a default of 60s) for the balancer to execute one
+ * round. If no round has been executed, throws an error.
+ */
+sh.awaitBalancerRound = function(timeout, interval) {
+    timeout = timeout || 60000;
+
     var initialStatus = sh._getBalancerStatus();
-    if (!initialStatus.inBalancerRound && !wait) {
-        return;
-    }
     var currentStatus;
     assert.soon(function() {
         currentStatus = sh._getBalancerStatus();
-        return (currentStatus.numBalancerRounds - initialStatus.numBalancerRounds) != 0;
+        assert.eq(currentStatus.mode, 'full', "Balancer is disabled");
+        if (!friendlyEqual(currentStatus.term, initialStatus.term)) {
+            // A new primary of the csrs has been elected
+            initialStatus = currentStatus;
+            return false;
+        }
+        assert.gte(currentStatus.numBalancerRounds,
+                   initialStatus.numBalancerRounds,
+                   'Number of balancer rounds moved back in time unexpectedly. Current status: ' +
+                       tojson(currentStatus) + ', initial status: ' + tojson(initialStatus));
+        return currentStatus.numBalancerRounds > initialStatus.numBalancerRounds;
     }, 'Latest balancer status: ' + tojson(currentStatus), timeout, interval);
 };
 
-- 
cgit v1.2.1