From 8aae16ebae4d30dca45556adf9cc8b735e647e14 Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Thu, 13 Jun 2024 15:24:52 +0400 Subject: [PATCH] chore: refactor autobuild/notify to use clock test --- cli/ssh.go | 4 +- coderd/autobuild/notify/notifier.go | 69 ++++++++++---------- coderd/autobuild/notify/notifier_test.go | 80 ++++++++++++------------ 3 files changed, 80 insertions(+), 73 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index aa8bdadb9d0dd..ac849649b9184 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -711,12 +711,12 @@ func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, work lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String())) conditionCtx, cancelCondition := context.WithCancel(ctx) condition := notifyCondition(conditionCtx, client, workspace.ID, lock) - stopFunc := notify.Notify(condition, workspacePollInterval, autostopNotifyCountdown...) + notifier := notify.New(condition, workspacePollInterval, autostopNotifyCountdown) return func() { // With many "ssh" processes running, `lock.TryLockContext` can be hanging until the context canceled. // Without this cancellation, a CLI process with failed remote-forward could be hanging indefinitely. cancelCondition() - stopFunc() + notifier.Close() } } diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go index e0db12af35475..d8226161507ef 100644 --- a/coderd/autobuild/notify/notifier.go +++ b/coderd/autobuild/notify/notifier.go @@ -5,9 +5,16 @@ import ( "sort" "sync" "time" + + "github.com/coder/coder/v2/clock" ) -// Notifier calls a Condition at most once for each count in countdown. +// Notifier triggers callbacks at given intervals until some event happens. The +// intervals (e.g. 10 minute warning, 5 minute warning) are given in the +// countdown. The Notifier periodically polls the condition to get the time of +// the event (the Condition's deadline) and the callback. The callback is +// called at most once per entry in the countdown, the first time the time to +// the deadline is shorter than the duration. type Notifier struct { ctx context.Context cancel context.CancelFunc @@ -17,12 +24,15 @@ type Notifier struct { condition Condition notifiedAt map[time.Duration]bool countdown []time.Duration + + // for testing + clock clock.Clock } -// Condition is a function that gets executed with a certain time. +// Condition is a function that gets executed periodically, and receives the +// current time as an argument. // - It should return the deadline for the notification, as well as a -// callback function to execute once the time to the deadline is -// less than one of the notify attempts. If deadline is the zero +// callback function to execute. If deadline is the zero // time, callback will not be executed. // - Callback is executed once for every time the difference between deadline // and the current time is less than an element of countdown. @@ -30,23 +40,19 @@ type Notifier struct { // the returned deadline to the minimum interval. type Condition func(now time.Time) (deadline time.Time, callback func()) -// Notify is a convenience function that initializes a new Notifier -// with the given condition, interval, and countdown. -// It is the responsibility of the caller to call close to stop polling. -func Notify(cond Condition, interval time.Duration, countdown ...time.Duration) (closeFunc func()) { - notifier := New(cond, countdown...) - ticker := time.NewTicker(interval) - go notifier.Poll(ticker.C) - return func() { - ticker.Stop() - _ = notifier.Close() +type Option func(*Notifier) + +// WithTestClock is used in tests to inject a mock Clock +func WithTestClock(clk clock.Clock) Option { + return func(n *Notifier) { + n.clock = clk } } // New returns a Notifier that calls cond once every time it polls. // - Duplicate values are removed from countdown, and it is sorted in // descending order. -func New(cond Condition, countdown ...time.Duration) *Notifier { +func New(cond Condition, interval time.Duration, countdown []time.Duration, opts ...Option) *Notifier { // Ensure countdown is sorted in descending order and contains no duplicates. ct := unique(countdown) sort.Slice(ct, func(i, j int) bool { @@ -61,38 +67,36 @@ func New(cond Condition, countdown ...time.Duration) *Notifier { countdown: ct, condition: cond, notifiedAt: make(map[time.Duration]bool), + clock: clock.NewReal(), } + for _, opt := range opts { + opt(n) + } + go n.poll(interval) return n } -// Poll polls once immediately, and then once for every value from ticker. +// poll polls once immediately, and then periodically according to the interval. // Poll exits when ticker is closed. -func (n *Notifier) Poll(ticker <-chan time.Time) { +func (n *Notifier) poll(interval time.Duration) { defer close(n.pollDone) // poll once immediately - n.pollOnce(time.Now()) - for { - select { - case <-n.ctx.Done(): - return - case t, ok := <-ticker: - if !ok { - return - } - n.pollOnce(t) - } - } + _ = n.pollOnce() + tkr := n.clock.TickerFunc(n.ctx, interval, n.pollOnce, "notifier", "poll") + _ = tkr.Wait() } -func (n *Notifier) Close() error { +func (n *Notifier) Close() { n.cancel() <-n.pollDone - return nil } -func (n *Notifier) pollOnce(tick time.Time) { +// pollOnce only returns an error so it matches the signature expected of TickerFunc +// nolint: revive // bare returns are fine here +func (n *Notifier) pollOnce() (_ error) { + tick := n.clock.Now() n.lock.Lock() defer n.lock.Unlock() @@ -113,6 +117,7 @@ func (n *Notifier) pollOnce(tick time.Time) { n.notifiedAt[tock] = true return } + return } func unique(ds []time.Duration) []time.Duration { diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go index 09e8158abaa99..d53b06c1a2133 100644 --- a/coderd/autobuild/notify/notifier_test.go +++ b/coderd/autobuild/notify/notifier_test.go @@ -1,34 +1,36 @@ package notify_test import ( - "sync" "testing" "time" "github.com/stretchr/testify/require" - "go.uber.org/atomic" "go.uber.org/goleak" + "github.com/coder/coder/v2/clock" "github.com/coder/coder/v2/coderd/autobuild/notify" + "github.com/coder/coder/v2/testutil" ) func TestNotifier(t *testing.T) { t.Parallel() - now := time.Now() + now := time.Date(2022, 5, 13, 0, 0, 0, 0, time.UTC) testCases := []struct { Name string Countdown []time.Duration - Ticks []time.Time + PollInterval time.Duration + NTicks int ConditionDeadline time.Time - NumConditions int64 - NumCallbacks int64 + NumConditions int + NumCallbacks int }{ { Name: "zero deadline", Countdown: durations(), - Ticks: fakeTicker(now, time.Second, 0), + PollInterval: time.Second, + NTicks: 0, ConditionDeadline: time.Time{}, NumConditions: 1, NumCallbacks: 0, @@ -36,7 +38,8 @@ func TestNotifier(t *testing.T) { { Name: "no calls", Countdown: durations(), - Ticks: fakeTicker(now, time.Second, 0), + PollInterval: time.Second, + NTicks: 0, ConditionDeadline: now, NumConditions: 1, NumCallbacks: 0, @@ -44,7 +47,8 @@ func TestNotifier(t *testing.T) { { Name: "exactly one call", Countdown: durations(time.Second), - Ticks: fakeTicker(now, time.Second, 1), + PollInterval: time.Second, + NTicks: 1, ConditionDeadline: now.Add(time.Second), NumConditions: 2, NumCallbacks: 1, @@ -52,7 +56,8 @@ func TestNotifier(t *testing.T) { { Name: "two calls", Countdown: durations(4*time.Second, 2*time.Second), - Ticks: fakeTicker(now, time.Second, 5), + PollInterval: time.Second, + NTicks: 5, ConditionDeadline: now.Add(5 * time.Second), NumConditions: 6, NumCallbacks: 2, @@ -60,7 +65,8 @@ func TestNotifier(t *testing.T) { { Name: "wrong order should not matter", Countdown: durations(2*time.Second, 4*time.Second), - Ticks: fakeTicker(now, time.Second, 5), + PollInterval: time.Second, + NTicks: 5, ConditionDeadline: now.Add(5 * time.Second), NumConditions: 6, NumCallbacks: 2, @@ -68,7 +74,8 @@ func TestNotifier(t *testing.T) { { Name: "ssh autostop notify", Countdown: durations(5*time.Minute, time.Minute), - Ticks: fakeTicker(now, 30*time.Second, 120), + PollInterval: 30 * time.Second, + NTicks: 120, ConditionDeadline: now.Add(30 * time.Minute), NumConditions: 121, NumCallbacks: 2, @@ -79,30 +86,33 @@ func TestNotifier(t *testing.T) { testCase := testCase t.Run(testCase.Name, func(t *testing.T) { t.Parallel() - ch := make(chan time.Time) - numConditions := atomic.NewInt64(0) - numCalls := atomic.NewInt64(0) + ctx := testutil.Context(t, testutil.WaitShort) + mClock := clock.NewMock(t) + mClock.Set(now).MustWait(ctx) + numConditions := 0 + numCalls := 0 cond := func(time.Time) (time.Time, func()) { - numConditions.Inc() + numConditions++ return testCase.ConditionDeadline, func() { - numCalls.Inc() + numCalls++ } } - var wg sync.WaitGroup - go func() { - defer wg.Done() - n := notify.New(cond, testCase.Countdown...) - defer n.Close() - n.Poll(ch) - }() - wg.Add(1) - for _, tick := range testCase.Ticks { - ch <- tick + + trap := mClock.Trap().TickerFunc("notifier", "poll") + defer trap.Close() + + n := notify.New(cond, testCase.PollInterval, testCase.Countdown, notify.WithTestClock(mClock)) + defer n.Close() + + trap.MustWait(ctx).Release() // ensure ticker started + for i := 0; i < testCase.NTicks; i++ { + interval, w := mClock.AdvanceNext() + w.MustWait(ctx) + require.Equal(t, testCase.PollInterval, interval) } - close(ch) - wg.Wait() - require.Equal(t, testCase.NumCallbacks, numCalls.Load()) - require.Equal(t, testCase.NumConditions, numConditions.Load()) + + require.Equal(t, testCase.NumCallbacks, numCalls) + require.Equal(t, testCase.NumConditions, numConditions) }) } } @@ -111,14 +121,6 @@ func durations(ds ...time.Duration) []time.Duration { return ds } -func fakeTicker(t time.Time, d time.Duration, n int) []time.Time { - var ts []time.Time - for i := 1; i <= n; i++ { - ts = append(ts, t.Add(time.Duration(n)*d)) - } - return ts -} - func TestMain(m *testing.M) { goleak.VerifyTestMain(m) }