From 2dd30ec036cccc4fe03339da705355e023eacfc1 Mon Sep 17 00:00:00 2001
From: Fullstop000 <fullstop1005@gmail.com>
Date: Thu, 17 Jan 2019 14:58:16 +0800
Subject: [PATCH] fix: correct learner checking

---
 src/raft.rs                          | 14 +++++++-------
 tests/integration_cases/test_raft.rs | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/raft.rs b/src/raft.rs
index fa91683a1..c768124b9 100644
--- a/src/raft.rs
+++ b/src/raft.rs
@@ -1203,13 +1203,13 @@ impl<T: Storage> Raft<T> {
         }
     }
 
-    fn handle_transfer_leader(&mut self, m: &Message, pr: &mut Progress) {
-        if self.is_learner {
-            debug!("{} is learner. Ignored transferring leadership", self.tag);
+    fn handle_transfer_leader(&mut self, m: &Message, prs: &mut ProgressSet) {
+        let from = m.get_from();
+        if prs.learner_ids().contains(&from) {
+            debug!("{} is learner. Ignored transferring leadership", from);
             return;
         }
-
-        let lead_transferee = m.get_from();
+        let lead_transferee = from;
         let last_lead_transferee = self.lead_transferee;
         if last_lead_transferee.is_some() {
             if last_lead_transferee.unwrap() == lead_transferee {
@@ -1244,6 +1244,7 @@ impl<T: Storage> Raft<T> {
         // so reset r.electionElapsed.
         self.election_elapsed = 0;
         self.lead_transferee = Some(lead_transferee);
+        let pr = prs.get_mut(from).unwrap();
         if pr.matched == self.raft_log.last_index() {
             self.send_timeout_now(lead_transferee);
             info!(
@@ -1323,8 +1324,7 @@ impl<T: Storage> Raft<T> {
                 );
             }
             MessageType::MsgTransferLeader => {
-                let pr = prs.get_mut(m.get_from()).unwrap();
-                self.handle_transfer_leader(m, pr);
+                self.handle_transfer_leader(m, &mut prs);
             }
             _ => {}
         }
diff --git a/tests/integration_cases/test_raft.rs b/tests/integration_cases/test_raft.rs
index ada967ca3..ab2f85cd3 100644
--- a/tests/integration_cases/test_raft.rs
+++ b/tests/integration_cases/test_raft.rs
@@ -3274,6 +3274,23 @@ fn test_leader_transfer_to_non_existing_node() {
     check_leader_transfer_state(&nt.peers[&1], StateRole::Leader, 1);
 }
 
+#[test]
+fn test_leader_transfer_to_learner() {
+    setup_for_test();
+    let mut leader_config = new_test_config(1, vec![1], 10, 1);
+    leader_config.learners = vec![2];
+    let leader = new_test_raft_with_config(&leader_config, new_storage());
+    let mut learner_config = new_test_config(2, vec![1], 10, 1);
+    learner_config.learners = vec![2];
+    let learner = new_test_raft_with_config(&learner_config, new_storage());
+    let mut nt = Network::new(vec![Some(leader), Some(learner)]);
+    nt.send(vec![new_message(1, 1, MessageType::MsgHup, 0)]);
+
+    // Transfer leadership to learner node, there will be noop.
+    nt.send(vec![new_message(2, 1, MessageType::MsgTransferLeader, 0)]);
+    check_leader_transfer_state(&nt.peers[&1], StateRole::Leader, 1);
+}
+
 #[test]
 fn test_leader_transfer_timeout() {
     setup_for_test();