SERVER-29908 Move the migration critical section out of MigrationSourceManager

1 files changed, 16 insertions, 2 deletions

diff --git a/src/mongo/db/s/collection_sharding_state.h b/src/mongo/db/s/collection_sharding_state.h
index 14b403cd93f..51ff7d26435 100644
--- a/src/mongo/db/s/collection_sharding_state.h
+++ b/src/mongo/db/s/collection_sharding_state.h
@@ -35,6 +35,7 @@
 #include "mongo/bson/bsonobj.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/s/metadata_manager.h"
+#include "mongo/db/s/sharding_migration_critical_section.h"
 
 namespace mongo {
 
@@ -153,9 +154,16 @@ public:
     std::vector<ScopedCollectionMetadata> overlappingMetadata(ChunkRange const& range) const;
 
     /**
-     * Returns the active migration source manager, if one is available.
+     * Methods to control the collection's critical section. Must be called with the collection X
+     * lock held.
      */
-    MigrationSourceManager* getMigrationSourceManager();
+    void enterCriticalSectionCatchUpPhase(OperationContext* opCtx);
+    void enterCriticalSectionCommitPhase(OperationContext* opCtx);
+    void exitCriticalSection(OperationContext* opCtx);
+
+    auto getCriticalSectionSignal(ShardingMigrationCriticalSection::Operation op) const {
+        return _critSec.getSignal(op);
+    }
 
     /**
      * Attaches a migration source manager to this collection's sharding state. Must be called with
@@ -164,6 +172,10 @@ public:
      */
     void setMigrationSourceManager(OperationContext* opCtx, MigrationSourceManager* sourceMgr);
 
+    auto getMigrationSourceManager() const {
+        return _sourceMgr;
+    }
+
     /**
      * Removes a migration source manager from this collection's sharding state. Must be called with
      * collection X lock. May not be called if there isn't a migration source manager installed
@@ -255,6 +267,8 @@ private:
     // Contains all the metadata associated with this collection.
     std::shared_ptr<MetadataManager> _metadataManager;
 
+    ShardingMigrationCriticalSection _critSec;
+
     // If this collection is serving as a source shard for chunk migration, this value will be
     // non-null. To write this value there needs to be X-lock on the collection in order to
     // synchronize with other callers, which read it.