Skip to content

Commit 8aae16e

Browse files
spikecurtisspikecurtis
committed
chore: refactor autobuild/notify to use clock test
1 parent 4b0b9b0 commit 8aae16e

File tree

3 files changed

+80
-73
lines changed

3 files changed

+80
-73
lines changed

‎cli/ssh.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -711,12 +711,12 @@ func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, work
711711
lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
712712
conditionCtx, cancelCondition := context.WithCancel(ctx)
713713
condition := notifyCondition(conditionCtx, client, workspace.ID, lock)
714-
stopFunc := notify.Notify(condition, workspacePollInterval, autostopNotifyCountdown...)
714+
notifier := notify.New(condition, workspacePollInterval, autostopNotifyCountdown)
715715
return func() {
716716
// With many "ssh" processes running, `lock.TryLockContext` can be hanging until the context canceled.
717717
// Without this cancellation, a CLI process with failed remote-forward could be hanging indefinitely.
718718
cancelCondition()
719-
stopFunc()
719+
notifier.Close()
720720
}
721721
}
722722

‎coderd/autobuild/notify/notifier.go

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,16 @@ import (
55
"sort"
66
"sync"
77
"time"
8+
9+
"github.com/coder/coder/v2/clock"
810
)
911

10-
// Notifier calls a Condition at most once for each count in countdown.
12+
// Notifier triggers callbacks at given intervals until some event happens. The
13+
// intervals (e.g. 10 minute warning, 5 minute warning) are given in the
14+
// countdown. The Notifier periodically polls the condition to get the time of
15+
// the event (the Condition's deadline) and the callback. The callback is
16+
// called at most once per entry in the countdown, the first time the time to
17+
// the deadline is shorter than the duration.
1118
type Notifier struct {
1219
ctx context.Context
1320
cancel context.CancelFunc
@@ -17,36 +24,35 @@ type Notifier struct {
1724
condition Condition
1825
notifiedAt map[time.Duration]bool
1926
countdown []time.Duration
27+
28+
// for testing
29+
clock clock.Clock
2030
}
2131

22-
// Condition is a function that gets executed with a certain time.
32+
// Condition is a function that gets executed periodically, and receives the
33+
// current time as an argument.
2334
// - It should return the deadline for the notification, as well as a
24-
// callback function to execute once the time to the deadline is
25-
// less than one of the notify attempts. If deadline is the zero
35+
// callback function to execute. If deadline is the zero
2636
// time, callback will not be executed.
2737
// - Callback is executed once for every time the difference between deadline
2838
// and the current time is less than an element of countdown.
2939
// - To enforce a minimum interval between consecutive callbacks, truncate
3040
// the returned deadline to the minimum interval.
3141
type Condition func(now time.Time) (deadline time.Time, callback func())
3242

33-
// Notify is a convenience function that initializes a new Notifier
34-
// with the given condition, interval, and countdown.
35-
// It is the responsibility of the caller to call close to stop polling.
36-
func Notify(cond Condition, interval time.Duration, countdown ...time.Duration) (closeFunc func()) {
37-
notifier := New(cond, countdown...)
38-
ticker := time.NewTicker(interval)
39-
go notifier.Poll(ticker.C)
40-
return func() {
41-
ticker.Stop()
42-
_ = notifier.Close()
43+
type Option func(*Notifier)
44+
45+
// WithTestClock is used in tests to inject a mock Clock
46+
func WithTestClock(clk clock.Clock) Option {
47+
return func(n *Notifier) {
48+
n.clock = clk
4349
}
4450
}
4551

4652
// New returns a Notifier that calls cond once every time it polls.
4753
// - Duplicate values are removed from countdown, and it is sorted in
4854
// descending order.
49-
func New(cond Condition, countdown ...time.Duration) *Notifier {
55+
func New(cond Condition, interval time.Duration, countdown []time.Duration, opts ...Option) *Notifier {
5056
// Ensure countdown is sorted in descending order and contains no duplicates.
5157
ct := unique(countdown)
5258
sort.Slice(ct, func(i, j int) bool {
@@ -61,38 +67,36 @@ func New(cond Condition, countdown ...time.Duration) *Notifier {
6167
countdown: ct,
6268
condition: cond,
6369
notifiedAt: make(map[time.Duration]bool),
70+
clock: clock.NewReal(),
6471
}
72+
for _, opt := range opts {
73+
opt(n)
74+
}
75+
go n.poll(interval)
6576

6677
return n
6778
}
6879

69-
// Poll polls once immediately, and then once for every value from ticker.
80+
// poll polls once immediately, and then periodically according to the interval.
7081
// Poll exits when ticker is closed.
71-
func (n *Notifier) Poll(ticker <-chan time.Time) {
82+
func (n *Notifier) poll(interval time.Duration) {
7283
defer close(n.pollDone)
7384

7485
// poll once immediately
75-
n.pollOnce(time.Now())
76-
for {
77-
select {
78-
case <-n.ctx.Done():
79-
return
80-
case t, ok := <-ticker:
81-
if !ok {
82-
return
83-
}
84-
n.pollOnce(t)
85-
}
86-
}
86+
_ = n.pollOnce()
87+
tkr := n.clock.TickerFunc(n.ctx, interval, n.pollOnce, "notifier", "poll")
88+
_ = tkr.Wait()
8789
}
8890

89-
func (n *Notifier) Close() error {
91+
func (n *Notifier) Close() {
9092
n.cancel()
9193
<-n.pollDone
92-
return nil
9394
}
9495

95-
func (n *Notifier) pollOnce(tick time.Time) {
96+
// pollOnce only returns an error so it matches the signature expected of TickerFunc
97+
// nolint: revive // bare returns are fine here
98+
func (n *Notifier) pollOnce() (_ error) {
99+
tick := n.clock.Now()
96100
n.lock.Lock()
97101
defer n.lock.Unlock()
98102

@@ -113,6 +117,7 @@ func (n *Notifier) pollOnce(tick time.Time) {
113117
n.notifiedAt[tock] = true
114118
return
115119
}
120+
return
116121
}
117122

118123
func unique(ds []time.Duration) []time.Duration {

‎coderd/autobuild/notify/notifier_test.go

Lines changed: 41 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,74 +1,81 @@
11
package notify_test
22

33
import (
4-
"sync"
54
"testing"
65
"time"
76

87
"github.com/stretchr/testify/require"
9-
"go.uber.org/atomic"
108
"go.uber.org/goleak"
119

10+
"github.com/coder/coder/v2/clock"
1211
"github.com/coder/coder/v2/coderd/autobuild/notify"
12+
"github.com/coder/coder/v2/testutil"
1313
)
1414

1515
func TestNotifier(t *testing.T) {
1616
t.Parallel()
1717

18-
now := time.Now()
18+
now := time.Date(2022, 5, 13, 0, 0, 0, 0, time.UTC)
1919

2020
testCases := []struct {
2121
Name string
2222
Countdown []time.Duration
23-
Ticks []time.Time
23+
PollInterval time.Duration
24+
NTicks int
2425
ConditionDeadline time.Time
25-
NumConditions int64
26-
NumCallbacks int64
26+
NumConditions int
27+
NumCallbacks int
2728
}{
2829
{
2930
Name: "zero deadline",
3031
Countdown: durations(),
31-
Ticks: fakeTicker(now, time.Second, 0),
32+
PollInterval: time.Second,
33+
NTicks: 0,
3234
ConditionDeadline: time.Time{},
3335
NumConditions: 1,
3436
NumCallbacks: 0,
3537
},
3638
{
3739
Name: "no calls",
3840
Countdown: durations(),
39-
Ticks: fakeTicker(now, time.Second, 0),
41+
PollInterval: time.Second,
42+
NTicks: 0,
4043
ConditionDeadline: now,
4144
NumConditions: 1,
4245
NumCallbacks: 0,
4346
},
4447
{
4548
Name: "exactly one call",
4649
Countdown: durations(time.Second),
47-
Ticks: fakeTicker(now, time.Second, 1),
50+
PollInterval: time.Second,
51+
NTicks: 1,
4852
ConditionDeadline: now.Add(time.Second),
4953
NumConditions: 2,
5054
NumCallbacks: 1,
5155
},
5256
{
5357
Name: "two calls",
5458
Countdown: durations(4*time.Second, 2*time.Second),
55-
Ticks: fakeTicker(now, time.Second, 5),
59+
PollInterval: time.Second,
60+
NTicks: 5,
5661
ConditionDeadline: now.Add(5 * time.Second),
5762
NumConditions: 6,
5863
NumCallbacks: 2,
5964
},
6065
{
6166
Name: "wrong order should not matter",
6267
Countdown: durations(2*time.Second, 4*time.Second),
63-
Ticks: fakeTicker(now, time.Second, 5),
68+
PollInterval: time.Second,
69+
NTicks: 5,
6470
ConditionDeadline: now.Add(5 * time.Second),
6571
NumConditions: 6,
6672
NumCallbacks: 2,
6773
},
6874
{
6975
Name: "ssh autostop notify",
7076
Countdown: durations(5*time.Minute, time.Minute),
71-
Ticks: fakeTicker(now, 30*time.Second, 120),
77+
PollInterval: 30 * time.Second,
78+
NTicks: 120,
7279
ConditionDeadline: now.Add(30 * time.Minute),
7380
NumConditions: 121,
7481
NumCallbacks: 2,
@@ -79,30 +86,33 @@ func TestNotifier(t *testing.T) {
7986
testCase := testCase
8087
t.Run(testCase.Name, func(t *testing.T) {
8188
t.Parallel()
82-
ch := make(chan time.Time)
83-
numConditions := atomic.NewInt64(0)
84-
numCalls := atomic.NewInt64(0)
89+
ctx := testutil.Context(t, testutil.WaitShort)
90+
mClock := clock.NewMock(t)
91+
mClock.Set(now).MustWait(ctx)
92+
numConditions := 0
93+
numCalls := 0
8594
cond := func(time.Time) (time.Time, func()) {
86-
numConditions.Inc()
95+
numConditions++
8796
return testCase.ConditionDeadline, func() {
88-
numCalls.Inc()
97+
numCalls++
8998
}
9099
}
91-
var wg sync.WaitGroup
92-
go func() {
93-
defer wg.Done()
94-
n := notify.New(cond, testCase.Countdown...)
95-
defer n.Close()
96-
n.Poll(ch)
97-
}()
98-
wg.Add(1)
99-
for _, tick := range testCase.Ticks {
100-
ch <- tick
100+
101+
trap := mClock.Trap().TickerFunc("notifier", "poll")
102+
defer trap.Close()
103+
104+
n := notify.New(cond, testCase.PollInterval, testCase.Countdown, notify.WithTestClock(mClock))
105+
defer n.Close()
106+
107+
trap.MustWait(ctx).Release() // ensure ticker started
108+
for i := 0; i < testCase.NTicks; i++ {
109+
interval, w := mClock.AdvanceNext()
110+
w.MustWait(ctx)
111+
require.Equal(t, testCase.PollInterval, interval)
101112
}
102-
close(ch)
103-
wg.Wait()
104-
require.Equal(t, testCase.NumCallbacks, numCalls.Load())
105-
require.Equal(t, testCase.NumConditions, numConditions.Load())
113+
114+
require.Equal(t, testCase.NumCallbacks, numCalls)
115+
require.Equal(t, testCase.NumConditions, numConditions)
106116
})
107117
}
108118
}
@@ -111,14 +121,6 @@ func durations(ds ...time.Duration) []time.Duration {
111121
return ds
112122
}
113123

114-
func fakeTicker(t time.Time, d time.Duration, n int) []time.Time {
115-
var ts []time.Time
116-
for i := 1; i <= n; i++ {
117-
ts = append(ts, t.Add(time.Duration(n)*d))
118-
}
119-
return ts
120-
}
121-
122124
func TestMain(m *testing.M) {
123125
goleak.VerifyTestMain(m)
124126
}

0 commit comments

Comments
 (0)