SERVER-65956 fix minor bugs in defragmentation source code/tests

(cherry picked from commit 4afa98e3a445ee7f6738c8717f7283c2ef6eeefd)

5 files changed, 20 insertions, 12 deletions

diff --git a/jstests/concurrency/fsm_libs/cluster.js b/jstests/concurrency/fsm_libs/cluster.js
index 5cf6c7facff..381131e2b54 100644
--- a/jstests/concurrency/fsm_libs/cluster.js
+++ b/jstests/concurrency/fsm_libs/cluster.js
@@ -319,6 +319,11 @@ var Cluster = function(options) {
         });
     };
 
+    this.getConfigPrimaryNode = function getConfigPrimaryNode() {
+        assert(initialized, 'cluster must be initialized first');
+        return st._configsvr.getPrimary();
+    };
+
     this.synchronizeMongosClusterTimes = function synchronizeMongosClusterTimes() {
         const contactConfigServerFn = ((mongosConn) => {
             // The admin database is hosted on the config server.
diff --git a/jstests/concurrency/fsm_workloads/collection_defragmentation.js b/jstests/concurrency/fsm_workloads/collection_defragmentation.js
index 9d651a3566b..90f7e3de9c8 100644
--- a/jstests/concurrency/fsm_workloads/collection_defragmentation.js
+++ b/jstests/concurrency/fsm_workloads/collection_defragmentation.js
@@ -239,7 +239,7 @@ var $config = (function() {
                 // Wait for defragmentation to complete and check final state
                 defragmentationUtil.waitForEndOfDefragmentation(mongos, fullNs);
                 defragmentationUtil.checkPostDefragmentationState(
-                    mongos, fullNs, maxChunkSizeMB, "key");
+                    cluster.getConfigPrimaryNode(), mongos, fullNs, maxChunkSizeMB, "key");
                 // Resume original throttling value
                 cluster.executeOnConfigNodes((db) => {
                     assert.commandWorked(db.adminCommand({
diff --git a/jstests/sharding/defragment_large_collection.js b/jstests/sharding/defragment_large_collection.js
index 379e6e4297a..20210762258 100644
--- a/jstests/sharding/defragment_large_collection.js
+++ b/jstests/sharding/defragment_large_collection.js
@@ -82,7 +82,8 @@ let runTest = function(numCollections, dbName) {
         const finalNumberChunks = findChunksUtil.countChunksForNs(st.s.getDB('config'), ns);
         jsTest.log("Finished defragmentation of collection " + coll + " with " + finalNumberChunks +
                    " chunks.");
-        defragmentationUtil.checkPostDefragmentationState(st.s, ns, maxChunkSizeMB, "key");
+        defragmentationUtil.checkPostDefragmentationState(
+            st.configRS.getPrimary(), st.s, ns, maxChunkSizeMB, "key");
     });
 
     st.printShardingStatus();
diff --git a/jstests/sharding/libs/defragmentation_util.js b/jstests/sharding/libs/defragmentation_util.js
index 8fcc8a35d99..52ac1333ef7 100644
--- a/jstests/sharding/libs/defragmentation_util.js
+++ b/jstests/sharding/libs/defragmentation_util.js
@@ -122,9 +122,11 @@ var defragmentationUtil = (function() {
                        tojson(rightChunk)} are mergeable with combined size ${combinedDataSize}`);
     };
 
-    let checkPostDefragmentationState = function(mongos, ns, maxChunkSizeMB, shardKey) {
+    let checkPostDefragmentationState = function(configSvr, mongos, ns, maxChunkSizeMB, shardKey) {
         const withAutoSplitActive =
-            !FeatureFlagUtil.isEnabled(mongos.getDB('admin'), 'NoMoreAutoSplitter');
+            !FeatureFlagUtil.isEnabled(configSvr.getDB('admin'), 'NoMoreAutoSplitter');
+        jsTest.log(`Chunk (auto)splitting functionalities assumed to be ${
+            withAutoSplitActive ? "ON" : "OFF"}`);
         const oversizedChunkThreshold = maxChunkSizeMB * 1024 * 1024 * 4 / 3;
         const chunks = findChunksUtil.findChunksByNs(mongos.getDB('config'), ns)
                            .sort({[shardKey]: 1})
@@ -160,7 +162,8 @@ var defragmentationUtil = (function() {
                                                        oversizedChunkThreshold);
                     } else {
                         assert(false,
-                               `Chunks ${leftChunk} and ${rightChunk} should have been merged`);
+                               `Chunks ${tojson(leftChunk)} and ${
+                                   tojson(rightChunk)} should have been merged`);
                     }
                 }
             }
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
index 5563223fb0a..0f1063d6d8c 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -1611,6 +1611,11 @@ std::unique_ptr<DefragmentationPhase> BalancerDefragmentationPolicyImpl::_transi
     DefragmentationPhaseEnum nextPhase,
     bool shouldPersistPhase) {
     std::unique_ptr<DefragmentationPhase> nextPhaseObject(nullptr);
+    if (nextPhase == DefragmentationPhaseEnum::kSplitChunks &&
+        feature_flags::gNoMoreAutoSplitter.isEnabled(serverGlobalParams.featureCompatibility)) {
+        nextPhase = DefragmentationPhaseEnum::kFinished;
+    }
+
     try {
         if (shouldPersistPhase) {
             _persistPhaseUpdate(opCtx, nextPhase, coll.getUuid());
@@ -1629,13 +1634,7 @@ std::unique_ptr<DefragmentationPhase> BalancerDefragmentationPolicyImpl::_transi
                 nextPhaseObject = MergeChunksPhase::build(opCtx, coll);
                 break;
             case DefragmentationPhaseEnum::kSplitChunks:
-                if (feature_flags::gNoMoreAutoSplitter.isEnabled(
-                        serverGlobalParams.featureCompatibility)) {
-                    _clearDefragmentationState(opCtx, coll.getUuid());
-                } else {
-                    nextPhaseObject = SplitChunksPhase::build(opCtx, coll);
-                }
-
+                nextPhaseObject = SplitChunksPhase::build(opCtx, coll);
                 break;
             case DefragmentationPhaseEnum::kFinished:
                 _clearDefragmentationState(opCtx, coll.getUuid());