Skip to content

Commit ad55c7a

Browse files
spikecurtisspikecurtis
committed
chore: add Now, Since, AfterFunc to clock; use clock for configmaps test
1 parent ddbd5b9 commit ad55c7a

File tree

8 files changed

+103
-43
lines changed

8 files changed

+103
-43
lines changed

‎clock/clock.go

+9
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,15 @@ type Clock interface {
1717
// NewTimer creates a new Timer that will send the current time on its channel after at least
1818
// duration d.
1919
NewTimer(d time.Duration, tags ...string) *Timer
20+
// AfterFunc waits for the duration to elapse and then calls f in its own goroutine. It returns
21+
// a Timer that can be used to cancel the call using its Stop method. The returned Timer's C
22+
// field is not used and will be nil.
23+
AfterFunc(d time.Duration, f func(), tags ...string) *Timer
24+
25+
// Now returns the current local time.
26+
Now(tags ...string) time.Time
27+
// Since returns the time elapsed since t. It is shorthand for Clock.Now().Sub(t).
28+
Since(t time.Time, tags ...string) time.Duration
2029
}
2130

2231
// Waiter can be waited on for an error.

‎clock/mock.go

+63-4
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,46 @@ func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {
7171
return t
7272
}
7373

74+
func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {
75+
if d < 0 {
76+
panic("duration must be positive or zero")
77+
}
78+
m.mu.Lock()
79+
defer m.mu.Unlock()
80+
m.matchCallLocked(&Call{
81+
fn: clockFunctionAfterFunc,
82+
Duration: d,
83+
Tags: tags,
84+
})
85+
t := &Timer{
86+
nxt: m.cur.Add(d),
87+
fn: f,
88+
mock: m,
89+
}
90+
m.addTimerLocked(t)
91+
return t
92+
}
93+
94+
func (m *Mock) Now(tags ...string) time.Time {
95+
m.mu.Lock()
96+
defer m.mu.Unlock()
97+
m.matchCallLocked(&Call{
98+
fn: clockFunctionNow,
99+
Tags: tags,
100+
})
101+
return m.cur
102+
}
103+
104+
func (m *Mock) Since(t time.Time, tags ...string) time.Duration {
105+
m.mu.Lock()
106+
defer m.mu.Unlock()
107+
m.matchCallLocked(&Call{
108+
fn: clockFunctionSince,
109+
Tags: tags,
110+
})
111+
return m.cur.Sub(t)
112+
}
113+
74114
func (m *Mock) addTimerLocked(t *Timer) {
75115
m.all = append(m.all, t)
76116
m.recomputeNextLocked()
@@ -142,9 +182,13 @@ func (m *Mock) matchCallLocked(c *Call) {
142182
m.mu.Lock()
143183
}
144184

145-
// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for tick functions of
146-
// tickers created using TickerFunc to complete. If multiple timers or tickers trigger
147-
// simultaneously, they are all run on separate go routines.
185+
// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for:
186+
//
187+
// - tick functions of tickers created using NewContextTicker
188+
// - functions passed to AfterFunc
189+
//
190+
// to complete. If multiple timers or tickers trigger simultaneously, they are all run on separate
191+
// go routines.
148192
type AdvanceWaiter struct {
149193
ch chan struct{}
150194
}
@@ -190,7 +234,7 @@ func (m *Mock) Advance(d time.Duration) AdvanceWaiter {
190234

191235
func (m *Mock) advanceLocked(d time.Duration) {
192236
if m.advancing {
193-
panic("multiple simultaneous calls to Advance not supported")
237+
panic("multiple simultaneous calls to Advance/Set not supported")
194238
}
195239
m.advancing = true
196240
defer func() {
@@ -262,6 +306,10 @@ func (t Trapper) NewTimer(tags ...string) *Trap {
262306
return t.mock.newTrap(clockFunctionNewTimer, tags)
263307
}
264308

309+
func (t Trapper) AfterFunc(tags ...string) *Trap {
310+
return t.mock.newTrap(clockFunctionAfterFunc, tags)
311+
}
312+
265313
func (t Trapper) TimerStop(tags ...string) *Trap {
266314
return t.mock.newTrap(clockFunctionTimerStop, tags)
267315
}
@@ -278,6 +326,14 @@ func (t Trapper) TickerFuncWait(tags ...string) *Trap {
278326
return t.mock.newTrap(clockFunctionTickerFuncWait, tags)
279327
}
280328

329+
func (t Trapper) Now(tags ...string) *Trap {
330+
return t.mock.newTrap(clockFunctionNow, tags)
331+
}
332+
333+
func (t Trapper) Since(tags ...string) *Trap {
334+
return t.mock.newTrap(clockFunctionSince, tags)
335+
}
336+
281337
func (m *Mock) Trap() Trapper {
282338
return Trapper{m}
283339
}
@@ -385,10 +441,13 @@ type clockFunction int
385441

386442
const (
387443
clockFunctionNewTimer clockFunction = iota
444+
clockFunctionAfterFunc
388445
clockFunctionTimerStop
389446
clockFunctionTimerReset
390447
clockFunctionTickerFunc
391448
clockFunctionTickerFuncWait
449+
clockFunctionNow
450+
clockFunctionSince
392451
)
393452

394453
type callArg func(c *Call)

‎clock/real.go

+13
Original file line numberDiff line numberDiff line change
@@ -55,4 +55,17 @@ func (realClock) NewTimer(d time.Duration, _ ...string) *Timer {
5555
return &Timer{C: rt.C, timer: rt}
5656
}
5757

58+
func (realClock) AfterFunc(d time.Duration, f func(), _ ...string) *Timer {
59+
rt := time.AfterFunc(d, f)
60+
return &Timer{C: rt.C, timer: rt}
61+
}
62+
63+
func (realClock) Now(_ ...string) time.Time {
64+
return time.Now()
65+
}
66+
67+
func (realClock) Since(t time.Time, _ ...string) time.Duration {
68+
return time.Since(t)
69+
}
70+
5871
var _ Clock = realClock{}

‎clock/timer.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@ func (t *Timer) fire(tt time.Time) {
1818
panic("mock timer fired at wrong time")
1919
}
2020
t.mock.removeTimer(t)
21-
t.c <- tt
2221
if t.fn != nil {
2322
t.fn()
23+
} else {
24+
t.c <- tt
2425
}
2526
}
2627

‎go.mod

-1
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,6 @@ require (
198198
require go.uber.org/mock v0.4.0
199199

200200
require (
201-
github.com/benbjohnson/clock v1.3.5
202201
github.com/coder/serpent v0.7.0
203202
github.com/gomarkdown/markdown v0.0.0-20231222211730-1d6d20845b47
204203
github.com/google/go-github/v61 v61.0.0

‎go.sum

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

‎tailnet/configmaps.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99
"sync"
1010
"time"
1111

12-
"github.com/benbjohnson/clock"
1312
"github.com/google/uuid"
1413
"go4.org/netipx"
1514
"tailscale.com/ipn/ipnstate"
@@ -25,6 +24,7 @@ import (
2524
"tailscale.com/wgengine/wgcfg/nmcfg"
2625

2726
"cdr.dev/slog"
27+
"github.com/coder/coder/v2/clock"
2828
"github.com/coder/coder/v2/tailnet/proto"
2929
)
3030

@@ -116,7 +116,7 @@ func newConfigMaps(logger slog.Logger, engine engineConfigurable, nodeID tailcfg
116116
}},
117117
},
118118
peers: make(map[uuid.UUID]*peerLifecycle),
119-
clock: clock.New(),
119+
clock: clock.NewReal(),
120120
}
121121
go c.configLoop()
122122
return c

‎tailnet/configmaps_internal_test.go

+14-33
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77
"testing"
88
"time"
99

10-
"github.com/benbjohnson/clock"
1110
"github.com/google/uuid"
1211
"github.com/stretchr/testify/assert"
1312
"github.com/stretchr/testify/require"
@@ -22,6 +21,7 @@ import (
2221

2322
"cdr.dev/slog"
2423
"cdr.dev/slog/sloggers/slogtest"
24+
"github.com/coder/coder/v2/clock"
2525
"github.com/coder/coder/v2/tailnet/proto"
2626
"github.com/coder/coder/v2/testutil"
2727
)
@@ -195,9 +195,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_neverConfigures(t *testing.
195195
discoKey := key.NewDisco()
196196
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
197197
defer uut.close()
198-
start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
199198
mClock := clock.NewMock()
200-
mClock.Set(start)
201199
uut.clock = mClock
202200

203201
p1ID := uuid.UUID{1}
@@ -241,9 +239,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_outOfOrder(t *testing.T) {
241239
discoKey := key.NewDisco()
242240
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
243241
defer uut.close()
244-
start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
245242
mClock := clock.NewMock()
246-
mClock.Set(start)
247243
uut.clock = mClock
248244

249245
p1ID := uuid.UUID{1}
@@ -314,9 +310,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake(t *testing.T) {
314310
discoKey := key.NewDisco()
315311
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
316312
defer uut.close()
317-
start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
318313
mClock := clock.NewMock()
319-
mClock.Set(start)
320314
uut.clock = mClock
321315

322316
p1ID := uuid.UUID{1}
@@ -387,9 +381,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
387381
discoKey := key.NewDisco()
388382
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
389383
defer uut.close()
390-
start := time.Date(2024, time.March, 29, 8, 0, 0, 0, time.UTC)
391384
mClock := clock.NewMock()
392-
mClock.Set(start)
393385
uut.clock = mClock
394386

395387
p1ID := uuid.UUID{1}
@@ -412,7 +404,7 @@ func TestConfigMaps_updatePeers_new_waitForHandshake_timeout(t *testing.T) {
412404
}
413405
uut.updatePeers(u1)
414406

415-
mClock.Add(5 * time.Second)
407+
mClock.Advance(5*time.Second).MustWait(ctx, t)
416408

417409
// it should now send the peer to the netmap
418410

@@ -574,9 +566,8 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
574566
discoKey := key.NewDisco()
575567
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
576568
defer uut.close()
577-
start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
578569
mClock := clock.NewMock()
579-
mClock.Set(start)
570+
start := mClock.Now()
580571
uut.clock = mClock
581572

582573
p1ID := uuid.UUID{1}
@@ -600,7 +591,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
600591
require.Len(t, r.wg.Peers, 1)
601592
_ = testutil.RequireRecvCtx(ctx, t, s1)
602593

603-
mClock.Add(5 * time.Second)
594+
mClock.Advance(5*time.Second).MustWait(ctx, t)
604595

605596
s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
606597

@@ -621,7 +612,7 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
621612
// latest handshake has advanced by a minute, so we don't remove the peer.
622613
lh := start.Add(time.Minute)
623614
s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
624-
mClock.Add(lostTimeout)
615+
mClock.Advance(lostTimeout).MustWait(ctx, t)
625616
_ = testutil.RequireRecvCtx(ctx, t, s3)
626617
select {
627618
case <-fEng.setNetworkMap:
@@ -630,18 +621,10 @@ func TestConfigMaps_updatePeers_lost(t *testing.T) {
630621
// OK!
631622
}
632623

633-
// Before we update the clock again, we need to be sure the timeout has
634-
// completed running. To do that, we check the new lastHandshake has been set
635-
require.Eventually(t, func() bool {
636-
uut.L.Lock()
637-
defer uut.L.Unlock()
638-
return uut.peers[p1ID].lastHandshake == lh
639-
}, testutil.WaitShort, testutil.IntervalFast)
640-
641624
// Advance the clock again by a minute, which should trigger the reprogrammed
642625
// timeout.
643626
s4 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, lh)
644-
mClock.Add(time.Minute)
627+
mClock.Advance(time.Minute).MustWait(ctx, t)
645628

646629
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
647630
r = testutil.RequireRecvCtx(ctx, t, fEng.reconfig)
@@ -667,9 +650,8 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
667650
discoKey := key.NewDisco()
668651
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
669652
defer uut.close()
670-
start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
671653
mClock := clock.NewMock()
672-
mClock.Set(start)
654+
start := mClock.Now()
673655
uut.clock = mClock
674656

675657
p1ID := uuid.UUID{1}
@@ -693,7 +675,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
693675
require.Len(t, r.wg.Peers, 1)
694676
_ = testutil.RequireRecvCtx(ctx, t, s1)
695677

696-
mClock.Add(5 * time.Second)
678+
mClock.Advance(5*time.Second).MustWait(ctx, t)
697679

698680
s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
699681

@@ -710,7 +692,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
710692
// OK!
711693
}
712694

713-
mClock.Add(5 * time.Second)
695+
mClock.Advance(5*time.Second).MustWait(ctx, t)
714696
s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
715697

716698
updates[0].Kind = proto.CoordinateResponse_PeerUpdate_NODE
@@ -727,7 +709,7 @@ func TestConfigMaps_updatePeers_lost_and_found(t *testing.T) {
727709

728710
// When we advance the clock, nothing happens because the timeout was
729711
// canceled
730-
mClock.Add(lostTimeout)
712+
mClock.Advance(lostTimeout).MustWait(ctx, t)
731713
select {
732714
case <-fEng.setNetworkMap:
733715
t.Fatal("should not reprogram")
@@ -753,9 +735,8 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
753735
discoKey := key.NewDisco()
754736
uut := newConfigMaps(logger, fEng, nodeID, nodePrivateKey, discoKey.Public())
755737
defer uut.close()
756-
start := time.Date(2024, time.January, 1, 8, 0, 0, 0, time.UTC)
757738
mClock := clock.NewMock()
758-
mClock.Set(start)
739+
start := mClock.Now()
759740
uut.clock = mClock
760741

761742
p1ID := uuid.UUID{1}
@@ -788,7 +769,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
788769
require.Len(t, r.wg.Peers, 2)
789770
_ = testutil.RequireRecvCtx(ctx, t, s1)
790771

791-
mClock.Add(5 * time.Second)
772+
mClock.Advance(5*time.Second).MustWait(ctx, t)
792773
uut.setAllPeersLost()
793774

794775
// No reprogramming yet, since we keep the peer around.
@@ -802,7 +783,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
802783
// When we advance the clock, even by a few ms, the timeout for peer 2 pops
803784
// because our status only includes a handshake for peer 1
804785
s2 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
805-
mClock.Add(time.Millisecond * 10)
786+
mClock.Advance(time.Millisecond*10).MustWait(ctx, t)
806787
_ = testutil.RequireRecvCtx(ctx, t, s2)
807788

808789
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)
@@ -812,7 +793,7 @@ func TestConfigMaps_setAllPeersLost(t *testing.T) {
812793

813794
// Finally, advance the clock until after the timeout
814795
s3 := expectStatusWithHandshake(ctx, t, fEng, p1Node.Key, start)
815-
mClock.Add(lostTimeout)
796+
mClock.Advance(lostTimeout).MustWait(ctx, t)
816797
_ = testutil.RequireRecvCtx(ctx, t, s3)
817798

818799
nm = testutil.RequireRecvCtx(ctx, t, fEng.setNetworkMap)

0 commit comments

Comments
 (0)