coder · spikecurtis · Jun 13, 2024 · Jun 13, 2024 · spikecurtis · Jun 13, 2024
diff --git 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
@@ -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,36 +24,35 @@ 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.
 //   - To enforce a minimum interval between consecutive callbacks, truncate
 //     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
@@ -1,74 +1,81 @@
 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,
 		},
 		{
 			Name:              "no calls",
 			Countdown:         durations(),
-			Ticks:             fakeTicker(now, time.Second, 0),
+			PollInterval:      time.Second,
+			NTicks:            0,
 			ConditionDeadline: now,
 			NumConditions:     1,
 			NumCallbacks:      0,
 		},
 		{
 			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,
 		},
 		{
 			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,
 		},
 		{
 			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,
 		},
 		{
 			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)
 }