feat: send native system notification on scheduled workspace shutdown (#1414)

* feat: send native system notification on scheduled workspace shutdown

This commit adds a fairly generic notification package and uses it
to notify users connected over SSH of pending workspace shutdowns.
Only one notification will be sent at most 5 minutes prior to the scheduled
shutdown, and only one CLI instance will send notifications if multiple
instances are running.

coderd/autobuild/notify/notifier.go

@@ -0,0 +1,100 @@
+package notify
+
+import (
+	"sort"
+	"sync"
+	"time"
+)
+
+// Notifier calls a Condition at most once for each count in countdown.
+type Notifier struct {
+	lock       sync.Mutex
+	condition  Condition
+	notifiedAt map[time.Duration]bool
+	countdown  []time.Duration
+}
+
+// Condition is a function that gets executed with a certain time.
+// - 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
+//   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) (close func()) {
+	notifier := New(cond, countdown...)
+	ticker := time.NewTicker(interval)
+	go notifier.Poll(ticker.C)
+	return ticker.Stop
+}
+
+// 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 {
+	// Ensure countdown is sorted in descending order and contains no duplicates.
+	ct := unique(countdown)
+	sort.Slice(ct, func(i, j int) bool {
+		return ct[i] < ct[j]
+	})
+
+	n := &Notifier{
+		countdown:  ct,
+		condition:  cond,
+		notifiedAt: make(map[time.Duration]bool),
+	}
+
+	return n
+}
+
+// Poll polls once immediately, and then once for every value from ticker.
+// Poll exits when ticker is closed.
+func (n *Notifier) Poll(ticker <-chan time.Time) {
+	// poll once immediately
+	n.pollOnce(time.Now())
+	for t := range ticker {
+		n.pollOnce(t)
+	}
+}
+
+func (n *Notifier) pollOnce(tick time.Time) {
+	n.lock.Lock()
+	defer n.lock.Unlock()
+
+	deadline, callback := n.condition(tick)
+	if deadline.IsZero() {
+		return
+	}
+
+	timeRemaining := deadline.Sub(tick)
+	for _, tock := range n.countdown {
+		if n.notifiedAt[tock] {
+			continue
+		}
+		if timeRemaining > tock {
+			continue
+		}
+		callback()
+		n.notifiedAt[tock] = true
+		return
+	}
+}
+
+func unique(ds []time.Duration) []time.Duration {
+	m := make(map[time.Duration]bool)
+	for _, d := range ds {
+		m[d] = true
+	}
+	var ks []time.Duration
+	for k := range m {
+		ks = append(ks, k)
+	}
+	return ks
+}

coderd/autobuild/notify/notifier_test.go

@@ -0,0 +1,123 @@
+package notify_test
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+	"go.uber.org/goleak"
+
+	"github.com/coder/coder/coderd/autobuild/notify"
+)
+
+func TestNotifier(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+
+	testCases := []struct {
+		Name              string
+		Countdown         []time.Duration
+		Ticks             []time.Time
+		ConditionDeadline time.Time
+		NumConditions     int64
+		NumCallbacks      int64
+	}{
+		{
+			Name:              "zero deadline",
+			Countdown:         durations(),
+			Ticks:             fakeTicker(now, time.Second, 0),
+			ConditionDeadline: time.Time{},
+			NumConditions:     1,
+			NumCallbacks:      0,
+		},
+		{
+			Name:              "no calls",
+			Countdown:         durations(),
+			Ticks:             fakeTicker(now, time.Second, 0),
+			ConditionDeadline: now,
+			NumConditions:     1,
+			NumCallbacks:      0,
+		},
+		{
+			Name:              "exactly one call",
+			Countdown:         durations(time.Second),
+			Ticks:             fakeTicker(now, time.Second, 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),
+			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),
+			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),
+			ConditionDeadline: now.Add(30 * time.Minute),
+			NumConditions:     121,
+			NumCallbacks:      2,
+		},
+	}
+
+	for _, testCase := range testCases {
+		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)
+			cond := func(time.Time) (time.Time, func()) {
+				numConditions.Inc()
+				return testCase.ConditionDeadline, func() {
+					numCalls.Inc()
+				}
+			}
+			var wg sync.WaitGroup
+			go func() {
+				n := notify.New(cond, testCase.Countdown...)
+				n.Poll(ch)
+				wg.Done()
+			}()
+			wg.Add(1)
+			for _, tick := range testCase.Ticks {
+				ch <- tick
+			}
+			close(ch)
+			wg.Wait()
+			require.Equal(t, testCase.NumCallbacks, numCalls.Load())
+			require.Equal(t, testCase.NumConditions, numConditions.Load())
+		})
+	}
+}
+
+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)
+}