Ensure known leader is cleared on state transition

raft.go

@@ -667,8 +667,8 @@ func (r *Raft) runCandidate() {
 			// Check if we've become the leader
 			if grantedVotes >= votesNeeded {
 				r.logger.Printf("[INFO] raft: Election won. Tally: %d", grantedVotes)
-				r.setLeader(r.localAddr)
 				r.setState(Leader)
+				r.setLeader(r.localAddr)
 				return
 			}
 
@@ -1567,6 +1567,14 @@ func (r *Raft) setCurrentTerm(t uint64) {
 	r.raftState.setCurrentTerm(t)
 }
 
+// setState is used to update the current state. Any state
+// transition causes the known leader to be cleared. This means
+// that leader should be set only after updating the state.
+func (r *Raft) setState(state RaftState) {
+	r.setLeader(nil)
+	r.raftState.setState(state)
+}
+
 // runSnapshots is a long running goroutine used to manage taking
 // new snapshots of the FSM. It runs in parallel to the FSM and
 // main goroutines, so that snapshots do not block normal operation.

raft_test.go

@@ -1279,6 +1279,14 @@ func TestRaft_LeaderLeaseExpire(t *testing.T) {
 	if last != follower.LastContact() {
 		t.Fatalf("unexpected further contact")
 	}
+
+	// Ensure both have cleared their leader
+	if l := leader.Leader(); l != nil {
+		t.Fatalf("bad: %v", l)
+	}
+	if l := follower.Leader(); l != nil {
+		t.Fatalf("bad: %v", l)
+	}
 }
 
 func TestRaft_Barrier(t *testing.T) {
@@ -1374,6 +1382,11 @@ func TestRaft_VerifyLeader_Fail(t *testing.T) {
 	if err := verify.Error(); err != ErrNotLeader && err != ErrLeadershipLost {
 		t.Fatalf("err: %v", err)
 	}
+
+	// Ensure the known leader is cleared
+	if l := leader.Leader(); l != nil {
+		t.Fatalf("bad: %v", l)
+	}
 }
 
 func TestRaft_VerifyLeader_ParitalConnect(t *testing.T) {