SERVER-69136 Tests should consider balancerCollectionStatus may report balancerCompliant too early

(cherry picked from commit cf6316677c0ddc98835789c8e659f356f6295bfa)

4 files changed, 13 insertions, 33 deletions

diff --git a/jstests/concurrency/fsm_workloads/collection_defragmentation.js b/jstests/concurrency/fsm_workloads/collection_defragmentation.js
index c84729399e5..6d5714caa69 100644
--- a/jstests/concurrency/fsm_workloads/collection_defragmentation.js
+++ b/jstests/concurrency/fsm_workloads/collection_defragmentation.js
@@ -236,11 +236,7 @@ var $config = (function() {
                 // Enable balancing and wait for balanced
                 assertAlways.commandWorked(mongos.getDB('config').collections.update(
                     {_id: fullNs}, {$set: {"noBalance": false}}));
-                assertAlways.soon(function() {
-                    let res = mongos.adminCommand({balancerCollectionStatus: fullNs});
-                    assertAlways.commandWorked(res);
-                    return res.balancerCompliant;
-                });
+                sh.awaitCollectionBalance(mongos.getCollection(fullNs));
                 // Begin defragmentation again
                 assertAlways.commandWorked(mongos.adminCommand({
                     configureCollectionBalancing: fullNs,
diff --git a/jstests/sharding/balancing_based_on_size.js b/jstests/sharding/balancing_based_on_size.js
index 4ee6334f8fa..21b91a5dd6e 100644
--- a/jstests/sharding/balancing_based_on_size.js
+++ b/jstests/sharding/balancing_based_on_size.js
@@ -13,19 +13,6 @@
 
 load("jstests/sharding/libs/find_chunks_util.js");
 
-function getCollSizeMB(ns, node) {
-    let res;
-    let collections = [{ns: ns}];
-    assert.soon(() => {
-        res = assert.commandWorkedOrFailedWithCode(
-            node.adminCommand({_shardsvrGetStatsForBalancing: 1, collections: collections}),
-            [ErrorCodes.NotYetInitialized]);
-        return res.ok;
-    });
-
-    return res['stats'][0]['collSize'];
-}
-
 const maxChunkSizeMB = 1;
 const st = new ShardingTest(
     {shards: 2, mongos: 1, other: {chunkSize: maxChunkSizeMB, enableBalancer: false}});
@@ -75,21 +62,11 @@ jsTestLog("Printing sharding status before starting balancer");
 st.printShardingStatus();
 st.startBalancer();
 
-assert.soon(function() {
-    return assert.commandWorked(st.s0.adminCommand({balancerCollectionStatus: ns}))
-        .balancerCompliant;
-}, 'Timed out waiting for the collection to be balanced', 60000 /* timeout */, 1000 /* interval */);
-
-// Check that the collection size diff between shards is small (2 * maxChunkSize)
-const collSizeOnShard0BeforeNoopRounds = getCollSizeMB(ns, st.shard0.rs.getPrimary());
-const collSizeOnShard1BeforeNoopRounds = getCollSizeMB(ns, st.shard1.rs.getPrimary());
+st.awaitCollectionBalance(coll, 60000 /* timeout */, 1000 /* interval */);
 const chunksBeforeNoopRound = findChunksUtil.findChunksByNs(st.config, ns).toArray();
-var errMsg = '[BEFORE NOOP ROUND] Data on shard0 = ' + collSizeOnShard0BeforeNoopRounds +
-    ' and data on shard 1 = ' + collSizeOnShard1BeforeNoopRounds +
-    ' - chunks before noop round = ' + JSON.stringify(chunksBeforeNoopRound);
-assert.lte(collSizeOnShard0BeforeNoopRounds - collSizeOnShard1BeforeNoopRounds,
-           3 * maxChunkSizeMB,
-           errMsg);
+
+// Check that the collection is balanced
+st.verifyCollectionIsBalanced(coll);
 
 // Wait for some more rounds and then check the balancer is not wrongly moving around data
 st.forEachConfigServer((conn) => {
diff --git a/jstests/sharding/libs/defragmentation_util.js b/jstests/sharding/libs/defragmentation_util.js
index f416bafa7f7..e9c71aa2de8 100644
--- a/jstests/sharding/libs/defragmentation_util.js
+++ b/jstests/sharding/libs/defragmentation_util.js
@@ -207,6 +207,13 @@ var defragmentationUtil = (function() {
         assert.soon(function() {
             let balancerStatus =
                 assert.commandWorked(mongos.adminCommand({balancerCollectionStatus: ns}));
+
+            if (balancerStatus.balancerCompliant) {
+                // As we can't rely on `balancerCompliant` due to orphan counter non atomic update,
+                // we need to ensure the collection is balanced by some extra checks
+                sh.awaitCollectionBalance(mongos.getCollection(ns));
+            }
+
             return balancerStatus.balancerCompliant ||
                 balancerStatus.firstComplianceViolation !== 'defragmentingChunks';
         });
diff --git a/src/mongo/shell/utils_sh.js b/src/mongo/shell/utils_sh.js
index fc7574c5cb1..a64292a6004 100644
--- a/src/mongo/shell/utils_sh.js
+++ b/src/mongo/shell/utils_sh.js
@@ -331,7 +331,7 @@ sh.awaitCollectionBalance = function(coll, timeout, interval) {
         {'$group': {'_id': null, 'totalNumOrphanDocs': {'$sum': '$storageStats.numOrphanDocs'}}}
     ];
 
-    var oldDb = db;
+    var oldDb = (typeof (db) === 'undefined' ? undefined : db);
     try {
         db = coll.getDB();