Merge branch 'master' into fix/afterfunc_goroutine

.github/workflows/build.yml

@@ -0,0 +1,37 @@
+name: Go
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+    branches: [master]
+
+jobs:
+
+  build:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        # Test on one "earliest" Go as well as the latest two major Go
+        # versions. If some change requires bumping the "earliest" Go versiion,
+        # that's fine - just include that in the commit description so that
+        # users are aware.
+        go: ["1.16.x", "1.19.x", "1.20.x"]
+
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Set up Go
+      uses: actions/setup-go@v2
+      with:
+        go-version: ${{ matrix.go }}
+
+    - name: Load cached dependencies
+      uses: actions/cache@v2
+      with:
+        path: ~/go/pkg/mod
+        key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
+
+    - name: Test
+      run: go test -v -race ./...

clock.go

@@ -165,9 +165,17 @@ func (m *Mock) After(d time.Duration) <-chan time.Time {
 // AfterFunc waits for the duration to elapse and then executes a function in its own goroutine.
 // A Timer is returned that can be stopped.
 func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
-	t := m.Timer(d)
-	t.C = nil
-	t.fn = f
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	ch := make(chan time.Time, 1)
+	t := &Timer{
+		c:       ch,
+		fn:      f,
+		mock:    m,
+		next:    m.now.Add(d),
+		stopped: false,
+	}
+	m.timers = append(m.timers, (*internalTimer)(t))
 	return t
 }
 
@@ -324,12 +332,13 @@ func (t *internalTimer) Tick(now time.Time) {
 
 // Ticker holds a channel that receives "ticks" at regular intervals.
 type Ticker struct {
-	C      <-chan time.Time
-	c      chan time.Time
-	ticker *time.Ticker  // realtime impl, if set
-	next   time.Time     // next tick time
-	mock   *Mock         // mock clock, if set
-	d      time.Duration // time between ticks
+	C       <-chan time.Time
+	c       chan time.Time
+	ticker  *time.Ticker  // realtime impl, if set
+	next    time.Time     // next tick time
+	mock    *Mock         // mock clock, if set
+	d       time.Duration // time between ticks
+	stopped bool          // True if stopped, false if running
 }
 
 // Stop turns off the ticker.
@@ -339,6 +348,7 @@ func (t *Ticker) Stop() {
 	} else {
 		t.mock.mu.Lock()
 		t.mock.removeClockTimer((*internalTicker)(t))
+		t.stopped = true
 		t.mock.mu.Unlock()
 	}
 }
@@ -353,6 +363,11 @@ func (t *Ticker) Reset(dur time.Duration) {
 	t.mock.mu.Lock()
 	defer t.mock.mu.Unlock()
 
+	if t.stopped {
+		t.mock.timers = append(t.mock.timers, (*internalTicker)(t))
+		t.stopped = false
+	}
+
 	t.d = dur
 	t.next = t.mock.now.Add(dur)
 }

clock_test.go

@@ -128,6 +128,26 @@ func TestClock_Ticker_Rst(t *testing.T) {
 	ticker.Stop()
 }
 
+// Ensure that the clock's ticker can stop and then be reset correctly.
+func TestClock_Ticker_Stop_Rst(t *testing.T) {
+	start := time.Now()
+	ticker := New().Ticker(20 * time.Millisecond)
+	<-ticker.C
+	ticker.Stop()
+	select {
+	case <-ticker.C:
+		t.Fatal("unexpected send")
+	case <-time.After(30 * time.Millisecond):
+	}
+	ticker.Reset(5 * time.Millisecond)
+	<-ticker.C
+	dur := time.Since(start)
+	if dur >= 60*time.Millisecond {
+		t.Fatal("took more than 60ms")
+	}
+	ticker.Stop()
+}
+
 // Ensure that the clock's timer waits correctly.
 func TestClock_Timer(t *testing.T) {
 	start := time.Now()
@@ -475,6 +495,52 @@ func TestMock_Ticker_Reset(t *testing.T) {
 	}
 }
 
+func TestMock_Ticker_Stop_Reset(t *testing.T) {
+	var n int32
+	clock := NewMock()
+
+	ticker := clock.Ticker(5 * time.Second)
+	defer ticker.Stop()
+
+	go func() {
+		for {
+			<-ticker.C
+			atomic.AddInt32(&n, 1)
+		}
+	}()
+	gosched()
+
+	// Move clock forward.
+	clock.Add(10 * time.Second)
+	if atomic.LoadInt32(&n) != 2 {
+		t.Fatalf("expected 2, got: %d", n)
+	}
+
+	ticker.Stop()
+
+	// Move clock forward again.
+	clock.Add(5 * time.Second)
+	if atomic.LoadInt32(&n) != 2 {
+		t.Fatalf("still expected 2, got: %d", n)
+	}
+
+	ticker.Reset(2 * time.Second)
+
+	// Advance the remaining 2 seconds
+	clock.Add(2 * time.Second)
+
+	if atomic.LoadInt32(&n) != 3 {
+		t.Fatalf("expected 3, got: %d", n)
+	}
+
+	// Advance another 2 seconds
+	clock.Add(2 * time.Second)
+
+	if atomic.LoadInt32(&n) != 4 {
+		t.Fatalf("expected 4, got: %d", n)
+	}
+}
+
 // Ensure that multiple tickers can be used together.
 func TestMock_Ticker_Multi(t *testing.T) {
 	var n int32
@@ -651,5 +717,44 @@ func TestMock_ReentrantDeadlock(t *testing.T) {
 	mockedClock.Add(15 * time.Second)
 }
 
+func TestMock_AddAfterFuncRace(t *testing.T) {
+	// start blocks the goroutines in this test
+	// until we're ready for them to race.
+	start := make(chan struct{})
+
+	var wg sync.WaitGroup
+
+	mockedClock := NewMock()
+
+	called := false
+	defer func() {
+		if !called {
+			t.Errorf("AfterFunc did not call the function")
+		}
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		<-start
+
+		mockedClock.AfterFunc(time.Millisecond, func() {
+			called = true
+		})
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		<-start
+
+		mockedClock.Add(time.Millisecond)
+		mockedClock.Add(time.Millisecond)
+	}()
+
+	close(start) // unblock the goroutines
+	wg.Wait()    // and wait for them
+}
+
 func warn(v ...interface{})              { fmt.Fprintln(os.Stderr, v...) }
 func warnf(msg string, v ...interface{}) { fmt.Fprintf(os.Stderr, msg+"\n", v...) }