coder · spikecurtis · Jun 20, 2024 · Jun 18, 2024
diff --git a/clock/clock.go b/clock/clock.go
@@ -10,9 +10,16 @@ import (
 )
 
 type Clock interface {
+	// NewTicker returns a new Ticker containing a channel that will send the current time on the
+	// channel after each tick. The period of the ticks is specified by the duration argument. The
+	// ticker will adjust the time interval or drop ticks to make up for slow receivers. The
+	// duration d must be greater than zero; if not, NewTicker will panic. Stop the ticker to
+	// release associated resources.
+	NewTicker(d time.Duration, tags ...string) *Ticker
 	// TickerFunc is a convenience function that calls f on the interval d until either the given
 	// context expires or f returns an error.  Callers may call Wait() on the returned Waiter to
-	// wait until this happens and obtain the error.
+	// wait until this happens and obtain the error. The duration d must be greater than zero; if
+	// not, TickerFunc will panic.
 	TickerFunc(ctx context.Context, d time.Duration, f func() error, tags ...string) Waiter
 	// NewTimer creates a new Timer that will send the current time on its channel after at least
 	// duration d.

diff --git a/clock/mock.go b/clock/mock.go
@@ -32,6 +32,9 @@ type event interface {
 }
 
 func (m *Mock) TickerFunc(ctx context.Context, d time.Duration, f func() error, tags ...string) Waiter {
+	if d <= 0 {
+		panic("TickerFunc called with negative or zero duration")
+	}
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	c := newCall(clockFunctionTickerFunc, tags, withDuration(d))
@@ -51,6 +54,28 @@ func (m *Mock) TickerFunc(ctx context.Context, d time.Duration, f func() error,
 	return t
 }
 
+func (m *Mock) NewTicker(d time.Duration, tags ...string) *Ticker {
+	if d <= 0 {
+		panic("NewTicker called with negative or zero duration")
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	c := newCall(clockFunctionNewTicker, tags, withDuration(d))
+	m.matchCallLocked(c)
+	defer close(c.complete)
+	// 1 element buffer follows standard library implementation
+	ticks := make(chan time.Time, 1)
+	t := &Ticker{
+		C:    ticks,
+		c:    ticks,
+		d:    d,
+		nxt:  m.cur.Add(d),
+		mock: m,
+	}
+	m.addEventLocked(t)
+	return t
+}
+
 func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -70,7 +95,7 @@ func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
 		go t.fire(t.mock.cur)
 		return t
 	}
-	m.addTimerLocked(t)
+	m.addEventLocked(t)
 	return t
 }
 
@@ -91,7 +116,7 @@ func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {
 		go t.fire(t.mock.cur)
 		return t
 	}
-	m.addTimerLocked(t)
+	m.addEventLocked(t)
 	return t
 }
 
@@ -122,8 +147,8 @@ func (m *Mock) Until(t time.Time, tags ...string) time.Duration {
 	return t.Sub(m.cur)
 }
 
-func (m *Mock) addTimerLocked(t *Timer) {
-	m.all = append(m.all, t)
+func (m *Mock) addEventLocked(e event) {
+	m.all = append(m.all, e)
 	m.recomputeNextLocked()
 }
 
@@ -152,20 +177,12 @@ func (m *Mock) removeTimer(t *Timer) {
 }
 
 func (m *Mock) removeTimerLocked(t *Timer) {
-	defer m.recomputeNextLocked()
 	t.stopped = true
-	var e event = t
-	for i := range m.all {
-		if m.all[i] == e {
-			m.all = append(m.all[:i], m.all[i+1:]...)
-			return
-		}
-	}
+	m.removeEventLocked(t)
 }
 
-func (m *Mock) removeTickerFuncLocked(ct *mockTickerFunc) {
+func (m *Mock) removeEventLocked(e event) {
 	defer m.recomputeNextLocked()
-	var e event = ct
 	for i := range m.all {
 		if m.all[i] == e {
 			m.all = append(m.all[:i], m.all[i+1:]...)
@@ -371,6 +388,18 @@ func (t Trapper) TickerFuncWait(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionTickerFuncWait, tags)
 }
 
+func (t Trapper) NewTicker(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionNewTicker, tags)
+}
+
+func (t Trapper) TickerStop(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionTickerStop, tags)
+}
+
+func (t Trapper) TickerReset(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionTickerReset, tags)
+}
+
 func (t Trapper) Now(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionNow, tags)
 }
@@ -459,7 +488,7 @@ func (m *mockTickerFunc) exitLocked(err error) {
 	}
 	m.done = true
 	m.err = err
-	m.mock.removeTickerFuncLocked(m)
+	m.mock.removeEventLocked(m)
 	m.cond.Broadcast()
 }
 
@@ -493,6 +522,9 @@ const (
 	clockFunctionTimerReset
 	clockFunctionTickerFunc
 	clockFunctionTickerFuncWait
+	clockFunctionNewTicker
+	clockFunctionTickerReset
+	clockFunctionTickerStop
 	clockFunctionNow
 	clockFunctionSince
 	clockFunctionUntil

diff --git a/clock/mock_test.go b/clock/mock_test.go
@@ -80,3 +80,90 @@ func TestAfterFunc_NegativeDuration(t *testing.T) {
 		t.Fatal("timer still running")
 	}
 }
+
+func TestNewTicker(t *testing.T) {
+	t.Parallel()
+	// nolint:gocritic // trying to avoid Coder-specific stuff with an eye toward spinning this out
+	ctx, cancel := context.WithTimeout(context.Background(), 1000*time.Second)
+	defer cancel()
+
+	mClock := clock.NewMock(t)
+	start := mClock.Now()
+	trapNT := mClock.Trap().NewTicker("new")
+	defer trapNT.Close()
+	trapStop := mClock.Trap().TickerStop("stop")
+	defer trapStop.Close()
+	trapReset := mClock.Trap().TickerReset("reset")
+	defer trapReset.Close()
+
+	tickers := make(chan *clock.Ticker, 1)
+	go func() {
+		tickers <- mClock.NewTicker(time.Hour, "new")
+	}()
+	c := trapNT.MustWait(ctx)
+	c.Release()
+	if c.Duration != time.Hour {
+		t.Fatalf("expected time.Hour, got: %v", c.Duration)
+	}
+	tkr := <-tickers
+
+	for i := 0; i < 3; i++ {
+		mClock.Advance(time.Hour).MustWait(ctx)
+	}
+
+	// should get first tick, rest dropped
+	tTime := start.Add(time.Hour)
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for ticker")
+	case tick := <-tkr.C:
+		if !tick.Equal(tTime) {
+			t.Fatalf("expected time %v, got %v", tTime, tick)
+		}
+	}
+
+	go tkr.Reset(time.Minute, "reset")
+	c = trapReset.MustWait(ctx)
+	mClock.Advance(time.Second).MustWait(ctx)
+	c.Release()
+	if c.Duration != time.Minute {
+		t.Fatalf("expected time.Minute, got: %v", c.Duration)
+	}
+	mClock.Advance(time.Minute).MustWait(ctx)
+
+	// tick should show present time, ensuring the 2 hour ticks got dropped when
+	// we didn't read from the channel.
+	tTime = mClock.Now()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for ticker")
+	case tick := <-tkr.C:
+		if !tick.Equal(tTime) {
+			t.Fatalf("expected time %v, got %v", tTime, tick)
+		}
+	}
+
+	go tkr.Stop("stop")
+	trapStop.MustWait(ctx).Release()
+	mClock.Advance(time.Hour).MustWait(ctx)
+	select {
+	case <-tkr.C:
+		t.Fatal("ticker still running")
+	default:
+		// OK
+	}
+
+	// Resetting after stop
+	go tkr.Reset(time.Minute, "reset")
+	trapReset.MustWait(ctx).Release()
+	mClock.Advance(time.Minute).MustWait(ctx)
+	tTime = mClock.Now()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for ticker")
+	case tick := <-tkr.C:
+		if !tick.Equal(tTime) {
+			t.Fatalf("expected time %v, got %v", tTime, tick)
+		}
+	}
+}
diff --git a/clock/real.go b/clock/real.go
@@ -11,6 +11,11 @@ func NewReal() Clock {
 	return realClock{}
 }
 
+func (realClock) NewTicker(d time.Duration, _ ...string) *Ticker {
+	tkr := time.NewTicker(d)
+	return &Ticker{ticker: tkr, C: tkr.C}
+}
+
 func (realClock) TickerFunc(ctx context.Context, d time.Duration, f func() error, _ ...string) Waiter {
 	ct := &realContextTicker{
 		ctx: ctx,

diff --git a/clock/ticker.go b/clock/ticker.go
@@ -0,0 +1,68 @@
+package clock
+
+import "time"
+
+type Ticker struct {
+	C <-chan time.Time
+	//nolint: revive
+	c       chan time.Time
+	ticker  *time.Ticker  // realtime impl, if set
+	d       time.Duration // period, if set
+	nxt     time.Time     // next tick time
+	mock    *Mock         // mock clock, if set
+	stopped bool          // true if the ticker is not running
+}
+
+func (t *Ticker) fire(tt time.Time) {
+	t.mock.mu.Lock()
+	defer t.mock.mu.Unlock()
+	if t.stopped {
+		return
+	}
+	for !t.nxt.After(t.mock.cur) {
+		t.nxt = t.nxt.Add(t.d)
+	}
+	t.mock.recomputeNextLocked()
+	select {
+	case t.c <- tt:
+	default:
+	}
+}
+
+func (t *Ticker) next() time.Time {
+	return t.nxt
+}
+
+func (t *Ticker) Stop(tags ...string) {
+	if t.ticker != nil {
+		t.ticker.Stop()
+		return
+	}
+	t.mock.mu.Lock()
+	defer t.mock.mu.Unlock()
+	c := newCall(clockFunctionTickerStop, tags)
+	t.mock.matchCallLocked(c)
+	defer close(c.complete)
+	t.mock.removeEventLocked(t)
+	t.stopped = true
+}
+
+func (t *Ticker) Reset(d time.Duration, tags ...string) {
+	if t.ticker != nil {
+		t.ticker.Reset(d)
+		return
+	}
+	t.mock.mu.Lock()
+	defer t.mock.mu.Unlock()
+	c := newCall(clockFunctionTickerReset, tags, withDuration(d))
+	t.mock.matchCallLocked(c)
+	defer close(c.complete)
+	t.nxt = t.mock.cur.Add(d)
+	t.d = d
+	if t.stopped {
+		t.stopped = false
+		t.mock.addEventLocked(t)
+	} else {
+		t.mock.recomputeNextLocked()
+	}
+}
diff --git a/clock/timer.go b/clock/timer.go
@@ -64,6 +64,6 @@ func (t *Timer) Reset(d time.Duration, tags ...string) bool {
 	t.mock.removeTimerLocked(t)
 	t.stopped = false
 	t.nxt = t.mock.cur.Add(d)
-	t.mock.addTimerLocked(t)
+	t.mock.addEventLocked(t)
 	return result
 }