Improve leadership expiration and add a test for removal of a follower

tinyraft.go

@@ -9,7 +9,9 @@
 // The only requirement is to guarantee preservation of entries confirmed by
 // all hosts participating in consensus.
 //
-// Supports leader expiration like in NuRaft:
+// Supports leader expiration similar to NuRaft, but even better: the leader
+// tracks followers successfully replying to pings and removes followers which
+// don't reply in LeadershipTimeout. NuRaft original:
 // https://github.com/eBay/NuRaft/blob/master/docs/leadership_expiration.md
 
 package tinyraft
@@ -79,6 +81,11 @@ type Raft struct {
     electionTimer  *time.Timer
     heartbeatTimer *time.Timer
 
+    // for leadership expiration
+    pongSent           bool
+    confirmedFollowers map[NodeID]struct{}
+    pongCheckTimer     *time.Timer
+
     voted int
     votes map[NodeID][]NodeID
 }
@@ -102,6 +109,7 @@ func (r *Raft) Shutdown() {
     r.mu.Lock()
     defer r.mu.Unlock()
     r.shutdown = true
+    r.setPongCheck(0, nil)
     r.noNextTerm()
     if r.heartbeatTimer != nil {
         r.heartbeatTimer.Stop()
@@ -110,6 +118,16 @@ func (r *Raft) Shutdown() {
     r.onChange = nil
 }
 
+func (r *Raft) setPongCheck(d time.Duration, f func()) {
+    if r.pongCheckTimer != nil {
+        r.pongCheckTimer.Stop()
+        r.pongCheckTimer = nil
+    }
+    if f != nil {
+        r.pongCheckTimer = time.AfterFunc(d, f)
+    }
+}
+
 func (r *Raft) noNextTerm() {
     if r.electionTimer != nil {
         r.electionTimer.Stop()
@@ -139,6 +157,8 @@ func (r *Raft) startUnlocked() {
     r.votes[r.config.NodeID] = []NodeID{r.config.NodeID}
     r.state = StateCandidate
     r.leader = r.config.NodeID
+    r.setPongCheck(0, nil)
+    r.confirmedFollowers = nil
     for _, node := range r.config.Nodes {
         if node != r.config.NodeID {
             r.send(node, RaftMsg{Type: MsgVoteRequest, Term: r.term, Leader: r.leader})
@@ -174,6 +194,10 @@ func (r *Raft) heartbeatUnlocked() {
     if r.state != StateLeader {
         return
     }
+    if r.config.LeadershipTimeout > 0 && r.confirmedFollowers == nil {
+        r.confirmedFollowers = make(map[NodeID]struct{})
+        r.setPongCheck(r.config.LeadershipTimeout, r.checkPongs)
+    }
     for _, node := range r.followers {
         if node != r.config.NodeID {
             r.send(node, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
@@ -197,6 +221,7 @@ func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
             r.term = msg.Term
             r.state = StateCandidate
             r.followers = nil
+            r.confirmedFollowers = nil
             r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
             r.changed()
         }
@@ -210,13 +235,14 @@ func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
                 r.state = StateLeader
                 r.leader = r.config.NodeID
                 r.followers = r.votes[r.config.NodeID]
+                r.confirmedFollowers = nil
                 // Send a heartbeat to confirm leadership
                 r.heartbeatUnlocked()
                 r.scheduleHeartbeat()
+                r.noNextTerm()
                 if r.config.LeadershipTimeout > 0 {
-                    r.nextTerm(r.config.LeadershipTimeout)
-                } else {
-                    r.noNextTerm()
+                    r.confirmedFollowers = make(map[NodeID]struct{})
+                    r.setPongCheck(r.config.LeadershipTimeout, r.checkPongs)
                 }
                 r.changed()
             } else {
@@ -238,16 +264,14 @@ func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
         if r.state == StateFollower && from == r.leader {
             r.markAliveUnlocked()
         }
-        if r.config.LeadershipTimeout > 0 {
-            r.send(from, RaftMsg{Type: MsgPong, Term: r.term, Leader: r.leader})
-        }
+        // Always respond to "ping" with "pong" if LeadershipTimeout is enabled
+        r.pongSent = false
+        r.sendPongUnlocked()
     } else if msg.Type == MsgPong && r.state == StateLeader {
         if msg.Leader != r.config.NodeID {
             r.startUnlocked()
-        } else if r.config.LeadershipTimeout > 0 {
-            r.nextTerm(r.config.LeadershipTimeout)
-        } else {
-            r.noNextTerm()
+        } else if r.confirmedFollowers != nil {
+            r.confirmedFollowers[from] = struct{}{}
         }
     }
 }
@@ -256,6 +280,48 @@ func (r *Raft) MarkAlive() {
     r.mu.Lock()
     defer r.mu.Unlock()
     r.markAliveUnlocked()
+    // Send "pong" to the leader at least once in HeartbeatTimeout when
+    // the user manually marks follower as live if LeadershipTimeout is enabled
+    r.sendPongUnlocked()
+}
+
+func (r *Raft) sendPong() {
+    r.mu.Lock()
+    defer r.mu.Unlock()
+    r.pongSent = false
+    r.sendPongUnlocked()
+}
+
+func (r *Raft) sendPongUnlocked() {
+    if r.config.LeadershipTimeout > 0 && r.state == StateFollower {
+        if !r.pongSent {
+            r.pongSent = true
+            r.send(r.leader, RaftMsg{Type: MsgPong, Term: r.term, Leader: r.leader})
+        }
+        r.setPongCheck(r.config.HeartbeatTimeout, r.sendPong)
+    }
+}
+
+func (r *Raft) checkPongs() {
+    r.mu.Lock()
+    defer r.mu.Unlock()
+    if len(r.confirmedFollowers)+1 < len(r.followers) {
+        if len(r.confirmedFollowers)+1 > len(r.config.Nodes)/2 {
+            // Still in quorum
+            r.followers = []NodeID{ r.config.NodeID }
+            for follower := range r.confirmedFollowers {
+                r.followers = append(r.followers, follower)
+            }
+            r.confirmedFollowers = nil
+            r.changed()
+        } else {
+            // Quorum is lost
+            r.startUnlocked()
+        }
+    } else {
+        r.confirmedFollowers = nil
+        r.setPongCheck(r.config.LeadershipTimeout, r.checkPongs)
+    }
 }
 
 func (r *Raft) markAliveUnlocked() {
@@ -304,7 +370,7 @@ func (r *Raft) SetNodes(nodes []NodeID) {
 
 // for fmt.Printf("%v")
 func (t MsgType) String() string {
-       return t.GoString()
+    return t.GoString()
 }
 
 // for fmt.Printf("%#v")

tinyraft_test.go

@@ -2,30 +2,35 @@ package tinyraft
 
 import (
     "fmt"
+    "sync/atomic"
     "testing"
     "time"
 )
 
-func addTestNode(nodes []*Raft, i int, partitions map[int]bool) {
-    from := NodeID(i + 1)
-    n := NewRaft(RaftConfig{
-        NodeID: from,
+func testRaftConfig() RaftConfig {
+    return RaftConfig{
         Nodes: []NodeID{ 1, 2, 3, 4, 5 },
         HeartbeatTimeout: 100 * time.Millisecond,
         ElectionTimeout: 500 * time.Millisecond,
         RandomTimeout: 500 * time.Millisecond,
         LeadershipTimeout: 1500 * time.Millisecond,
-        SendFunction: func(to NodeID, msg RaftMsg) {
-            intTo := to.(int)
-            intFrom := from.(int)
-            if !partitions[(intTo-1)*100 + intFrom-1] && nodes[intTo-1] != nil {
-                go func() {
-                    fmt.Printf("received from %v to %v: %#v\n", from, to, msg)
-                    nodes[intTo-1].OnReceive(from, msg)
-                }()
-            }
-        },
-    })
+    }
+}
+
+func addTestNode(cfg RaftConfig, partitions map[int]bool, nodes []*Raft, i int) {
+    from := NodeID(i + 1)
+    cfg.NodeID = from
+    cfg.SendFunction = func(to NodeID, msg RaftMsg) {
+        intTo := to.(int)
+        intFrom := from.(int)
+        if !partitions[(intTo-1)*100 + intFrom-1] && nodes[intTo-1] != nil {
+            go func() {
+                fmt.Printf("received from %v to %v: %#v\n", from, to, msg)
+                nodes[intTo-1].OnReceive(from, msg)
+            }()
+        }
+    }
+    n := NewRaft(cfg)
     n.OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
         fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", from, state, term, leader, followers)
     })
@@ -68,9 +73,10 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 
 func TestStartThenRemoveNode(t *testing.T) {
     partitions := map[int]bool{}
+    cfg := testRaftConfig()
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        addTestNode(nodes, i, partitions)
+        addTestNode(cfg, partitions, nodes, i)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -141,8 +147,9 @@ func TestStartThenRemoveNode(t *testing.T) {
 func TestAddNode(t *testing.T) {
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
+    cfg := testRaftConfig()
     for i := 0; i < 5; i++ {
-        addTestNode(nodes, i, partitions)
+        addTestNode(cfg, partitions, nodes, i)
     }
     for i := 0; i < 5; i++ {
         nodes[i].Start()
@@ -151,7 +158,7 @@ func TestAddNode(t *testing.T) {
     time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 5)
     // Add a node
-    addTestNode(nodes, 5, partitions)
+    addTestNode(cfg, partitions, nodes, 5)
     nodes[5].Start()
     nodes[0].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[1].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
@@ -170,8 +177,9 @@ func TestAddNode(t *testing.T) {
 func TestLeadershipExpiration(t *testing.T) {
     partitions := map[int]bool{}
     nodes := make([]*Raft, 5)
+    cfg := testRaftConfig()
     for i := 0; i < 5; i++ {
-        addTestNode(nodes, i, partitions)
+        addTestNode(cfg, partitions, nodes, i)
     }
     for i := 0; i < 5; i++ {
         nodes[i].Start()
@@ -202,3 +210,82 @@ func TestLeadershipExpiration(t *testing.T) {
         node.Shutdown()
     }
 }
+
+func TestLeadershipExpireFollower(t *testing.T) {
+    partitions := map[int]bool{}
+    nodes := make([]*Raft, 5)
+    cfg := testRaftConfig()
+    for i := 0; i < 5; i++ {
+        addTestNode(cfg, partitions, nodes, i)
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, 5)
+    // Break network on a follower
+    leader := nodes[0].Leader().(int)
+    follower := leader % 5
+    fmt.Printf("Stopping a follower (%v)\n", follower)
+    nodes[follower].Shutdown()
+    nodes[follower] = nil
+    // Check that the leader sees a follower list change in 2 * LeadershipTimeout
+    time.Sleep(3000 * time.Millisecond)
+    if len(nodes[leader-1].Followers()) >= len(nodes) {
+        t.Fatalf("follower is not removed on leader")
+    }
+    // Cleanup
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
+}
+
+func TestManualHeartbeat(t *testing.T) {
+    partitions := map[int]bool{}
+    nodes := make([]*Raft, 5)
+    cfg := testRaftConfig()
+    cfg.DisableHeartbeat = true
+    for i := 0; i < 5; i++ {
+        addTestNode(cfg, partitions, nodes, i)
+    }
+    // Count leader changes
+    prevLeader1 := NoLeader
+    leaderChanges := int32(0)
+    nodes[1].OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
+        fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", 2, state, term, leader, followers)
+        if state != StateCandidate && leader != NoLeader && leader != prevLeader1 {
+            prevLeader1 = leader
+            atomic.AddInt32(&leaderChanges, 1)
+        }
+    })
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    for i := 0; i < 80; i++ {
+        time.Sleep(50 * time.Millisecond)
+        for i := 0; i < 5; i++ {
+            // Pretend we're doing replication and followers receive some replication messages
+            if nodes[i].State() == StateFollower {
+                nodes[i].MarkAlive()
+            }
+        }
+        if i == 10 {
+            atomic.StoreInt32(&leaderChanges, 0)
+        }
+    }
+    // Check that 5 nodes are in quorum after 4000ms
+    checkQuorum(t, nodes, 5, 5)
+    // And the leader didn't change after initial election
+    if atomic.LoadInt32(&leaderChanges) > 0 {
+        t.Fatalf("leader timed out and changed")
+    }
+    // Cleanup
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
+}