SERVER-54030 Validate state transitions in TenantMigrationAccessBlocker

(cherry picked from commit 4c4b33ea47c92de59f7cb1d15f0aabd5e5512845)

5 files changed, 165 insertions, 66 deletions

diff --git a/src/mongo/db/repl/tenant_migration_donor_access_blocker.cpp b/src/mongo/db/repl/tenant_migration_donor_access_blocker.cpp
index cb8ac0a7466..9a45a4e3621 100644
--- a/src/mongo/db/repl/tenant_migration_donor_access_blocker.cpp
+++ b/src/mongo/db/repl/tenant_migration_donor_access_blocker.cpp
@@ -83,15 +83,15 @@ TenantMigrationDonorAccessBlocker::TenantMigrationDonorAccessBlocker(
 Status TenantMigrationDonorAccessBlocker::checkIfCanWrite() {
     stdx::lock_guard<Latch> lg(_mutex);
 
-    switch (_state) {
-        case State::kAllow:
-        case State::kAborted:
+    switch (_state.getState()) {
+        case BlockerState::State::kAllow:
+        case BlockerState::State::kAborted:
             return Status::OK();
-        case State::kBlockWrites:
-        case State::kBlockWritesAndReads:
+        case BlockerState::State::kBlockWrites:
+        case BlockerState::State::kBlockWritesAndReads:
             return {TenantMigrationConflictInfo(_tenantId, shared_from_this()),
                     "Write must block until this tenant migration commits or aborts"};
-        case State::kReject:
+        case BlockerState::State::kReject:
             return {ErrorCodes::TenantMigrationCommitted,
                     "Write must be re-routed to the new owner of this tenant"};
         default:
@@ -153,7 +153,7 @@ SharedSemiFuture<void> TenantMigrationDonorAccessBlocker::getCanReadFuture(Opera
     stdx::lock_guard<Latch> lk(_mutex);
     if (!readTimestamp) {
         if (!MONGO_unlikely(tenantMigrationDonorAllowsNonTimestampedReads.shouldFail()) &&
-            _state == State::kReject &&
+            _state.isReject() &&
             commandDenyListAfterMigration.find(command) != commandDenyListAfterMigration.end()) {
             LOGV2_DEBUG(5505100,
                         1,
@@ -168,14 +168,14 @@ SharedSemiFuture<void> TenantMigrationDonorAccessBlocker::getCanReadFuture(Opera
         }
     }
 
-    auto canRead = _state == State::kAllow || _state == State::kAborted ||
-        _state == State::kBlockWrites || *readTimestamp < *_blockTimestamp;
+    auto canRead = _state.isAllow() || _state.isAborted() || _state.isBlockWrites() ||
+        *readTimestamp < *_blockTimestamp;
 
     if (canRead) {
         return SharedSemiFuture<void>();
     }
 
-    if (_state == State::kReject) {
+    if (_state.isReject()) {
         return SharedSemiFuture<void>(
             Status(ErrorCodes::TenantMigrationCommitted,
                    "Read must be re-routed to the new owner of this tenant"));
@@ -188,7 +188,7 @@ SharedSemiFuture<void> TenantMigrationDonorAccessBlocker::getCanReadFuture(Opera
 Status TenantMigrationDonorAccessBlocker::checkIfLinearizableReadWasAllowed(
     OperationContext* opCtx) {
     stdx::lock_guard<Latch> lg(_mutex);
-    if (_state == State::kReject) {
+    if (_state.isReject()) {
         return {ErrorCodes::TenantMigrationCommitted,
                 "Read must be re-routed to the new owner of this tenant"};
     }
@@ -197,16 +197,16 @@ Status TenantMigrationDonorAccessBlocker::checkIfLinearizableReadWasAllowed(
 
 Status TenantMigrationDonorAccessBlocker::checkIfCanBuildIndex() {
     stdx::lock_guard<Latch> lg(_mutex);
-    switch (_state) {
-        case State::kAllow:
-        case State::kBlockWrites:
-        case State::kBlockWritesAndReads:
+    switch (_state.getState()) {
+        case BlockerState::State::kAllow:
+        case BlockerState::State::kBlockWrites:
+        case BlockerState::State::kBlockWritesAndReads:
             return {TenantMigrationConflictInfo(_tenantId, shared_from_this(), kIndexBuild),
                     "Index build must block until tenant migration is committed or aborted."};
-        case State::kReject:
+        case BlockerState::State::kReject:
             return {ErrorCodes::TenantMigrationCommitted,
                     "Index build must be re-routed to the new owner of this tenant"};
-        case State::kAborted:
+        case BlockerState::State::kAborted:
             return Status::OK();
     }
     MONGO_UNREACHABLE;
@@ -217,12 +217,11 @@ void TenantMigrationDonorAccessBlocker::startBlockingWrites() {
 
     LOGV2(5093800, "Tenant migration starting to block writes", "tenantId"_attr = _tenantId);
 
-    invariant(_state == State::kAllow);
     invariant(!_blockTimestamp);
     invariant(!_commitOpTime);
     invariant(!_abortOpTime);
 
-    _state = State::kBlockWrites;
+    _state.transitionTo(BlockerState::State::kBlockWrites);
 }
 
 void TenantMigrationDonorAccessBlocker::startBlockingReadsAfter(const Timestamp& blockTimestamp) {
@@ -233,23 +232,21 @@ void TenantMigrationDonorAccessBlocker::startBlockingReadsAfter(const Timestamp&
           "tenantId"_attr = _tenantId,
           "blockTimestamp"_attr = blockTimestamp);
 
-    invariant(_state == State::kBlockWrites);
     invariant(!_blockTimestamp);
     invariant(!_commitOpTime);
     invariant(!_abortOpTime);
 
-    _state = State::kBlockWritesAndReads;
+    _state.transitionTo(BlockerState::State::kBlockWritesAndReads);
     _blockTimestamp = blockTimestamp;
 }
 
 void TenantMigrationDonorAccessBlocker::rollBackStartBlocking() {
     stdx::lock_guard<Latch> lg(_mutex);
 
-    invariant(_state == State::kBlockWrites || _state == State::kBlockWritesAndReads);
     invariant(!_commitOpTime);
     invariant(!_abortOpTime);
 
-    _state = State::kAllow;
+    _state.transitionTo(BlockerState::State::kAllow);
     _blockTimestamp.reset();
     _transitionOutOfBlockingPromise.setFrom(Status::OK());
 }
@@ -259,7 +256,7 @@ void TenantMigrationDonorAccessBlocker::setCommitOpTime(OperationContext* opCtx,
     {
         stdx::lock_guard<Latch> lg(_mutex);
 
-        invariant(_state == State::kBlockWritesAndReads);
+        invariant(_state.isBlockWritesAndReads());
         invariant(!_commitOpTime);
         invariant(!_abortOpTime);
 
@@ -331,12 +328,11 @@ void TenantMigrationDonorAccessBlocker::onMajorityCommitPointUpdate(repl::OpTime
 
 void TenantMigrationDonorAccessBlocker::_onMajorityCommitCommitOpTime(
     stdx::unique_lock<Latch>& lk) {
-    invariant(_state == State::kBlockWritesAndReads);
     invariant(_blockTimestamp);
     invariant(_commitOpTime);
     invariant(!_abortOpTime);
 
-    _state = State::kReject;
+    _state.transitionTo(BlockerState::State::kReject);
     Status error{ErrorCodes::TenantMigrationCommitted,
                  "Write or read must be re-routed to the new owner of this tenant"};
     _completionPromise.setError(error);
@@ -349,12 +345,10 @@ void TenantMigrationDonorAccessBlocker::_onMajorityCommitCommitOpTime(
 }
 
 void TenantMigrationDonorAccessBlocker::_onMajorityCommitAbortOpTime(stdx::unique_lock<Latch>& lk) {
-    invariant(_state != State::kReject);
     invariant(!_commitOpTime);
-    invariant(_state != State::kAborted);
     invariant(_abortOpTime);
 
-    _state = State::kAborted;
+    _state.transitionTo(BlockerState::State::kAborted);
     _transitionOutOfBlockingPromise.setFrom(Status::OK());
     _completionPromise.setError({ErrorCodes::TenantMigrationAborted, "Tenant migration aborted"});
 
@@ -369,7 +363,7 @@ void TenantMigrationDonorAccessBlocker::appendInfoForServerStatus(BSONObjBuilder
     invariant(!_commitOpTime || !_abortOpTime);
 
     BSONObjBuilder tenantBuilder;
-    tenantBuilder.append("state", _stateToString(_state));
+    tenantBuilder.append("state", _state.toString());
     if (_blockTimestamp) {
         tenantBuilder.append("blockTimestamp", _blockTimestamp.get());
     }
@@ -383,23 +377,6 @@ void TenantMigrationDonorAccessBlocker::appendInfoForServerStatus(BSONObjBuilder
     builder->append(_tenantId, tenantBuilder.obj());
 }
 
-std::string TenantMigrationDonorAccessBlocker::_stateToString(State state) const {
-    switch (state) {
-        case State::kAllow:
-            return "allow";
-        case State::kBlockWrites:
-            return "blockWrites";
-        case State::kBlockWritesAndReads:
-            return "blockWritesAndReads";
-        case State::kReject:
-            return "reject";
-        case State::kAborted:
-            return "aborted";
-        default:
-            MONGO_UNREACHABLE;
-    }
-}
-
 BSONObj TenantMigrationDonorAccessBlocker::getDebugInfo() const {
     return BSON("tenantId" << _tenantId << "recipientConnectionString" << _recipientConnString);
 }
@@ -421,4 +398,69 @@ void TenantMigrationDonorAccessBlocker::Stats::report(BSONObjBuilder* builder) c
     builder->append("numTenantMigrationAbortedErrors", numTenantMigrationAbortedErrors.load());
 }
 
+std::string TenantMigrationDonorAccessBlocker::BlockerState::toString(State state) {
+    switch (state) {
+        case State::kAllow:
+            return "allow";
+        case State::kBlockWrites:
+            return "blockWrites";
+        case State::kBlockWritesAndReads:
+            return "blockWritesAndReads";
+        case State::kReject:
+            return "reject";
+        case State::kAborted:
+            return "aborted";
+        default:
+            MONGO_UNREACHABLE;
+    }
+}
+
+bool TenantMigrationDonorAccessBlocker::BlockerState::_isLegalTransition(State oldState,
+                                                                         State newState) {
+    switch (oldState) {
+        case State::kAllow:
+            switch (newState) {
+                case State::kBlockWrites:
+                case State::kAborted:
+                    return true;
+                default:
+                    return false;
+            }
+            MONGO_UNREACHABLE;
+        case State::kBlockWrites:
+            switch (newState) {
+                case State::kAllow:
+                case State::kBlockWritesAndReads:
+                case State::kAborted:
+                    return true;
+                default:
+                    return false;
+            }
+            MONGO_UNREACHABLE;
+        case State::kBlockWritesAndReads:
+            switch (newState) {
+                case State::kAllow:
+                case State::kReject:
+                case State::kAborted:
+                    return true;
+                default:
+                    return false;
+            }
+            MONGO_UNREACHABLE;
+        case State::kReject:
+            return false;
+        case State::kAborted:
+            return false;
+    }
+    MONGO_UNREACHABLE;
+}
+
+void TenantMigrationDonorAccessBlocker::BlockerState::transitionTo(State newState) {
+    invariant(BlockerState::_isLegalTransition(_state, newState),
+              str::stream() << "Current state: " << toString(_state)
+                            << ", Illegal attempted next state: " << toString(newState));
+
+    _state = newState;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/repl/tenant_migration_donor_access_blocker.h b/src/mongo/db/repl/tenant_migration_donor_access_blocker.h
index 54ffb5123b0..7cc22eca265 100644
--- a/src/mongo/db/repl/tenant_migration_donor_access_blocker.h
+++ b/src/mongo/db/repl/tenant_migration_donor_access_blocker.h
@@ -245,8 +245,47 @@ private:
     /**
      * The access states of an mtab.
      */
-    enum class State { kAllow, kBlockWrites, kBlockWritesAndReads, kReject, kAborted };
-    std::string _stateToString(State state) const;
+    class BlockerState {
+    public:
+        enum class State { kAllow, kBlockWrites, kBlockWritesAndReads, kReject, kAborted };
+
+        void transitionTo(State newState);
+
+        State getState() const {
+            return _state;
+        }
+
+        bool isAllow() const {
+            return _state == State::kAllow;
+        }
+
+        bool isBlockWrites() const {
+            return _state == State::kBlockWrites;
+        }
+
+        bool isBlockWritesAndReads() const {
+            return _state == State::kBlockWritesAndReads;
+        }
+
+        bool isReject() const {
+            return _state == State::kReject;
+        }
+
+        bool isAborted() const {
+            return _state == State::kAborted;
+        }
+
+        std::string toString() const {
+            return toString(_state);
+        }
+
+        static std::string toString(State state);
+
+    private:
+        static bool _isLegalTransition(State oldState, State newState);
+
+        State _state = State::kAllow;
+    };
 
     /**
      * Encapsulates runtime statistics on blocked reads and writes, and tenant migration errors
@@ -279,7 +318,7 @@ private:
     // Protects the state below.
     mutable Mutex _mutex = MONGO_MAKE_LATCH("TenantMigrationDonorAccessBlocker::_mutex");
 
-    State _state{State::kAllow};
+    BlockerState _state;
 
     boost::optional<Timestamp> _blockTimestamp;
     boost::optional<repl::OpTime> _commitOpTime;
diff --git a/src/mongo/db/repl/tenant_migration_recipient_access_blocker.cpp b/src/mongo/db/repl/tenant_migration_recipient_access_blocker.cpp
index d164ca37330..37848c63b05 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_access_blocker.cpp
+++ b/src/mongo/db/repl/tenant_migration_recipient_access_blocker.cpp
@@ -105,11 +105,11 @@ SharedSemiFuture<void> TenantMigrationRecipientAccessBlocker::getCanReadFuture(
     }();
 
     stdx::lock_guard<Latch> lk(_mutex);
-    if (_state == State::kReject) {
+    if (_state.isReject()) {
         return SharedSemiFuture<void>(Status(
             ErrorCodes::SnapshotTooOld, "Tenant read is not allowed before migration completes"));
     }
-    invariant(_state == State::kRejectBefore);
+    invariant(_state.isRejectBefore());
     invariant(_rejectBeforeTimestamp);
     if (atClusterTime && *atClusterTime < *_rejectBeforeTimestamp) {
         return SharedSemiFuture<void>(Status(
@@ -168,7 +168,7 @@ void TenantMigrationRecipientAccessBlocker::appendInfoForServerStatus(
     stdx::lock_guard<Latch> lg(_mutex);
 
     BSONObjBuilder tenantBuilder;
-    tenantBuilder.append("state", _stateToString(_state));
+    tenantBuilder.append("state", _state.toString());
     if (_rejectBeforeTimestamp) {
         tenantBuilder.append("rejectBeforeTimestamp", _rejectBeforeTimestamp.get());
     }
@@ -176,8 +176,8 @@ void TenantMigrationRecipientAccessBlocker::appendInfoForServerStatus(
     builder->append(_tenantId, tenantBuilder.obj());
 }
 
-std::string TenantMigrationRecipientAccessBlocker::_stateToString(State state) const {
-    switch (state) {
+std::string TenantMigrationRecipientAccessBlocker::BlockerState::toString() const {
+    switch (_state) {
         case State::kReject:
             return "reject";
         case State::kRejectBefore:
@@ -197,7 +197,7 @@ BSONObj TenantMigrationRecipientAccessBlocker::getDebugInfo() const {
 
 void TenantMigrationRecipientAccessBlocker::startRejectingReadsBefore(const Timestamp& timestamp) {
     stdx::lock_guard<Latch> lk(_mutex);
-    _state = State::kRejectBefore;
+    _state.transitionToRejectBefore();
     if (!_rejectBeforeTimestamp || timestamp > *_rejectBeforeTimestamp) {
         LOGV2(5358100,
               "Tenant migration recipient starting to reject reads before timestamp",
diff --git a/src/mongo/db/repl/tenant_migration_recipient_access_blocker.h b/src/mongo/db/repl/tenant_migration_recipient_access_blocker.h
index e06a6a82bf5..7decbe897b9 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_access_blocker.h
+++ b/src/mongo/db/repl/tenant_migration_recipient_access_blocker.h
@@ -72,11 +72,6 @@ class TenantMigrationRecipientAccessBlocker
     : public std::enable_shared_from_this<TenantMigrationRecipientAccessBlocker>,
       public TenantMigrationAccessBlocker {
 public:
-    /**
-     * The access states of an mtab.
-     */
-    enum class State { kReject, kRejectBefore };
-
     TenantMigrationRecipientAccessBlocker(ServiceContext* serviceContext,
                                           UUID migrationId,
                                           std::string tenantId,
@@ -126,12 +121,35 @@ public:
     //
     void startRejectingReadsBefore(const Timestamp& timestamp);
 
-    State getState() const {
-        return _state;
+    bool inStateReject() const {
+        return _state.isReject();
     }
 
 private:
-    std::string _stateToString(State state) const;
+    /**
+     * The access states of an mtab.
+     */
+    class BlockerState {
+    public:
+        void transitionToRejectBefore() {
+            _state = State::kRejectBefore;
+        }
+
+        bool isReject() const {
+            return _state == State::kReject;
+        }
+
+        bool isRejectBefore() const {
+            return _state == State::kRejectBefore;
+        }
+
+        std::string toString() const;
+
+    private:
+        enum class State { kReject, kRejectBefore };
+
+        State _state = State::kReject;
+    };
 
     ServiceContext* _serviceContext;
     const UUID _migrationId;
@@ -141,7 +159,7 @@ private:
     // Protects the state below.
     mutable Mutex _mutex = MONGO_MAKE_LATCH("TenantMigrationRecipientAccessBlocker::_mutex");
 
-    State _state{State::kReject};
+    BlockerState _state;
 
     boost::optional<Timestamp> _rejectBeforeTimestamp;
 
diff --git a/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp b/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
index d84d706e6b0..e16004be680 100644
--- a/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
+++ b/src/mongo/db/repl/tenant_migration_recipient_op_observer.cpp
@@ -96,7 +96,7 @@ void TenantMigrationRecipientOpObserver::onUpdate(OperationContext* opCtx,
                 opCtx->getServiceContext(), recipientStateDoc.getTenantId());
 
             if (recipientStateDoc.getExpireAt() && mtab) {
-                if (mtab->getState() == TenantMigrationRecipientAccessBlocker::State::kReject) {
+                if (mtab->inStateReject()) {
                     // The TenantMigrationRecipientAccessBlocker entry needs to be removed to
                     // re-allow reads and future migrations with the same tenantId as this migration
                     // has already been aborted and forgotten.