SERVER-47532 ShardServerProcessInterface: Convert usage of getCollectionDescription_DEPRECATED to getCollectionDescription

17 files changed, 137 insertions, 83 deletions

diff --git a/src/mongo/db/index_builds_coordinator_mongod.cpp b/src/mongo/db/index_builds_coordinator_mongod.cpp
index f028d2e2dc2..2a2fd13169e 100644
--- a/src/mongo/db/index_builds_coordinator_mongod.cpp
+++ b/src/mongo/db/index_builds_coordinator_mongod.cpp
@@ -199,7 +199,7 @@ IndexBuildsCoordinatorMongod::startIndexBuild(OperationContext* opCtx,
 
     const auto nss = CollectionCatalog::get(opCtx).resolveNamespaceStringOrUUID(opCtx, nssOrUuid);
 
-    const auto& oss = OperationShardingState::get(opCtx);
+    auto& oss = OperationShardingState::get(opCtx);
     const auto shardVersion = oss.getShardVersion(nss);
     const auto dbVersion = oss.getDbVersion(dbName);
 
diff --git a/src/mongo/db/pipeline/document_source_graph_lookup.cpp b/src/mongo/db/pipeline/document_source_graph_lookup.cpp
index a84d1773a1f..5f6214400df 100644
--- a/src/mongo/db/pipeline/document_source_graph_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_graph_lookup.cpp
@@ -48,14 +48,8 @@
 namespace mongo {
 
 namespace {
-void assertIsValidCollectionState(const boost::intrusive_ptr<ExpressionContext>& expCtx) {
-    if (expCtx->mongoProcessInterface->isSharded(expCtx->opCtx, expCtx->ns)) {
-        const bool foreignShardedAllowed =
-            getTestCommandsEnabled() && internalQueryAllowShardedLookup.load();
-        if (!foreignShardedAllowed) {
-            uasserted(31428, "Cannot run $graphLookup with sharded foreign collection");
-        }
-    }
+bool foreignShardedLookupAllowed() {
+    return getTestCommandsEnabled() && internalQueryAllowShardedLookup.load();
 }
 }  // namespace
 
@@ -187,8 +181,12 @@ void DocumentSourceGraphLookUp::doDispose() {
 void DocumentSourceGraphLookUp::doBreadthFirstSearch() {
     long long depth = 0;
     bool shouldPerformAnotherQuery;
-    assertIsValidCollectionState(_fromExpCtx);
     do {
+        if (!foreignShardedLookupAllowed()) {
+            // Enforce that the foreign collection must be unsharded for $graphLookup.
+            _fromExpCtx->mongoProcessInterface->setExpectedShardVersion(
+                _fromExpCtx->opCtx, _fromExpCtx->ns, ChunkVersion::UNSHARDED());
+        }
         shouldPerformAnotherQuery = false;
 
         // Check whether each key in the frontier exists in the cache or needs to be queried.
@@ -359,7 +357,19 @@ void DocumentSourceGraphLookUp::performSearch() {
         _frontierUsageBytes += startingValue.getApproximateSize();
     }
 
-    doBreadthFirstSearch();
+    try {
+        doBreadthFirstSearch();
+    } catch (const ExceptionForCat<ErrorCategory::StaleShardVersionError>& ex) {
+        // If lookup on a sharded collection is disallowed and the foreign collection is sharded,
+        // throw a custom exception.
+        if (auto staleInfo = ex.extraInfo<StaleConfigInfo>()) {
+            uassert(31428,
+                    "Cannot run $graphLookup with sharded foreign collection",
+                    foreignShardedLookupAllowed() || !staleInfo->getVersionWanted() ||
+                        staleInfo->getVersionWanted() == ChunkVersion::UNSHARDED());
+        }
+        throw;
+    }
 }
 
 DocumentSource::GetModPathsReturn DocumentSourceGraphLookUp::getModifiedPaths() const {
diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index af45df02f2f..95c33761454 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -234,16 +234,6 @@ BSONObj buildEqualityOrQuery(const std::string& fieldName, const BSONArray& valu
     return orBuilder.obj();
 }
 
-void assertIsValidCollectionState(const boost::intrusive_ptr<ExpressionContext>& expCtx) {
-    if (expCtx->mongoProcessInterface->isSharded(expCtx->opCtx, expCtx->ns)) {
-        const bool foreignShardedAllowed =
-            getTestCommandsEnabled() && internalQueryAllowShardedLookup.load();
-        if (!foreignShardedAllowed) {
-            uasserted(51069, "Cannot run $lookup with sharded foreign collection");
-        }
-    }
-}
-
 void lookupPipeValidator(const Pipeline& pipeline) {
     const auto& sources = pipeline.getSources();
     std::for_each(sources.begin(), sources.end(), [](auto& src) {
@@ -253,6 +243,10 @@ void lookupPipeValidator(const Pipeline& pipeline) {
                 src->constraints().isAllowedInLookupPipeline());
     });
 }
+
+bool foreignShardedLookupAllowed() {
+    return getTestCommandsEnabled() && internalQueryAllowShardedLookup.load();
+}
 }  // namespace
 
 DocumentSource::GetNextResult DocumentSourceLookUp::doGetNext() {
@@ -278,7 +272,20 @@ DocumentSource::GetNextResult DocumentSourceLookUp::doGetNext() {
         _resolvedPipeline.back() = matchStage;
     }
 
-    auto pipeline = buildPipeline(inputDoc);
+    std::unique_ptr<Pipeline, PipelineDeleter> pipeline;
+    try {
+        pipeline = buildPipeline(inputDoc);
+    } catch (const ExceptionForCat<ErrorCategory::StaleShardVersionError>& ex) {
+        // If lookup on a sharded collection is disallowed and the foreign collection is sharded,
+        // throw a custom exception.
+        if (auto staleInfo = ex.extraInfo<StaleConfigInfo>()) {
+            uassert(51069,
+                    "Cannot run $lookup with sharded foreign collection",
+                    foreignShardedLookupAllowed() || !staleInfo->getVersionWanted() ||
+                        staleInfo->getVersionWanted() == ChunkVersion::UNSHARDED());
+        }
+        throw;
+    }
 
     std::vector<Value> results;
     long long objsize = 0;
@@ -308,11 +315,15 @@ std::unique_ptr<Pipeline, PipelineDeleter> DocumentSourceLookUp::buildPipeline(
     // Copy all 'let' variables into the foreign pipeline's expression context.
     _variables.copyToExpCtx(_variablesParseState, _fromExpCtx.get());
 
-    assertIsValidCollectionState(_fromExpCtx);
-
     // Resolve the 'let' variables to values per the given input document.
     resolveLetVariables(inputDoc, &_fromExpCtx->variables);
 
+    if (!foreignShardedLookupAllowed()) {
+        // Enforce that the foreign collection must be unsharded for lookup.
+        _fromExpCtx->mongoProcessInterface->setExpectedShardVersion(
+            _fromExpCtx->opCtx, _fromExpCtx->ns, ChunkVersion::UNSHARDED());
+    }
+
     // If we don't have a cache, build and return the pipeline immediately.
     if (!_cache || _cache->isAbandoned()) {
         MakePipelineOptions pipelineOpts;
diff --git a/src/mongo/db/pipeline/document_source_merge.cpp b/src/mongo/db/pipeline/document_source_merge.cpp
index 256f23f01a2..2dd399d6fce 100644
--- a/src/mongo/db/pipeline/document_source_merge.cpp
+++ b/src/mongo/db/pipeline/document_source_merge.cpp
@@ -477,7 +477,8 @@ StageConstraints DocumentSourceMerge::constraints(Pipeline::SplitState pipeState
     // either choice will work correctly, we are simply applying a heuristic optimization.
     return {StreamType::kStreaming,
             PositionRequirement::kLast,
-            pExpCtx->mongoProcessInterface->isSharded(pExpCtx->opCtx, _outputNs)
+            pExpCtx->inMongos &&
+                    pExpCtx->mongoProcessInterface->isSharded(pExpCtx->opCtx, _outputNs)
                 ? HostTypeRequirement::kAnyShard
                 : HostTypeRequirement::kPrimaryShard,
             DiskUseRequirement::kWritesPersistentData,
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index d99123a1404..43915889af4 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -2622,6 +2622,7 @@ class MergeWithShardedCollection : public ShardMergerBase {
         };
 
         auto expCtx = ShardMergerBase::createExpressionContext(request);
+        expCtx->inMongos = true;
         expCtx->mongoProcessInterface = std::make_shared<ProcessInterface>();
         return expCtx;
     }
diff --git a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
index f7d8f8e0a2e..440648566a6 100644
--- a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
@@ -418,6 +418,16 @@ public:
                                               const NamespaceString& nss,
                                               ChunkVersion targetCollectionVersion) const = 0;
 
+    /**
+     * Sets the expected shard version for the given namespace. Invariants if the caller attempts to
+     * change an existing shard version, or if the shard version for this namespace has already been
+     * checked by the commands infrastructure. Used by $lookup and $graphLookup to enforce the
+     * constraint that the foreign collection must be unsharded.
+     */
+    virtual void setExpectedShardVersion(OperationContext* opCtx,
+                                         const NamespaceString& nss,
+                                         boost::optional<ChunkVersion> chunkVersion) = 0;
+
     virtual std::unique_ptr<ResourceYielder> getResourceYielder() const = 0;
 
     /**
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp b/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
index cce89f5177d..72868642cf5 100644
--- a/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
@@ -278,8 +278,9 @@ std::vector<GenericCursor> MongosProcessInterface::getIdleCursors(
 }
 
 bool MongosProcessInterface::isSharded(OperationContext* opCtx, const NamespaceString& nss) {
-    auto routingInfo = Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss);
-    return routingInfo.isOK() && routingInfo.getValue().cm();
+    auto routingInfo =
+        uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss));
+    return static_cast<bool>(routingInfo.cm());
 }
 
 bool MongosProcessInterface::fieldsHaveSupportingUniqueIndex(
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
index 25f5dfa2f9c..8ce9c181fde 100644
--- a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
@@ -211,6 +211,12 @@ public:
         MONGO_UNREACHABLE;
     }
 
+    void setExpectedShardVersion(OperationContext* opCtx,
+                                 const NamespaceString& nss,
+                                 boost::optional<ChunkVersion> chunkVersion) override {
+        MONGO_UNREACHABLE;
+    }
+
     std::unique_ptr<ResourceYielder> getResourceYielder() const override {
         return nullptr;
     }
diff --git a/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.h b/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.h
index 0e8422d9e8e..fcb015c25b4 100644
--- a/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.h
@@ -111,6 +111,12 @@ public:
                                         const NamespaceString& ns,
                                         const std::vector<BSONObj>& indexSpecs) override;
 
+    void setExpectedShardVersion(OperationContext* opCtx,
+                                 const NamespaceString& nss,
+                                 boost::optional<ChunkVersion> chunkVersion) override {
+        // Do nothing on a non-shardsvr mongoD.
+    }
+
 protected:
     // This constructor is marked as protected in order to prevent instantiation since this
     // interface is designed to have a concrete process interface for each possible
diff --git a/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.cpp b/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.cpp
index f97b06e49b3..1e8cfcd38fc 100644
--- a/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.cpp
@@ -42,6 +42,7 @@
 #include "mongo/db/pipeline/sharded_agg_helpers.h"
 #include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/db/s/collection_sharding_state.h"
+#include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/s/cluster_commands_helpers.h"
@@ -53,11 +54,9 @@ namespace mongo {
 using namespace fmt::literals;
 
 bool ShardServerProcessInterface::isSharded(OperationContext* opCtx, const NamespaceString& nss) {
-    Lock::DBLock dbLock(opCtx, nss.db(), MODE_IS);
-    Lock::CollectionLock collLock(opCtx, nss, MODE_IS);
-    return CollectionShardingState::get(opCtx, nss)
-        ->getCollectionDescription_DEPRECATED()
-        .isSharded();
+    auto routingInfo =
+        uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss));
+    return static_cast<bool>(routingInfo.cm());
 }
 
 void ShardServerProcessInterface::checkRoutingInfoEpochOrThrow(
@@ -350,4 +349,17 @@ ShardServerProcessInterface::attachCursorSourceToPipeline(Pipeline* ownedPipelin
     return sharded_agg_helpers::attachCursorToPipeline(ownedPipeline, allowTargetingShards);
 }
 
+void ShardServerProcessInterface::setExpectedShardVersion(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    boost::optional<ChunkVersion> chunkVersion) {
+    auto& oss = OperationShardingState::get(opCtx);
+    if (oss.hasShardVersion(nss)) {
+        invariant(oss.getShardVersion(nss) == chunkVersion);
+    } else {
+        OperationShardingState::get(opCtx).initializeClientRoutingVersions(
+            nss, chunkVersion, boost::none);
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.h b/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.h
index 76cf5887fbb..6b0ecc26bea 100644
--- a/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/shardsvr_process_interface.h
@@ -43,10 +43,13 @@ public:
     using CommonMongodProcessInterface::CommonMongodProcessInterface;
 
     /**
-     * Note: Information returned can be stale. Caller should always attach shardVersion when
-     * sending request against nss based on this information.
+     * Note: Cannot be called while holding a lock. Refreshes from the config servers if the
+     * metadata for the given namespace does not exist. Otherwise, will not automatically refresh,
+     * so the answer may be stale or become stale after calling. Caller should always attach
+     * shardVersion when sending request against nss based on this information.
      */
     bool isSharded(OperationContext* opCtx, const NamespaceString& nss) final;
+
     void checkRoutingInfoEpochOrThrow(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                       const NamespaceString& nss,
                                       ChunkVersion targetCollectionVersion) const final;
@@ -119,6 +122,10 @@ public:
      */
     std::unique_ptr<Pipeline, PipelineDeleter> attachCursorSourceToPipeline(
         Pipeline* pipeline, bool allowTargetingShards) final;
+
+    void setExpectedShardVersion(OperationContext* opCtx,
+                                 const NamespaceString& nss,
+                                 boost::optional<ChunkVersion> chunkVersion) final;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
index 7ebd3beba13..9648cce4bdc 100644
--- a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
@@ -259,5 +259,11 @@ public:
 
         return {*fieldPaths, targetCollectionVersion};
     }
+
+    void setExpectedShardVersion(OperationContext* opCtx,
+                                 const NamespaceString& nss,
+                                 boost::optional<ChunkVersion> chunkVersion) override {
+        // Do nothing.
+    }
 };
 }  // namespace mongo
diff --git a/src/mongo/db/s/collection_sharding_runtime.cpp b/src/mongo/db/s/collection_sharding_runtime.cpp
index cb8ad8e1b32..a8bf443d7fb 100644
--- a/src/mongo/db/s/collection_sharding_runtime.cpp
+++ b/src/mongo/db/s/collection_sharding_runtime.cpp
@@ -73,11 +73,9 @@ const auto kUnshardedCollection = std::make_shared<UnshardedCollection>();
 
 boost::optional<ChunkVersion> getOperationReceivedVersion(OperationContext* opCtx,
                                                           const NamespaceString& nss) {
-    auto& oss = OperationShardingState::get(opCtx);
-
     // If there is a version attached to the OperationContext, use it as the received version.
-    if (oss.hasShardVersion()) {
-        return oss.getShardVersion(nss);
+    if (OperationShardingState::isOperationVersioned(opCtx)) {
+        return OperationShardingState::get(opCtx).getShardVersion(nss);
     }
 
     // There is no shard version information on the 'opCtx'. This means that the operation
@@ -128,7 +126,7 @@ ScopedCollectionDescription CollectionShardingRuntime::getCollectionDescription(
     // consider all collections as unsharded. Also, return unsharded if no shard version or db
     // version is present on the context.
     if (!ShardingState::get(_serviceContext)->enabled() ||
-        (!oss.hasShardVersion() && !oss.hasDbVersion())) {
+        (!OperationShardingState::isOperationVersioned(opCtx) && !oss.hasDbVersion())) {
         return {kUnshardedCollection};
     }
 
diff --git a/src/mongo/db/s/collection_sharding_runtime_test.cpp b/src/mongo/db/s/collection_sharding_runtime_test.cpp
index 08d6ad3edad..a58d4375a92 100644
--- a/src/mongo/db/s/collection_sharding_runtime_test.cpp
+++ b/src/mongo/db/s/collection_sharding_runtime_test.cpp
@@ -55,11 +55,12 @@ CollectionMetadata makeShardedMetadata(OperationContext* opCtx, UUID uuid = UUID
         kTestNss, uuid, kShardKeyPattern, nullptr, false, epoch, {std::move(chunk)});
     std::shared_ptr<ChunkManager> cm = std::make_shared<ChunkManager>(rt, boost::none);
 
-    auto& oss = OperationShardingState::get(opCtx);
-    if (!oss.hasShardVersion()) {
+    if (!OperationShardingState::isOperationVersioned(opCtx)) {
         const auto version = cm->getVersion(ShardId("0"));
         BSONObjBuilder builder;
         version.appendToCommand(&builder);
+
+        auto& oss = OperationShardingState::get(opCtx);
         oss.initializeClientRoutingVersionsFromCommand(kTestNss, builder.obj());
     }
     return CollectionMetadata(std::move(cm), ShardId("0"));
diff --git a/src/mongo/db/s/operation_sharding_state.cpp b/src/mongo/db/s/operation_sharding_state.cpp
index 58351127f90..f7c47e6231a 100644
--- a/src/mongo/db/s/operation_sharding_state.cpp
+++ b/src/mongo/db/s/operation_sharding_state.cpp
@@ -47,7 +47,6 @@ const Milliseconds kMaxWaitForMovePrimaryCriticalSection = Minutes(5);
 
 // The name of the field in which the client attaches its database version.
 constexpr auto kDbVersionField = "databaseVersion"_sd;
-
 }  // namespace
 
 OperationShardingState::OperationShardingState() = default;
@@ -61,7 +60,7 @@ OperationShardingState& OperationShardingState::get(OperationContext* opCtx) {
 }
 
 bool OperationShardingState::isOperationVersioned(OperationContext* opCtx) {
-    return get(opCtx).hasShardVersion();
+    return !(get(opCtx)._shardVersions.empty());
 }
 
 void OperationShardingState::setAllowImplicitCollectionCreation(
@@ -79,12 +78,9 @@ bool OperationShardingState::allowImplicitCollectionCreation() const {
 
 void OperationShardingState::initializeClientRoutingVersionsFromCommand(NamespaceString nss,
                                                                         const BSONObj& cmdObj) {
-    invariant(_shardVersions.empty());
-    invariant(_databaseVersions.empty());
-
+    // TODO SERVER-48618 Enforce that the nss shoud not be empty.
     boost::optional<ChunkVersion> shardVersion;
     boost::optional<DatabaseVersion> dbVersion;
-
     const auto shardVersionElem = cmdObj.getField(ChunkVersion::kShardVersionField);
     if (!shardVersionElem.eoo()) {
         shardVersion = uassertStatusOK(ChunkVersion::parseFromCommand(cmdObj));
@@ -108,30 +104,28 @@ void OperationShardingState::initializeClientRoutingVersions(
     NamespaceString nss,
     const boost::optional<ChunkVersion>& shardVersion,
     const boost::optional<DatabaseVersion>& dbVersion) {
-    invariant(_shardVersions.empty());
-    invariant(_databaseVersions.empty());
-
+    // TODO SERVER-48618 Enforce that the nss shoud not be empty. For now, all empty
+    // NamespaceStrings will be mapped under the empty string "".
     if (shardVersion) {
+        invariant(_shardVersionsChecked.find(nss.ns()) == _shardVersionsChecked.end(), nss.ns());
         _shardVersions[nss.ns()] = *shardVersion;
     }
     if (dbVersion) {
-        // Unforunately this is a bit ugly; it's because a command comes with a shardVersion or
-        // databaseVersion, and the assumption is that those versions are applied to whatever is
-        // returned by the Command's parseNs(), which can either be a full namespace or just a db.
-        _databaseVersions[nss.db().empty() ? nss.ns() : nss.db()] = *dbVersion;
+        invariant(_databaseVersions.find(nss.db()) == _databaseVersions.end());
+        _databaseVersions[nss.db()] = *dbVersion;
     }
 }
 
-bool OperationShardingState::hasShardVersion() const {
-    return _globalUnshardedShardVersion || !_shardVersions.empty();
+bool OperationShardingState::hasShardVersion(const NamespaceString& nss) const {
+    // TODO SERVER-48618 Enforce that the nss shoud not be empty. For now, all empty
+    // NamespaceStrings will be treated as the same namespace.
+    return _shardVersions.find(nss.ns()) != _shardVersions.end();
 }
 
-boost::optional<ChunkVersion> OperationShardingState::getShardVersion(
-    const NamespaceString& nss) const {
-    if (_globalUnshardedShardVersion) {
-        return ChunkVersion::UNSHARDED();
-    }
-
+boost::optional<ChunkVersion> OperationShardingState::getShardVersion(const NamespaceString& nss) {
+    // TODO SERVER-48618 Enforce that the nss shoud not be empty. For now, all empty
+    // NamespaceStrings will be treated as the same namespace.
+    _shardVersionsChecked.insert(nss.ns());
     const auto it = _shardVersions.find(nss.ns());
     if (it != _shardVersions.end()) {
         return it->second;
@@ -153,11 +147,6 @@ boost::optional<DatabaseVersion> OperationShardingState::getDbVersion(
     return it->second;
 }
 
-void OperationShardingState::setGlobalUnshardedShardVersion() {
-    invariant(_shardVersions.empty());
-    _globalUnshardedShardVersion = true;
-}
-
 bool OperationShardingState::waitForMigrationCriticalSectionSignal(OperationContext* opCtx) {
     // Must not block while holding a lock
     invariant(!opCtx->lockState()->isLocked());
diff --git a/src/mongo/db/s/operation_sharding_state.h b/src/mongo/db/s/operation_sharding_state.h
index 77dd848d9c4..558202e6d4a 100644
--- a/src/mongo/db/s/operation_sharding_state.h
+++ b/src/mongo/db/s/operation_sharding_state.h
@@ -99,25 +99,26 @@ public:
     void initializeClientRoutingVersionsFromCommand(NamespaceString nss, const BSONObj& cmdObj);
 
     /**
-     * Stores the given shardVersion and databaseVersion for the given namespace.
+     * Stores the given shardVersion and databaseVersion for the given namespace. Note: The shard
+     * version for the given namespace stored in the OperationShardingState can be overwritten if it
+     * has not been checked yet.
      */
     void initializeClientRoutingVersions(NamespaceString nss,
                                          const boost::optional<ChunkVersion>& shardVersion,
                                          const boost::optional<DatabaseVersion>& dbVersion);
 
     /**
-     * Returns whether or not there is a shard version associated with this operation.
+     * Returns whether or not there is a shard version for the namespace associated with this
+     * operation.
      */
-    bool hasShardVersion() const;
+    bool hasShardVersion(const NamespaceString& nss) const;
 
     /**
      * Returns the shard version (i.e. maximum chunk version) of a namespace being used by the
      * operation. Documents in chunks which did not belong on this shard at this shard version
      * will be filtered out.
-     *
-     * Returns ChunkVersion::UNSHARDED() if setGlobalUnshardedShardVersion has been called.
      */
-    boost::optional<ChunkVersion> getShardVersion(const NamespaceString& nss) const;
+    boost::optional<ChunkVersion> getShardVersion(const NamespaceString& nss);
 
     /**
      * Returns true if the client sent a databaseVersion for any namespace.
@@ -131,12 +132,6 @@ public:
     boost::optional<DatabaseVersion> getDbVersion(const StringData dbName) const;
 
     /**
-     * Makes the OperationShardingState behave as if an UNSHARDED shardVersion was sent for every
-     * possible namespace.
-     */
-    void setGlobalUnshardedShardVersion();
-
-    /**
      * This call is a no op if there isn't a currently active migration critical section. Otherwise
      * it will wait for the critical section to complete up to the remaining operation time.
      *
@@ -189,15 +184,15 @@ private:
     // Specifies whether the request is allowed to create database/collection implicitly
     bool _allowImplicitCollectionCreation{true};
 
-    // Should be set to true if all collections accessed are expected to be unsharded.
-    bool _globalUnshardedShardVersion = false;
-
     // The OperationShardingState class supports storing shardVersions for multiple namespaces (and
     // databaseVersions for multiple databases), even though client code has not been written yet to
     // *send* multiple shardVersions or databaseVersions.
     StringMap<ChunkVersion> _shardVersions;
     StringMap<DatabaseVersion> _databaseVersions;
 
+    // Stores shards that have undergone a version check.
+    StringDataSet _shardVersionsChecked;
+
     // This value will only be non-null if version check during the operation execution failed due
     // to stale version and there was a migration for that namespace, which was in critical section.
     std::shared_ptr<Notification<void>> _migrationCriticalSectionSignal;
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 20c80bc7d83..15681bc8522 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -1131,7 +1131,7 @@ void execCommandDatabase(OperationContext* opCtx,
             oss.initializeClientRoutingVersionsFromCommand(invocation->ns(), request.body);
 
             auto const shardingState = ShardingState::get(opCtx);
-            if (oss.hasShardVersion() || oss.hasDbVersion()) {
+            if (OperationShardingState::isOperationVersioned(opCtx) || oss.hasDbVersion()) {
                 uassertStatusOK(shardingState->canAcceptShardedCommands());
             }