SERVER-62617 Optimise the logic to obtain the shard version while generating defragmentation actions

2 files changed, 142 insertions, 159 deletions

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 72afa835b31..7d640a06170 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -56,24 +56,11 @@ const std::string kProgress("progress");
 const std::string kNoPhase("none");
 const std::string kRemainingChunksToProcess("remainingChunksToProcess");
 
-// TODO (SERVER-62617) Avoid access to disk on each invocation
-ChunkVersion getShardVersion(OperationContext* opCtx, const ShardId& shardId, const UUID& uuid) {
-    auto coll = Grid::get(opCtx)->catalogClient()->getCollection(opCtx, uuid);
-    auto chunkVector = uassertStatusOK(Grid::get(opCtx)->catalogClient()->getChunks(
-        opCtx,
-        BSON(ChunkType::collectionUUID()
-             << coll.getUuid() << ChunkType::shard(shardId.toString())) /*query*/,
-        BSON(ChunkType::lastmod << -1) /*sort*/,
-        1 /*limit*/,
-        nullptr /*opTime*/,
-        coll.getEpoch(),
-        coll.getTimestamp(),
-        repl::ReadConcernLevel::kLocalReadConcern,
-        boost::none));
-    uassert(ErrorCodes::BadValue,
-            "No chunks or chunk version in collection",
-            !chunkVector.empty() && chunkVector.front().isVersionSet());
-    return chunkVector.front().getVersion();
+ChunkVersion getShardVersion(OperationContext* opCtx,
+                             const ShardId& shardId,
+                             const NamespaceString& nss) {
+    auto cm = Grid::get(opCtx)->catalogCache()->getShardedCollectionRoutingInfo(opCtx, nss);
+    return cm.getVersion(shardId);
 }
 
 std::vector<ChunkType> getCollectionChunks(OperationContext* opCtx, const CollectionType& coll) {
@@ -107,7 +94,8 @@ bool isRetriableForDefragmentation(const Status& error) {
             error == ErrorCodes::StaleShardVersion || error == ErrorCodes::StaleConfig);
 }
 
-void handleActionResult(const NamespaceString& nss,
+void handleActionResult(OperationContext* opCtx,
+                        const NamespaceString& nss,
                         const UUID& uuid,
                         const DefragmentationPhaseEnum currentPhase,
                         const Status& status,
@@ -118,7 +106,24 @@ void handleActionResult(const NamespaceString& nss,
         onSuccess();
         return;
     }
+
+    if (status.isA<ErrorCategory::StaleShardVersionError>()) {
+        if (auto staleInfo = status.extraInfo<StaleConfigInfo>()) {
+            Grid::get(opCtx)
+                ->catalogCache()
+                ->invalidateShardOrEntireCollectionEntryForShardedCollection(
+                    nss, staleInfo->getVersionWanted(), staleInfo->getShardId());
+        }
+    }
+
     if (isRetriableForDefragmentation(status)) {
+        LOGV2_DEBUG(6261701,
+                    1,
+                    "Hit retriable error while defragmenting collection",
+                    "namespace"_attr = nss,
+                    "uuid"_attr = uuid,
+                    "currentPhase"_attr = currentPhase,
+                    "error"_attr = redact(status));
         onRetriableError();
     } else {
         LOGV2_ERROR(6258601,
@@ -126,7 +131,7 @@ void handleActionResult(const NamespaceString& nss,
                     "namespace"_attr = nss,
                     "uuid"_attr = uuid,
                     "currentPhase"_attr = currentPhase,
-                    "error"_attr = status);
+                    "error"_attr = redact(status));
         onNonRetriableError();
     }
 }
@@ -199,7 +204,7 @@ public:
         if (!_pendingActionsByShards.empty()) {
             // TODO (SERVER-61635) improve fairness if needed
             auto& [shardId, pendingActions] = *_pendingActionsByShards.begin();
-            auto shardVersion = getShardVersion(opCtx, shardId, _uuid);
+            auto shardVersion = getShardVersion(opCtx, shardId, _nss);
 
             if (pendingActions.rangesWithoutDataSize.size() > pendingActions.rangesToMerge.size()) {
                 const auto& rangeToMeasure = pendingActions.rangesWithoutDataSize.back();
@@ -241,6 +246,7 @@ public:
                     auto& mergeResponse = stdx::get<Status>(response);
                     auto& shardingPendingActions = _pendingActionsByShards[mergeAction.shardId];
                     handleActionResult(
+                        opCtx,
                         _nss,
                         _uuid,
                         getType(),
@@ -258,6 +264,7 @@ public:
                 [&](const DataSizeInfo& dataSizeAction) {
                     auto& dataSizeResponse = stdx::get<StatusWith<DataSizeResponse>>(response);
                     handleActionResult(
+                        opCtx,
                         _nss,
                         _uuid,
                         getType(),
@@ -395,7 +402,7 @@ public:
         _outstandingMerges.push_back(_actionableMerges.front());
         _actionableMerges.pop_front();
         const auto& nextRequest = _outstandingMerges.back();
-        auto version = getShardVersion(opCtx, nextRequest.getDestinationShard(), _uuid);
+        auto version = getShardVersion(opCtx, nextRequest.getDestinationShard(), _nss);
         return boost::optional<DefragmentationAction>(
             nextRequest.asMergeInfo(_uuid, _nss, version));
     }
@@ -436,7 +443,7 @@ public:
             targetSibling->busyInOperation = true;
             usedShards->insert(nextSmallChunk->shard);
             usedShards->insert(targetSibling->shard);
-            auto smallChunkVersion = getShardVersion(opCtx, nextSmallChunk->shard, _uuid);
+            auto smallChunkVersion = getShardVersion(opCtx, nextSmallChunk->shard, _nss);
             _outstandingMigrations.emplace_back(nextSmallChunk, targetSibling);
             return _outstandingMigrations.back().asMigrateInfo(_uuid, _nss, smallChunkVersion);
         }
@@ -467,6 +474,11 @@ public:
                     }
 
                     if (migrationResponse.isOK()) {
+                        Grid::get(opCtx)
+                            ->catalogCache()
+                            ->invalidateShardOrEntireCollectionEntryForShardedCollection(
+                                _nss, boost::none, moveRequest.getDestinationShard());
+
                         auto transferredAmount = moveRequest.getMovedDataSizeBytes();
                         _shardInfos.at(moveRequest.getSourceShard()).currentSizeBytes -=
                             transferredAmount;
@@ -540,6 +552,12 @@ public:
                     auto onSuccess = [&] {
                         // The sequence is complete; update the state of the merged chunk...
                         auto& mergedChunk = mergeRequest.chunkToMergeWith;
+
+                        Grid::get(opCtx)
+                            ->catalogCache()
+                            ->invalidateShardOrEntireCollectionEntryForShardedCollection(
+                                _nss, boost::none, mergedChunk->shard);
+
                         auto& chunkToDelete = mergeRequest.chunkToMove;
                         mergedChunk->range = mergeRequest.asMergedRange();
                         mergedChunk->estimatedSizeBytes += chunkToDelete->estimatedSizeBytes;
@@ -570,7 +588,8 @@ public:
                     };
 
                     if (!_aborted) {
-                        handleActionResult(_nss,
+                        handleActionResult(opCtx,
+                                           _nss,
                                            _uuid,
                                            getType(),
                                            mergeResponse,
@@ -1011,7 +1030,7 @@ public:
 
         auto& [shardId, unmergedRanges] = *_unmergedRangesByShard.begin();
         invariant(!unmergedRanges.empty());
-        auto shardVersion = getShardVersion(opCtx, shardId, _uuid);
+        auto shardVersion = getShardVersion(opCtx, shardId, _nss);
         const auto& rangeToMerge = unmergedRanges.back();
         boost::optional<DefragmentationAction> nextAction = boost::optional<DefragmentationAction>(
             MergeInfo(shardId, _nss, _uuid, shardVersion, rangeToMerge));
@@ -1045,7 +1064,8 @@ public:
                             mergeAction.chunkRange);
                     };
                     auto onNonretriableError = [this] { _abort(getType()); };
-                    handleActionResult(_nss,
+                    handleActionResult(opCtx,
+                                       _nss,
                                        _uuid,
                                        getType(),
                                        mergeResponse,
@@ -1159,7 +1179,7 @@ public:
         boost::optional<DefragmentationAction> nextAction = boost::none;
         if (!_pendingActionsByShards.empty()) {
             auto& [shardId, pendingActions] = *_pendingActionsByShards.begin();
-            auto shardVersion = getShardVersion(opCtx, shardId, _uuid);
+            auto shardVersion = getShardVersion(opCtx, shardId, _nss);
 
             if (!pendingActions.rangesToSplit.empty()) {
                 const auto& [rangeToSplit, splitPoints] = pendingActions.rangesToSplit.back();
@@ -1226,6 +1246,7 @@ public:
                 [&](const AutoSplitVectorInfo& autoSplitVectorAction) {
                     auto& splitVectorResponse = stdx::get<StatusWith<SplitPoints>>(response);
                     handleActionResult(
+                        opCtx,
                         _nss,
                         _uuid,
                         getType(),
@@ -1257,6 +1278,7 @@ public:
                 [&](const SplitInfoWithKeyPattern& splitAction) {
                     auto& splitResponse = stdx::get<Status>(response);
                     handleActionResult(
+                        opCtx,
                         _nss,
                         _uuid,
                         getType(),
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
index a477213504c..2863bd91dac 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_test.cpp
@@ -73,31 +73,31 @@ protected:
 
     BalancerDefragmentationPolicyTest() : _clusterStats(), _defragmentationPolicy(&_clusterStats) {}
 
-    CollectionType makeConfigCollectionEntry(
-        boost::optional<DefragmentationPhaseEnum> phase = boost::none,
-        boost::optional<int64_t> maxChunkSizeBytes = boost::none) {
-        CollectionType shardedCollection(kNss, OID::gen(), Timestamp(1, 1), Date_t::now(), kUuid);
-        shardedCollection.setKeyPattern(kShardKeyPattern);
-        shardedCollection.setDefragmentCollection(true);
-        shardedCollection.setDefragmentationPhase(phase);
-        if (maxChunkSizeBytes) {
-            shardedCollection.setMaxChunkSizeBytes(maxChunkSizeBytes.get());
-        }
-        ASSERT_OK(insertToConfigCollection(
-            operationContext(), CollectionType::ConfigNS, shardedCollection.toBSON()));
-        return shardedCollection;
-    }
-
     CollectionType setupCollectionWithPhase(
         const std::vector<ChunkType>& chunkList,
-        boost::optional<DefragmentationPhaseEnum> startingPhase) {
+        boost::optional<DefragmentationPhaseEnum> startingPhase = boost::none,
+        boost::optional<int64_t> maxChunkSizeBytes = boost::none) {
+
         setupShards(kShardList);
         setupCollection(kNss, kShardKeyPattern, chunkList);
-        auto updateClause = startingPhase.has_value()
-            ? BSON("$set" << BSON(CollectionType::kDefragmentCollectionFieldName
-                                  << true << CollectionType::kDefragmentationPhaseFieldName
-                                  << DefragmentationPhase_serializer(*startingPhase)))
-            : BSON("$set" << BSON(CollectionType::kDefragmentCollectionFieldName << true));
+
+        const auto updateClause = [&] {
+            BSONObjBuilder builder;
+            BSONObjBuilder setObj(builder.subobjStart("$set"));
+            setObj.append(CollectionType::kDefragmentCollectionFieldName, true);
+
+            if (startingPhase) {
+                setObj.append(CollectionType::kDefragmentationPhaseFieldName,
+                              DefragmentationPhase_serializer(*startingPhase));
+            }
+
+            if (maxChunkSizeBytes) {
+                setObj.append(CollectionType::kMaxChunkSizeBytesFieldName, *maxChunkSizeBytes);
+            }
+            setObj.done();
+            return builder.obj();
+        }();
+
         ASSERT_OK(updateToConfigCollection(operationContext(),
                                            CollectionType::ConfigNS,
                                            BSON(CollectionType::kUuidFieldName << kUuid),
@@ -108,19 +108,15 @@ protected:
             ->getCollection(operationContext(), kUuid);
     }
 
-    void makeConfigChunkEntry(const boost::optional<long long>& estimatedSize = boost::none) {
-        ChunkType chunk(kUuid, ChunkRange(kKeyAtZero, kKeyAtTen), kCollectionVersion, kShardId0);
+    ChunkType makeConfigChunkEntry(const boost::optional<long long>& estimatedSize = boost::none) {
+        ChunkType chunk(kUuid, ChunkRange(kKeyAtMin, kKeyAtMax), kCollectionVersion, kShardId0);
         chunk.setEstimatedSizeBytes(estimatedSize);
-        ASSERT_OK(insertToConfigCollection(
-            operationContext(), ChunkType::ConfigNS, chunk.toConfigBSON()));
+        return chunk;
     }
 
-    void makeMergeableConfigChunkEntries() {
-        auto opCtx = operationContext();
-        ChunkType chunk(kUuid, ChunkRange(kKeyAtZero, kKeyAtTen), kCollectionVersion, kShardId0);
-        ChunkType chunk2(kUuid, ChunkRange(kKeyAtTen, kKeyAtTwenty), kCollectionVersion, kShardId0);
-        ASSERT_OK(insertToConfigCollection(opCtx, ChunkType::ConfigNS, chunk.toConfigBSON()));
-        ASSERT_OK(insertToConfigCollection(opCtx, ChunkType::ConfigNS, chunk2.toConfigBSON()));
+    std::vector<ChunkType> makeMergeableConfigChunkEntries() {
+        return {ChunkType(kUuid, ChunkRange(kKeyAtMin, kKeyAtTen), kCollectionVersion, kShardId0),
+                ChunkType(kUuid, ChunkRange(kKeyAtTen, kKeyAtMax), kCollectionVersion, kShardId0)};
     }
 
     BSONObj getConfigCollectionEntry() {
@@ -181,18 +177,6 @@ TEST_F(BalancerDefragmentationPolicyTest, TestGetNextActionIsNotReadyWhenNotDefr
     ASSERT_FALSE(future.isReady());
 }
 
-TEST_F(BalancerDefragmentationPolicyTest, TestAddEmptyCollectionDoesNotTriggerDefragmentation) {
-    auto coll = makeConfigCollectionEntry();
-    setDefaultClusterStats();
-    _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
-
-    ASSERT_FALSE(_defragmentationPolicy.isDefragmentingCollection(coll.getUuid()));
-    verifyExpectedDefragmentationPhaseOndisk(boost::none);
-
-    auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
-    ASSERT_FALSE(future.isReady());
-}
-
 TEST_F(BalancerDefragmentationPolicyTest, TestAddCollectionWhenCollectionRemovedFailsGracefully) {
     CollectionType coll(kNss, OID::gen(), Timestamp(1, 1), Date_t::now(), kUuid);
     coll.setKeyPattern(kShardKeyPattern);
@@ -211,8 +195,7 @@ TEST_F(BalancerDefragmentationPolicyTest, TestAddCollectionWhenCollectionRemoved
 // Phase 1 tests.
 
 TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAddSingleChunkCollectionTriggersDataSize) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     ASSERT_FALSE(future.isReady());
     ASSERT_FALSE(_defragmentationPolicy.isDefragmentingCollection(coll.getUuid()));
@@ -228,15 +211,14 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAddSingleChunkCollectionTr
     // 3. with the expected content
     // TODO refactor chunk builder
     ASSERT_EQ(coll.getNss(), dataSizeAction.nss);
-    ASSERT_BSONOBJ_EQ(kKeyAtZero, dataSizeAction.chunkRange.getMin());
-    ASSERT_BSONOBJ_EQ(kKeyAtTen, dataSizeAction.chunkRange.getMax());
+    ASSERT_BSONOBJ_EQ(kKeyAtMin, dataSizeAction.chunkRange.getMin());
+    ASSERT_BSONOBJ_EQ(kKeyAtMax, dataSizeAction.chunkRange.getMax());
 }
 
 TEST_F(BalancerDefragmentationPolicyTest,
        AddSingleChunkCollectionWithKnownDataSizeCompletesDefragmentationWithNoOperationIssued) {
-    auto coll = makeConfigCollectionEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(1024)});
     setDefaultClusterStats();
-    makeConfigChunkEntry(1024);
 
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
 
@@ -270,8 +252,7 @@ TEST_F(BalancerDefragmentationPolicyTest,
 
 TEST_F(BalancerDefragmentationPolicyTest,
        TestPhaseOneAcknowledgeFinalDataSizeActionCompletesPhase) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     setDefaultClusterStats();
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
@@ -282,7 +263,7 @@ TEST_F(BalancerDefragmentationPolicyTest,
 
     // 1. The outcome of the data size has been stored in the expected document...
     auto chunkQuery = BSON(ChunkType::collectionUUID()
-                           << kUuid << ChunkType::min(kKeyAtZero) << ChunkType::max(kKeyAtTen));
+                           << kUuid << ChunkType::min(kKeyAtMin) << ChunkType::max(kKeyAtMax));
     auto configChunkDoc =
         findOneOnConfigCollection(operationContext(), ChunkType::ConfigNS, chunkQuery).getValue();
     ASSERT_EQ(configChunkDoc.getIntField(ChunkType::estimatedSizeBytes.name()), 2000);
@@ -295,8 +276,7 @@ TEST_F(BalancerDefragmentationPolicyTest,
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestRetriableFailedDataSizeActionGetsReissued) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     DataSizeInfo failingDataSizeAction = stdx::get<DataSizeInfo>(future.get());
@@ -326,8 +306,7 @@ TEST_F(BalancerDefragmentationPolicyTest, TestRetriableFailedDataSizeActionGetsR
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestRemoveCollectionEndsDefragmentation) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     DataSizeInfo dataSizeAction = stdx::get<DataSizeInfo>(future.get());
@@ -347,8 +326,7 @@ TEST_F(BalancerDefragmentationPolicyTest, TestRemoveCollectionEndsDefragmentatio
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneUserCancellationBeginsPhase3) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
 
     // Collection should be in phase 1
@@ -365,8 +343,7 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneUserCancellationBeginsPhas
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestNonRetriableErrorRebuildsCurrentPhase) {
-    auto coll = makeConfigCollectionEntry();
-    makeConfigChunkEntry();
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     DataSizeInfo failingDataSizeAction = stdx::get<DataSizeInfo>(future.get());
@@ -386,20 +363,15 @@ TEST_F(BalancerDefragmentationPolicyTest, TestNonRetriableErrorRebuildsCurrentPh
     // 3. with the expected content
     // TODO refactor chunk builder
     ASSERT_EQ(coll.getNss(), dataSizeAction.nss);
-    ASSERT_BSONOBJ_EQ(kKeyAtZero, dataSizeAction.chunkRange.getMin());
-    ASSERT_BSONOBJ_EQ(kKeyAtTen, dataSizeAction.chunkRange.getMax());
+    ASSERT_BSONOBJ_EQ(kKeyAtMin, dataSizeAction.chunkRange.getMin());
+    ASSERT_BSONOBJ_EQ(kKeyAtMax, dataSizeAction.chunkRange.getMax());
 }
 
 TEST_F(BalancerDefragmentationPolicyTest,
        TestNonRetriableErrorWaitsForAllOutstandingActionsToComplete) {
-    auto coll = makeConfigCollectionEntry();
-    ChunkType chunk1(kUuid, ChunkRange(kKeyAtZero, kKeyAtTen), kCollectionVersion, kShardId0);
-    ChunkType chunk2(
-        kUuid, ChunkRange(BSON("x" << 11), kKeyAtTwenty), kCollectionVersion, kShardId0);
-    ASSERT_OK(
-        insertToConfigCollection(operationContext(), ChunkType::ConfigNS, chunk1.toConfigBSON()));
-    ASSERT_OK(
-        insertToConfigCollection(operationContext(), ChunkType::ConfigNS, chunk2.toConfigBSON()));
+    auto coll = setupCollectionWithPhase(
+        {ChunkType{kUuid, ChunkRange(kKeyAtMin, kKeyAtTen), kCollectionVersion, kShardId0},
+         ChunkType{kUuid, ChunkRange(BSON("x" << 11), kKeyAtMax), kCollectionVersion, kShardId0}});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     DataSizeInfo failingDataSizeAction = stdx::get<DataSizeInfo>(future.get());
@@ -431,8 +403,7 @@ TEST_F(BalancerDefragmentationPolicyTest,
 
 TEST_F(BalancerDefragmentationPolicyTest,
        TestPhaseOneAcknowledgeMergeChunkActionsTriggersDataSizeOnResultingRange) {
-    auto coll = makeConfigCollectionEntry();
-    makeMergeableConfigChunkEntries();
+    auto coll = setupCollectionWithPhase({makeMergeableConfigChunkEntries()});
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto mergeChunksAction = stdx::get<MergeInfo>(future.get());
@@ -455,8 +426,7 @@ TEST_F(BalancerDefragmentationPolicyTest,
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneFailedMergeChunksActionGetsReissued) {
-    auto coll = makeConfigCollectionEntry();
-    makeMergeableConfigChunkEntries();
+    auto coll = setupCollectionWithPhase(makeMergeableConfigChunkEntries());
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto failingMergeChunksAction = stdx::get<MergeInfo>(future.get());
@@ -482,15 +452,14 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneFailedMergeChunksActionGet
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAcknowledgeSuccessfulMergeAction) {
-    auto coll = makeConfigCollectionEntry();
+    auto coll = setupCollectionWithPhase(makeMergeableConfigChunkEntries());
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     ASSERT_FALSE(future.isReady());
-    makeMergeableConfigChunkEntries();
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     ASSERT_TRUE(future.isReady());
     MergeInfo mergeInfoAction = stdx::get<MergeInfo>(future.get());
-    ASSERT_BSONOBJ_EQ(mergeInfoAction.chunkRange.getMin(), kKeyAtZero);
-    ASSERT_BSONOBJ_EQ(mergeInfoAction.chunkRange.getMax(), kKeyAtTwenty);
+    ASSERT_BSONOBJ_EQ(mergeInfoAction.chunkRange.getMin(), kKeyAtMin);
+    ASSERT_BSONOBJ_EQ(mergeInfoAction.chunkRange.getMax(), kKeyAtMax);
     future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     ASSERT_FALSE(future.isReady());
     _defragmentationPolicy.acknowledgeMergeResult(
@@ -506,22 +475,26 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAllConsecutive) {
     // Set up collection with all mergeable chunks
     std::vector<ChunkType> chunkList;
     for (int i = 0; i < 5; i++) {
+        const auto minKey = (i == 0) ? kKeyAtMin : BSON("x" << i);
+        const auto maxKey = BSON("x" << i + 1);
         ChunkType chunk(
             kUuid,
-            ChunkRange(BSON("x" << i), BSON("x" << i + 1)),
+            ChunkRange(minKey, maxKey),
             ChunkVersion(1, i, kCollectionVersion.epoch(), kCollectionVersion.getTimestamp()),
             kShardId0);
         chunkList.push_back(chunk);
     }
     for (int i = 5; i < 10; i++) {
+        const auto minKey = BSON("x" << i);
+        const auto maxKey = (i == 9) ? kKeyAtMax : BSON("x" << i + 1);
         ChunkType chunk(
             kUuid,
-            ChunkRange(BSON("x" << i), BSON("x" << i + 1)),
+            ChunkRange(minKey, maxKey),
             ChunkVersion(1, i, kCollectionVersion.epoch(), kCollectionVersion.getTimestamp()),
             kShardId1);
         chunkList.push_back(chunk);
     }
-    auto coll = setupCollectionWithPhase(chunkList, boost::none);
+    auto coll = setupCollectionWithPhase(chunkList, boost::none, boost::none);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     // Test
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
@@ -532,16 +505,16 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAllConsecutive) {
     // (Note: there is no guarantee on the order provided by the stream)
     MergeInfo mergeAction = stdx::get<MergeInfo>(future.get());
     MergeInfo mergeAction2 = stdx::get<MergeInfo>(future2.get());
-    if (mergeAction.chunkRange.getMin().woCompare(BSON("x" << 0)) == 0) {
-        ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMin(), BSON("x" << 0));
+    if (mergeAction.chunkRange.getMin().woCompare(kKeyAtMin) == 0) {
+        ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMin(), kKeyAtMin);
         ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMax(), BSON("x" << 5));
         ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMin(), BSON("x" << 5));
-        ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMax(), BSON("x" << 10));
+        ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMax(), kKeyAtMax);
     } else {
-        ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMin(), BSON("x" << 0));
+        ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMin(), kKeyAtMin);
         ASSERT_BSONOBJ_EQ(mergeAction2.chunkRange.getMax(), BSON("x" << 5));
         ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMin(), BSON("x" << 5));
-        ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMax(), BSON("x" << 10));
+        ASSERT_BSONOBJ_EQ(mergeAction.chunkRange.getMax(), kKeyAtMax);
     }
     auto future3 = _defragmentationPolicy.getNextStreamingAction(operationContext());
     ASSERT_FALSE(future3.isReady());
@@ -550,15 +523,17 @@ TEST_F(BalancerDefragmentationPolicyTest, TestPhaseOneAllConsecutive) {
 TEST_F(BalancerDefragmentationPolicyTest, PhaseOneNotConsecutive) {
     std::vector<ChunkType> chunkList;
     for (int i = 0; i < 10; i++) {
+        const auto minKey = (i == 0) ? kKeyAtMin : BSON("x" << i);
+        const auto maxKey = (i == 9) ? kKeyAtMax : BSON("x" << i + 1);
         ShardId chosenShard = (i == 5) ? kShardId1 : kShardId0;
         ChunkType chunk(
             kUuid,
-            ChunkRange(BSON("x" << i), BSON("x" << i + 1)),
+            ChunkRange(minKey, maxKey),
             ChunkVersion(1, i, kCollectionVersion.epoch(), kCollectionVersion.getTimestamp()),
             chosenShard);
         chunkList.push_back(chunk);
     }
-    auto coll = setupCollectionWithPhase(chunkList, boost::none);
+    auto coll = setupCollectionWithPhase(chunkList, boost::none, boost::none);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     // Three actions (in an unspecified order) should be immediately available.
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
@@ -575,12 +550,12 @@ TEST_F(BalancerDefragmentationPolicyTest, PhaseOneNotConsecutive) {
         stdx::visit(
             visit_helper::Overloaded{
                 [&](const MergeInfo& mergeAction) {
-                    if (mergeAction.chunkRange.getMin().woCompare(BSON("x" << 0)) == 0 &&
+                    if (mergeAction.chunkRange.getMin().woCompare(kKeyAtMin) == 0 &&
                         mergeAction.chunkRange.getMax().woCompare(BSON("x" << 5)) == 0) {
                         ++timesLowerRangeMergeFound;
                     }
                     if (mergeAction.chunkRange.getMin().woCompare(BSON("x" << 6)) == 0 &&
-                        mergeAction.chunkRange.getMax().woCompare(BSON("x" << 10)) == 0) {
+                        mergeAction.chunkRange.getMax().woCompare(kKeyAtMax) == 0) {
                         ++timesUpperRangeMergeFound;
                     }
                 },
@@ -609,26 +584,10 @@ TEST_F(BalancerDefragmentationPolicyTest, PhaseOneNotConsecutive) {
 
 // Phase 2 tests.
 
-TEST_F(BalancerDefragmentationPolicyTest,
-       TestPhaseTwoEmptyCollectionDoesNotTriggerDefragmentation) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kMoveAndMergeChunks);
-    setDefaultClusterStats();
-    _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
-
-    ASSERT_FALSE(_defragmentationPolicy.isDefragmentingCollection(coll.getUuid()));
-    verifyExpectedDefragmentationPhaseOndisk(boost::none);
-    stdx::unordered_set<ShardId> usedShards;
-    auto pendingMigrations =
-        _defragmentationPolicy.selectChunksToMove(operationContext(), &usedShards);
-    ASSERT_TRUE(pendingMigrations.empty());
-    auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
-    ASSERT_FALSE(future.isReady());
-}
-
 TEST_F(BalancerDefragmentationPolicyTest, TestPhaseTwoMissingDataSizeRestartsPhase1) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kMoveAndMergeChunks);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry()},
+                                         DefragmentationPhaseEnum::kMoveAndMergeChunks);
     setDefaultClusterStats();
-    makeConfigChunkEntry();
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
 
     // Should be in phase 1
@@ -754,7 +713,8 @@ TEST_F(BalancerDefragmentationPolicyTest,
                                           chunkOnShard3,
                                           secondChunkOnShard0,
                                           secondChunkOnShard1},
-                                         DefragmentationPhaseEnum::kMoveAndMergeChunks);
+                                         DefragmentationPhaseEnum::kMoveAndMergeChunks,
+                                         boost::none);
     setDefaultClusterStats();
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
 
@@ -768,14 +728,14 @@ TEST_F(BalancerDefragmentationPolicyTest,
 }
 
 /** Phase 3 tests. By passing in DefragmentationPhaseEnum::kSplitChunks to
- * makeConfigCollectionEntry, the persisted collection entry will have
+ * setupCollectionWithPhase, the persisted collection entry will have
  * kDefragmentationPhaseFieldName set to kSplitChunks and defragmentation will be started with
  * phase 3.
  */
 
 TEST_F(BalancerDefragmentationPolicyTest, DefragmentationBeginsWithPhase3FromPersistedSetting) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     // Defragmentation does not start until refreshCollectionDefragmentationStatus is called
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     ASSERT_FALSE(future.isReady());
@@ -788,8 +748,8 @@ TEST_F(BalancerDefragmentationPolicyTest, DefragmentationBeginsWithPhase3FromPer
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, SingleLargeChunkCausesAutoSplitAndSplitActions) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
 
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
@@ -799,13 +759,14 @@ TEST_F(BalancerDefragmentationPolicyTest, SingleLargeChunkCausesAutoSplitAndSpli
     AutoSplitVectorInfo splitVectorAction = stdx::get<AutoSplitVectorInfo>(future.get());
     // with the expected content
     ASSERT_EQ(coll.getNss(), splitVectorAction.nss);
-    ASSERT_BSONOBJ_EQ(kKeyAtZero, splitVectorAction.minKey);
-    ASSERT_BSONOBJ_EQ(kKeyAtTen, splitVectorAction.maxKey);
+    ASSERT_BSONOBJ_EQ(kKeyAtMin, splitVectorAction.minKey);
+    ASSERT_BSONOBJ_EQ(kKeyAtMax, splitVectorAction.maxKey);
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, CollectionMaxChunkSizeIsUsedForPhase3) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks, 1024);
-    makeConfigChunkEntry(2 * 1024);  // > 1KB should trigger AutoSplitVector
+    // One chunk > 1KB should trigger AutoSplitVector
+    auto coll = setupCollectionWithPhase(
+        {makeConfigChunkEntry(2 * 1024)}, DefragmentationPhaseEnum::kSplitChunks, 1024);
 
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
 
@@ -816,13 +777,13 @@ TEST_F(BalancerDefragmentationPolicyTest, CollectionMaxChunkSizeIsUsedForPhase3)
     AutoSplitVectorInfo splitVectorAction = stdx::get<AutoSplitVectorInfo>(future.get());
     // with the expected content
     ASSERT_EQ(coll.getNss(), splitVectorAction.nss);
-    ASSERT_BSONOBJ_EQ(kKeyAtZero, splitVectorAction.minKey);
-    ASSERT_BSONOBJ_EQ(kKeyAtTen, splitVectorAction.maxKey);
+    ASSERT_BSONOBJ_EQ(kKeyAtMin, splitVectorAction.minKey);
+    ASSERT_BSONOBJ_EQ(kKeyAtMax, splitVectorAction.maxKey);
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestRetryableFailedAutoSplitActionGetsReissued) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     AutoSplitVectorInfo failingAutoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
@@ -850,8 +811,8 @@ TEST_F(BalancerDefragmentationPolicyTest, TestRetryableFailedAutoSplitActionGets
 
 TEST_F(BalancerDefragmentationPolicyTest,
        TestAcknowledgeAutoSplitActionTriggersSplitOnResultingRange) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
@@ -877,8 +838,8 @@ TEST_F(BalancerDefragmentationPolicyTest,
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestAutoSplitWithNoSplitPointsDoesNotTriggerSplit) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
@@ -893,8 +854,8 @@ TEST_F(BalancerDefragmentationPolicyTest, TestAutoSplitWithNoSplitPointsDoesNotT
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestMoreThan16MBSplitPointsTriggersSplitAndAutoSplit) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
@@ -934,8 +895,8 @@ TEST_F(BalancerDefragmentationPolicyTest, TestMoreThan16MBSplitPointsTriggersSpl
 }
 
 TEST_F(BalancerDefragmentationPolicyTest, TestFailedSplitChunkActionGetsReissued) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());
@@ -967,8 +928,8 @@ TEST_F(BalancerDefragmentationPolicyTest, TestFailedSplitChunkActionGetsReissued
 
 TEST_F(BalancerDefragmentationPolicyTest,
        TestAcknowledgeLastSuccessfulSplitActionEndsDefragmentation) {
-    auto coll = makeConfigCollectionEntry(DefragmentationPhaseEnum::kSplitChunks);
-    makeConfigChunkEntry(kPhase3DefaultChunkSize);
+    auto coll = setupCollectionWithPhase({makeConfigChunkEntry(kPhase3DefaultChunkSize)},
+                                         DefragmentationPhaseEnum::kSplitChunks);
     _defragmentationPolicy.refreshCollectionDefragmentationStatus(operationContext(), coll);
     auto future = _defragmentationPolicy.getNextStreamingAction(operationContext());
     auto autoSplitAction = stdx::get<AutoSplitVectorInfo>(future.get());