SERVER-46386 Change MetadataManager::_findNewestOverlappingMetadata to work with refined shard keys

1 files changed, 55 insertions, 3 deletions

diff --git a/src/mongo/db/s/metadata_manager.cpp b/src/mongo/db/s/metadata_manager.cpp
index f17fd52d8e0..dc5adf82c97 100644
--- a/src/mongo/db/s/metadata_manager.cpp
+++ b/src/mongo/db/s/metadata_manager.cpp
@@ -52,6 +52,59 @@ namespace mongo {
 namespace {
 using TaskExecutor = executor::TaskExecutor;
 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());
+
+    // 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;
+        }
+    }();
+
+    return metadata->rangeOverlapsChunk(chunkRangeToCompareToMetadata);
+}
+
 }  // namespace
 
 class RangePreserver : public ScopedCollectionDescription::Impl {
@@ -427,15 +480,14 @@ auto MetadataManager::_findNewestOverlappingMetadata(WithLock, ChunkRange const&
     invariant(!_metadata.empty());
 
     auto it = _metadata.rbegin();
-    if ((*it)->metadata && (*it)->metadata->rangeOverlapsChunk(range)) {
+    if (metadataOverlapsRange((*it)->metadata, range)) {
         return (*it).get();
     }
 
     ++it;
     for (; it != _metadata.rend(); ++it) {
         auto& tracker = *it;
-        if (tracker->usageCounter && tracker->metadata &&
-            tracker->metadata->rangeOverlapsChunk(range)) {
+        if (tracker->usageCounter && metadataOverlapsRange(tracker->metadata, range)) {
             return tracker.get();
         }
     }