SERVER-35723 Make the collection critical section into an RAII class

4 files changed, 44 insertions, 17 deletions

diff --git a/src/mongo/db/s/collection_sharding_state.cpp b/src/mongo/db/s/collection_sharding_state.cpp
index fc925d5d153..ebbfc83891b 100644
--- a/src/mongo/db/s/collection_sharding_state.cpp
+++ b/src/mongo/db/s/collection_sharding_state.cpp
@@ -387,4 +387,21 @@ boost::optional<ChunkRange> CollectionShardingState::getNextOrphanRange(BSONObj
     return _metadataManager->getNextOrphanRange(from);
 }
 
+CollectionCriticalSection::CollectionCriticalSection(OperationContext* opCtx, NamespaceString ns)
+    : _nss(std::move(ns)), _opCtx(opCtx) {
+    AutoGetCollection autoColl(_opCtx, _nss, MODE_IX, MODE_X);
+    CollectionShardingState::get(opCtx, _nss)->enterCriticalSectionCatchUpPhase(_opCtx);
+}
+
+CollectionCriticalSection::~CollectionCriticalSection() {
+    UninterruptibleLockGuard noInterrupt(_opCtx->lockState());
+    AutoGetCollection autoColl(_opCtx, _nss, MODE_IX, MODE_X);
+    CollectionShardingState::get(_opCtx, _nss)->exitCriticalSection(_opCtx);
+}
+
+void CollectionCriticalSection::enterCommitPhase() {
+    AutoGetCollection autoColl(_opCtx, _nss, MODE_IX, MODE_X);
+    CollectionShardingState::get(_opCtx, _nss)->enterCriticalSectionCommitPhase(_opCtx);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/collection_sharding_state.h b/src/mongo/db/s/collection_sharding_state.h
index 3af06506280..501ae2477c7 100644
--- a/src/mongo/db/s/collection_sharding_state.h
+++ b/src/mongo/db/s/collection_sharding_state.h
@@ -202,4 +202,26 @@ private:
         -> boost::optional<Date_t>;
 };
 
+/**
+ * RAII-style class, which obtains a reference to the critical section for the
+ * specified collection.
+ */
+class CollectionCriticalSection {
+    MONGO_DISALLOW_COPYING(CollectionCriticalSection);
+
+public:
+    CollectionCriticalSection(OperationContext* opCtx, NamespaceString ns);
+    ~CollectionCriticalSection();
+
+    /**
+     * Enters the commit phase of the critical section and blocks reads.
+     */
+    void enterCommitPhase();
+
+private:
+    NamespaceString _nss;
+
+    OperationContext* _opCtx;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 259b10ebc6c..7391648af31 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -317,17 +317,7 @@ Status MigrationSourceManager::enterCriticalSection(OperationContext* opCtx) {
         return status;
     }
 
-    {
-        // The critical section must be entered with collection X lock in order to ensure there are
-        // no writes which could have entered and passed the version check just before we entered
-        // the crticial section, but managed to complete after we left it.
-        UninterruptibleLockGuard noInterrupt(opCtx->lockState());
-        AutoGetCollection autoColl(opCtx, getNss(), MODE_IX, MODE_X);
-
-        // IMPORTANT: After this line, the critical section is in place and needs to be signaled
-        CollectionShardingState::get(opCtx, _args.getNss())
-            ->enterCriticalSectionCatchUpPhase(opCtx);
-    }
+    _critSec.emplace(opCtx, _args.getNss());
 
     _state = kCriticalSection;
 
@@ -419,11 +409,7 @@ Status MigrationSourceManager::commitChunkMetadataOnConfig(OperationContext* opC
 
     // Read operations must begin to wait on the critical section just before we send the commit
     // operation to the config server
-    {
-        UninterruptibleLockGuard noInterrupt(opCtx->lockState());
-        AutoGetCollection autoColl(opCtx, getNss(), MODE_IX, MODE_X);
-        CollectionShardingState::get(opCtx, _args.getNss())->enterCriticalSectionCommitPhase(opCtx);
-    }
+    _critSec->enterCommitPhase();
 
     Timer t;
 
@@ -687,7 +673,7 @@ void MigrationSourceManager::_cleanup(OperationContext* opCtx) {
         auto css = CollectionShardingState::get(opCtx, getNss());
 
         invariant(this == std::exchange(msmForCss(css), nullptr));
-        css->exitCriticalSection(opCtx);
+        _critSec.reset();
         return std::move(_cloneDriver);
     }();
 
diff --git a/src/mongo/db/s/migration_source_manager.h b/src/mongo/db/s/migration_source_manager.h
index c15daec931d..e4e91d8b0fb 100644
--- a/src/mongo/db/s/migration_source_manager.h
+++ b/src/mongo/db/s/migration_source_manager.h
@@ -234,6 +234,8 @@ private:
 
     // The statistics about a chunk migration to be included in moveChunk.commit
     BSONObj _recipientCloneCounts;
+
+    boost::optional<CollectionCriticalSection> _critSec;
 };
 
 }  // namespace mongo