SERVER-47975 Optimize ScopedShardVersionCriticalSection in order to avoid convoy on SSV after a shardVersion change

9 files changed, 168 insertions, 183 deletions

diff --git a/src/mongo/db/s/SConscript b/src/mongo/db/s/SConscript
index c932e48c2e2..3cb6a57100b 100644
--- a/src/mongo/db/s/SConscript
+++ b/src/mongo/db/s/SConscript
@@ -67,7 +67,6 @@ env.Library(
         'range_deletion_util.cpp',
         'read_only_catalog_cache_loader.cpp',
         'scoped_operation_completion_sharding_actions.cpp',
-        'scoped_shard_version_critical_section.cpp',
         'session_catalog_migration_destination.cpp',
         'session_catalog_migration_source.cpp',
         'shard_filtering_metadata_refresh.cpp',
diff --git a/src/mongo/db/s/collection_sharding_runtime.cpp b/src/mongo/db/s/collection_sharding_runtime.cpp
index 56b9c141c76..51861f21ddb 100644
--- a/src/mongo/db/s/collection_sharding_runtime.cpp
+++ b/src/mongo/db/s/collection_sharding_runtime.cpp
@@ -166,15 +166,13 @@ void CollectionShardingRuntime::checkShardVersionOrThrow(OperationContext* opCtx
     (void)_getMetadataWithVersionCheckAt(opCtx, boost::none);
 }
 
-void CollectionShardingRuntime::enterCriticalSectionCatchUpPhase(OperationContext* opCtx) {
-    invariant(opCtx->lockState()->isCollectionLockedForMode(_nss, MODE_X));
-    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(opCtx, this);
+void CollectionShardingRuntime::enterCriticalSectionCatchUpPhase(OperationContext* opCtx,
+                                                                 const CSRLock&) {
     _critSec.enterCriticalSectionCatchUpPhase();
 }
 
-void CollectionShardingRuntime::enterCriticalSectionCommitPhase(OperationContext* opCtx) {
-    invariant(opCtx->lockState()->isCollectionLockedForMode(_nss, MODE_X));
-    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(opCtx, this);
+void CollectionShardingRuntime::enterCriticalSectionCommitPhase(OperationContext* opCtx,
+                                                                const CSRLock&) {
     _critSec.enterCriticalSectionCommitPhase();
 }
 
@@ -428,7 +426,8 @@ CollectionCriticalSection::CollectionCriticalSection(OperationContext* opCtx, Na
                                opCtx->getServiceContext()->getPreciseClockSource()->now() +
                                    Milliseconds(migrationLockAcquisitionMaxWaitMS.load()));
     auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
-    csr->enterCriticalSectionCatchUpPhase(_opCtx);
+    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(opCtx, csr);
+    csr->enterCriticalSectionCatchUpPhase(_opCtx, csrLock);
 }
 
 CollectionCriticalSection::~CollectionCriticalSection() {
@@ -446,7 +445,8 @@ void CollectionCriticalSection::enterCommitPhase() {
                                _opCtx->getServiceContext()->getPreciseClockSource()->now() +
                                    Milliseconds(migrationLockAcquisitionMaxWaitMS.load()));
     auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
-    csr->enterCriticalSectionCommitPhase(_opCtx);
+    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(_opCtx, csr);
+    csr->enterCriticalSectionCommitPhase(_opCtx, csrLock);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/collection_sharding_runtime.h b/src/mongo/db/s/collection_sharding_runtime.h
index 31d901dbae2..3551bf39eb5 100644
--- a/src/mongo/db/s/collection_sharding_runtime.h
+++ b/src/mongo/db/s/collection_sharding_runtime.h
@@ -167,8 +167,8 @@ public:
      *
      * In these methods, the CSRLock ensures concurrent access to the critical section.
      */
-    void enterCriticalSectionCatchUpPhase(OperationContext* opCtx);
-    void enterCriticalSectionCommitPhase(OperationContext* opCtx);
+    void enterCriticalSectionCatchUpPhase(OperationContext* opCtx, const CSRLock&);
+    void enterCriticalSectionCommitPhase(OperationContext* opCtx, const CSRLock&);
 
     /**
      * Method to control the collection's critical secion. Method listed below must be called with
diff --git a/src/mongo/db/s/scoped_shard_version_critical_section.cpp b/src/mongo/db/s/scoped_shard_version_critical_section.cpp
deleted file mode 100644
index 6be946fc1b0..00000000000
--- a/src/mongo/db/s/scoped_shard_version_critical_section.cpp
+++ /dev/null
@@ -1,76 +0,0 @@
-/**
- *    Copyright (C) 2018-present MongoDB, Inc.
- *
- *    This program is free software: you can redistribute it and/or modify
- *    it under the terms of the Server Side Public License, version 1,
- *    as published by MongoDB, Inc.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    Server Side Public License for more details.
- *
- *    You should have received a copy of the Server Side Public License
- *    along with this program. If not, see
- *    <http://www.mongodb.com/licensing/server-side-public-license>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the Server Side Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#include "mongo/platform/basic.h"
-
-#include "mongo/db/s/scoped_shard_version_critical_section.h"
-
-#include "mongo/db/catalog_raii.h"
-#include "mongo/db/s/shard_filtering_metadata_refresh.h"
-
-namespace mongo {
-
-ScopedShardVersionCriticalSection::ScopedShardVersionCriticalSection(OperationContext* opCtx,
-                                                                     NamespaceString nss)
-    : _nss(std::move(nss)), _opCtx(opCtx) {
-
-    {
-        // Enters the commit phase of the critical section and blocks reads
-        AutoGetCollection autoColl(_opCtx,
-                                   _nss,
-                                   MODE_X,
-                                   AutoGetCollection::ViewMode::kViewsForbidden,
-                                   _opCtx->getServiceContext()->getPreciseClockSource()->now() +
-                                       Milliseconds(migrationLockAcquisitionMaxWaitMS.load()));
-        auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
-        csr->enterCriticalSectionCatchUpPhase(_opCtx);
-    }
-
-    forceShardFilteringMetadataRefresh(_opCtx, _nss, true);
-}
-
-ScopedShardVersionCriticalSection::~ScopedShardVersionCriticalSection() {
-    UninterruptibleLockGuard noInterrupt(_opCtx->lockState());
-    AutoGetCollection autoColl(_opCtx, _nss, MODE_IX);
-    auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
-    csr->exitCriticalSection(_opCtx);
-}
-
-void ScopedShardVersionCriticalSection::enterCommitPhase() {
-    AutoGetCollection autoColl(_opCtx,
-                               _nss,
-                               MODE_X,
-                               AutoGetCollection::ViewMode::kViewsForbidden,
-                               _opCtx->getServiceContext()->getPreciseClockSource()->now() +
-                                   Milliseconds(migrationLockAcquisitionMaxWaitMS.load()));
-    auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
-    csr->enterCriticalSectionCommitPhase(_opCtx);
-}
-
-}  // namespace mongo
diff --git a/src/mongo/db/s/scoped_shard_version_critical_section.h b/src/mongo/db/s/scoped_shard_version_critical_section.h
deleted file mode 100644
index d5ec579313e..00000000000
--- a/src/mongo/db/s/scoped_shard_version_critical_section.h
+++ /dev/null
@@ -1,57 +0,0 @@
-/**
- *    Copyright (C) 2018-present MongoDB, Inc.
- *
- *    This program is free software: you can redistribute it and/or modify
- *    it under the terms of the Server Side Public License, version 1,
- *    as published by MongoDB, Inc.
- *
- *    This program is distributed in the hope that it will be useful,
- *    but WITHOUT ANY WARRANTY; without even the implied warranty of
- *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *    Server Side Public License for more details.
- *
- *    You should have received a copy of the Server Side Public License
- *    along with this program. If not, see
- *    <http://www.mongodb.com/licensing/server-side-public-license>.
- *
- *    As a special exception, the copyright holders give permission to link the
- *    code of portions of this program with the OpenSSL library under certain
- *    conditions as described in each individual source file and distribute
- *    linked combinations including the program with the OpenSSL library. You
- *    must comply with the Server Side Public License in all respects for
- *    all of the code used other than as permitted herein. If you modify file(s)
- *    with this exception, you may extend this exception to your version of the
- *    file(s), but you are not obligated to do so. If you do not wish to do so,
- *    delete this exception statement from your version. If you delete this
- *    exception statement from all source files in the program, then also delete
- *    it in the license file.
- */
-
-#pragma once
-
-#include "mongo/db/namespace_string.h"
-#include "mongo/db/s/collection_sharding_runtime.h"
-
-namespace mongo {
-
-/**
- * RAII-style class that enters the migration critical section and refresh the filtering
- * metadata for the specified collection. The critical section is released when this object
- * goes out of scope.
- */
-class ScopedShardVersionCriticalSection {
-    ScopedShardVersionCriticalSection(const ScopedShardVersionCriticalSection&) = delete;
-    ScopedShardVersionCriticalSection& operator=(const ScopedShardVersionCriticalSection&) = delete;
-
-public:
-    ScopedShardVersionCriticalSection(OperationContext* opCtx, NamespaceString nss);
-    ~ScopedShardVersionCriticalSection();
-
-    void enterCommitPhase();
-
-private:
-    const NamespaceString _nss;
-    OperationContext* _opCtx;
-};
-
-}  // namespace mongo
diff --git a/src/mongo/db/s/set_shard_version_command.cpp b/src/mongo/db/s/set_shard_version_command.cpp
index a7c71f7f8b2..ae45d23e844 100644
--- a/src/mongo/db/s/set_shard_version_command.cpp
+++ b/src/mongo/db/s/set_shard_version_command.cpp
@@ -239,10 +239,16 @@ public:
 
         // Step 6
 
-        // Note: The forceRefresh flag controls whether we make sure to do our
-        // own refresh or if we're okay with joining another thread
-        const auto status = onShardVersionMismatchNoExcept(
-            opCtx, nss, requestedVersion, forceRefresh /*forceRefreshFromThisThread*/);
+        // Note: The forceRefresh flag controls whether we make sure to do our own refresh or if
+        // we're okay with joining another thread
+        const auto status = [&] {
+            try {
+                forceShardFilteringMetadataRefresh(opCtx, nss, forceRefresh);
+                return Status::OK();
+            } catch (const DBException& ex) {
+                return ex.toStatus();
+            }
+        }();
 
         {
             // Avoid using AutoGetCollection() as it returns the InvalidViewDefinition error code
diff --git a/src/mongo/db/s/shard_filtering_metadata_refresh.cpp b/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
index 5db26090bcd..60b02fe846a 100644
--- a/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
+++ b/src/mongo/db/s/shard_filtering_metadata_refresh.cpp
@@ -36,6 +36,7 @@
 #include "mongo/db/catalog/database_holder.h"
 #include "mongo/db/catalog_raii.h"
 #include "mongo/db/commands/feature_compatibility_version.h"
+#include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/s/collection_sharding_runtime.h"
 #include "mongo/db/s/database_sharding_state.h"
@@ -53,16 +54,15 @@ MONGO_FAIL_POINT_DEFINE(skipDatabaseVersionMetadataRefresh);
 MONGO_FAIL_POINT_DEFINE(skipShardFilteringMetadataRefresh);
 
 namespace {
-
 void onShardVersionMismatch(OperationContext* opCtx,
                             const NamespaceString& nss,
-                            ChunkVersion shardVersionReceived,
-                            bool forceRefreshFromThisThread) {
+                            boost::optional<ChunkVersion> shardVersionReceived) {
     invariant(!opCtx->lockState()->isLocked());
     invariant(!opCtx->getClient()->isInDirectClient());
-
     invariant(ShardingState::get(opCtx)->canAcceptShardedCommands());
 
+    ShardingStatistics::get(opCtx).countStaleConfigErrors.addAndFetch(1);
+
     LOGV2_DEBUG(22061,
                 2,
                 "Metadata refresh requested for {namespace} at shard version "
@@ -71,37 +71,66 @@ void onShardVersionMismatch(OperationContext* opCtx,
                 "namespace"_attr = nss.ns(),
                 "shardVersionReceived"_attr = shardVersionReceived);
 
-    ShardingStatistics::get(opCtx).countStaleConfigErrors.addAndFetch(1);
-
-    // Ensure any ongoing migrations have completed before trying to do the refresh. This wait is
-    // just an optimization so that mongos does not exhaust its maximum number of StaleShardVersion
-    // retry attempts while the migration is being committed.
-    OperationShardingState::get(opCtx).waitForMigrationCriticalSectionSignal(opCtx);
-
-    {
-        // 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_IS);
-        Lock::CollectionLock collLock(opCtx, nss, MODE_IS);
-        const auto collDescr =
-            CollectionShardingRuntime::get(opCtx, nss)->getCurrentMetadataIfKnown();
-        if (collDescr) {
-            const auto currentShardVersion = collDescr->getShardVersion();
-            if (currentShardVersion.epoch() == shardVersionReceived.epoch() &&
-                currentShardVersion.majorVersion() >= shardVersionReceived.majorVersion()) {
-                // Don't need to remotely reload if we're in the same epoch and the requested
-                // version is smaller than the one we know about. This means that the remote side is
-                // behind.
-                return;
+    bool runRecover;
+    while (true) {
+        std::shared_ptr<Notification<void>> critSecSignal;
+
+        {
+            AutoGetDb autoDb(opCtx, nss.db(), MODE_IS);
+            Lock::CollectionLock collLock(opCtx, nss, MODE_IS);
+
+            auto* const csr = CollectionShardingRuntime::get(opCtx, nss);
+            critSecSignal =
+                csr->getCriticalSectionSignal(opCtx, ShardingMigrationCriticalSection::kWrite);
+            if (!critSecSignal) {
+                const auto collDesc = csr->getCurrentMetadataIfKnown();
+                if (collDesc) {
+                    if (shardVersionReceived) {
+                        const auto currentShardVersion = collDesc->getShardVersion();
+                        // Don't need to remotely reload if we're in the same epoch and the
+                        // requested version is smaller than the one we know about. This means that
+                        // the remote side is behind.
+                        if (currentShardVersion.epoch() == shardVersionReceived->epoch() &&
+                            currentShardVersion.majorVersion() >=
+                                shardVersionReceived->majorVersion())
+                            return;
+                    }
+
+                    runRecover = false;
+                    break;
+                } else {
+                    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(opCtx, csr);
+                    critSecSignal = csr->getCriticalSectionSignal(
+                        opCtx, ShardingMigrationCriticalSection::kWrite);
+                    if (!critSecSignal) {
+                        CollectionShardingRuntime::get(opCtx, nss)
+                            ->enterCriticalSectionCatchUpPhase(opCtx, csrLock);
+                        runRecover = true;
+                        break;
+                    }
+                }
             }
         }
+
+        invariant(critSecSignal);
+        critSecSignal->get(opCtx);
     }
 
-    if (MONGO_unlikely(skipShardFilteringMetadataRefresh.shouldFail())) {
-        return;
+    ON_BLOCK_EXIT([&] {
+        if (runRecover) {
+            UninterruptibleLockGuard noInterrupt(opCtx->lockState());
+            AutoGetDb autoDb(opCtx, nss.db(), MODE_IX);
+            Lock::CollectionLock collLock(opCtx, nss, MODE_IX);
+            CollectionShardingRuntime::get(opCtx, nss)->exitCriticalSection(opCtx);
+        }
+    });
+
+    if (runRecover) {
+        // TODO (SERVER-47985): Invoke recovery of the shardVersion after a (possible) failed
+        // migration
     }
 
-    forceShardFilteringMetadataRefresh(opCtx, nss, forceRefreshFromThisThread);
+    forceShardFilteringMetadataRefresh(opCtx, nss, !shardVersionReceived);
 }
 
 void onDbVersionMismatch(OperationContext* opCtx,
@@ -133,12 +162,74 @@ void onDbVersionMismatch(OperationContext* opCtx,
 
 }  // namespace
 
+ScopedShardVersionCriticalSection::ScopedShardVersionCriticalSection(OperationContext* opCtx,
+                                                                     NamespaceString nss)
+    : _opCtx(opCtx), _nss(std::move(nss)) {
+
+    while (true) {
+        std::shared_ptr<Notification<void>> critSecSignal;
+
+        {
+            AutoGetDb autoDb(_opCtx, _nss.db(), MODE_IS);
+            Lock::CollectionLock collLock(_opCtx, _nss, MODE_S);
+
+            auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
+            critSecSignal =
+                csr->getCriticalSectionSignal(_opCtx, ShardingMigrationCriticalSection::kWrite);
+            if (!critSecSignal) {
+                auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(_opCtx, csr);
+                critSecSignal =
+                    csr->getCriticalSectionSignal(_opCtx, ShardingMigrationCriticalSection::kWrite);
+                if (!critSecSignal) {
+                    CollectionShardingRuntime::get(_opCtx, _nss)
+                        ->enterCriticalSectionCatchUpPhase(_opCtx, csrLock);
+                    break;
+                }
+            }
+        }
+
+        invariant(critSecSignal);
+        critSecSignal->get(_opCtx);
+    }
+
+    // TODO (SERVER-47985): Invoke recovery of the shardVersion after a (possible) failed migration
+
+    forceShardFilteringMetadataRefresh(_opCtx, _nss, true);
+}
+
+ScopedShardVersionCriticalSection::~ScopedShardVersionCriticalSection() {
+    UninterruptibleLockGuard noInterrupt(_opCtx->lockState());
+    AutoGetCollection autoColl(_opCtx, _nss, MODE_IX);
+    auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
+    csr->exitCriticalSection(_opCtx);
+}
+
+void ScopedShardVersionCriticalSection::enterCommitPhase() {
+    AutoGetCollection autoColl(_opCtx,
+                               _nss,
+                               MODE_IS,
+                               AutoGetCollection::ViewMode::kViewsForbidden,
+                               _opCtx->getServiceContext()->getPreciseClockSource()->now() +
+                                   Milliseconds(migrationLockAcquisitionMaxWaitMS.load()));
+    auto* const csr = CollectionShardingRuntime::get(_opCtx, _nss);
+    auto csrLock = CollectionShardingRuntime::CSRLock::lockExclusive(_opCtx, csr);
+    csr->enterCriticalSectionCommitPhase(_opCtx, csrLock);
+}
+
+CatalogCacheLoader& getCatalogCacheLoaderForFiltering(ServiceContext* serviceContext) {
+    return CatalogCacheLoader::get(serviceContext);
+}
+
+CatalogCacheLoader& getCatalogCacheLoaderForFiltering(OperationContext* opCtx) {
+    return getCatalogCacheLoaderForFiltering(opCtx->getServiceContext());
+}
+
+
 Status onShardVersionMismatchNoExcept(OperationContext* opCtx,
                                       const NamespaceString& nss,
-                                      ChunkVersion shardVersionReceived,
-                                      bool forceRefreshFromThisThread) noexcept {
+                                      ChunkVersion shardVersionReceived) noexcept {
     try {
-        onShardVersionMismatch(opCtx, nss, shardVersionReceived, forceRefreshFromThisThread);
+        onShardVersionMismatch(opCtx, nss, shardVersionReceived);
         return Status::OK();
     } catch (const DBException& ex) {
         LOGV2(22062,
@@ -156,6 +247,10 @@ ChunkVersion forceShardFilteringMetadataRefresh(OperationContext* opCtx,
     invariant(!opCtx->lockState()->isLocked());
     invariant(!opCtx->getClient()->isInDirectClient());
 
+    if (MONGO_unlikely(skipShardFilteringMetadataRefresh.shouldFail())) {
+        uasserted(ErrorCodes::InternalError, "skipShardFilteringMetadataRefresh failpoint");
+    }
+
     auto* const shardingState = ShardingState::get(opCtx);
     invariant(shardingState->canAcceptShardedCommands());
 
diff --git a/src/mongo/db/s/shard_filtering_metadata_refresh.h b/src/mongo/db/s/shard_filtering_metadata_refresh.h
index f378c60622f..0576403b0fd 100644
--- a/src/mongo/db/s/shard_filtering_metadata_refresh.h
+++ b/src/mongo/db/s/shard_filtering_metadata_refresh.h
@@ -56,8 +56,7 @@ class OperationContext;
  */
 Status onShardVersionMismatchNoExcept(OperationContext* opCtx,
                                       const NamespaceString& nss,
-                                      ChunkVersion shardVersionReceived,
-                                      bool forceRefreshFromThisThread = false) noexcept;
+                                      ChunkVersion shardVersionReceived) noexcept;
 
 /**
  * Unconditionally causes the shard's filtering metadata to be refreshed from the config server and
@@ -84,4 +83,24 @@ Status onDbVersionMismatchNoExcept(
 
 void forceDatabaseRefresh(OperationContext* opCtx, const StringData dbName);
 
+/**
+ * RAII-style class that enters the migration critical section and refresh the filtering
+ * metadata for the specified collection. The critical section is released when this object
+ * goes out of scope.
+ */
+class ScopedShardVersionCriticalSection {
+    ScopedShardVersionCriticalSection(const ScopedShardVersionCriticalSection&) = delete;
+    ScopedShardVersionCriticalSection& operator=(const ScopedShardVersionCriticalSection&) = delete;
+
+public:
+    ScopedShardVersionCriticalSection(OperationContext* opCtx, NamespaceString nss);
+    ~ScopedShardVersionCriticalSection();
+
+    void enterCommitPhase();
+
+private:
+    OperationContext* const _opCtx;
+    const NamespaceString _nss;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/shardsvr_shard_collection.cpp b/src/mongo/db/s/shardsvr_shard_collection.cpp
index 081033eccf1..c1a6cca5e22 100644
--- a/src/mongo/db/s/shardsvr_shard_collection.cpp
+++ b/src/mongo/db/s/shardsvr_shard_collection.cpp
@@ -48,7 +48,6 @@
 #include "mongo/db/s/collection_sharding_runtime.h"
 #include "mongo/db/s/config/initial_split_policy.h"
 #include "mongo/db/s/config/sharding_catalog_manager.h"
-#include "mongo/db/s/scoped_shard_version_critical_section.h"
 #include "mongo/db/s/shard_filtering_metadata_refresh.h"
 #include "mongo/db/s/shard_key_util.h"
 #include "mongo/db/s/sharding_logging.h"