SERVER-29908 Move OpObserver and MigrationSourceManager logic out of CollectionShardingState.

Makes MigrationSourceManager a decoration on CollectionShardingState,
not a member and moves the op observer behavior from
CollectionShardingState to free functions in shard_server_op_observer.h/cpp.

9 files changed, 120 insertions, 158 deletions

diff --git a/src/mongo/db/s/active_migrations_registry.cpp b/src/mongo/db/s/active_migrations_registry.cpp
index 6c07939d7d7..0e64ee167b2 100644
--- a/src/mongo/db/s/active_migrations_registry.cpp
+++ b/src/mongo/db/s/active_migrations_registry.cpp
@@ -120,9 +120,9 @@ BSONObj ActiveMigrationsRegistry::getActiveMigrationStatusReport(OperationContex
         // Lock the collection so nothing changes while we're getting the migration report.
         AutoGetCollection autoColl(opCtx, nss.get(), MODE_IS);
 
-        auto css = CollectionShardingState::get(opCtx, nss.get());
-        if (css->getMigrationSourceManager()) {
-            return css->getMigrationSourceManager()->getMigrationStatusReport();
+        if (auto msm =
+                MigrationSourceManager::get(CollectionShardingState::get(opCtx, nss.get()))) {
+            return msm->getMigrationStatusReport();
         }
     }
 
diff --git a/src/mongo/db/s/collection_sharding_state.cpp b/src/mongo/db/s/collection_sharding_state.cpp
index 5bafa1f37fd..a19c93d0e64 100644
--- a/src/mongo/db/s/collection_sharding_state.cpp
+++ b/src/mongo/db/s/collection_sharding_state.cpp
@@ -37,7 +37,6 @@
 #include "mongo/db/catalog_raii.h"
 #include "mongo/db/client.h"
 #include "mongo/db/operation_context.h"
-#include "mongo/db/s/migration_source_manager.h"
 #include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/s/sharded_connection_info.h"
 #include "mongo/db/s/sharding_state.h"
@@ -162,10 +161,6 @@ CollectionShardingState::CollectionShardingState(ServiceContext* sc, NamespaceSt
       _metadataManager(std::make_shared<MetadataManager>(
           sc, _nss, getRangeDeleterExecutorHolder(sc).getOrCreateExecutor())) {}
 
-CollectionShardingState::~CollectionShardingState() {
-    invariant(!_sourceMgr);
-}
-
 CollectionShardingState* CollectionShardingState::get(OperationContext* opCtx,
                                                       const NamespaceString& nss) {
     return CollectionShardingState::get(opCtx, nss.ns());
@@ -238,22 +233,6 @@ void CollectionShardingState::exitCriticalSection(OperationContext* opCtx) {
     _critSec.exitCriticalSection();
 }
 
-void CollectionShardingState::setMigrationSourceManager(OperationContext* opCtx,
-                                                        MigrationSourceManager* sourceMgr) {
-    invariant(opCtx->lockState()->isCollectionLockedForMode(_nss.ns(), MODE_X));
-    invariant(sourceMgr);
-    invariant(!_sourceMgr);
-
-    _sourceMgr = sourceMgr;
-}
-
-void CollectionShardingState::clearMigrationSourceManager(OperationContext* opCtx) {
-    invariant(opCtx->lockState()->isCollectionLockedForMode(_nss.ns(), MODE_X));
-    invariant(_sourceMgr);
-
-    _sourceMgr = nullptr;
-}
-
 void CollectionShardingState::checkShardVersionOrThrow(OperationContext* opCtx) {
     std::string errmsg;
     ChunkVersion received;
@@ -330,50 +309,6 @@ boost::optional<ChunkRange> CollectionShardingState::getNextOrphanRange(BSONObj
     return _metadataManager->getNextOrphanRange(from);
 }
 
-void CollectionShardingState::onInsertOp(OperationContext* opCtx,
-                                         const BSONObj& insertedDoc,
-                                         const repl::OpTime& opTime) {
-    dassert(opCtx->lockState()->isCollectionLockedForMode(_nss.ns(), MODE_IX));
-
-    checkShardVersionOrThrow(opCtx);
-
-    if (_sourceMgr) {
-        _sourceMgr->getCloner()->onInsertOp(opCtx, insertedDoc, opTime);
-    }
-}
-
-void CollectionShardingState::onUpdateOp(OperationContext* opCtx,
-                                         const BSONObj& updatedDoc,
-                                         const repl::OpTime& opTime,
-                                         const repl::OpTime& prePostImageOpTime) {
-    dassert(opCtx->lockState()->isCollectionLockedForMode(_nss.ns(), MODE_IX));
-
-    checkShardVersionOrThrow(opCtx);
-
-    if (_sourceMgr) {
-        _sourceMgr->getCloner()->onUpdateOp(opCtx, updatedDoc, opTime, prePostImageOpTime);
-    }
-}
-
-auto CollectionShardingState::makeDeleteState(OperationContext* opCtx, BSONObj const& doc)
-    -> DeleteState {
-    return {getMetadata(opCtx).extractDocumentKey(doc).getOwned(),
-            _sourceMgr && _sourceMgr->getCloner()->isDocumentInMigratingChunk(doc)};
-}
-
-void CollectionShardingState::onDeleteOp(OperationContext* opCtx,
-                                         const DeleteState& deleteState,
-                                         const repl::OpTime& opTime,
-                                         const repl::OpTime& preImageOpTime) {
-    dassert(opCtx->lockState()->isCollectionLockedForMode(_nss.ns(), MODE_IX));
-
-    checkShardVersionOrThrow(opCtx);
-
-    if (_sourceMgr && deleteState.isMigrating) {
-        _sourceMgr->getCloner()->onDeleteOp(opCtx, deleteState.documentKey, opTime, preImageOpTime);
-    }
-}
-
 bool CollectionShardingState::_checkShardVersionOk(OperationContext* opCtx,
                                                    std::string* errmsg,
                                                    ChunkVersion* expectedShardVersion,
diff --git a/src/mongo/db/s/collection_sharding_state.h b/src/mongo/db/s/collection_sharding_state.h
index 6aa2463dc9d..2df1699c450 100644
--- a/src/mongo/db/s/collection_sharding_state.h
+++ b/src/mongo/db/s/collection_sharding_state.h
@@ -36,10 +36,10 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/s/metadata_manager.h"
 #include "mongo/db/s/sharding_migration_critical_section.h"
+#include "mongo/util/decorable.h"
 
 namespace mongo {
 
-class MigrationSourceManager;
 class OperationContext;
 
 /**
@@ -50,7 +50,7 @@ class OperationContext;
  * Synchronization rules: In order to look-up this object in the instance's sharding map, one must
  * have some lock on the respective collection.
  */
-class CollectionShardingState {
+class CollectionShardingState : public Decorable<CollectionShardingState> {
     MONGO_DISALLOW_COPYING(CollectionShardingState);
 
 public:
@@ -60,22 +60,6 @@ public:
      * Instantiates a new per-collection sharding state as unsharded.
      */
     CollectionShardingState(ServiceContext* sc, NamespaceString nss);
-    ~CollectionShardingState();
-
-    /**
-     * Details of documents being removed from a sharded collection.
-     */
-    struct DeleteState {
-        // Contains the fields of the document that are in the collection's shard key, and "_id".
-        BSONObj documentKey;
-
-        // True if the document being deleted belongs to a chunk which, while still in the shard,
-        // is being migrated out. (Not to be confused with "fromMigrate", which tags operations
-        // that are steps in performing the migration.)
-        bool isMigrating;
-    };
-
-    DeleteState makeDeleteState(OperationContext* opCtx, BSONObj const& doc);
 
     /**
      * Obtains the sharding state for the specified collection. If it does not exist, it will be
@@ -168,24 +152,6 @@ public:
     }
 
     /**
-     * Attaches a migration source manager to this collection's sharding state. Must be called with
-     * collection X lock. May not be called if there is a migration source manager already
-     * installed. Must be followed by a call to clearMigrationSourceManager.
-     */
-    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
-     * already through a previous call to setMigrationSourceManager.
-     */
-    void clearMigrationSourceManager(OperationContext* opCtx);
-
-    /**
      * Checks whether the shard version in the context is compatible with the shard version of the
      * collection locally and if not throws StaleConfigException populated with the expected and
      * actual versions.
@@ -226,24 +192,6 @@ public:
      */
     boost::optional<ChunkRange> getNextOrphanRange(BSONObj const& startingFrom);
 
-    /**
-     * Replication oplog OpObserver hooks. Informs the sharding system of changes that may be
-     * relevant to ongoing operations.
-     *
-     * The global exclusive lock is expected to be held by the caller of any of these functions.
-     */
-    void onInsertOp(OperationContext* opCtx,
-                    const BSONObj& insertedDoc,
-                    const repl::OpTime& opTime);
-    void onUpdateOp(OperationContext* opCtx,
-                    const BSONObj& updatedDoc,
-                    const repl::OpTime& opTime,
-                    const repl::OpTime& prePostImageOpTime);
-    void onDeleteOp(OperationContext* opCtx,
-                    const DeleteState& deleteState,
-                    const repl::OpTime& opTime,
-                    const repl::OpTime& preImageOpTime);
-
 private:
     /**
      * Checks whether the shard version of the operation matches that of the collection.
@@ -271,13 +219,6 @@ private:
 
     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.
-    //
-    // NOTE: The value is not owned by this class.
-    MigrationSourceManager* _sourceMgr{nullptr};
-
     // for access to _metadataManager
     friend auto CollectionRangeDeleter::cleanUpNextRange(OperationContext*,
                                                          NamespaceString const&,
diff --git a/src/mongo/db/s/collection_sharding_state_test.cpp b/src/mongo/db/s/collection_sharding_state_test.cpp
index 72f01c98f53..d416840861c 100644
--- a/src/mongo/db/s/collection_sharding_state_test.cpp
+++ b/src/mongo/db/s/collection_sharding_state_test.cpp
@@ -31,6 +31,7 @@
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/s/collection_sharding_state.h"
+#include "mongo/db/s/shard_server_op_observer.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/s/type_shard_identity.h"
 #include "mongo/s/shard_server_test_fixture.h"
@@ -156,7 +157,7 @@ TEST_F(DeleteStateTest, MakeDeleteStateUnsharded) {
 
     // First, check that an order for deletion from an unsharded collection (where css has not been
     // "refreshed" with chunk metadata) extracts just the "_id" field:
-    auto deleteState = css->makeDeleteState(operationContext(), doc);
+    auto deleteState = ShardObserverDeleteState::make(operationContext(), css, doc);
     ASSERT_BSONOBJ_EQ(deleteState.documentKey,
                       BSON("_id"
                            << "hello"));
@@ -181,7 +182,7 @@ TEST_F(DeleteStateTest, MakeDeleteStateShardedWithoutIdInShardKey) {
                     << true);
 
     // Verify the shard key is extracted, in correct order, followed by the "_id" field.
-    auto deleteState = css->makeDeleteState(operationContext(), doc);
+    auto deleteState = ShardObserverDeleteState::make(operationContext(), css, doc);
     ASSERT_BSONOBJ_EQ(deleteState.documentKey,
                       BSON("key" << 100 << "key3"
                                  << "abc"
@@ -207,7 +208,7 @@ TEST_F(DeleteStateTest, MakeDeleteStateShardedWithIdInShardKey) {
                            << 100);
 
     // Verify the shard key is extracted with "_id" in the right place.
-    auto deleteState = css->makeDeleteState(operationContext(), doc);
+    auto deleteState = ShardObserverDeleteState::make(operationContext(), css, doc);
     ASSERT_BSONOBJ_EQ(deleteState.documentKey,
                       BSON("key" << 100 << "_id"
                                  << "hello"
@@ -231,7 +232,7 @@ TEST_F(DeleteStateTest, MakeDeleteStateShardedWithIdHashInShardKey) {
                            << 100);
 
     // Verify the shard key is extracted with "_id" in the right place, not hashed.
-    auto deleteState = css->makeDeleteState(operationContext(), doc);
+    auto deleteState = ShardObserverDeleteState::make(operationContext(), css, doc);
     ASSERT_BSONOBJ_EQ(deleteState.documentKey,
                       BSON("_id"
                            << "hello"));
diff --git a/src/mongo/db/s/migration_chunk_cloner_source_legacy_commands.cpp b/src/mongo/db/s/migration_chunk_cloner_source_legacy_commands.cpp
index cecca24da1b..e56347589a2 100644
--- a/src/mongo/db/s/migration_chunk_cloner_source_legacy_commands.cpp
+++ b/src/mongo/db/s/migration_chunk_cloner_source_legacy_commands.cpp
@@ -70,15 +70,16 @@ public:
                 str::stream() << "Collection " << nss->ns() << " does not exist",
                 _autoColl->getCollection());
 
-        auto css = CollectionShardingState::get(opCtx, *nss);
-        uassert(ErrorCodes::IllegalOperation,
-                str::stream() << "No active migrations were found for collection " << nss->ns(),
-                css->getMigrationSourceManager());
-
-        // It is now safe to access the cloner
-        _chunkCloner = dynamic_cast<MigrationChunkClonerSourceLegacy*>(
-            css->getMigrationSourceManager()->getCloner());
-        invariant(_chunkCloner);
+        if (auto msm = MigrationSourceManager::get(CollectionShardingState::get(opCtx, *nss))) {
+            // It is now safe to access the cloner
+            _chunkCloner = dynamic_cast<MigrationChunkClonerSourceLegacy*>(msm->getCloner());
+            invariant(_chunkCloner);
+
+        } else {
+            uasserted(ErrorCodes::IllegalOperation,
+                      str::stream() << "No active migrations were found for collection "
+                                    << nss->ns());
+        }
 
         // Ensure the session ids are correct
         uassert(ErrorCodes::IllegalOperation,
diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 253496d2832..e68dc4c443f 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -68,6 +68,8 @@ using namespace shardmetadatautil;
 
 namespace {
 
+const auto msmForCss = CollectionShardingState::declareDecoration<MigrationSourceManager*>();
+
 // Wait at most this much time for the recipient to catch up sufficiently so critical section can be
 // entered
 const Hours kMaxWaitToEnterCriticalSectionTimeout(6);
@@ -129,6 +131,10 @@ MONGO_FP_DECLARE(failMigrationCommit);
 MONGO_FP_DECLARE(hangBeforeLeavingCriticalSection);
 MONGO_FP_DECLARE(migrationCommitNetworkError);
 
+MigrationSourceManager* MigrationSourceManager::get(CollectionShardingState& css) {
+    return msmForCss(css);
+}
+
 MigrationSourceManager::MigrationSourceManager(OperationContext* opCtx,
                                                MoveChunkRequest request,
                                                ConnectionString donorConnStr,
@@ -250,7 +256,7 @@ Status MigrationSourceManager::startClone(OperationContext* opCtx) {
         _cloneDriver = stdx::make_unique<MigrationChunkClonerSourceLegacy>(
             _args, metadata->getKeyPattern(), _donorConnStr, _recipientHost);
 
-        css->setMigrationSourceManager(opCtx, this);
+        invariant(nullptr == std::exchange(msmForCss(css), this));
     }
 
     Status startCloneStatus = _cloneDriver->startClone(opCtx);
@@ -686,18 +692,13 @@ void MigrationSourceManager::_cleanup(OperationContext* opCtx) {
     invariant(_state != kDone);
 
     auto cloneDriver = [&]() {
-        // Unregister from the collection's sharding state
+        // Unregister from the collection's sharding state and exit the migration critical section.
         UninterruptibleLockGuard noInterrupt(opCtx->lockState());
         AutoGetCollection autoColl(opCtx, getNss(), MODE_IX, MODE_X);
         auto css = CollectionShardingState::get(opCtx, getNss());
 
-        // The migration source manager is not visible anymore after it is unregistered from the
-        // collection
-        css->clearMigrationSourceManager(opCtx);
-
-        // Leave the critical section.
-        CollectionShardingState::get(opCtx, _args.getNss())->exitCriticalSection(opCtx);
-
+        invariant(this == std::exchange(msmForCss(css), nullptr));
+        css->exitCriticalSection(opCtx);
         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 9c37802441e..c15daec931d 100644
--- a/src/mongo/db/s/migration_source_manager.h
+++ b/src/mongo/db/s/migration_source_manager.h
@@ -70,6 +70,11 @@ class MigrationSourceManager {
     MONGO_DISALLOW_COPYING(MigrationSourceManager);
 
 public:
+    static MigrationSourceManager* get(CollectionShardingState& css);
+    static MigrationSourceManager* get(CollectionShardingState* css) {
+        return get(*css);
+    }
+
     /**
      * Instantiates a new migration source manager with the specified migration parameters. Must be
      * called with the distributed lock acquired in advance (not asserted).
diff --git a/src/mongo/db/s/shard_server_op_observer.cpp b/src/mongo/db/s/shard_server_op_observer.cpp
index e10df7ab070..37691cab462 100644
--- a/src/mongo/db/s/shard_server_op_observer.cpp
+++ b/src/mongo/db/s/shard_server_op_observer.cpp
@@ -36,6 +36,7 @@
 #include "mongo/db/catalog_raii.h"
 #include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/database_sharding_state.h"
+#include "mongo/db/s/migration_source_manager.h"
 #include "mongo/db/s/shard_identity_rollback_notifier.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/s/type_shard_identity.h"
@@ -49,10 +50,7 @@
 namespace mongo {
 namespace {
 
-using DeleteState = CollectionShardingState::DeleteState;
-
-const OperationContext::Decoration<DeleteState> getDeleteState =
-    OperationContext::declareDecoration<DeleteState>();
+const auto getDeleteState = OperationContext::declareDecoration<ShardObserverDeleteState>();
 
 bool isStandaloneOrPrimary(OperationContext* opCtx) {
     auto replCoord = repl::ReplicationCoordinator::get(opCtx);
@@ -327,9 +325,8 @@ void ShardServerOpObserver::onUpdate(OperationContext* opCtx, const OplogUpdateE
 void ShardServerOpObserver::aboutToDelete(OperationContext* opCtx,
                                           NamespaceString const& nss,
                                           BSONObj const& doc) {
-    auto& deleteState = getDeleteState(opCtx);
-    auto* css = CollectionShardingState::get(opCtx, nss.ns());
-    deleteState = css->makeDeleteState(opCtx, doc);
+    auto css = CollectionShardingState::get(opCtx, nss.ns());
+    getDeleteState(opCtx) = ShardObserverDeleteState::make(opCtx, css, doc);
 }
 
 void ShardServerOpObserver::onDelete(OperationContext* opCtx,
@@ -396,4 +393,47 @@ repl::OpTime ShardServerOpObserver::onDropCollection(OperationContext* opCtx,
     return {};
 }
 
+void shardObserveInsertOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const BSONObj& insertedDoc,
+                          const repl::OpTime& opTime) {
+    css->checkShardVersionOrThrow(opCtx);
+    auto msm = MigrationSourceManager::get(css);
+    if (msm) {
+        msm->getCloner()->onInsertOp(opCtx, insertedDoc, opTime);
+    }
+}
+
+void shardObserveUpdateOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const BSONObj& updatedDoc,
+                          const repl::OpTime& opTime,
+                          const repl::OpTime& prePostImageOpTime) {
+    css->checkShardVersionOrThrow(opCtx);
+    auto msm = MigrationSourceManager::get(css);
+    if (msm) {
+        msm->getCloner()->onUpdateOp(opCtx, updatedDoc, opTime, prePostImageOpTime);
+    }
+}
+
+void shardObserveDeleteOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const ShardObserverDeleteState& deleteState,
+                          const repl::OpTime& opTime,
+                          const repl::OpTime& preImageOpTime) {
+    css->checkShardVersionOrThrow(opCtx);
+    auto msm = MigrationSourceManager::get(css);
+    if (msm && deleteState.isMigrating) {
+        msm->getCloner()->onDeleteOp(opCtx, deleteState.documentKey, opTime, preImageOpTime);
+    }
+}
+
+ShardObserverDeleteState ShardObserverDeleteState::make(OperationContext* opCtx,
+                                                        CollectionShardingState* css,
+                                                        const BSONObj& docToDelete) {
+    auto msm = MigrationSourceManager::get(css);
+    return {css->getMetadata(opCtx).extractDocumentKey(docToDelete).getOwned(),
+            msm && msm->getCloner()->isDocumentInMigratingChunk(docToDelete)};
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/s/shard_server_op_observer.h b/src/mongo/db/s/shard_server_op_observer.h
index 4bccaacc3d9..5f5e6723516 100644
--- a/src/mongo/db/s/shard_server_op_observer.h
+++ b/src/mongo/db/s/shard_server_op_observer.h
@@ -30,6 +30,7 @@
 
 #include "mongo/base/disallow_copying.h"
 #include "mongo/db/op_observer.h"
+#include "mongo/db/s/collection_sharding_state.h"
 
 namespace mongo {
 
@@ -127,4 +128,41 @@ public:
     void onReplicationRollback(OperationContext* opCtx, const RollbackObserverInfo& rbInfo) {}
 };
 
+
+// Replication oplog OpObserver hooks. Informs the sharding system of changes that may be
+// relevant to ongoing operations.
+//
+// The global lock is expected to be held in mode IX by the caller of any of these functions.
+
+/**
+ * Details of documents being removed from a sharded collection.
+ */
+struct ShardObserverDeleteState {
+    static ShardObserverDeleteState make(OperationContext* opCtx,
+                                         CollectionShardingState* css,
+                                         const BSONObj& docToDelete);
+    // Contains the fields of the document that are in the collection's shard key, and "_id".
+    BSONObj documentKey;
+
+    // True if the document being deleted belongs to a chunk which, while still in the shard,
+    // is being migrated out. (Not to be confused with "fromMigrate", which tags operations
+    // that are steps in performing the migration.)
+    bool isMigrating;
+};
+
+void shardObserveInsertOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const BSONObj& insertedDoc,
+                          const repl::OpTime& opTime);
+void shardObserveUpdateOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const BSONObj& updatedDoc,
+                          const repl::OpTime& opTime,
+                          const repl::OpTime& prePostImageOpTime);
+void shardObserveDeleteOp(OperationContext* opCtx,
+                          CollectionShardingState* css,
+                          const ShardObserverDeleteState& deleteState,
+                          const repl::OpTime& opTime,
+                          const repl::OpTime& preImageOpTime);
+
 }  // namespace mongo