SERVER-46370 Maintain _receivingChunks list correctly after shard key refine

5 files changed, 104 insertions, 22 deletions

diff --git a/src/mongo/db/s/metadata_manager.cpp b/src/mongo/db/s/metadata_manager.cpp
index dc5adf82c97..73736707a67 100644
--- a/src/mongo/db/s/metadata_manager.cpp
+++ b/src/mongo/db/s/metadata_manager.cpp
@@ -57,13 +57,8 @@ using CallbackArgs = TaskExecutor::CallbackArgs;
  * Returns whether the given metadata object has a chunk owned by this shard that overlaps the
  * input range.
  */
-bool metadataOverlapsRange(const boost::optional<CollectionMetadata>& metadata,
-                           const ChunkRange& range) {
-    if (!metadata) {
-        return false;
-    }
-
-    auto metadataShardKeyPattern = KeyPattern(metadata->getKeyPattern());
+bool metadataOverlapsRange(const CollectionMetadata& metadata, const ChunkRange& range) {
+    auto metadataShardKeyPattern = KeyPattern(metadata.getKeyPattern());
 
     // If the input range is shorter than the range in the ChunkManager inside
     // 'metadata', we must extend its bounds to get a correct comparison. If the input
@@ -102,7 +97,15 @@ bool metadataOverlapsRange(const boost::optional<CollectionMetadata>& metadata,
         }
     }();
 
-    return metadata->rangeOverlapsChunk(chunkRangeToCompareToMetadata);
+    return metadata.rangeOverlapsChunk(chunkRangeToCompareToMetadata);
+}
+
+bool metadataOverlapsRange(const boost::optional<CollectionMetadata>& metadata,
+                           const ChunkRange& range) {
+    if (!metadata) {
+        return false;
+    }
+    return metadataOverlapsRange(metadata.get(), range);
 }
 
 }  // namespace
@@ -240,7 +243,7 @@ void MetadataManager::setFilteringMetadata(CollectionMetadata remoteMetadata) {
     for (auto it = _receivingChunks.begin(); it != _receivingChunks.end();) {
         const ChunkRange receivingRange(it->first, it->second);
 
-        if (!remoteMetadata.rangeOverlapsChunk(receivingRange)) {
+        if (!metadataOverlapsRange(remoteMetadata, receivingRange)) {
             ++it;
             continue;
         }
diff --git a/src/mongo/db/s/metadata_manager.h b/src/mongo/db/s/metadata_manager.h
index bd2c0189742..27049e16490 100644
--- a/src/mongo/db/s/metadata_manager.h
+++ b/src/mongo/db/s/metadata_manager.h
@@ -107,6 +107,13 @@ public:
     void toBSONPending(BSONArrayBuilder& bb) const;
 
     /**
+     * Returns the number of items in the _receivingChunks list. Useful for unit tests.
+     */
+    size_t numberOfReceivingChunks() {
+        return _receivingChunks.size();
+    }
+
+    /**
      * Appends information on all the chunk ranges in rangesToClean to builder.
      */
     void append(BSONObjBuilder* builder) const;
diff --git a/src/mongo/db/s/metadata_manager_test.cpp b/src/mongo/db/s/metadata_manager_test.cpp
index 58c663c4a4f..57947948fcf 100644
--- a/src/mongo/db/s/metadata_manager_test.cpp
+++ b/src/mongo/db/s/metadata_manager_test.cpp
@@ -73,20 +73,20 @@ protected:
     /**
      * Returns an instance of CollectionMetadata which has no chunks owned by 'thisShard'.
      */
-    static CollectionMetadata makeEmptyMetadata() {
+    static CollectionMetadata makeEmptyMetadata(
+        const KeyPattern& shardKeyPattern = kShardKeyPattern,
+        const ChunkRange& range = ChunkRange{BSON(kPattern << MINKEY), BSON(kPattern << MAXKEY)},
+        UUID uuid = UUID::gen()) {
         const OID epoch = OID::gen();
 
         auto rt = RoutingTableHistory::makeNew(
             kNss,
-            UUID::gen(),
-            kShardKeyPattern,
+            uuid,
+            shardKeyPattern,
             nullptr,
             false,
             epoch,
-            {ChunkType{kNss,
-                       ChunkRange{BSON(kPattern << MINKEY), BSON(kPattern << MAXKEY)},
-                       ChunkVersion(1, 0, epoch),
-                       kOtherShard}});
+            {ChunkType{kNss, range, ChunkVersion(1, 0, epoch), kOtherShard}});
 
         std::shared_ptr<ChunkManager> cm = std::make_shared<ChunkManager>(rt, boost::none);
 
@@ -103,11 +103,16 @@ protected:
      */
     static CollectionMetadata cloneMetadataPlusChunk(const ScopedCollectionDescription& collDesc,
                                                      const ChunkRange& range) {
+        return cloneMetadataPlusChunk(collDesc.get(), range);
+    }
+
+    static CollectionMetadata cloneMetadataPlusChunk(const CollectionMetadata& collMetadata,
+                                                     const ChunkRange& range) {
         const BSONObj& minKey = range.getMin();
         const BSONObj& maxKey = range.getMax();
-        ASSERT(!rangeMapOverlaps(collDesc->getChunks(), minKey, maxKey));
+        ASSERT(!rangeMapOverlaps(collMetadata.getChunks(), minKey, maxKey));
 
-        auto cm = collDesc->getChunkManager();
+        auto cm = collMetadata.getChunkManager();
 
         const auto chunkToSplit = cm->findIntersectingChunkWithSimpleCollation(minKey);
         ASSERT_BSONOBJ_GTE(minKey, chunkToSplit.getMin());
@@ -179,6 +184,63 @@ TEST_F(MetadataManagerTest, CleanUpForMigrateIn) {
     ASSERT_EQ(0UL, _manager->numberOfRangesToCleanStillInUse());
 }
 
+TEST_F(MetadataManagerTest,
+       ChunkInReceivingChunksListIsRemovedAfterShardKeyRefineIfMigrationSucceeded) {
+    _manager->setFilteringMetadata(makeEmptyMetadata());
+
+    // Simulate receiving a range. This will add an item to _receivingChunks.
+    ChunkRange range(BSON("key" << 0), BSON("key" << 10));
+    auto notif1 = _manager->beginReceive(range);
+
+    ASSERT_EQ(_manager->numberOfReceivingChunks(), 1);
+
+    // Simulate a situation in which the migration completes, and then the shard key is refined,
+    // before this shard discovers the updated metadata.
+    auto uuid = _manager->getActiveMetadata(boost::none)->getChunkManager()->getUUID().get();
+    ChunkRange refinedRange(BSON("key" << 0 << "other" << MINKEY),
+                            BSON("key" << 10 << "other" << MINKEY));
+    auto refinedMetadata = makeEmptyMetadata(BSON(kPattern << 1 << "other" << 1),
+                                             ChunkRange(BSON("key" << MINKEY << "other" << MINKEY),
+                                                        BSON("key" << MAXKEY << "other" << MAXKEY)),
+                                             uuid);
+
+    // Set the updated chunk map on the MetadataManager.
+    _manager->setFilteringMetadata(cloneMetadataPlusChunk(refinedMetadata, refinedRange));
+    // Because the refined range overlaps with the received range (pre-refine), this should remove
+    // the item in _receivingChunks.
+    ASSERT_EQ(_manager->numberOfReceivingChunks(), 0);
+}
+
+TEST_F(MetadataManagerTest,
+       ChunkInReceivingChunksListIsNotRemovedAfterShardKeyRefineIfNonOverlappingRangeIsReceived) {
+    _manager->setFilteringMetadata(makeEmptyMetadata());
+
+    // Simulate receiving a range. This will add an item to _receivingChunks.
+    ChunkRange range(BSON("key" << 0), BSON("key" << 10));
+    auto notif1 = _manager->beginReceive(range);
+    ASSERT_EQ(_manager->numberOfReceivingChunks(), 1);
+
+    // Simulate a situation in which the shard key is refined and this shard discovers
+    // updated metadata where it owns some range that does not overlap with the range being migrated
+    // in.
+    auto uuid = _manager->getActiveMetadata(boost::none)->getChunkManager()->getUUID().get();
+    ChunkRange refinedNonOverlappingRange(BSON("key" << -10 << "other" << MINKEY),
+                                          BSON("key" << 0 << "other" << MINKEY));
+
+    auto refinedMetadata = makeEmptyMetadata(BSON(kPattern << 1 << "other" << 1),
+                                             ChunkRange(BSON("key" << MINKEY << "other" << MINKEY),
+                                                        BSON("key" << MAXKEY << "other" << MAXKEY)),
+                                             uuid);
+
+    // Set the updated chunk map on the MetadataManager.
+    _manager->setFilteringMetadata(
+        cloneMetadataPlusChunk(refinedMetadata, refinedNonOverlappingRange));
+
+    // Because the refined range does not overlap with the received range (pre-refine), this should
+    // NOT remove the item in _receivingChunks.
+    ASSERT_EQ(_manager->numberOfReceivingChunks(), 1);
+}
+
 TEST_F(MetadataManagerTest, TrackOrphanedDataCleanupBlocksOnScheduledRangeDeletions) {
     ChunkRange cr1(BSON("key" << 0), BSON("key" << 10));
 
diff --git a/src/mongo/db/s/migration_destination_manager.cpp b/src/mongo/db/s/migration_destination_manager.cpp
index bea8ac3aa30..19de4193563 100644
--- a/src/mongo/db/s/migration_destination_manager.cpp
+++ b/src/mongo/db/s/migration_destination_manager.cpp
@@ -809,6 +809,7 @@ void MigrationDestinationManager::_migrateThread() {
 
     stdx::lock_guard<Latch> lk(_mutex);
     _sessionId.reset();
+    _collUuid.reset();
     _scopedReceiveChunk.reset();
     _isActiveCV.notify_all();
 }
@@ -892,6 +893,9 @@ void MigrationDestinationManager::_migrateDriver(OperationContext* outerOpCtx) {
             // Synchronously delete any data which might have been left orphaned in the range
             // being moved, and wait for completion
 
+            // Needed for _forgetPending to make sure the collection has the same UUID at the end of
+            // an aborted migration as at the beginning. Must be set before calling _notePending.
+            _collUuid = donorCollectionOptionsAndIndexes.uuid;
             auto cleanupCompleteFuture = _notePending(outerOpCtx, range);
             auto cleanupStatus = cleanupCompleteFuture.getNoThrow(outerOpCtx);
             // Wait for the range deletion to report back. Swallow
@@ -1356,8 +1360,9 @@ SharedSemiFuture<void> MigrationDestinationManager::_notePending(OperationContex
     auto* const css = CollectionShardingRuntime::get(opCtx, _nss);
     const auto optMetadata = css->getCurrentMetadataIfKnown();
 
-    // This can currently happen because drops aren't synchronized with in-migrations. The idea for
-    // checking this here is that in the future we shouldn't have this problem.
+    // This can currently happen because drops and shard key refine operations aren't guaranteed to
+    // be synchronized with in-migrations. The idea for checking this here is that in the future we
+    // shouldn't have this problem.
     if (!optMetadata || !(*optMetadata)->isSharded() ||
         (*optMetadata)->getCollVersion().epoch() != _epoch) {
         return Status{ErrorCodes::StaleShardVersion,
@@ -1388,10 +1393,14 @@ void MigrationDestinationManager::_forgetPending(OperationContext* opCtx, ChunkR
 
     // This can currently happen because drops aren't synchronized with in-migrations. The idea for
     // checking this here is that in the future we shouldn't have this problem.
+    //
+    // _collUuid will always be set if _notePending was called, so if it is not set, there is no
+    // need to do anything. If it is set, we use it to ensure that the collection UUID has not
+    // changed since the beginning of migration.
     if (!optMetadata || !(*optMetadata)->isSharded() ||
-        (*optMetadata)->getCollVersion().epoch() != _epoch) {
+        (_collUuid && !(*optMetadata)->uuidMatches(*_collUuid))) {
         LOGV2(22009,
-              "No need to forget pending chunk {range} because the epoch for {nss_ns} changed",
+              "No need to forget pending chunk {range} because the uuid for {nss_ns} changed",
               "range"_attr = redact(range.toString()),
               "nss_ns"_attr = _nss.ns());
         return;
diff --git a/src/mongo/db/s/migration_destination_manager.h b/src/mongo/db/s/migration_destination_manager.h
index ea322655509..403800427fd 100644
--- a/src/mongo/db/s/migration_destination_manager.h
+++ b/src/mongo/db/s/migration_destination_manager.h
@@ -210,6 +210,7 @@ private:
     LogicalSessionId _lsid;
     TxnNumber _txnNumber;
     NamespaceString _nss;
+    boost::optional<UUID> _collUuid;
     ConnectionString _fromShardConnString;
     ShardId _fromShard;
     ShardId _toShard;