SERVER-46487 Add fast path to getShardIdsForRange() when chunk range is [MinKey, MaxKey]

(cherry picked from commit d9d92a15561dc84223d302475e09d0fa91309550)

2 files changed, 46 insertions, 5 deletions

diff --git a/src/mongo/s/chunk_manager.cpp b/src/mongo/s/chunk_manager.cpp
index f0b7042b398..37e591e7d41 100644
--- a/src/mongo/s/chunk_manager.cpp
+++ b/src/mongo/s/chunk_manager.cpp
@@ -50,12 +50,19 @@ namespace {
 // Used to generate sequence numbers to assign to each newly created RoutingTableHistory
 AtomicUInt32 nextCMSequenceNumber(0);
 
-void checkAllElementsAreOfType(BSONType type, const BSONObj& o) {
-    for (auto&& element : o) {
-        uassert(ErrorCodes::ConflictingOperationInProgress,
-                str::stream() << "Not all elements of " << o << " are of type " << typeName(type),
-                element.type() == type);
+bool allElementsAreOfType(BSONType type, const BSONObj& obj) {
+    for (auto&& elem : obj) {
+        if (elem.type() != type) {
+            return false;
+        }
     }
+    return true;
+}
+
+void checkAllElementsAreOfType(BSONType type, const BSONObj& o) {
+    uassert(ErrorCodes::ConflictingOperationInProgress,
+            str::stream() << "Not all elements of " << o << " are of type " << typeName(type),
+            allElementsAreOfType(type, o));
 }
 
 std::string extractKeyStringInternal(const BSONObj& shardKeyValue, Ordering ordering) {
@@ -193,6 +200,16 @@ void ChunkManager::getShardIdsForQuery(OperationContext* opCtx,
 void ChunkManager::getShardIdsForRange(const BSONObj& min,
                                        const BSONObj& max,
                                        std::set<ShardId>* shardIds) const {
+    // If our range is [MinKey, MaxKey], we can simply return all shard ids right away. However,
+    // this optimization does not apply when we are reading from a snapshot because _shardVersions
+    // contains shards with chunks and is built based on the last refresh. Therefore, it is
+    // possible for _shardVersions to have fewer entries if a shard no longer owns chunks when it
+    // used to at _clusterTime.
+    if (!_clusterTime && allElementsAreOfType(MinKey, min) && allElementsAreOfType(MaxKey, max)) {
+        getAllShardIds(shardIds);
+        return;
+    }
+
     const auto bounds = _rt->overlappingRanges(min, max, true);
     for (auto it = bounds.first; it != bounds.second; ++it) {
         shardIds->insert(it->second->getShardIdAt(_clusterTime));
diff --git a/src/mongo/s/chunk_manager_refresh_bm.cpp b/src/mongo/s/chunk_manager_refresh_bm.cpp
index 62ea6f33e5f..89f630e91a8 100644
--- a/src/mongo/s/chunk_manager_refresh_bm.cpp
+++ b/src/mongo/s/chunk_manager_refresh_bm.cpp
@@ -269,6 +269,24 @@ void BM_GetShardIdsForRange(benchmark::State& state,
 }
 
 template <typename CollectionMetadataBuilderFn>
+void BM_GetShardIdsForRangeMinKeyToMaxKey(benchmark::State& state,
+                                          CollectionMetadataBuilderFn makeCollectionMetadata) {
+    const int nShards = state.range(0);
+    const int nChunks = state.range(1);
+
+    auto cm = makeCollectionMetadata(nShards, nChunks);
+    auto min = BSON("_id" << MINKEY);
+    auto max = BSON("_id" << MAXKEY);
+
+    for (auto keepRunning : state) {
+        std::set<ShardId> shardIds;
+        cm->getChunkManager()->getShardIdsForRange(min, max, &shardIds);
+    }
+
+    state.SetItemsProcessed(state.iterations());
+}
+
+template <typename CollectionMetadataBuilderFn>
 void BM_KeyBelongsToMe(benchmark::State& state,
                        CollectionMetadataBuilderFn makeCollectionMetadata) {
     const int nShards = state.range(0);
@@ -333,6 +351,12 @@ MONGO_INITIALIZER(RegisterBenchmarks)(InitializerContext* context) {
             BM_GetShardIdsForRange, Pessimal, makeChunkManagerWithPessimalBalancedDistribution),
         REGISTER_BENCHMARK_CAPTURE(
             BM_GetShardIdsForRange, Optimal, makeChunkManagerWithOptimalBalancedDistribution),
+        REGISTER_BENCHMARK_CAPTURE(BM_GetShardIdsForRangeMinKeyToMaxKey,
+                                   Pessimal,
+                                   makeChunkManagerWithPessimalBalancedDistribution),
+        REGISTER_BENCHMARK_CAPTURE(BM_GetShardIdsForRangeMinKeyToMaxKey,
+                                   Optimal,
+                                   makeChunkManagerWithOptimalBalancedDistribution),
         REGISTER_BENCHMARK_CAPTURE(
             BM_KeyBelongsToMe, Pessimal, makeChunkManagerWithPessimalBalancedDistribution),
         REGISTER_BENCHMARK_CAPTURE(