SERVER-40258 Relax locking requirements for sharding metadata refresh on shards

(cherry picked from commit 0d07bf5e7a72a5bce3f7d7d681a71d7ecfe7eb8c)

5 files changed, 36 insertions, 23 deletions

diff --git a/src/mongo/db/s/collection_sharding_runtime.cpp b/src/mongo/db/s/collection_sharding_runtime.cpp
index 1fc98f41876..dcd52ea2ebc 100644
--- a/src/mongo/db/s/collection_sharding_runtime.cpp
+++ b/src/mongo/db/s/collection_sharding_runtime.cpp
@@ -96,7 +96,8 @@ void CollectionShardingRuntime::setFilteringMetadata(OperationContext* opCtx,
                                                      CollectionMetadata newMetadata) {
     invariant(!newMetadata.isSharded() || !isNamespaceAlwaysUnsharded(_nss),
               str::stream() << "Namespace " << _nss.ns() << " must never be sharded.");
-    invariant(opCtx->lockState()->isCollectionLockedForMode(_nss, MODE_X));
+
+    auto csrLock = CollectionShardingState::CSRLock::lockExclusive(opCtx, this);
 
     _metadataManager->setFilteringMetadata(std::move(newMetadata));
 }
diff --git a/src/mongo/db/s/collection_sharding_state.cpp b/src/mongo/db/s/collection_sharding_state.cpp
index e0dd987600d..feb519090e3 100644
--- a/src/mongo/db/s/collection_sharding_state.cpp
+++ b/src/mongo/db/s/collection_sharding_state.cpp
@@ -160,12 +160,8 @@ void CollectionShardingState::report(OperationContext* opCtx, BSONObjBuilder* bu
 }
 
 ScopedCollectionMetadata CollectionShardingState::getOrphansFilter(OperationContext* opCtx) {
-    const auto receivedShardVersion = getOperationReceivedVersion(opCtx, _nss);
-    if (!receivedShardVersion)
-        return {kUnshardedCollection};
-
     const auto atClusterTime = repl::ReadConcernArgs::get(opCtx).getArgsAtClusterTime();
-    auto optMetadata = _getMetadata(atClusterTime);
+    auto optMetadata = _getMetadataWithVersionCheckAt(opCtx, atClusterTime);
 
     if (!optMetadata)
         return {kUnshardedCollection};
@@ -199,26 +195,34 @@ boost::optional<ChunkVersion> CollectionShardingState::getCurrentShardVersionIfK
 }
 
 void CollectionShardingState::checkShardVersionOrThrow(OperationContext* opCtx) {
+    (void)_getMetadataWithVersionCheckAt(opCtx, boost::none);
+}
+
+boost::optional<ScopedCollectionMetadata> CollectionShardingState::_getMetadataWithVersionCheckAt(
+    OperationContext* opCtx, const boost::optional<mongo::LogicalTime>& atClusterTime) {
     const auto optReceivedShardVersion = getOperationReceivedVersion(opCtx, _nss);
 
     if (!optReceivedShardVersion)
-        return;
+        return ScopedCollectionMetadata(kUnshardedCollection);
 
     const auto& receivedShardVersion = *optReceivedShardVersion;
     if (ChunkVersion::isIgnoredVersion(receivedShardVersion)) {
-        return;
+        return boost::none;
     }
 
     // An operation with read concern 'available' should never have shardVersion set.
     invariant(repl::ReadConcernArgs::get(opCtx).getLevel() !=
               repl::ReadConcernLevel::kAvailableReadConcern);
 
-    const auto metadata = getCurrentMetadata();
-    const auto wantedShardVersion =
-        metadata->isSharded() ? metadata->getShardVersion() : ChunkVersion::UNSHARDED();
+    auto csrLock = CSRLock::lockShared(opCtx, this);
+
+    auto metadata = _getMetadata(atClusterTime);
+    auto wantedShardVersion = ChunkVersion::UNSHARDED();
+    if (metadata && (*metadata)->isSharded()) {
+        wantedShardVersion = (*metadata)->getShardVersion();
+    }
 
     auto criticalSectionSignal = [&] {
-        auto csrLock = CSRLock::lockShared(opCtx, this);
         return _critSec.getSignal(opCtx->lockState()->isWriteLocked()
                                       ? ShardingMigrationCriticalSection::kWrite
                                       : ShardingMigrationCriticalSection::kRead);
@@ -235,7 +239,7 @@ void CollectionShardingState::checkShardVersionOrThrow(OperationContext* opCtx)
     }
 
     if (receivedShardVersion.isWriteCompatibleWith(wantedShardVersion)) {
-        return;
+        return metadata;
     }
 
     //
diff --git a/src/mongo/db/s/collection_sharding_state.h b/src/mongo/db/s/collection_sharding_state.h
index d62b010bec6..906c366b8fb 100644
--- a/src/mongo/db/s/collection_sharding_state.h
+++ b/src/mongo/db/s/collection_sharding_state.h
@@ -91,8 +91,9 @@ public:
      * metadata object.
      *
      * The intended users of this method are callers which need to perform orphan filtering. Use
-     * 'getCurrentMetadata' for all other cases, where just sharding-related properties of the
-     * collection are necessary (e.g., isSharded or the shard key).
+     * 'getCurrentMetadata' for other cases, like obtaining information about sharding-related
+     * properties of the collection are necessary that won't change under collection IX/IS lock
+     * (e.g., isSharded or the shard key).
      *
      * The returned object is safe to access even after the collection lock has been dropped.
      */
@@ -155,6 +156,13 @@ protected:
 private:
     friend CSRLock;
 
+    /**
+     * Returns the latest version of collection metadata with filtering configured for
+     * atClusterTime if specified.
+     */
+    boost::optional<ScopedCollectionMetadata> _getMetadataWithVersionCheckAt(
+        OperationContext* opCtx, const boost::optional<mongo::LogicalTime>& atClusterTime);
+
     // Object-wide ResourceMutex to protect changes to the CollectionShardingRuntime or objects
     // held within. Use only the CollectionShardingRuntimeLock to lock this mutex.
     Lock::ResourceMutex _stateChangeMutex;
diff --git a/src/mongo/db/s/op_observer_sharding_impl.cpp b/src/mongo/db/s/op_observer_sharding_impl.cpp
index 1ac9e743879..ec0fcb370e9 100644
--- a/src/mongo/db/s/op_observer_sharding_impl.cpp
+++ b/src/mongo/db/s/op_observer_sharding_impl.cpp
@@ -42,9 +42,10 @@ namespace {
 const auto getIsMigrating = OperationContext::declareDecoration<bool>();
 
 /**
- * Write operations do shard version checking, but do not perform orphan document filtering. Because
- * of this, if an update operation runs as part of a 'readConcern:snapshot' transaction, it might
- * get routed to a shard which no longer owns the chunk being written to. In such cases, throw a
+ * Write operations do shard version checking, but if an update operation runs as part of a
+ * 'readConcern:snapshot' transaction, the router could have used the metadata at the snapshot
+ * time and yet set the latest shard version on the request. This is why the write can get routed
+ * to a shard which no longer owns the chunk being written to. In such cases, throw a
  * MigrationConflict exception to indicate that the transaction needs to be rolled-back and
  * restarted.
  */
diff --git a/src/mongo/db/s/shard_filtering_metadata_refresh.cpp b/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
index bd55c8d6b9e..2d8b2359378 100644
--- a/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
+++ b/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
@@ -157,7 +157,7 @@ ChunkVersion forceShardFilteringMetadataRefresh(OperationContext* opCtx,
         // No chunk manager, so unsharded. Avoid using AutoGetCollection() as it returns the
         // InvalidViewDefinition error code if an invalid view is in the 'system.views' collection.
         AutoGetDb autoDb(opCtx, nss.db(), MODE_IX);
-        Lock::CollectionLock collLock(opCtx, nss, MODE_X);
+        Lock::CollectionLock collLock(opCtx, nss, MODE_IX);
         CollectionShardingRuntime::get(opCtx, nss)
             ->setFilteringMetadata(opCtx, CollectionMetadata());
 
@@ -186,11 +186,10 @@ ChunkVersion forceShardFilteringMetadataRefresh(OperationContext* opCtx,
         }
     }
 
-    // Exclusive collection lock needed since we're now changing the metadata. Avoid using
-    // AutoGetCollection() as it returns the InvalidViewDefinition error code if an invalid view is
-    // in the 'system.views' collection.
+    // Avoid using AutoGetCollection() as it returns the InvalidViewDefinition error code if an
+    // invalid view is in the 'system.views' collection.
     AutoGetDb autoDb(opCtx, nss.db(), MODE_IX);
-    Lock::CollectionLock collLock(opCtx, nss, MODE_X);
+    Lock::CollectionLock collLock(opCtx, nss, MODE_IX);
     auto* const css = CollectionShardingRuntime::get(opCtx, nss);
 
     {