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

@@ -19,6 +19,8 @@ import (
     "math/rand"
     "sync"
     "time"
+
+    "github.com/jonboulle/clockwork"
 )
 
 type MsgType int8
@@ -82,6 +84,10 @@ type Raft struct {
 }
 
 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
     }
@@ -110,7 +116,7 @@ func NewRaft(config RaftConfig) *Raft {
                 case <-r.shutdownCh:
                     return
                 case dur, ok = <-r.electionTimerCh:
-                case <-time.After(dur):
+                case <-clock.After(dur):
                     r.Start()
                 }
             }
@@ -121,7 +127,7 @@ func NewRaft(config RaftConfig) *Raft {
             select {
             case <-r.shutdownCh:
                 return
-            case <-time.After(r.config.HeartbeatTimeout):
+            case <-clock.After(r.config.HeartbeatTimeout):
                 r.heartbeat()
             }
         }

tinyraft_test.go

@@ -4,11 +4,13 @@ import (
     "fmt"
     "testing"
     "time"
+
+    "github.com/jonboulle/clockwork"
 )
 
-func addTestNode(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(RaftConfig{
+    n := newRaft(RaftConfig{
         NodeID: from,
         Nodes: []NodeID{ 1, 2, 3, 4, 5 },
         HeartbeatTimeout: 100 * time.Millisecond,
@@ -23,7 +25,7 @@ func addTestNode(nodes []*Raft, i int, partitions map[int]bool) {
                 }()
             }
         },
-    })
+    }, 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)
     })
@@ -65,10 +67,11 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 }
 
 func TestStartThenRemoveNode(t *testing.T) {
+    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        addTestNode(nodes, i, partitions)
+        addTestNode(clock, nodes, i, partitions)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -83,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
@@ -103,7 +109,9 @@ func TestStartThenRemoveNode(t *testing.T) {
     nodes[leader-1].Shutdown()
     nodes[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++ {
@@ -114,7 +122,9 @@ func TestStartThenRemoveNode(t *testing.T) {
     nodes[leader-1].Shutdown()
     nodes[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++ {
@@ -125,7 +135,9 @@ func TestStartThenRemoveNode(t *testing.T) {
     nodes[leader-1].Shutdown()
     nodes[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)
@@ -137,19 +149,22 @@ func TestStartThenRemoveNode(t *testing.T) {
 }
 
 func TestAddNode(t *testing.T) {
+    clock := clockwork.NewFakeClock()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
     for i := 0; i < 5; i++ {
-        addTestNode(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(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})
@@ -158,7 +173,9 @@ 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)
     for _, node := range nodes {
         node.Shutdown()