SERVER-52742 Ensure both donorFields and recipientFields exists on the original collection entry after a resharding operation commits

4 files changed, 96 insertions, 44 deletions

diff --git a/src/mongo/db/s/resharding/resharding_coordinator_service.cpp b/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
index 8ca360a992d..94fd81363a6 100644
--- a/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
+++ b/src/mongo/db/s/resharding/resharding_coordinator_service.cpp
@@ -168,6 +168,32 @@ void writeToCoordinatorStateNss(OperationContext* opCtx,
     }
 }
 
+/**
+ * Extracts the ShardId from each Donor/RecipientShardEntry in participantShardEntries.
+ */
+template <class T>
+std::vector<ShardId> extractShardIds(const std::vector<T>& participantShardEntries) {
+    std::vector<ShardId> shardIds(participantShardEntries.size());
+    std::transform(participantShardEntries.begin(),
+                   participantShardEntries.end(),
+                   shardIds.begin(),
+                   [](auto& shardEntry) { return shardEntry.getId(); });
+    return shardIds;
+}
+
+/**
+ * Creates reshardingFields.recipientFields for the resharding operation. Note: these should not
+ * change once the operation has begun.
+ */
+TypeCollectionRecipientFields constructRecipientFields(
+    const ReshardingCoordinatorDocument& coordinatorDoc) {
+    auto donorShardIds = extractShardIds(coordinatorDoc.getDonorShards());
+    TypeCollectionRecipientFields recipientFields(
+        std::move(donorShardIds), coordinatorDoc.getExistingUUID(), coordinatorDoc.getNss());
+    emplaceFetchTimestampIfExists(recipientFields, coordinatorDoc.getFetchTimestamp());
+    return recipientFields;
+}
+
 BSONObj createReshardingFieldsUpdateForOriginalNss(
     OperationContext* opCtx,
     const ReshardingCoordinatorDocument& coordinatorDoc,
@@ -193,15 +219,17 @@ BSONObj createReshardingFieldsUpdateForOriginalNss(
             // the new sharded collection. Set 'uuid' to the reshardingUUID, 'key' to the new shard
             // key, 'lastmodEpoch' to newCollectionEpoch, and 'timestamp' to
             // newCollectionTimestamp (if newCollectionTimestamp has a value; i.e. when the
-            // shardingFullDDLSupport feature flag is enabled). Also update the 'state' field in the
-            // 'reshardingFields' section
+            // shardingFullDDLSupport feature flag is enabled). Also update the 'state' field and
+            // add the 'recipientFields' to the 'reshardingFields' section.
+            auto recipientFields = constructRecipientFields(coordinatorDoc);
             BSONObj setFields =
                 BSON("uuid" << coordinatorDoc.get_id() << "key"
                             << coordinatorDoc.getReshardingKey().toBSON() << "lastmodEpoch"
                             << newCollectionEpoch.get() << "lastmod"
                             << opCtx->getServiceContext()->getPreciseClockSource()->now()
                             << "reshardingFields.state"
-                            << CoordinatorState_serializer(coordinatorDoc.getState()).toString());
+                            << CoordinatorState_serializer(coordinatorDoc.getState()).toString()
+                            << "reshardingFields.recipientFields" << recipientFields.toBSON());
             if (newCollectionTimestamp.has_value()) {
                 setFields = setFields.addFields(BSON("timestamp" << newCollectionTimestamp.get()));
             }
@@ -396,19 +424,6 @@ void updateChunkAndTagsDocsForTempNss(OperationContext* opCtx,
         opCtx, TagsType::ConfigNS, tagsRequest, txnNumber);
 }
 
-/**
- * Extracts the ShardId from each Donor/RecipientShardEntry in participantShardEntries.
- */
-template <class T>
-std::vector<ShardId> extractShardIds(const std::vector<T>& participantShardEntries) {
-    std::vector<ShardId> shardIds(participantShardEntries.size());
-    std::transform(participantShardEntries.begin(),
-                   participantShardEntries.end(),
-                   shardIds.begin(),
-                   [](auto& shardEntry) { return shardEntry.getId(); });
-    return shardIds;
-}
-
 //
 // Helper methods for ensuring donors/ recipients are able to notice when certain state transitions
 // occur.
@@ -530,12 +545,8 @@ CollectionType createTempReshardingCollectionType(
     TypeCollectionReshardingFields tempEntryReshardingFields(coordinatorDoc.get_id());
     tempEntryReshardingFields.setState(coordinatorDoc.getState());
 
-    auto donorShardIds = extractShardIds(coordinatorDoc.getDonorShards());
-
-    TypeCollectionRecipientFields recipient(
-        std::move(donorShardIds), coordinatorDoc.getExistingUUID(), coordinatorDoc.getNss());
-    emplaceFetchTimestampIfExists(recipient, coordinatorDoc.getFetchTimestamp());
-    tempEntryReshardingFields.setRecipientFields(recipient);
+    auto recipientFields = constructRecipientFields(coordinatorDoc);
+    tempEntryReshardingFields.setRecipientFields(std::move(recipientFields));
     collType.setReshardingFields(std::move(tempEntryReshardingFields));
     return collType;
 }
diff --git a/src/mongo/db/s/resharding/resharding_coordinator_test.cpp b/src/mongo/db/s/resharding/resharding_coordinator_test.cpp
index af86c07a826..11eb9dde0d6 100644
--- a/src/mongo/db/s/resharding/resharding_coordinator_test.cpp
+++ b/src/mongo/db/s/resharding/resharding_coordinator_test.cpp
@@ -286,13 +286,16 @@ protected:
     }
 
     void readOriginalCollectionCatalogEntryAndAssertReshardingFieldsMatchExpected(
-        OperationContext* opCtx, CollectionType expectedCollType, bool doneState) {
+        OperationContext* opCtx,
+        CollectionType expectedCollType,
+        const ReshardingCoordinatorDocument& expectedCoordinatorDoc) {
         DBDirectClient client(opCtx);
         CollectionType onDiskEntry(
             client.findOne(CollectionType::ConfigNS.ns(), Query(BSON("_id" << _originalNss.ns()))));
 
         auto expectedReshardingFields = expectedCollType.getReshardingFields();
-        if (doneState ||
+        auto expectedCoordinatorState = expectedCoordinatorDoc.getState();
+        if (expectedCoordinatorState == CoordinatorStateEnum::kDone ||
             (expectedReshardingFields &&
              expectedReshardingFields->getState() >= CoordinatorStateEnum::kCommitted &&
              expectedReshardingFields->getState() != CoordinatorStateEnum::kError)) {
@@ -316,8 +319,19 @@ protected:
                 expectedReshardingFields->getDonorFields()->getReshardingKey().toBSON()),
             0);
 
-        // 'recipientFields' should only in the entry for the temporary collection.
-        ASSERT(!onDiskReshardingFields.getRecipientFields());
+        // Check the reshardingFields.recipientFields.
+        if (expectedCoordinatorState != CoordinatorStateEnum::kError) {
+            // Don't bother checking the recipientFields if the coordinator state is already kError.
+            if (expectedCoordinatorState < CoordinatorStateEnum::kCommitted) {
+                // Until CoordinatorStateEnum::kCommitted, recipientsFields only live on the
+                // temporaryNss entry in config.collections.
+                ASSERT(!onDiskReshardingFields.getRecipientFields());
+            } else {
+                // The entry for the temporaryNss has been removed, recipientFields are appended to
+                // the originalCollection's reshardingFields.
+                ASSERT(onDiskReshardingFields.getRecipientFields());
+            }
+        }
     }
 
     void readTemporaryCollectionCatalogEntryAndAssertReshardingFieldsMatchExpected(
@@ -417,9 +431,7 @@ protected:
             std::move(collectionEpoch),
             opCtx->getServiceContext()->getPreciseClockSource()->now());
         readOriginalCollectionCatalogEntryAndAssertReshardingFieldsMatchExpected(
-            opCtx,
-            originalCollType,
-            expectedCoordinatorDoc.getState() == CoordinatorStateEnum::kDone);
+            opCtx, originalCollType, expectedCoordinatorDoc);
 
         // Check the resharding fields in the config.collections entry for the temp collection. If
         // the expected state is >= kCommitted, the entry for the temp collection should have been
@@ -528,7 +540,7 @@ protected:
             _finalEpoch,
             opCtx->getServiceContext()->getPreciseClockSource()->now());
         readOriginalCollectionCatalogEntryAndAssertReshardingFieldsMatchExpected(
-            opCtx, collType, true);
+            opCtx, collType, expectedCoordinatorDoc);
     }
 
     void assertChunkVersionDidNotIncreaseAfterStateTransition(
diff --git a/src/mongo/db/s/resharding/resharding_donor_recipient_common.cpp b/src/mongo/db/s/resharding/resharding_donor_recipient_common.cpp
index e4150f9589e..55f928fd7ea 100644
--- a/src/mongo/db/s/resharding/resharding_donor_recipient_common.cpp
+++ b/src/mongo/db/s/resharding/resharding_donor_recipient_common.cpp
@@ -31,6 +31,8 @@
 
 #include "mongo/db/s/resharding/resharding_donor_recipient_common.h"
 
+#include <fmt/format.h>
+
 namespace mongo {
 namespace resharding {
 
@@ -38,6 +40,8 @@ using DonorStateMachine = ReshardingDonorService::DonorStateMachine;
 using RecipientStateMachine = ReshardingRecipientService::RecipientStateMachine;
 
 namespace {
+using namespace fmt::literals;
+
 std::vector<DonorShardMirroringEntry> createDonorShardMirroringEntriesFromDonorShardIds(
     const std::vector<ShardId>& shardIds) {
     std::vector<DonorShardMirroringEntry> donorShards(shardIds.size());
@@ -170,19 +174,42 @@ void processReshardingFieldsForCollection(OperationContext* opCtx,
                                           const NamespaceString& nss,
                                           const CollectionMetadata& metadata,
                                           const ReshardingFields& reshardingFields) {
+    auto coordinatorState = reshardingFields.getState();
+    if (coordinatorState != CoordinatorStateEnum::kError) {
+        if (coordinatorState < CoordinatorStateEnum::kCommitted) {
+            // The reshardingFields are either (1) on the originalNss because this shard is a donor
+            // or (2) temporaryNss because this shard is a recipient. Until the cordinator is in
+            // state CoordinatorStateEnum::kCommitted, reshardingFields should contain either
+            // recipientFields or donorFields, not both.
+            uassert(
+                5274201,
+                fmt::format("reshardingFields must contain either donorFields or recipientFields "
+                            "(and not both) when the "
+                            "coordinator is in state {}. Got reshardingFields {}",
+                            CoordinatorState_serializer(reshardingFields.getState()),
+                            reshardingFields.toBSON().toString()),
+                reshardingFields.getDonorFields().is_initialized() ^
+                    reshardingFields.getRecipientFields().is_initialized());
+        } else {
+            // Once the entry in config.collections for the temporaryNss is removed, recipientFields
+            // and donorFields should both exist in the reshardingFields.
+            uassert(
+                5274202,
+                fmt::format("reshardingFields must contain both donorFields and recipientFields "
+                            "when the coordinator's state is greater than or equal to "
+                            "CoordinatorStateEnum::kCommitted. Got reshardingFields {}",
+                            reshardingFields.toBSON().toString()),
+                reshardingFields.getDonorFields() && reshardingFields.getRecipientFields());
+        }
+    }
+
     if (reshardingFields.getDonorFields()) {
-        invariant(!reshardingFields.getRecipientFields());
         processReshardingFieldsForDonorCollection(opCtx, nss, metadata, reshardingFields);
-        return;
     }
 
     if (reshardingFields.getRecipientFields()) {
-        invariant(!reshardingFields.getDonorFields());
         processReshardingFieldsForRecipientCollection(opCtx, metadata, reshardingFields);
-        return;
     }
-
-    MONGO_UNREACHABLE
 }
 
 }  // namespace resharding
diff --git a/src/mongo/db/s/resharding/resharding_donor_recipient_common_test.cpp b/src/mongo/db/s/resharding/resharding_donor_recipient_common_test.cpp
index 429dca541d2..fdb30128ea3 100644
--- a/src/mongo/db/s/resharding/resharding_donor_recipient_common_test.cpp
+++ b/src/mongo/db/s/resharding/resharding_donor_recipient_common_test.cpp
@@ -351,8 +351,10 @@ TEST_F(ReshardingDonorRecipientCommonTest,
     appendRecipientFieldsToReshardingFields(
         reshardingFields, kShardIds, kExistingUUID, kOriginalNss, kFetchTimestamp);
 
-    resharding::processReshardingFieldsForCollection(
-        opCtx, kTemporaryReshardingNss, metadata, reshardingFields);
+    ASSERT_THROWS_CODE(resharding::processReshardingFieldsForCollection(
+                           opCtx, kTemporaryReshardingNss, metadata, reshardingFields),
+                       DBException,
+                       5274202);
 
     auto recipientStateMachine =
         resharding::tryGetReshardingStateMachine<ReshardingRecipientService,
@@ -363,19 +365,19 @@ TEST_F(ReshardingDonorRecipientCommonTest,
     ASSERT(recipientStateMachine == boost::none);
 }
 
-DEATH_TEST_F(ReshardingDonorRecipientCommonTest,
-             ProcessReshardingFieldsWithoutDonorOrRecipientFields,
-             "invariant") {
+TEST_F(ReshardingDonorRecipientCommonTest, ProcessReshardingFieldsWithoutDonorOrRecipientFields) {
     OperationContext* opCtx = operationContext();
     auto metadata = makeShardedMetadataForTemporaryReshardingCollection(opCtx);
 
     auto reshardingFields =
         createCommonReshardingFields(kReshardingUUID, CoordinatorStateEnum::kCloning);
 
-    resharding::processReshardingFieldsForCollection(
-        opCtx, kTemporaryReshardingNss, metadata, reshardingFields);
+    ASSERT_THROWS_CODE(resharding::processReshardingFieldsForCollection(
+                           opCtx, kTemporaryReshardingNss, metadata, reshardingFields),
+                       DBException,
+                       5274201);
 }
 
 }  // namespace
 
-}  // namespace mongo
\ No newline at end of file
+}  // namespace mongo