SERVER-48198 Account for extended range bounds when recovering migration decision

3 files changed, 55 insertions, 40 deletions

diff --git a/src/mongo/db/s/metadata_manager.cpp b/src/mongo/db/s/metadata_manager.cpp
index b0dadd35e17..953098ac44f 100644
--- a/src/mongo/db/s/metadata_manager.cpp
+++ b/src/mongo/db/s/metadata_manager.cpp
@@ -38,6 +38,7 @@
 #include "mongo/bson/util/builder.h"
 #include "mongo/db/query/internal_plans.h"
 #include "mongo/db/range_arithmetic.h"
+#include "mongo/db/s/migration_util.h"
 #include "mongo/db/s/range_deletion_util.h"
 #include "mongo/db/s/sharding_runtime_d_params_gen.h"
 #include "mongo/logv2/log.h"
@@ -56,45 +57,8 @@ using CallbackArgs = TaskExecutor::CallbackArgs;
  * input range.
  */
 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
-    // range is longer than the range in the ChunkManager, we likewise must shorten it.
-    // We make sure to match what's in the ChunkManager instead of the other way around,
-    // since the ChunkManager only stores ranges and compares overlaps using a string version of the
-    // key, rather than a BSONObj. This logic is necessary because the _metadata list can
-    // contain ChunkManagers with different shard keys if the shard key has been refined.
-    //
-    // Note that it's safe to use BSONObj::nFields() (which returns the number of top level
-    // fields in the BSONObj) to compare the two, since shard key refine operations can only add
-    // top-level fields.
-    //
-    // Using extractFieldsUndotted to shorten the input range is correct because the ChunkRange and
-    // the shard key pattern will both already store nested shard key fields as top-level dotted
-    // fields, and extractFieldsUndotted uses the top-level fields verbatim rather than treating
-    // dots as accessors for subfields.
-    auto chunkRangeToCompareToMetadata = [&] {
-        auto metadataShardKeyPatternBson = metadataShardKeyPattern.toBSON();
-        auto numFieldsInMetadataShardKey = metadataShardKeyPatternBson.nFields();
-        auto numFieldsInInputRangeShardKey = range.getMin().nFields();
-        if (numFieldsInInputRangeShardKey < numFieldsInMetadataShardKey) {
-            auto extendedRangeMin = metadataShardKeyPattern.extendRangeBound(
-                range.getMin(), false /* makeUpperInclusive */);
-            auto extendedRangeMax = metadataShardKeyPattern.extendRangeBound(
-                range.getMax(), false /* makeUpperInclusive */);
-            return ChunkRange(extendedRangeMin, extendedRangeMax);
-        } else if (numFieldsInInputRangeShardKey > numFieldsInMetadataShardKey) {
-            auto shortenedRangeMin =
-                range.getMin().extractFieldsUndotted(metadataShardKeyPatternBson);
-            auto shortenedRangeMax =
-                range.getMax().extractFieldsUndotted(metadataShardKeyPatternBson);
-            return ChunkRange(shortenedRangeMin, shortenedRangeMax);
-        } else {
-            return range;
-        }
-    }();
-
+    auto chunkRangeToCompareToMetadata =
+        migrationutil::extendOrTruncateBoundsForMetadata(metadata, range);
     return metadata.rangeOverlapsChunk(chunkRangeToCompareToMetadata);
 }
 
diff --git a/src/mongo/db/s/migration_util.cpp b/src/mongo/db/s/migration_util.cpp
index ce7a20dea64..fa8c0f89365 100644
--- a/src/mongo/db/s/migration_util.cpp
+++ b/src/mongo/db/s/migration_util.cpp
@@ -49,6 +49,7 @@
 #include "mongo/db/repl/repl_client_info.h"
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/db/s/active_migrations_registry.h"
+#include "mongo/db/s/collection_metadata.h"
 #include "mongo/db/s/collection_sharding_runtime.h"
 #include "mongo/db/s/migration_coordinator.h"
 #include "mongo/db/s/shard_filtering_metadata_refresh.h"
@@ -219,6 +220,44 @@ BSONObj makeMigrationStatusDocument(const NamespaceString& nss,
     return builder.obj();
 }
 
+ChunkRange extendOrTruncateBoundsForMetadata(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
+    // range is longer than the range in the ChunkManager, we likewise must shorten it.
+    // We make sure to match what's in the ChunkManager instead of the other way around,
+    // since the ChunkManager only stores ranges and compares overlaps using a string version of the
+    // key, rather than a BSONObj. This logic is necessary because the _metadata list can
+    // contain ChunkManagers with different shard keys if the shard key has been refined.
+    //
+    // Note that it's safe to use BSONObj::nFields() (which returns the number of top level
+    // fields in the BSONObj) to compare the two, since shard key refine operations can only add
+    // top-level fields.
+    //
+    // Using extractFieldsUndotted to shorten the input range is correct because the ChunkRange and
+    // the shard key pattern will both already store nested shard key fields as top-level dotted
+    // fields, and extractFieldsUndotted uses the top-level fields verbatim rather than treating
+    // dots as accessors for subfields.
+    auto metadataShardKeyPatternBson = metadataShardKeyPattern.toBSON();
+    auto numFieldsInMetadataShardKey = metadataShardKeyPatternBson.nFields();
+    auto numFieldsInInputRangeShardKey = range.getMin().nFields();
+    if (numFieldsInInputRangeShardKey < numFieldsInMetadataShardKey) {
+        auto extendedRangeMin = metadataShardKeyPattern.extendRangeBound(
+            range.getMin(), false /* makeUpperInclusive */);
+        auto extendedRangeMax = metadataShardKeyPattern.extendRangeBound(
+            range.getMax(), false /* makeUpperInclusive */);
+        return ChunkRange(extendedRangeMin, extendedRangeMax);
+    } else if (numFieldsInInputRangeShardKey > numFieldsInMetadataShardKey) {
+        auto shortenedRangeMin = range.getMin().extractFieldsUndotted(metadataShardKeyPatternBson);
+        auto shortenedRangeMax = range.getMax().extractFieldsUndotted(metadataShardKeyPatternBson);
+        return ChunkRange(shortenedRangeMin, shortenedRangeMax);
+    } else {
+        return range;
+    }
+}
+
 Query overlappingRangeQuery(const ChunkRange& range, const UUID& uuid) {
     return QUERY(RangeDeletionTask::kCollectionUuidFieldName
                  << uuid << RangeDeletionTask::kRangeFieldName + "." + ChunkRange::kMinKey << LT
@@ -897,7 +936,11 @@ void resumeMigrationCoordinationsOnStepUp(OperationContext* opCtx) {
                         return true;
                     }
 
-                    if (refreshedMetadata->keyBelongsToMe(doc.getRange().getMin())) {
+                    // Note this should only extend the range boundaries (if there has been a shard
+                    // key refine since the migration began) and never truncate them.
+                    auto chunkRangeToCompareToMetadata =
+                        extendOrTruncateBoundsForMetadata(*refreshedMetadata, doc.getRange());
+                    if (refreshedMetadata->keyBelongsToMe(chunkRangeToCompareToMetadata.getMin())) {
                         coordinator.setMigrationDecision(MigrationCoordinator::Decision::kAborted);
                     } else {
                         coordinator.setMigrationDecision(
diff --git a/src/mongo/db/s/migration_util.h b/src/mongo/db/s/migration_util.h
index dbc51495e56..9754f139e7d 100644
--- a/src/mongo/db/s/migration_util.h
+++ b/src/mongo/db/s/migration_util.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include "mongo/db/repl/optime.h"
+#include "mongo/db/s/collection_metadata.h"
 #include "mongo/db/s/migration_coordinator_document_gen.h"
 #include "mongo/db/s/persistent_task_store.h"
 #include "mongo/db/s/range_deletion_task_gen.h"
@@ -66,6 +67,13 @@ BSONObj makeMigrationStatusDocument(const NamespaceString& nss,
                                     const BSONObj& max);
 
 /**
+ * Returns a chunk range with extended or truncated boundaries to match the number of fields in the
+ * given metadata's shard key pattern.
+ */
+ChunkRange extendOrTruncateBoundsForMetadata(const CollectionMetadata& metadata,
+                                             const ChunkRange& range);
+
+/**
  * Returns an executor to be used to run commands related to submitting tasks to the range deleter.
  * The executor is initialized on the first call to this function. Uses a shared_ptr
  * because a shared_ptr is required to work with ExecutorFutures.