Fake timer for tests?

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

@@ -15,12 +15,11 @@
 package tinyraft
 
 import (
+    "context"
     "fmt"
     "math/rand"
     "sync"
     "time"
-
-    "github.com/jonboulle/clockwork"
 )
 
 type MsgType int8
@@ -61,15 +60,17 @@ type RaftConfig struct {
     HeartbeatTimeout  time.Duration
     LeadershipTimeout time.Duration
     SendFunction      func(to NodeID, msg RaftMsg)
+    FakeTimer         func(forHeartbeat bool, duration time.Duration) <-chan time.Time
 }
 
 type Raft struct {
+    ctx      context.Context
     config   RaftConfig
     send     func(to NodeID, msg RaftMsg)
     onChange func(state NodeState, term uint64, leader NodeID, followers []NodeID)
+    timerFn  func(forHeartbeat bool, duration time.Duration) <-chan time.Time
 
-    mu         sync.Mutex
-    shutdownCh chan struct{}
+    mu sync.Mutex
 
     term      uint64
     state     NodeState
@@ -83,23 +84,26 @@ type Raft struct {
     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
     }
     if len(config.Nodes) == 0 {
         config.Nodes = append(config.Nodes, config.NodeID)
     }
+    timerFn := config.FakeTimer
+    if timerFn == nil {
+        timerFn = func(forHeartbeat bool, duration time.Duration) <-chan time.Time {
+            return time.After(duration)
+        }
+    }
     r := &Raft{
+        ctx:     ctx,
         config:  config,
         send:    config.SendFunction,
+        timerFn: timerFn,
         term:    config.InitialTerm,
     }
-    r.shutdownCh = make(chan struct{})
     r.electionTimerCh = make(chan time.Duration, 1)
     go func() {
         ok := true
@@ -107,16 +111,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 <-r.timerFn(false, dur):
                     r.Start()
                 }
             }
@@ -125,9 +129,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 <-r.timerFn(true, r.config.HeartbeatTimeout):
                 r.heartbeat()
             }
         }
@@ -135,16 +139,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
 }
@@ -160,9 +154,6 @@ 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()
 }
 
@@ -204,8 +195,7 @@ func (r *Raft) changed() {
 func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
     r.mu.Lock()
     defer r.mu.Unlock()
-    if r.shutdownCh == nil {
-    } else if msg.Type == MsgVoteRequest {
+    if msg.Type == MsgVoteRequest {
         if msg.Term > r.term && msg.Leader != NoLeader {
             r.leader = msg.Leader
             r.term = msg.Term
@@ -316,12 +306,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:
@@ -336,10 +320,6 @@ 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,16 +1,17 @@
-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, nodes []*Raft, i int, partitions map[int]bool) {
     from := NodeID(i + 1)
-    n := newRaft(RaftConfig{
+    n := NewRaft(ctx, RaftConfig{
         NodeID: from,
         Nodes: []NodeID{ 1, 2, 3, 4, 5 },
         HeartbeatTimeout: 100 * time.Millisecond,
@@ -25,7 +26,7 @@ func addTestNode(clock clockwork.Clock, nodes []*Raft, i int, partitions map[int
                 }()
             }
         },
-    }, 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)
     })
@@ -67,11 +68,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, nodes, i, partitions)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -86,10 +89,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 +106,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 +118,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,39 +130,36 @@ 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, 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, 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})
@@ -173,11 +168,7 @@ 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()
-    }
+    cancel()
 }