From d1301de8713ca8796e3144caa1c86ef4237fb42b Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Tue, 20 Feb 2024 16:41:41 +0000
Subject: [PATCH] Do not reset the electionElapsed if the node doesn't grant
 vote

If the local node receives a MsgVote message with higher
term but it doesn't grant the vote; it turns into a follower,
but it shouldn't reset the electionElapsed, to ensure it
has higher priority to start a campaign in the next round
of election. If we reject a node, it's highly likely we
will reject it again if it immediately campaigns again.
So it may waste a long time to elect a leader if we reset
the electionElapsed.

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>
---
 raft.go      | 24 ++++++++++++++++++++++++
 raft_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/raft.go b/raft.go
index 3357ae45..60e6eaf5 100644
--- a/raft.go
+++ b/raft.go
@@ -1058,12 +1058,24 @@ func (r *raft) poll(id uint64, t pb.MessageType, v bool) (granted int, rejected
 }
 
 func (r *raft) Step(m pb.Message) error {
+	var (
+		bakElectionElapsed           int
+		bakRandomizedElectionTimeout int
+		higherTerm                   bool
+	)
+
 	// Handle the message term, which may result in our stepping down to a follower.
 	switch {
 	case m.Term == 0:
 		// local message
 	case m.Term > r.Term:
+		higherTerm = true
+
 		if m.Type == pb.MsgVote || m.Type == pb.MsgPreVote {
+			// backup the electionElapsed and randomizedElectionTimeout, we
+			// will restore the values if we don't grant the Vote.
+			bakElectionElapsed, bakRandomizedElectionTimeout = r.electionElapsed, r.randomizedElectionTimeout
+
 			force := bytes.Equal(m.Context, []byte(campaignTransfer))
 			inLease := r.checkQuorum && r.lead != None && r.electionElapsed < r.electionTimeout
 			if !force && inLease {
@@ -1221,6 +1233,18 @@ func (r *raft) Step(m pb.Message) error {
 				r.Vote = m.From
 			}
 		} else {
+			// If the local node receives a `MsgVote` message with higher
+			// term, but it doesn't grant the vote; it turns into a follower,
+			// but it shouldn't reset the electionElapsed, to ensure it
+			// has higher priority to start a campaign in the next round
+			// of election. If we reject a node, it's highly likely we
+			// will reject it again if it immediately campaigns again.
+			// So it may waste a long time to elect a leader if we reset
+			// the electionElapsed.
+			if higherTerm && m.Type == pb.MsgVote {
+				r.electionElapsed = bakElectionElapsed
+				r.randomizedElectionTimeout = bakRandomizedElectionTimeout
+			}
 			r.logger.Infof("%x [logterm: %d, index: %d, vote: %x] rejected %s from %x [logterm: %d, index: %d] at term %d",
 				r.id, lastID.term, lastID.index, r.Vote, m.Type, m.From, candLastID.term, candLastID.index, r.Term)
 			r.send(pb.Message{To: m.From, Term: r.Term, Type: voteRespMsgType(m.Type), Reject: true})
diff --git a/raft_test.go b/raft_test.go
index 5a258e56..5e2e4a13 100644
--- a/raft_test.go
+++ b/raft_test.go
@@ -1241,6 +1241,39 @@ func TestPastElectionTimeout(t *testing.T) {
 	}
 }
 
+func TestElectionElapsedOnRejectVote(t *testing.T) {
+	testCases := []struct {
+		electionElapsed           int
+		randomizedElectionTimeout int
+	}{
+		{18, 28},
+		{12, 30},
+		{7, 15},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			storage := newTestMemoryStorage(withPeers(1, 2))
+			storage.Append(index(1).terms(1, 2, 3, 4, 5))
+			r := newTestRaft(1, 10, 1, storage)
+
+			lastEntryId := r.raftLog.lastEntryID()
+			r.Term = lastEntryId.term
+
+			r.electionElapsed = tc.electionElapsed
+			r.randomizedElectionTimeout = tc.randomizedElectionTimeout
+
+			// ensure the MsgVote message has a higher term, but with stale data,
+			// so that the vote will be rejected.
+			err := r.Step(pb.Message{From: 2, To: 1, Term: lastEntryId.term + 1, Type: pb.MsgVote, LogTerm: lastEntryId.index - 1, Index: lastEntryId.index - 1})
+			require.NoError(t, err)
+
+			require.Equal(t, tc.electionElapsed, r.electionElapsed)
+			require.Equal(t, tc.randomizedElectionTimeout, r.randomizedElectionTimeout)
+		})
+	}
+}
+
 // TestStepIgnoreOldTermMsg to ensure that the Step function ignores the message
 // from old term and does not pass it to the actual stepX function.
 func TestStepIgnoreOldTermMsg(t *testing.T) {