chore: add Now, Since, AfterFunc to clock; use clock for configmaps test (#13476)

* chore: add Now, Since, AfterFunc to clock; use clock for configmaps test

* chore: update flake.nix vendor hash

Signed-off-by: Spike Curtis <spike@coder.com>

---------

Signed-off-by: Spike Curtis <spike@coder.com>

clock/clock.go

@@ -17,6 +17,15 @@ type Clock interface {
 	// NewTimer creates a new Timer that will send the current time on its channel after at least
 	// duration d.
 	NewTimer(d time.Duration, tags ...string) *Timer
+	// AfterFunc waits for the duration to elapse and then calls f in its own goroutine. It returns
+	// a Timer that can be used to cancel the call using its Stop method. The returned Timer's C
+	// field is not used and will be nil.
+	AfterFunc(d time.Duration, f func(), tags ...string) *Timer
+
+	// Now returns the current local time.
+	Now(tags ...string) time.Time
+	// Since returns the time elapsed since t. It is shorthand for Clock.Now().Sub(t).
+	Since(t time.Time, tags ...string) time.Duration
 }
 
 // Waiter can be waited on for an error.

clock/mock.go

@@ -71,6 +71,46 @@ func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
 	return t
 }
 
+func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {
+	if d < 0 {
+		panic("duration must be positive or zero")
+	}
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.matchCallLocked(&Call{
+		fn:       clockFunctionAfterFunc,
+		Duration: d,
+		Tags:     tags,
+	})
+	t := &Timer{
+		nxt:  m.cur.Add(d),
+		fn:   f,
+		mock: m,
+	}
+	m.addTimerLocked(t)
+	return t
+}
+
+func (m *Mock) Now(tags ...string) time.Time {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.matchCallLocked(&Call{
+		fn:   clockFunctionNow,
+		Tags: tags,
+	})
+	return m.cur
+}
+
+func (m *Mock) Since(t time.Time, tags ...string) time.Duration {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.matchCallLocked(&Call{
+		fn:   clockFunctionSince,
+		Tags: tags,
+	})
+	return m.cur.Sub(t)
+}
+
 func (m *Mock) addTimerLocked(t *Timer) {
 	m.all = append(m.all, t)
 	m.recomputeNextLocked()
@@ -142,9 +182,13 @@ func (m *Mock) matchCallLocked(c *Call) {
 	m.mu.Lock()
 }
 
-// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for tick functions of
-// tickers created using TickerFunc to complete. If multiple timers or tickers trigger
-// simultaneously, they are all run on separate go routines.
+// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for:
+//
+//   - tick functions of tickers created using NewContextTicker
+//   - functions passed to AfterFunc
+//
+// to complete. If multiple timers or tickers trigger simultaneously, they are all run on separate
+// go routines.
 type AdvanceWaiter struct {
 	ch chan struct{}
 }
@@ -191,7 +235,7 @@ func (m *Mock) Advance(d time.Duration) AdvanceWaiter {
 
 func (m *Mock) advanceLocked(d time.Duration) {
 	if m.advancing {
-		panic("multiple simultaneous calls to Advance not supported")
+		panic("multiple simultaneous calls to Advance/Set not supported")
 	}
 	m.advancing = true
 	defer func() {
@@ -263,6 +307,10 @@ func (t Trapper) NewTimer(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionNewTimer, tags)
 }
 
+func (t Trapper) AfterFunc(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionAfterFunc, tags)
+}
+
 func (t Trapper) TimerStop(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionTimerStop, tags)
 }
@@ -279,6 +327,14 @@ func (t Trapper) TickerFuncWait(tags ...string) *Trap {
 	return t.mock.newTrap(clockFunctionTickerFuncWait, tags)
 }
 
+func (t Trapper) Now(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionNow, tags)
+}
+
+func (t Trapper) Since(tags ...string) *Trap {
+	return t.mock.newTrap(clockFunctionSince, tags)
+}
+
 func (m *Mock) Trap() Trapper {
 	return Trapper{m}
 }
@@ -386,10 +442,13 @@ type clockFunction int
 
 const (
 	clockFunctionNewTimer clockFunction = iota
+	clockFunctionAfterFunc
 	clockFunctionTimerStop
 	clockFunctionTimerReset
 	clockFunctionTickerFunc
 	clockFunctionTickerFuncWait
+	clockFunctionNow
+	clockFunctionSince
 )
 
 type callArg func(c *Call)

clock/real.go

@@ -55,4 +55,17 @@ func (realClock) NewTimer(d time.Duration, _ ...string) *Timer {
 	return &Timer{C: rt.C, timer: rt}
 }
 
+func (realClock) AfterFunc(d time.Duration, f func(), _ ...string) *Timer {
+	rt := time.AfterFunc(d, f)
+	return &Timer{C: rt.C, timer: rt}
+}
+
+func (realClock) Now(_ ...string) time.Time {
+	return time.Now()
+}
+
+func (realClock) Since(t time.Time, _ ...string) time.Duration {
+	return time.Since(t)
+}
+
 var _ Clock = realClock{}

clock/timer.go

@@ -18,9 +18,10 @@ func (t *Timer) fire(tt time.Time) {
 		panic("mock timer fired at wrong time")
 	}
 	t.mock.removeTimer(t)
-	t.c <- tt
 	if t.fn != nil {
 		t.fn()
+	} else {
+		t.c <- tt
 	}
 }
 

flake.nix

@@ -97,7 +97,7 @@
             name = "coder-${osArch}";
             # Updated with ./scripts/update-flake.sh`.
             # This should be updated whenever go.mod changes!
-            vendorHash = "sha256-PDA+Yd/tYvMTJfw8zyperTYnSBijvECc6XjxqnYgtkw=";
+            vendorHash = "sha256-BVgjKZeVogPb/kyCtZw/R898TI4YGnu9oWzzxHSVIyY=";
             proxyVendor = true;
             src = ./.;
             nativeBuildInputs = with pkgs; [ getopt openssl zstd ];

go.mod

@@ -198,7 +198,6 @@ require (
 require go.uber.org/mock v0.4.0
 
 require (
-	github.com/benbjohnson/clock v1.3.5
 	github.com/coder/serpent v0.7.0
 	github.com/gomarkdown/markdown v0.0.0-20231222211730-1d6d20845b47
 	github.com/google/go-github/v61 v61.0.0

go.sum

@@ -126,8 +126,6 @@ github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiE
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
-github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o=
-github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bep/clocks v0.5.0 h1:hhvKVGLPQWRVsBP/UB7ErrHYIO42gINVbvqxvYTPVps=

tailnet/configmaps.go

@@ -9,7 +9,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/benbjohnson/clock"
 	"github.com/google/uuid"
 	"go4.org/netipx"
 	"tailscale.com/ipn/ipnstate"
@@ -25,6 +24,7 @@ import (
 	"tailscale.com/wgengine/wgcfg/nmcfg"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/clock"
 	"github.com/coder/coder/v2/tailnet/proto"
 )
 
@@ -116,7 +116,7 @@ func newConfigMaps(logger slog.Logger, engine engineConfigurable, nodeID tailcfg
 			}},
 		},
 		peers: make(map[uuid.UUID]*peerLifecycle),
-		clock: clock.New(),
+		clock: clock.NewReal(),
 	}
 	go c.configLoop()
 	return c

tailnet/configmaps_internal_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/benbjohnson/clock"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -22,6 +21,7 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/clock"
 	"github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -195,9 +195,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_neverConfigures(t *testing.
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -241,9 +239,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_outOfOrder(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -314,9 +310,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -387,9 +381,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -412,7 +404,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
 	}
 	uut.updatePeers(u1)
 
-	mClock.Add(5 * time.Second)
+	mClock.Advance(5*time.Second).MustWait(ctx, t)
 
 	// it should now send the peer to the netmap
 
@@ -574,9 +566,8 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
+	start := mClock.Now()
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -600,7 +591,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
 	require.Len(t, r.wg.Peers, 1)
 	_ = testutil.RequireRecvCtx(ctx, t, s1)
 
-	mClock.Add(5 * time.Second)
+	mClock.Advance(5*time.Second).MustWait(ctx, t)
 
 	s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
 
@@ -621,7 +612,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
 	// latest handshake has advanced by a minute, so we don't remove the peer.
 	lh := start.Add(time.Minute)
 	s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
-	mClock.Add(lostTimeout)
+	mClock.Advance(lostTimeout).MustWait(ctx, t)
 	_ = testutil.RequireRecvCtx(ctx, t, s3)
 	select {
 	case <-fEng.setNetworkMap:
@@ -630,18 +621,10 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
 		// OK!
 	}
 
-	// Before we update the clock again, we need to be sure the timeout has
-	// completed running. To do that, we check the new lastHandshake has been set
-	require.Eventually(t, func() bool {
-		uut.L.Lock()
-		defer uut.L.Unlock()
-		return uut.peers[p1ID].lastHandshake == lh
-	}, testutil.WaitShort, testutil.IntervalFast)
-
 	// Advance the clock again by a minute, which should trigger the reprogrammed
 	// timeout.
 	s4 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
-	mClock.Add(time.Minute)
+	mClock.Advance(time.Minute).MustWait(ctx, t)
 
 	nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
 	r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
@@ -667,9 +650,8 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
+	start := mClock.Now()
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -693,7 +675,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
 	require.Len(t, r.wg.Peers, 1)
 	_ = testutil.RequireRecvCtx(ctx, t, s1)
 
-	mClock.Add(5 * time.Second)
+	mClock.Advance(5*time.Second).MustWait(ctx, t)
 
 	s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
 
@@ -710,7 +692,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
 		// OK!
 	}
 
-	mClock.Add(5 * time.Second)
+	mClock.Advance(5*time.Second).MustWait(ctx, t)
 	s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
 
 	updates[0].Kind = proto.CoordinateResponse_PeerUpdate_NODE
@@ -727,7 +709,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
 
 	// When we advance the clock, nothing happens because the timeout was
 	// canceled
-	mClock.Add(lostTimeout)
+	mClock.Advance(lostTimeout).MustWait(ctx, t)
 	select {
 	case <-fEng.setNetworkMap:
 		t.Fatal("should not reprogram")
@@ -753,9 +735,8 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
 	discoKey := key.NewDisco()
 	uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
 	defer uut.close()
-	start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
 	mClock := clock.NewMock()
-	mClock.Set(start)
+	start := mClock.Now()
 	uut.clock = mClock
 
 	p1ID := uuid.UUID{1}
@@ -788,7 +769,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
 	require.Len(t, r.wg.Peers, 2)
 	_ = testutil.RequireRecvCtx(ctx, t, s1)
 
-	mClock.Add(5 * time.Second)
+	mClock.Advance(5*time.Second).MustWait(ctx, t)
 	uut.setAllPeersLost()
 
 	// No reprogramming yet, since we keep the peer around.
@@ -802,7 +783,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
 	// When we advance the clock, even by a few ms, the timeout for peer 2 pops
 	// because our status only includes a handshake for peer 1
 	s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
-	mClock.Add(time.Millisecond * 10)
+	mClock.Advance(time.Millisecond*10).MustWait(ctx, t)
 	_ = testutil.RequireRecvCtx(ctx, t, s2)
 
 	nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
@@ -812,7 +793,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
 
 	// Finally, advance the clock until after the timeout
 	s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
-	mClock.Add(lostTimeout)
+	mClock.Advance(lostTimeout).MustWait(ctx, t)
 	_ = testutil.RequireRecvCtx(ctx, t, s3)
 
 	nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)