Jun 13, 2024 · Jun 13, 2024 · Jun 13, 2024
diff --git a/cli/ssh.go b/cli/ssh.go
 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
 "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
 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 {
 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 thenonce for every value from ticker.
 //poll polls once immediately, and thenperiodically according to the interval.
 // Poll exits when ticker is closed.
 func (n *Notifier)Poll(ticker <-chantime.Time) {
 func (n *Notifier)poll(intervaltime.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()

 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
 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
 NumConditionsint64
 NumCallbacksint64
 NumConditionsint
 NumCallbacksint
 }{
 {
 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,
 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)
 })
 }
 }
 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)
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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()
		}
		}

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand All		@@ -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 {
Expand All		@@ -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 thenonce for every value from ticker.
		//poll polls once immediately, and thenperiodically according to the interval.
		// Poll exits when ticker is closed.
		func (n *Notifier)Poll(ticker <-chantime.Time) {
		func (n *Notifier)poll(intervaltime.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()

Expand All		@@ -113,6 +117,7 @@ func (n *Notifier) pollOnce(tick time.Time) {
		n.notifiedAt[tock] = true
		return
		}
		return
		}

		func unique(ds []time.Duration) []time.Duration {
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -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
		NumConditionsint64
		NumCallbacksint64
		NumConditionsint
		NumCallbacksint
		}{
		{
		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(4time.Second, 2time.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(2time.Second, 4time.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,
Expand All		@@ -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
Copy link ContributorAuthor spikecurtisJun 13, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. atomic was never necessary here, since polling is serialized.
		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)
		})
		}
		}
Expand All		@@ -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)
		}