SERVER-48938: Allow primary-elect to complete drain mode even if it is stepping down unconditionally

(cherry picked from commit f11b0351c33d2888607eebd1748d524d241fc9ba)

6 files changed, 17 insertions, 35 deletions

diff --git a/jstests/replsets/step_down_during_draining2.js b/jstests/replsets/step_down_during_draining2.js
index db37db7aada..5b972cfb51a 100644
--- a/jstests/replsets/step_down_during_draining2.js
+++ b/jstests/replsets/step_down_during_draining2.js
@@ -119,16 +119,6 @@ assert.soon(function() {
     return secondary.getDB("foo").foo.find().itcount() == numDocuments;
 });
 
-// Even though it finished draining its buffer, it shouldn't be able to exit drain mode due to
-// pending stepdown.
-assert.commandFailedWithCode(
-    secondary.adminCommand({
-        replSetTest: 1,
-        waitForDrainFinish: 5000,
-    }),
-    ErrorCodes.ExceededTimeLimit,
-    'replSetTest waitForDrainFinish should time out when in the middle of stepping down');
-
 jsTestLog("Checking that node is PRIMARY but not master");
 assert.eq(ReplSetTest.State.PRIMARY, secondary.adminCommand({replSetGetStatus: 1}).myState);
 assert(!secondary.adminCommand('ismaster').ismaster);
diff --git a/src/mongo/db/repl/replication_coordinator_impl.cpp b/src/mongo/db/repl/replication_coordinator_impl.cpp
index 6e839f37be2..a5753199a6b 100644
--- a/src/mongo/db/repl/replication_coordinator_impl.cpp
+++ b/src/mongo/db/repl/replication_coordinator_impl.cpp
@@ -1067,12 +1067,7 @@ void ReplicationCoordinatorImpl::signalDrainComplete(OperationContext* opCtx,
         OpTime firstOpTime = _externalState->onTransitionToPrimary(opCtx);
         lk.lock();
 
-        auto status = _topCoord->completeTransitionToPrimary(firstOpTime);
-        if (status.code() == ErrorCodes::PrimarySteppedDown) {
-            log() << "Transition to primary failed" << causedBy(status);
-            return;
-        }
-        invariant(status);
+        _topCoord->completeTransitionToPrimary(firstOpTime);
     }
 
     // Must calculate the commit level again because firstOpTimeOfMyTerm wasn't set when we logged
diff --git a/src/mongo/db/repl/replication_coordinator_impl_test.cpp b/src/mongo/db/repl/replication_coordinator_impl_test.cpp
index 9b347f968a1..34038462e90 100644
--- a/src/mongo/db/repl/replication_coordinator_impl_test.cpp
+++ b/src/mongo/db/repl/replication_coordinator_impl_test.cpp
@@ -1658,6 +1658,9 @@ TEST_F(ReplCoordTest, DrainCompletionMidStepDown) {
     ASSERT(updateTermEvh.isValid());
     ASSERT(termUpdated == TopologyCoordinator::UpdateTermResult::kTriggerStepDown);
 
+    // Set 'firstOpTimeOfMyTerm' to have term 1, so that the node will see that the noop entry has
+    // the correct term at the end of signalDrainComplete.
+    getExternalState()->setFirstOpTimeOfMyTerm(OpTime(Timestamp(100, 1), 1));
     // Now signal that replication applier is finished draining its buffer.
     const auto opCtx = makeOperationContext();
     getReplCoord()->signalDrainComplete(opCtx.get(), getReplCoord()->getTerm());
diff --git a/src/mongo/db/repl/topology_coordinator.cpp b/src/mongo/db/repl/topology_coordinator.cpp
index 98b88dad64c..3a563899472 100644
--- a/src/mongo/db/repl/topology_coordinator.cpp
+++ b/src/mongo/db/repl/topology_coordinator.cpp
@@ -2553,24 +2553,22 @@ bool TopologyCoordinator::canCompleteTransitionToPrimary(long long termWhenDrain
     }
     // Allow completing the transition to primary even when in the middle of a stepdown attempt,
     // in case the stepdown attempt fails.
-    if (_leaderMode != LeaderMode::kLeaderElect && _leaderMode != LeaderMode::kAttemptingStepDown) {
+    if (_leaderMode != LeaderMode::kLeaderElect && _leaderMode != LeaderMode::kAttemptingStepDown &&
+        _leaderMode != LeaderMode::kSteppingDown) {
         return false;
     }
 
     return true;
 }
 
-Status TopologyCoordinator::completeTransitionToPrimary(const OpTime& firstOpTimeOfTerm) {
-    if (!canCompleteTransitionToPrimary(firstOpTimeOfTerm.getTerm())) {
-        return Status(ErrorCodes::PrimarySteppedDown,
-                      "By the time this node was ready to complete its transition to PRIMARY it "
-                      "was no longer eligible to do so");
-    }
+void TopologyCoordinator::completeTransitionToPrimary(const OpTime& firstOpTimeOfTerm) {
+    invariant(canCompleteTransitionToPrimary(firstOpTimeOfTerm.getTerm()));
+
     if (_leaderMode == LeaderMode::kLeaderElect) {
         _setLeaderMode(LeaderMode::kMaster);
     }
+
     _firstOpTimeOfMyTerm = firstOpTimeOfTerm;
-    return Status::OK();
 }
 
 void TopologyCoordinator::adjustMaintenanceCountBy(int inc) {
diff --git a/src/mongo/db/repl/topology_coordinator.h b/src/mongo/db/repl/topology_coordinator.h
index ba17d534016..773442cc684 100644
--- a/src/mongo/db/repl/topology_coordinator.h
+++ b/src/mongo/db/repl/topology_coordinator.h
@@ -277,9 +277,8 @@ public:
      * "firstOpTimeOfTerm" is a floor on the OpTimes this node will be allowed to consider committed
      * for this tenure as primary. This prevents entries from before our election from counting as
      * committed in our view, until our election (the "firstOpTimeOfTerm" op) has been committed.
-     * Returns PrimarySteppedDown if this node is no longer eligible to begin accepting writes.
      */
-    Status completeTransitionToPrimary(const OpTime& firstOpTimeOfTerm);
+    void completeTransitionToPrimary(const OpTime& firstOpTimeOfTerm);
 
     /**
      * Adjusts the maintenance mode count by "inc".
diff --git a/src/mongo/db/repl/topology_coordinator_v1_test.cpp b/src/mongo/db/repl/topology_coordinator_v1_test.cpp
index d17dc192d57..df13a508780 100644
--- a/src/mongo/db/repl/topology_coordinator_v1_test.cpp
+++ b/src/mongo/db/repl/topology_coordinator_v1_test.cpp
@@ -114,7 +114,7 @@ protected:
         getTopoCoord().changeMemberState_forTest(MemberState::RS_PRIMARY, electionTimestamp);
         getTopoCoord().setCurrentPrimary_forTest(_selfIndex, electionTimestamp);
         OpTime dummyOpTime(Timestamp(1, 1), getTopoCoord().getTerm());
-        ASSERT_OK(getTopoCoord().completeTransitionToPrimary(dummyOpTime));
+        getTopoCoord().completeTransitionToPrimary(dummyOpTime);
     }
 
     void setMyOpTime(const OpTime& opTime, Date_t wallTime = Date_t()) {
@@ -5927,7 +5927,7 @@ TEST_F(HeartbeatResponseTestV1, NodeDoesNotStepDownSelfWhenRemoteNodeWasElectedL
     ASSERT_NO_ACTION(nextAction.getAction());
 }
 
-TEST_F(HeartbeatResponseTestV1, NodeWillNotTransitionToPrimaryAfterHearingAboutNewerTerm) {
+TEST_F(HeartbeatResponseTestV1, NodeWillCompleteTransitionToPrimaryAfterHearingAboutNewerTerm) {
     auto initialTerm = getTopoCoord().getTerm();
     OpTime firstOpTimeOfTerm(Timestamp(1, 1), initialTerm);
 
@@ -5936,17 +5936,14 @@ TEST_F(HeartbeatResponseTestV1, NodeWillNotTransitionToPrimaryAfterHearingAboutN
                                              firstOpTimeOfTerm.getTimestamp());
     getTopoCoord().setCurrentPrimary_forTest(getSelfIndex());
 
-    // At first transition to primary is OK
-    ASSERT(getTopoCoord().canCompleteTransitionToPrimary(initialTerm));
+    // Verify that transition to primary is OK.
+    ASSERT_TRUE(getTopoCoord().canCompleteTransitionToPrimary(initialTerm));
 
     // Now mark ourselves as mid-stepdown, as if we had heard about a new term.
     getTopoCoord().prepareForUnconditionalStepDown();
 
-    ASSERT_FALSE(getTopoCoord().canCompleteTransitionToPrimary(initialTerm));
-
-    // Check that transitioning to primary fails now that the term has been updated.
-    ASSERT_EQUALS(ErrorCodes::PrimarySteppedDown,
-                  getTopoCoord().completeTransitionToPrimary(firstOpTimeOfTerm));
+    // Verify that the transition to primary can still complete.
+    ASSERT_TRUE(getTopoCoord().canCompleteTransitionToPrimary(initialTerm));
 }
 
 TEST_F(HeartbeatResponseTestV1,