SERVER-44830 Modify StaleConfigException handlers to handle a single shard invalidation

6 files changed, 39 insertions, 9 deletions

diff --git a/src/mongo/s/commands/cluster_drop_cmd.cpp b/src/mongo/s/commands/cluster_drop_cmd.cpp
index 6bc58c2d0e8..3effbb64055 100644
--- a/src/mongo/s/commands/cluster_drop_cmd.cpp
+++ b/src/mongo/s/commands/cluster_drop_cmd.cpp
@@ -76,8 +76,7 @@ public:
         // Invalidate the routing table cache entry for this collection so that we reload it the
         // next time it is accessed, even if sending the command to the config server fails due
         // to e.g. a NetworkError.
-        ON_BLOCK_EXIT(
-            [opCtx, nss] { Grid::get(opCtx)->catalogCache()->invalidateShardedCollection(nss); });
+        ON_BLOCK_EXIT([opCtx, nss] { Grid::get(opCtx)->catalogCache()->onEpochChange(nss); });
 
         auto configShard = Grid::get(opCtx)->shardRegistry()->getConfigShard();
         auto cmdResponse = uassertStatusOK(configShard->runCommandWithFixedRetryAttempts(
diff --git a/src/mongo/s/commands/cluster_merge_chunks_cmd.cpp b/src/mongo/s/commands/cluster_merge_chunks_cmd.cpp
index c686efce22e..26b9435f91a 100644
--- a/src/mongo/s/commands/cluster_merge_chunks_cmd.cpp
+++ b/src/mongo/s/commands/cluster_merge_chunks_cmd.cpp
@@ -175,7 +175,8 @@ public:
             Shard::RetryPolicy::kNotIdempotent));
         uassertStatusOK(response.commandStatus);
 
-        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo));
+        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo),
+                                                              firstChunk.getShardId());
         CommandHelpers::filterCommandReplyForPassthrough(response.response, &result);
 
         return true;
diff --git a/src/mongo/s/commands/cluster_move_chunk_cmd.cpp b/src/mongo/s/commands/cluster_move_chunk_cmd.cpp
index 434ba36fb28..2336790882d 100644
--- a/src/mongo/s/commands/cluster_move_chunk_cmd.cpp
+++ b/src/mongo/s/commands/cluster_move_chunk_cmd.cpp
@@ -197,7 +197,9 @@ public:
                                                         cmdObj["waitForDelete"].trueValue(),
                                                     forceJumbo));
 
-        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo));
+        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo),
+                                                              chunk->getShardId());
+        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo), to->getId());
 
         result.append("millis", t.millis());
         return true;
diff --git a/src/mongo/s/commands/cluster_shard_collection_cmd.cpp b/src/mongo/s/commands/cluster_shard_collection_cmd.cpp
index d5c23dcacc7..502fdefde1c 100644
--- a/src/mongo/s/commands/cluster_shard_collection_cmd.cpp
+++ b/src/mongo/s/commands/cluster_shard_collection_cmd.cpp
@@ -120,8 +120,7 @@ public:
 
         // Invalidate the routing table cache entry for this collection so that we reload the
         // collection the next time it's accessed, even if we receive a failure, e.g. NetworkError.
-        ON_BLOCK_EXIT(
-            [opCtx, nss] { Grid::get(opCtx)->catalogCache()->invalidateShardedCollection(nss); });
+        ON_BLOCK_EXIT([opCtx, nss] { Grid::get(opCtx)->catalogCache()->onEpochChange(nss); });
 
         auto configShard = Grid::get(opCtx)->shardRegistry()->getConfigShard();
         auto cmdResponse = uassertStatusOK(configShard->runCommandWithFixedRetryAttempts(
diff --git a/src/mongo/s/commands/cluster_split_cmd.cpp b/src/mongo/s/commands/cluster_split_cmd.cpp
index 114d0f27bba..8c7467a75ec 100644
--- a/src/mongo/s/commands/cluster_split_cmd.cpp
+++ b/src/mongo/s/commands/cluster_split_cmd.cpp
@@ -270,7 +270,8 @@ public:
 
         // This invalidation is only necessary so that auto-split can begin to track statistics for
         // the chunks produced after the split instead of the single original chunk.
-        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo));
+        Grid::get(opCtx)->catalogCache()->onStaleShardVersion(std::move(routingInfo),
+                                                              chunk->getShardId());
 
         return true;
     }
diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index cca60108d5d..c50fd22055f 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -582,7 +582,21 @@ void runCommand(OperationContext* opCtx,
                     ShardConnection::checkMyConnectionVersions(opCtx, staleNs.ns());
                 }
 
-                Grid::get(opCtx)->catalogCache()->invalidateShardedCollection(staleNs);
+                if (auto staleInfo = ex.extraInfo<StaleConfigInfo>()) {
+                    Grid::get(opCtx)
+                        ->catalogCache()
+                        ->invalidateShardOrEntireCollectionEntryForShardedCollection(
+                            opCtx,
+                            staleNs,
+                            staleInfo->getVersionWanted(),
+                            staleInfo->getVersionReceived(),
+                            staleInfo->getShardId());
+                } else {
+                    // If we don't have the stale config info and therefore don't know the shard's
+                    // id, we have to force all further targetting requests for the namespace to
+                    // block on a refresh.
+                    Grid::get(opCtx)->catalogCache()->onEpochChange(staleNs);
+                }
 
                 // Retry logic specific to transactions. Throws and aborts the transaction if the
                 // error cannot be retried on.
@@ -1097,7 +1111,21 @@ void Strategy::explainFind(OperationContext* opCtx,
                 ShardConnection::checkMyConnectionVersions(opCtx, staleNs.ns());
             }
 
-            Grid::get(opCtx)->catalogCache()->invalidateShardedCollection(staleNs);
+            if (auto staleInfo = ex.extraInfo<StaleConfigInfo>()) {
+                Grid::get(opCtx)
+                    ->catalogCache()
+                    ->invalidateShardOrEntireCollectionEntryForShardedCollection(
+                        opCtx,
+                        staleNs,
+                        staleInfo->getVersionWanted(),
+                        staleInfo->getVersionReceived(),
+                        staleInfo->getShardId());
+            } else {
+                // If we don't have the stale config info and therefore don't know the shard's id,
+                // we have to force all further targetting requests for the namespace to block on
+                // a refresh.
+                Grid::get(opCtx)->catalogCache()->onEpochChange(staleNs);
+            }
 
             if (canRetry) {
                 continue;