SERVER-37591 Change MigrationSourceManager cleanup to only remove the MSM from the CSR when necessary

1 files changed, 12 insertions, 2 deletions

diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 87594ccad6d..799f062f219 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -259,6 +259,7 @@ Status MigrationSourceManager::startClone(OperationContext* opCtx) {
             _args, metadata->getKeyPattern(), _donorConnStr, _recipientHost);
 
         invariant(nullptr == std::exchange(msmForCsr(css), this));
+        _state = kCloning;
     }
 
     Status startCloneStatus = _cloneDriver->startClone(opCtx);
@@ -266,7 +267,6 @@ Status MigrationSourceManager::startClone(OperationContext* opCtx) {
         return startCloneStatus;
     }
 
-    _state = kCloning;
     scopedGuard.Dismiss();
     return Status::OK();
 }
@@ -695,7 +695,17 @@ void MigrationSourceManager::_cleanup(OperationContext* opCtx) {
         AutoGetCollection autoColl(opCtx, getNss(), MODE_IX, MODE_X);
         auto* const css = CollectionShardingRuntime::get(opCtx, getNss());
 
-        invariant(this == std::exchange(msmForCsr(css), nullptr));
+        // In the kCreated state there should be no state to clean up, but we can verify this
+        // just to be safe.
+        if (_state == kCreated) {
+            // Verify that we did not set the MSM on the CSR.
+            invariant(!msmForCsr(css));
+            // Verify that the clone driver was not initialized.
+            invariant(!_cloneDriver);
+        } else {
+            auto oldMsmOnCsr = std::exchange(msmForCsr(css), nullptr);
+            invariant(this == oldMsmOnCsr);
+        }
         _critSec.reset();
         return std::move(_cloneDriver);
     }();