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

go.mod

@@ -1,5 +1,3 @@
 module git.yourcmc.ru/vitalif/tinyraft
 
 go 1.16
-
-require github.com/jonboulle/clockwork v0.4.0

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
@@ -19,8 +21,6 @@ import (
     "math/rand"
     "sync"
     "time"
-
-    "github.com/jonboulle/clockwork"
 )
 
 type MsgType int8
@@ -42,9 +42,9 @@ const (
     StateFollower
 )
 
-type NodeID uint64
+type NodeID interface{}
 
-const NoLeader = NodeID(0)
+var NoLeader = NodeID(nil)
 
 type RaftMsg struct {
     Type   MsgType
@@ -60,6 +60,8 @@ type RaftConfig struct {
     RandomTimeout     time.Duration
     HeartbeatTimeout  time.Duration
     LeadershipTimeout time.Duration
+    // disable automatic heartbeats if you do them yourself and call MarkAlive() manually
+    DisableHeartbeat  bool
     SendFunction      func(to NodeID, msg RaftMsg)
 }
 
@@ -68,26 +70,27 @@ type Raft struct {
     send     func(to NodeID, msg RaftMsg)
     onChange func(state NodeState, term uint64, leader NodeID, followers []NodeID)
 
-    mu         sync.Mutex
-    shutdownCh chan struct{}
+    mu       sync.Mutex
+    shutdown bool
 
     term      uint64
     state     NodeState
     leader    NodeID
     followers []NodeID
 
-    electionTimer   *time.Timer
-    electionTimerCh chan time.Duration
+    electionTimer  *time.Timer
+    heartbeatTimer *time.Timer
+
+    // for leadership expiration
+    pongSent           bool
+    confirmedFollowers map[NodeID]struct{}
+    pongCheckTimer     *time.Timer
 
     voted int
     votes map[NodeID][]NodeID
 }
 
 func NewRaft(config RaftConfig) *Raft {
-    return newRaft(config, clockwork.NewRealClock())
-}
-
-func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
     if config.ElectionTimeout == 0 || config.NodeID == 0 || config.SendFunction == nil {
         return nil
     }
@@ -99,68 +102,49 @@ func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
         send:    config.SendFunction,
         term:    config.InitialTerm,
     }
-    r.shutdownCh = make(chan struct{})
-    r.electionTimerCh = make(chan time.Duration, 1)
-    go func() {
-        ok := true
-        dur := time.Duration(0)
-        for ok {
-            if dur == 0 {
-                select {
-                case <-r.shutdownCh:
-                    return
-                case dur, ok = <-r.electionTimerCh:
-                }
-            } else {
-                select {
-                case <-r.shutdownCh:
-                    return
-                case dur, ok = <-r.electionTimerCh:
-                case <-clock.After(dur):
-                    r.Start()
-                }
-            }
-        }
-    }()
-    go func() {
-        for {
-            select {
-            case <-r.shutdownCh:
-                return
-            case <-clock.After(r.config.HeartbeatTimeout):
-                r.heartbeat()
-            }
-        }
-    }()
     return r
 }
 
 func (r *Raft) Shutdown() {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.shutdownCh != nil {
-        close(r.shutdownCh)
-        r.shutdownCh = nil
-        r.onChange = nil
+    r.shutdown = true
+    r.setPongCheck(0, nil)
+    r.noNextTerm()
+    if r.heartbeatTimer != nil {
+        r.heartbeatTimer.Stop()
+        r.heartbeatTimer = nil
+    }
+    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() {
-    r.electionTimerCh <- 0
+    if r.electionTimer != nil {
+        r.electionTimer.Stop()
+        r.electionTimer = nil
+    }
 }
 
 func (r *Raft) nextTerm(after time.Duration) {
-    if r.config.RandomTimeout == 0 && after == 0 {
-        r.electionTimerCh <- time.Millisecond
-        return
-    }
-    r.electionTimerCh <- after + time.Duration(float64(r.config.RandomTimeout) * rand.Float64())
+    r.noNextTerm()
+    after += time.Duration(float64(r.config.RandomTimeout) * rand.Float64())
+    r.electionTimer = time.AfterFunc(after, r.Start)
 }
 
 func (r *Raft) Start() {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.shutdownCh == nil {
+    if r.shutdown {
         return
     }
     r.startUnlocked()
@@ -173,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})
@@ -186,11 +172,35 @@ func (r *Raft) startUnlocked() {
 func (r *Raft) heartbeat() {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.state == StateLeader {
-        for _, node := range r.votes[r.config.NodeID] {
-            if node != r.config.NodeID {
-                r.send(node, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
-            }
+    r.heartbeatUnlocked()
+    r.scheduleHeartbeat()
+}
+
+func (r *Raft) scheduleHeartbeat() {
+    if r.config.DisableHeartbeat {
+        return
+    }
+    if r.state != StateLeader {
+        return
+    }
+    if r.heartbeatTimer != nil {
+        r.heartbeatTimer.Stop()
+        r.heartbeatTimer = nil
+    }
+    r.heartbeatTimer = time.AfterFunc(r.config.HeartbeatTimeout, r.heartbeat)
+}
+
+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})
         }
     }
 }
@@ -204,13 +214,14 @@ func (r *Raft) changed() {
 func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.shutdownCh == nil {
+    if r.shutdown {
     } else if msg.Type == MsgVoteRequest {
         if msg.Term > r.term && msg.Leader != NoLeader {
             r.leader = msg.Leader
             r.term = msg.Term
             r.state = StateCandidate
             r.followers = nil
+            r.confirmedFollowers = nil
             r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
             r.changed()
         }
@@ -224,16 +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]
-                for _, node := range r.followers {
-                    if node != r.config.NodeID {
-                        // Send a heartbeat to confirm leadership
-                        r.send(node, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
-                    }
-                }
+                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 {
@@ -255,13 +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 {
-        r.noNextTerm()
         if msg.Leader != r.config.NodeID {
             r.startUnlocked()
+        } else if r.confirmedFollowers != nil {
+            r.confirmedFollowers[from] = struct{}{}
         }
     }
 }
@@ -270,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() {
@@ -318,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"
-
-    "github.com/jonboulle/clockwork"
 )
 
-func addTestNode(clock clockwork.Clock, 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: 0,
-        SendFunction: func(to NodeID, msg RaftMsg) {
-            if !partitions[int((to-1)*100 + from-1)] && nodes[to-1] != nil {
-                go func() {
-                    fmt.Printf("received from %v to %v: %#v\n", from, to, msg)
-                    nodes[to-1].OnReceive(from, msg)
-                }()
-            }
-        },
-    }, clock)
+        LeadershipTimeout: 1500 * time.Millisecond,
+    }
+}
+
+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)
     })
@@ -67,11 +72,11 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 }
 
 func TestStartThenRemoveNode(t *testing.T) {
-    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
+    cfg := testRaftConfig()
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        addTestNode(clock, nodes, i, partitions)
+        addTestNode(cfg, partitions, nodes, i)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -86,10 +91,7 @@ func TestStartThenRemoveNode(t *testing.T) {
         nodes[i].Start()
     }
     // Check that 5 nodes are in quorum after 2000ms
-    //time.Sleep(2000 * time.Millisecond)
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 5)
     leaderChanges := 0
     ok := true
@@ -105,39 +107,33 @@ func TestStartThenRemoveNode(t *testing.T) {
         t.Fatalf("leaderChanges = %v (expected < 3)", leaderChanges)
     }
     // Stop the leader
-    leader := nodes[0].Leader()
+    leader := nodes[0].Leader().(int)
     nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
     // Check quorum after 2000ms
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 4)
     // Stop the leader again
     for i := 0; i < 5; i++ {
         if nodes[i] != nil {
-            leader = nodes[i].Leader()
+            leader = nodes[i].Leader().(int)
         }
     }
     nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
     // Check quorum after 2000ms
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 3)
     // Stop the leader again
     for i := 0; i < 5; i++ {
         if nodes[i] != nil {
-            leader = nodes[i].Leader()
+            leader = nodes[i].Leader().(int)
         }
     }
     nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
     // Check for the absence of quorum
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkNoQuorum(t, nodes, 5)
     // Cleanup
     close(ch1)
@@ -149,22 +145,20 @@ func TestStartThenRemoveNode(t *testing.T) {
 }
 
 func TestAddNode(t *testing.T) {
-    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
+    cfg := testRaftConfig()
     for i := 0; i < 5; i++ {
-        addTestNode(clock, nodes, i, partitions)
+        addTestNode(cfg, partitions, nodes, i)
     }
     for i := 0; i < 5; i++ {
         nodes[i].Start()
     }
     // Check that 5 nodes are in quorum after 2000ms
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 5)
     // Add a node
-    addTestNode(clock, 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})
@@ -173,11 +167,125 @@ func TestAddNode(t *testing.T) {
     nodes[4].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[5].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     // Check that 6 nodes are in quorum after 2000ms
-    for i := 0; i < 40; i++ {
-        clock.Advance(50 * time.Millisecond)
-    }
+    time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 6, 6)
     for _, node := range nodes {
         node.Shutdown()
     }
 }
+
+func TestLeadershipExpiration(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 the leader
+    leader := nodes[0].Leader().(int)
+    fmt.Printf("Stopping the leader's (%v) network\n", leader)
+    for i := 0; i < 5; i++ {
+        nodes[i].mu.Lock()
+    }
+    for i := 0; i < 5; i++ {
+        partitions[(leader-1)*100 + i] = true
+        partitions[i*100 + leader-1] = true
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].mu.Unlock()
+    }
+    // Check that the leader loses leadership after 2 * LeadershipTimeout
+    time.Sleep(3000 * time.Millisecond)
+    if nodes[leader-1].State() != StateCandidate {
+        t.Fatalf("leadership expiration doesn't work")
+    }
+    // Cleanup
+    for _, node := range nodes {
+        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()
+        }
+    }
+}