Revert "SERVER-67733 ShardingTest::awaitBalancerRound() doesn't work in case of CSRS stepdowns"

This reverts commit 0c74d8baaef7d7c6d0e84401c1066b658f8c7bb2.

3 files changed, 30 insertions, 14 deletions

diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 6822de95f8a..f10a96cee4d 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -467,7 +467,6 @@ void Balancer::report(OperationContext* opCtx, BSONObjBuilder* builder) {
     builder->append("mode", BalancerSettingsType::kBalancerModes[mode]);
     builder->append("inBalancerRound", _inBalancerRound);
     builder->append("numBalancerRounds", _numBalancerRounds);
-    builder->append("term", repl::ReplicationCoordinator::get(opCtx)->getTerm());
 }
 
 void Balancer::_consumeActionStreamLoop() {
diff --git a/src/mongo/shell/shardingtest.js b/src/mongo/shell/shardingtest.js
index c304ad8de7c..465ee647492 100644
--- a/src/mongo/shell/shardingtest.js
+++ b/src/mongo/shell/shardingtest.js
@@ -602,9 +602,36 @@ var ShardingTest = function(params) {
     /**
      * 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) {
-        return this.waitForBalancer(true, timeoutMs || 60000);
+    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);
     };
 
     /**
diff --git a/src/mongo/shell/utils_sh.js b/src/mongo/shell/utils_sh.js
index 92f0c3a91a4..b4bd7175096 100644
--- a/src/mongo/shell/utils_sh.js
+++ b/src/mongo/shell/utils_sh.js
@@ -266,17 +266,7 @@ sh.waitForBalancer = function(wait, timeout, interval) {
     var currentStatus;
     assert.soon(function() {
         currentStatus = sh._getBalancerStatus();
-        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;
+        return (currentStatus.numBalancerRounds - initialStatus.numBalancerRounds) != 0;
     }, 'Latest balancer status: ' + tojson(currentStatus), timeout, interval);
 };