WIP Pre-voting (test is still unstable)

tinyraft.go

@@ -9,7 +9,8 @@
 // The only requirement is to guarantee preservation of entries confirmed by
 // all hosts participating in consensus.
 //
-// Supports leader expiration like in NuRaft:
+// Supports pre-vote protocol and leader expiration like in NuRaft:
+// https://github.com/eBay/NuRaft/blob/master/docs/prevote_protocol.md
 // https://github.com/eBay/NuRaft/blob/master/docs/leadership_expiration.md
 
 package tinyraft
@@ -30,6 +31,10 @@ const (
     MsgVote
     MsgPing
     MsgPong
+    // Following 3 are required only for the "pre-vote" feature
+    MsgPreVoteRequest
+    MsgPreVote
+    MsgFollow
 )
 
 type NodeState int8
@@ -59,6 +64,8 @@ type RaftConfig struct {
     RandomTimeout     time.Duration
     HeartbeatTimeout  time.Duration
     LeadershipTimeout time.Duration
+    EnablePrevote     bool
+    NeverLeader       bool
     SendFunction      func(to NodeID, msg RaftMsg)
 }
 
@@ -78,6 +85,10 @@ type Raft struct {
     electionTimer   *time.Timer
     electionTimerCh chan time.Duration
 
+    preVoting   bool
+    seenLeader  bool
+    preVoteFail int
+
     voted int
     votes map[NodeID][]NodeID
 }
@@ -142,6 +153,18 @@ func (r *Raft) nextTerm(after time.Duration) {
     r.electionTimerCh <- after + time.Duration(float64(r.config.RandomTimeout) * rand.Float64())
 }
 
+func (r *Raft) prevote() {
+    r.votes = make(map[NodeID][]NodeID)
+    r.preVoting = true
+    r.preVoteFail = 0
+    for _, node := range r.config.Nodes {
+        if node != r.config.NodeID {
+            r.send(node, RaftMsg{Type: MsgPreVoteRequest, Term: r.term})
+        }
+    }
+    r.nextTerm(r.config.ElectionTimeout)
+}
+
 func (r *Raft) Start() {
     r.mu.Lock()
     defer r.mu.Unlock()
@@ -149,6 +172,15 @@ func (r *Raft) Start() {
 }
 
 func (r *Raft) startUnlocked() {
+    r.seenLeader = false
+    if r.config.EnablePrevote && !r.preVoting {
+        r.prevote()
+        return
+    }
+    r.preVoting = false
+    if r.config.NeverLeader {
+        return
+    }
     r.term++
     r.voted = 1
     r.votes = make(map[NodeID][]NodeID)
@@ -186,11 +218,43 @@ func (r *Raft) changed() {
 func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if msg.Type == MsgVoteRequest {
+    if msg.Type == MsgPreVoteRequest {
+        leader := r.leader
+        if r.state != StateLeader && !r.seenLeader {
+            leader = NoLeader
+        }
+        r.send(from, RaftMsg{Type: MsgPreVote, Term: r.term, Leader: leader})
+    } else if msg.Type == MsgPreVote && msg.Term == r.term && r.preVoting {
+        if msg.Leader != NoLeader {
+            r.votes[msg.Leader] = append(r.votes[msg.Leader], from)
+            n := len(r.votes[msg.Leader])
+            if n == 1 + len(r.config.Nodes)/2 {
+                r.preVoting = false
+                r.preVoteFail = 0
+                if msg.Leader != r.leader {
+                    // Join the same leader. It may not know about us if we just started.
+                    r.leader = msg.Leader
+                    r.state = StateFollower
+                    r.seenLeader = false
+                    r.send(msg.Leader, RaftMsg{Type: MsgFollow, Term: r.term, Leader: msg.Leader})
+                    r.changed()
+                }
+                r.nextTerm(r.config.ElectionTimeout)
+            }
+        } else {
+            r.preVoteFail++
+        }
+        if r.preVoteFail > len(r.config.Nodes)/2 {
+            r.noNextTerm()
+            r.preVoteFail = 0
+            r.startUnlocked()
+        }
+    } 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.seenLeader = false
             r.followers = nil
             r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
             r.changed()
@@ -244,6 +308,12 @@ func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
         if msg.Leader != r.config.NodeID {
             r.startUnlocked()
         }
+    } else if msg.Type == MsgFollow && r.term == msg.Term && r.config.NodeID == msg.Leader && r.state == StateLeader {
+        // MsgFollow is sent to the leader by a node who failed prevote
+        if !r.hasNode(r.followers, from) {
+            r.followers = append(r.followers, from)
+            r.changed()
+        }
     }
 }
 
@@ -254,6 +324,7 @@ func (r *Raft) MarkAlive() {
 }
 
 func (r *Raft) markAliveUnlocked() {
+    r.seenLeader = true
     r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
 }
 
@@ -307,6 +378,12 @@ func (t MsgType) GoString() string {
         return "Ping"
     case MsgPong:
         return "Pong"
+    case MsgPreVoteRequest:
+        return "PreVoteRequest"
+    case MsgPreVote:
+        return "PreVote"
+    case MsgFollow:
+        return "Follow"
     }
     return fmt.Sprintf("unknown MsgType(%d)", t)
 }

tinyraft_test.go

@@ -9,24 +9,25 @@ import (
     . "git.yourcmc.ru/vitalif/tinyraft"
 )
 
-func addTestNode(ctx context.Context, nodes []*Raft, i int, partitions map[int]bool) {
+func addTestNode(ctx context.Context, cfg RaftConfig, nodes []*Raft, i int, partitions map[int]bool) {
     from := NodeID(i + 1)
-    n := NewRaft(ctx, RaftConfig{
-        NodeID: from,
-        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)
-                }()
-            }
-        },
-    })
+    cfg.NodeID = from
+    cfg.HeartbeatTimeout = 100 * time.Millisecond
+    if cfg.ElectionTimeout == 0 {
+        cfg.ElectionTimeout = 500 * time.Millisecond
+    }
+    cfg.RandomTimeout = cfg.ElectionTimeout
+    cfg.LeadershipTimeout = 0
+    cfg.EnablePrevote = true
+    cfg.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)
+            }()
+        }
+    }
+    n := NewRaft(ctx, 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)
     })
@@ -74,7 +75,7 @@ func TestStartThenRemoveNode(t *testing.T) {
     for i := 0; i < 5; i++ {
         ctx, cancel := context.WithCancel(context.Background())
         cancels = append(cancels, cancel)
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5}}, nodes, i, partitions)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -150,7 +151,7 @@ func TestAddNode(t *testing.T) {
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
     for i := 0; i < 5; i++ {
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5}}, nodes, i, partitions)
     }
     for i := 0; i < 5; i++ {
         nodes[i].Start()
@@ -159,16 +160,66 @@ func TestAddNode(t *testing.T) {
     time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 5, 5)
     // Add a node
-    addTestNode(ctx, nodes, 5, partitions)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5,6}}, nodes, 5, partitions)
     nodes[5].Start()
     nodes[0].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[1].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[2].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[3].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     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
     time.Sleep(2000 * time.Millisecond)
     checkQuorum(t, nodes, 6, 6)
     cancel()
 }
+
+func TestPrevote(t *testing.T) {
+    ctx, cancel := context.WithCancel(context.Background())
+    // The original example for pre-vote protocol from here:
+    // https://github.com/eBay/NuRaft/blob/master/docs/prevote_protocol.md
+    // only has partitions between nodes 1-5, 2-4, 3-4.
+    // But that setup has a problem of leader jumping between nodes 3 and 4
+    // which can't be prevented by prevotes because 4 only has connections
+    // to 2 nodes so they can't form a majority to reject the vote.
+    partitions := map[int]bool{}
+    for _, p := range [][]int{{1,5},{2,3},{3,4},{2,4}} {
+        partitions[(p[0]-1)*100 + (p[1]-1)] = true
+        partitions[(p[1]-1)*100 + (p[0]-1)] = true
+    }
+    nodes := make([]*Raft, 5)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5}}, nodes, 0, partitions)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5},NeverLeader:true}, nodes, 1, partitions)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5},NeverLeader:true}, nodes, 2, partitions)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5},NeverLeader:true}, nodes, 3, partitions)
+    addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5}}, nodes, 4, partitions)
+    ch1 := make(chan NodeID, 100)
+    nodes[2].OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
+        fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", 1, state, term, leader, followers)
+        if state != StateCandidate && leader != NoLeader {
+            ch1 <- leader
+        }
+    })
+    nodes[0].Start()
+    nodes[1].Start()
+    nodes[2].Start()
+    nodes[3].Start()
+    nodes[4].Start()
+    // First leader should be 1 or 5
+    var firstLeader NodeID
+    select {
+    case firstLeader = <-ch1:
+        if firstLeader != 1 && firstLeader != 5 {
+            t.Fatalf("Leader not 1 nor 5")
+        }
+    case <-time.After(2000 * time.Millisecond):
+        t.Fatalf("No leader in 500ms")
+    }
+    // Then leader should not change
+    select {
+    case <-ch1:
+        t.Fatalf("Leader changed in 5000ms")
+    case <-time.After(5000 * time.Millisecond):
+    }
+    close(ch1)
+    cancel()
+}