From 93b08502e4c5987da5ae5c0b0bfca19ef8f0fff6 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiang.li@coreos.com>
Date: Sat, 24 May 2014 21:08:06 -0700
Subject: [PATCH] raft: check voteFor

---
 raft/raft.go      | 15 ++++++++++++---
 raft/raft_test.go | 39 ++++++++++++++++++++++-----------------
 2 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/raft/raft.go b/raft/raft.go
index c8fda32d8..4c55b2bc8 100644
--- a/raft/raft.go
+++ b/raft/raft.go
@@ -288,6 +288,8 @@ func (sm *stateMachine) Step(m Message) {
 		case msgApp:
 			sm.becomeFollower(sm.term, m.From)
 			handleAppendEntries()
+		case msgVote:
+			sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
 		case msgVoteResp:
 			gr := sm.poll(m.From, m.Index >= 0)
 			switch sm.q() {
@@ -303,11 +305,18 @@ func (sm *stateMachine) Step(m Message) {
 		case msgApp:
 			handleAppendEntries()
 		case msgVote:
-			if sm.log.isUpToDate(m.Index, m.LogTerm) {
+			switch sm.vote {
+			case m.From:
 				sm.send(Message{To: m.From, Type: msgVoteResp, Index: sm.log.lastIndex()})
-			} else {
-				sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
+				return
+			case none:
+				if sm.log.isUpToDate(m.Index, m.LogTerm) {
+					sm.vote = m.From
+					sm.send(Message{To: m.From, Type: msgVoteResp, Index: sm.log.lastIndex()})
+					return
+				}
 			}
+			sm.send(Message{To: m.From, Type: msgVoteResp, Index: -1})
 		}
 	}
 }
diff --git a/raft/raft_test.go b/raft/raft_test.go
index 8592f9a2d..d518abc59 100644
--- a/raft/raft_test.go
+++ b/raft/raft_test.go
@@ -346,32 +346,37 @@ func TestCommit(t *testing.T) {
 func TestVote(t *testing.T) {
 	tests := []struct {
 		i, term int
+		voteFor int
 		w       int
 	}{
-		{0, 0, -1},
-		{0, 1, -1},
-		{0, 2, -1},
-		{0, 3, 2},
+		{0, 0, none, -1},
+		{0, 1, none, -1},
+		{0, 2, none, -1},
+		{0, 3, none, 2},
 
-		{1, 0, -1},
-		{1, 1, -1},
-		{1, 2, -1},
-		{1, 3, 2},
+		{1, 0, none, -1},
+		{1, 1, none, -1},
+		{1, 2, none, -1},
+		{1, 3, none, 2},
 
-		{2, 0, -1},
-		{2, 1, -1},
-		{2, 2, 2},
-		{2, 3, 2},
+		{2, 0, none, -1},
+		{2, 1, none, -1},
+		{2, 2, none, 2},
+		{2, 3, none, 2},
 
-		{3, 0, -1},
-		{3, 1, -1},
-		{3, 2, 2},
-		{3, 3, 2},
+		{3, 0, none, -1},
+		{3, 1, none, -1},
+		{3, 2, none, 2},
+		{3, 3, none, 2},
+
+		{3, 2, 0, 2},
+		{3, 2, 1, -1},
 	}
 
 	for i, tt := range tests {
 		called := false
-		sm := &nsm{stateMachine{log: &log{ents: []Entry{{}, {Term: 2}, {Term: 2}}}}, nil}
+		sm := &nsm{stateMachine{vote: tt.voteFor, log: &log{ents: []Entry{{}, {Term: 2}, {Term: 2}}}}, nil}
+
 		sm.next = stepperFunc(func(m Message) {
 			called = true
 			if m.Index != tt.w {