WIP fake clock in tests

go.mod

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

tinyraft.go

@@ -15,11 +15,12 @@
 package tinyraft
 
 import (
-    "context"
     "fmt"
     "math/rand"
     "sync"
     "time"
+
+    "github.com/jonboulle/clockwork"
 )
 
 type MsgType int8
@@ -63,12 +64,12 @@ type RaftConfig struct {
 }
 
 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
+    mu         sync.Mutex
+    shutdownCh chan struct{}
 
     term      uint64
     state     NodeState
@@ -82,7 +83,11 @@ type Raft struct {
     votes map[NodeID][]NodeID
 }
 
-func NewRaft(ctx context.Context, config RaftConfig) *Raft {
+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
     }
@@ -90,11 +95,11 @@ func NewRaft(ctx context.Context, config RaftConfig) *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
@@ -102,16 +107,16 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
         for ok {
             if dur == 0 {
                 select {
-                case <-r.ctx.Done():
+                case <-r.shutdownCh:
                     return
                 case dur, ok = <-r.electionTimerCh:
                 }
             } else {
                 select {
-                case <-r.ctx.Done():
+                case <-r.shutdownCh:
                     return
                 case dur, ok = <-r.electionTimerCh:
-                case <-time.After(dur):
+                case <-clock.After(dur):
                     r.Start()
                 }
             }
@@ -120,9 +125,9 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
     go func() {
         for {
             select {
-            case <-r.ctx.Done():
+            case <-r.shutdownCh:
                 return
-            case <-time.After(r.config.HeartbeatTimeout):
+            case <-clock.After(r.config.HeartbeatTimeout):
                 r.heartbeat()
             }
         }
@@ -130,6 +135,16 @@ func NewRaft(ctx context.Context, config RaftConfig) *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
 }
@@ -145,6 +160,9 @@ func (r *Raft) nextTerm(after time.Duration) {
 func (r *Raft) Start() {
     r.mu.Lock()
     defer r.mu.Unlock()
+    if r.shutdownCh == nil {
+        return
+    }
     r.startUnlocked()
 }
 
@@ -186,7 +204,8 @@ func (r *Raft) changed() {
 func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if msg.Type == MsgVoteRequest {
+    if r.shutdownCh == nil {
+    } else if msg.Type == MsgVoteRequest {
         if msg.Term > r.term && msg.Leader != NoLeader {
             r.leader = msg.Leader
             r.term = msg.Term
@@ -297,6 +316,12 @@ 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:
@@ -311,6 +336,10 @@ func (t MsgType) GoString() string {
     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,17 +1,16 @@
-package tinyraft_test
+package tinyraft
 
 import (
-    "context"
     "fmt"
     "testing"
     "time"
 
-    . "git.yourcmc.ru/vitalif/tinyraft"
+    "github.com/jonboulle/clockwork"
 )
 
-func addTestNode(ctx context.Context, nodes []*Raft, i int, partitions map[int]bool) {
+func addTestNode(clock clockwork.Clock, nodes []*Raft, i int, partitions map[int]bool) {
     from := NodeID(i + 1)
-    n := NewRaft(ctx, RaftConfig{
+    n := newRaft(RaftConfig{
         NodeID: from,
         Nodes: []NodeID{ 1, 2, 3, 4, 5 },
         HeartbeatTimeout: 100 * time.Millisecond,
@@ -26,7 +25,7 @@ func addTestNode(ctx context.Context, nodes []*Raft, i int, partitions map[int]b
                 }()
             }
         },
-    })
+    }, clock)
     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,13 +67,11 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 }
 
 func TestStartThenRemoveNode(t *testing.T) {
-    var cancels []func()
+    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        ctx, cancel := context.WithCancel(context.Background())
-        cancels = append(cancels, cancel)
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(clock, nodes, i, partitions)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -89,7 +86,10 @@ func TestStartThenRemoveNode(t *testing.T) {
         nodes[i].Start()
     }
     // Check that 5 nodes are in quorum after 2000ms
-    time.Sleep(2000 * time.Millisecond)
+    //time.Sleep(2000 * time.Millisecond)
+    for i := 0; i < 40; i++ {
+        clock.Advance(50 * time.Millisecond)
+    }
     checkQuorum(t, nodes, 5, 5)
     leaderChanges := 0
     ok := true
@@ -106,11 +106,12 @@ func TestStartThenRemoveNode(t *testing.T) {
     }
     // Stop the leader
     leader := nodes[0].Leader()
-    cancels[leader-1]()
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check quorum after 2000ms
-    time.Sleep(2000 * time.Millisecond)
+    for i := 0; i < 40; i++ {
+        clock.Advance(50 * time.Millisecond)
+    }
     checkQuorum(t, nodes, 5, 4)
     // Stop the leader again
     for i := 0; i < 5; i++ {
@@ -118,11 +119,12 @@ func TestStartThenRemoveNode(t *testing.T) {
             leader = nodes[i].Leader()
         }
     }
-    cancels[leader-1]()
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check quorum after 2000ms
-    time.Sleep(2000 * time.Millisecond)
+    for i := 0; i < 40; i++ {
+        clock.Advance(50 * time.Millisecond)
+    }
     checkQuorum(t, nodes, 5, 3)
     // Stop the leader again
     for i := 0; i < 5; i++ {
@@ -130,36 +132,39 @@ func TestStartThenRemoveNode(t *testing.T) {
             leader = nodes[i].Leader()
         }
     }
-    cancels[leader-1]()
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check for the absence of quorum
-    time.Sleep(2000 * time.Millisecond)
+    for i := 0; i < 40; i++ {
+        clock.Advance(50 * time.Millisecond)
+    }
     checkNoQuorum(t, nodes, 5)
     // Cleanup
     close(ch1)
-    for _, cancel := range cancels {
-        if cancel != nil {
-            cancel()
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
         }
     }
 }
 
 func TestAddNode(t *testing.T) {
-    ctx, cancel := context.WithCancel(context.Background())
+    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
     for i := 0; i < 5; i++ {
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(clock, nodes, i, partitions)
     }
     for i := 0; i < 5; i++ {
         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)
+    }
     checkQuorum(t, nodes, 5, 5)
     // Add a node
-    addTestNode(ctx, nodes, 5, partitions)
+    addTestNode(clock, 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})
@@ -168,7 +173,11 @@ 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
-    time.Sleep(2000 * time.Millisecond)
+    for i := 0; i < 40; i++ {
+        clock.Advance(50 * time.Millisecond)
+    }
     checkQuorum(t, nodes, 6, 6)
-    cancel()
+    for _, node := range nodes {
+        node.Shutdown()
+    }
 }