From 32051c152e429a9003d162d00f9a9ff466ad09ce Mon Sep 17 00:00:00 2001
From: Simon Fell <fellforce@gmail.com>
Date: Tue, 19 Jan 2016 10:51:10 -0800
Subject: [PATCH] allow a node to vote for the current leader, fixes #74

---
 raft.go      | 17 +++++++++--------
 raft_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/raft.go b/raft.go
index 84d85e21057..07ddb808e8d 100644
--- a/raft.go
+++ b/raft.go
@@ -1407,10 +1407,11 @@ func (r *Raft) requestVote(rpc RPC, req *RequestVoteRequest) {
 		rpc.Respond(resp, rpcErr)
 	}()
 
-	// Check if we have an existing leader
-	if leader := r.Leader(); leader != "" {
-		r.logger.Printf("[WARN] raft: Rejecting vote from %v since we have a leader: %v",
-			r.trans.DecodePeer(req.Candidate), leader)
+	// Check if we have an existing leader [who's not the candidate]
+	candidate := r.trans.DecodePeer(req.Candidate)
+	if leader := r.Leader(); leader != "" && leader != candidate {
+		r.logger.Printf("[WARN] raft: Rejecting vote request from %v since we have a leader: %v",
+			candidate, leader)
 		return
 	}
 
@@ -1452,14 +1453,14 @@ func (r *Raft) requestVote(rpc RPC, req *RequestVoteRequest) {
 	// Reject if their term is older
 	lastIdx, lastTerm := r.getLastEntry()
 	if lastTerm > req.LastLogTerm {
-		r.logger.Printf("[WARN] raft: Rejecting vote from %v since our last term is greater (%d, %d)",
-			r.trans.DecodePeer(req.Candidate), lastTerm, req.LastLogTerm)
+		r.logger.Printf("[WARN] raft: Rejecting vote request from %v since our last term is greater (%d, %d)",
+			candidate, lastTerm, req.LastLogTerm)
 		return
 	}
 
 	if lastIdx > req.LastLogIndex {
-		r.logger.Printf("[WARN] raft: Rejecting vote from %v since our last index is greater (%d, %d)",
-			r.trans.DecodePeer(req.Candidate), lastIdx, req.LastLogIndex)
+		r.logger.Printf("[WARN] raft: Rejecting vote request from %v since our last index is greater (%d, %d)",
+			candidate, lastIdx, req.LastLogIndex)
 		return
 	}
 
diff --git a/raft_test.go b/raft_test.go
index bbb39f9343e..69bd07451fa 100644
--- a/raft_test.go
+++ b/raft_test.go
@@ -206,6 +206,15 @@ func (c *cluster) Disconnect(a string) {
 	}
 }
 
+func (c *cluster) IndexOf(r *Raft) int {
+	for i, n := range c.rafts {
+		if n == r {
+			return i
+		}
+	}
+	return -1
+}
+
 func (c *cluster) EnsureLeader(t *testing.T, expect string) {
 	limit := time.Now().Add(400 * time.Millisecond)
 CHECK:
@@ -1589,7 +1598,7 @@ func TestRaft_NotifyCh(t *testing.T) {
 		if !v {
 			t.Fatalf("should become leader")
 		}
-	case <-time.After(conf.HeartbeatTimeout * 3):
+	case <-time.After(conf.HeartbeatTimeout * 6):
 		t.Fatalf("timeout becoming leader")
 	}
 
@@ -1602,7 +1611,38 @@ func TestRaft_NotifyCh(t *testing.T) {
 		if v {
 			t.Fatalf("should step down as leader")
 		}
-	case <-time.After(conf.HeartbeatTimeout * 3):
+	case <-time.After(conf.HeartbeatTimeout * 6):
 		t.Fatalf("timeout becoming leader")
 	}
 }
+
+func TestRaft_Voting(t *testing.T) {
+	c := MakeCluster(3, t, nil)
+	defer c.Close()
+	followers := c.Followers()
+	ldr := c.Leader()
+	ldrT := c.trans[c.IndexOf(ldr)]
+
+	reqVote := RequestVoteRequest{
+		Term:         42,
+		Candidate:    ldrT.EncodePeer(ldr.localAddr),
+		LastLogIndex: ldr.LastIndex(),
+		LastLogTerm:  1,
+	}
+	// a follower that thinks there's a leader should vote for that leader.
+	var resp RequestVoteResponse
+	if err := ldrT.RequestVote(followers[0].localAddr, &reqVote, &resp); err != nil {
+		t.Fatalf("RequestVote RPC failed %v", err)
+	}
+	if !resp.Granted {
+		t.Fatalf("expected vote to be granted, but wasn't %+v", resp)
+	}
+	// a follow that thinks there's a leader shouldn't vote for a different candidate
+	reqVote.Candidate = ldrT.EncodePeer(followers[0].localAddr)
+	if err := ldrT.RequestVote(followers[1].localAddr, &reqVote, &resp); err != nil {
+		t.Fatalf("RequestVote RPC failed %v", err)
+	}
+	if resp.Granted {
+		t.Fatalf("expected vote not to be granted, but was %+v", resp)
+	}
+}