From 10c30572b0b4b8680c998370851cc28bc299798f Mon Sep 17 00:00:00 2001
From: Ibrahim Kettaneh <ibrahim.kettaneh@cockroachlabs.com>
Date: Wed, 27 Nov 2024 21:19:56 -0500
Subject: [PATCH] raft: make PreCandidate ignore MsgDefortifyLeader

Right now if a PreCandidate receives a message from a leader, it becomes
a follower again since there is an active leader. However,
MsgDefortifyLeader should be a special case because it doesn't really
say that there is an active leader since it's only sent by non-leaders
when they step down.

Fixes: #136087

Release note: None
---
 pkg/raft/raft.go      |  4 +++-
 pkg/raft/raft_test.go | 54 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/pkg/raft/raft.go b/pkg/raft/raft.go
index 07fa04cca178..0851f3090623 100644
--- a/pkg/raft/raft.go
+++ b/pkg/raft/raft.go
@@ -2127,9 +2127,11 @@ func stepLeader(r *raft, m pb.Message) error {
 // stepCandidate is shared by StateCandidate and StatePreCandidate; the difference is
 // whether they respond to MsgVoteResp or MsgPreVoteResp.
 func stepCandidate(r *raft, m pb.Message) error {
-	if IsMsgFromLeader(m.Type) {
+	if IsMsgFromLeader(m.Type) && m.Type != pb.MsgDeFortifyLeader {
 		// If this is a message from a leader of r.Term, transition to a follower
 		// with the sender of the message as the leader, then process the message.
+		// One exception is MsgDeFortifyLeader where it doesn't mean that there is
+		// currently an active leader for this term.
 		assertTrue(m.Term == r.Term, "message term should equal current term")
 		r.becomeFollower(m.Term, m.From)
 		return r.step(r, m) // stepFollower
diff --git a/pkg/raft/raft_test.go b/pkg/raft/raft_test.go
index d1bae66d5289..5be8a2c26d5d 100644
--- a/pkg/raft/raft_test.go
+++ b/pkg/raft/raft_test.go
@@ -2488,6 +2488,60 @@ func testDisruptiveFollowerPreVote(t *testing.T, storeLivenessEnabled bool) {
 	require.Equal(t, pb.StateLeader, n1.state)
 }
 
+// TestPreCandidateIgnoresDefortification tests that a pre-candidate ignores
+// MsgDefortifyLeader and doesn't become a follower again.
+func TestPreCandidateIgnoresDefortification(t *testing.T) {
+	var fabric *raftstoreliveness.LivenessFabric
+	var n1, n2 *raft
+
+	fabric = raftstoreliveness.NewLivenessFabricWithPeers(1, 2)
+	n1 = newTestRaft(1, 10, 1, newTestMemoryStorage(withPeers(1, 2)),
+		withStoreLiveness(fabric.GetStoreLiveness(1)))
+	n2 = newTestRaft(2, 10, 1, newTestMemoryStorage(withPeers(1, 2)),
+		withStoreLiveness(fabric.GetStoreLiveness(2)))
+
+	n1.checkQuorum = true
+	n2.checkQuorum = true
+	n1.preVote = true
+	n2.preVote = true
+
+	n1.becomeFollower(1, None)
+	n2.becomeFollower(1, None)
+
+	nt := newNetworkWithConfigAndLivenessFabric(nil, fabric, n1, n2)
+
+	nt.send(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
+
+	// Check raft states.
+	require.Equal(t, pb.StateLeader, n1.state)
+	require.Equal(t, pb.StateFollower, n2.state)
+
+	// The term is 2 for both nodes.
+	require.Equal(t, uint64(2), n1.Term)
+	require.Equal(t, uint64(2), n2.Term)
+
+	// Withdraw 2's support for 1. This allows 2 to pre-campaign since it's not
+	// supporting a fortified leader.
+	nt.livenessFabric.WithdrawSupportFor(2, 1)
+
+	// Isolate 1 so that it doesn't receive the MsgVoteRequest from 2, and
+	// therefore it doesn't vote for it.
+	nt.isolate(1)
+	nt.send(pb.Message{From: 2, To: 2, Type: pb.MsgHup})
+
+	// 2 is now a pre-candidate.
+	require.Equal(t, pb.StatePreCandidate, n2.state)
+
+	// 2 should remain a PreCandidate even if it receives a MsgDefortifyLeader.
+	nt.send(pb.Message{From: 1, To: 2, Term: 2, Type: pb.MsgDeFortifyLeader})
+	require.Equal(t, pb.StatePreCandidate, n2.state)
+
+	// However, receiving another message from a leader would cause 2 to become
+	// follower again.
+	nt.send(pb.Message{From: 1, To: 2, Term: 2, Type: pb.MsgApp})
+	require.Equal(t, pb.StateFollower, n2.state)
+}
+
 func TestLeaderAppResp(t *testing.T) {
 	// initial progress: match = 0; next = 3
 	tests := []struct {