WIP Pre-voting (test is still unstable)

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,18 +9,18 @@
 // 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
 
 import (
+    "context"
     "fmt"
     "math/rand"
     "sync"
     "time"
-
-    "github.com/jonboulle/clockwork"
 )
 
 type MsgType int8
@@ -31,6 +31,10 @@ const (
     MsgVote
     MsgPing
     MsgPong
+    // Following 3 are required only for the "pre-vote" feature
+    MsgPreVoteRequest
+    MsgPreVote
+    MsgFollow
 )
 
 type NodeState int8
@@ -60,16 +64,18 @@ type RaftConfig struct {
     RandomTimeout     time.Duration
     HeartbeatTimeout  time.Duration
     LeadershipTimeout time.Duration
+    EnablePrevote     bool
+    NeverLeader       bool
     SendFunction      func(to NodeID, msg RaftMsg)
 }
 
 type Raft struct {
+    ctx      context.Context
     config   RaftConfig
     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
 
     term      uint64
     state     NodeState
@@ -79,15 +85,15 @@ type Raft struct {
     electionTimer   *time.Timer
     electionTimerCh chan time.Duration
 
+    preVoting   bool
+    seenLeader  bool
+    preVoteFail int
+
     voted int
     votes map[NodeID][]NodeID
 }
 
-func NewRaft(config RaftConfig) *Raft {
-    return newRaft(config, clockwork.NewRealClock())
-}
-
-func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
+func NewRaft(ctx context.Context, config RaftConfig) *Raft {
     if config.ElectionTimeout == 0 || config.NodeID == 0 || config.SendFunction == nil {
         return nil
     }
@@ -95,11 +101,11 @@ func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
         config.Nodes = append(config.Nodes, config.NodeID)
     }
     r := &Raft{
+        ctx:     ctx,
         config:  config,
         send:    config.SendFunction,
         term:    config.InitialTerm,
     }
-    r.shutdownCh = make(chan struct{})
     r.electionTimerCh = make(chan time.Duration, 1)
     go func() {
         ok := true
@@ -107,16 +113,16 @@ func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
         for ok {
             if dur == 0 {
                 select {
-                case <-r.shutdownCh:
+                case <-r.ctx.Done():
                     return
                 case dur, ok = <-r.electionTimerCh:
                 }
             } else {
                 select {
-                case <-r.shutdownCh:
+                case <-r.ctx.Done():
                     return
                 case dur, ok = <-r.electionTimerCh:
-                case <-clock.After(dur):
+                case <-time.After(dur):
                     r.Start()
                 }
             }
@@ -125,9 +131,9 @@ func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
     go func() {
         for {
             select {
-            case <-r.shutdownCh:
+            case <-r.ctx.Done():
                 return
-            case <-clock.After(r.config.HeartbeatTimeout):
+            case <-time.After(r.config.HeartbeatTimeout):
                 r.heartbeat()
             }
         }
@@ -135,16 +141,6 @@ func newRaft(config RaftConfig, clock clockwork.Clock) *Raft {
     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
-    }
-}
-
 func (r *Raft) noNextTerm() {
     r.electionTimerCh <- 0
 }
@@ -157,16 +153,34 @@ 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()
-    if r.shutdownCh == nil {
-        return
-    }
     r.startUnlocked()
 }
 
 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)
@@ -204,12 +218,43 @@ func (r *Raft) changed() {
 func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.shutdownCh == nil {
+    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()
@@ -263,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()
+        }
     }
 }
 
@@ -273,6 +324,7 @@ func (r *Raft) MarkAlive() {
 }
 
 func (r *Raft) markAliveUnlocked() {
+    r.seenLeader = true
     r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
 }
 
@@ -316,12 +368,6 @@ func (r *Raft) SetNodes(nodes []NodeID) {
     }
 }
 
-// for fmt.Printf("%v")
-func (t MsgType) String() string {
-       return t.GoString()
-}
-
-// for fmt.Printf("%#v")
 func (t MsgType) GoString() string {
     switch t {
     case MsgVoteRequest:
@@ -332,14 +378,16 @@ 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)
 }
 
-func (t NodeState) String() string {
-       return t.GoString()
-}
-
 func (t NodeState) GoString() string {
     switch t {
     case StateCandidate:

tinyraft_test.go

@@ -1,31 +1,33 @@
-package tinyraft
+package tinyraft_test
 
 import (
+    "context"
     "fmt"
     "testing"
     "time"
 
-    "github.com/jonboulle/clockwork"
+    . "git.yourcmc.ru/vitalif/tinyraft"
 )
 
-func addTestNode(clock clockwork.Clock, 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(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)
-                }()
-            }
-        },
-    }, clock)
+    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)
     })
@@ -67,11 +69,13 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 }
 
 func TestStartThenRemoveNode(t *testing.T) {
-    clock := clockwork.NewFakeClock()
+    var cancels []func()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        addTestNode(clock, nodes, i, partitions)
+        ctx, cancel := context.WithCancel(context.Background())
+        cancels = append(cancels, cancel)
+        addTestNode(ctx, RaftConfig{Nodes:[]NodeID{1,2,3,4,5}}, nodes, i, partitions)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -86,10 +90,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
@@ -106,12 +107,11 @@ func TestStartThenRemoveNode(t *testing.T) {
     }
     // Stop the leader
     leader := nodes[0].Leader()
-    nodes[leader-1].Shutdown()
+    cancels[leader-1]()
     nodes[leader-1] = nil
+    cancels[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++ {
@@ -119,12 +119,11 @@ func TestStartThenRemoveNode(t *testing.T) {
             leader = nodes[i].Leader()
         }
     }
-    nodes[leader-1].Shutdown()
+    cancels[leader-1]()
     nodes[leader-1] = nil
+    cancels[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++ {
@@ -132,52 +131,95 @@ func TestStartThenRemoveNode(t *testing.T) {
             leader = nodes[i].Leader()
         }
     }
-    nodes[leader-1].Shutdown()
+    cancels[leader-1]()
     nodes[leader-1] = nil
+    cancels[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)
-    for _, node := range nodes {
-        if node != nil {
-            node.Shutdown()
+    for _, cancel := range cancels {
+        if cancel != nil {
+            cancel()
         }
     }
 }
 
 func TestAddNode(t *testing.T) {
-    clock := clockwork.NewFakeClock()
+    ctx, cancel := context.WithCancel(context.Background())
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
     for i := 0; i < 5; i++ {
-        addTestNode(clock, 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()
     }
     // 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(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
-    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()
-    }
+    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()
 }