SERVER-43848 find/update/delete w/o shard key predicate under txn with snapshot read can miss documents

(cherry picked from commit 305bbe0ed709ffc88916093cfbd716bfb8fea60b)

2 files changed, 43 insertions, 4 deletions

diff --git a/src/mongo/s/chunk_manager.cpp b/src/mongo/s/chunk_manager.cpp
index e7ea2996dd9..1b61b34be59 100644
--- a/src/mongo/s/chunk_manager.cpp
+++ b/src/mongo/s/chunk_manager.cpp
@@ -173,8 +173,12 @@ void ChunkManager::getShardIdsForQuery(OperationContext* opCtx,
     for (BoundList::const_iterator it = ranges.begin(); it != ranges.end(); ++it) {
         getShardIdsForRange(it->first /*min*/, it->second /*max*/, shardIds);
 
-        // once we know we need to visit all shards no need to keep looping
-        if (shardIds->size() == _rt->_shardVersions.size()) {
+        // Once we know we need to visit all shards no need to keep looping.
+        // 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 && shardIds->size() == _rt->_shardVersions.size()) {
             break;
         }
     }
@@ -195,8 +199,11 @@ void ChunkManager::getShardIdsForRange(const BSONObj& min,
         shardIds->insert(it->second->getShardIdAt(_clusterTime));
 
         // No need to iterate through the rest of the ranges, because we already know we need to use
-        // all shards.
-        if (shardIds->size() == _rt->_shardVersions.size()) {
+        // all shards. 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 && shardIds->size() == _rt->_shardVersions.size()) {
             break;
         }
     }
diff --git a/src/mongo/s/chunk_manager_query_test.cpp b/src/mongo/s/chunk_manager_query_test.cpp
index 3513a3401b4..5aa8b8808d3 100644
--- a/src/mongo/s/chunk_manager_query_test.cpp
+++ b/src/mongo/s/chunk_manager_query_test.cpp
@@ -488,5 +488,37 @@ TEST_F(ChunkManagerQueryTest, SimpleCollationNumbersMultiShard) {
         {ShardId("0")});
 }
 
+TEST_F(ChunkManagerQueryTest, SnapshotQueryWithMoreShardsThanLatestMetadata) {
+    const auto epoch = OID::gen();
+    ChunkVersion version(1, 0, epoch);
+
+    ChunkType chunk0(kNss, {BSON("x" << MINKEY), BSON("x" << 0)}, version, ShardId("0"));
+
+    version.incMajor();
+    ChunkType chunk1(kNss, {BSON("x" << 0), BSON("x" << MAXKEY)}, version, ShardId("1"));
+
+    auto oldRoutingTable = RoutingTableHistory::makeNew(
+        kNss, boost::none, BSON("x" << 1), nullptr, false, epoch, {chunk0, chunk1});
+
+    // Simulate move chunk {x: 0} to shard 0. Effectively moving all remaining chunks to shard 0.
+    version.incMajor();
+    chunk1.setVersion(version);
+    chunk1.setShard(chunk0.getShard());
+    chunk1.setHistory({ChunkHistory(Timestamp(20, 0), ShardId("0")),
+                       ChunkHistory(Timestamp(1, 0), ShardId("1"))});
+
+    auto newRoutingTable = oldRoutingTable->makeUpdated({chunk1});
+    ChunkManager chunkManager(newRoutingTable, Timestamp(5, 0));
+
+    std::set<ShardId> shardIds;
+    chunkManager.getShardIdsForRange(BSON("x" << MINKEY), BSON("x" << MAXKEY), &shardIds);
+    ASSERT_EQ(2u, shardIds.size());
+
+    shardIds.clear();
+    chunkManager.getShardIdsForQuery(
+        operationContext(), BSON("x" << BSON("$gt" << -20)), {}, &shardIds);
+    ASSERT_EQ(2u, shardIds.size());
+}
+
 }  // namespace
 }  // namespace mongo