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.
This commit is contained in:
Cian Johnston 2022-05-13 18:09:04 +01:00 committed by GitHub
parent 4ab7a41f08
commit b2760b1faf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 309 additions and 0 deletions

View File

@ -2,10 +2,15 @@ package cli
import (
"context"
"fmt"
"io"
"os"
"path/filepath"
"strings"
"time"
"github.com/gen2brain/beeep"
"github.com/gofrs/flock"
"github.com/google/uuid"
"github.com/mattn/go-isatty"
"github.com/spf13/cobra"
@ -15,10 +20,15 @@ import (
"github.com/coder/coder/cli/cliflag"
"github.com/coder/coder/cli/cliui"
"github.com/coder/coder/coderd/autobuild/notify"
"github.com/coder/coder/coderd/autobuild/schedule"
"github.com/coder/coder/coderd/database"
"github.com/coder/coder/codersdk"
)
var autostopPollInterval = 30 * time.Second
var autostopNotifyCountdown = []time.Duration{5 * time.Minute}
func ssh() *cobra.Command {
var (
stdio bool
@ -108,6 +118,9 @@ func ssh() *cobra.Command {
}
defer conn.Close()
stopPolling := tryPollWorkspaceAutostop(cmd.Context(), client, workspace)
defer stopPolling()
if stdio {
rawSSH, err := conn.SSH()
if err != nil {
@ -179,3 +192,57 @@ func ssh() *cobra.Command {
return cmd
}
// Attempt to poll workspace autostop. We write a per-workspace lockfile to
// avoid spamming the user with notifications in case of multiple instances
// of the CLI running simultaneously.
func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace) (stop func()) {
lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
condition := notifyCondition(ctx, client, workspace.ID, lock)
return notify.Notify(condition, autostopPollInterval, autostopNotifyCountdown...)
}
// Notify the user if the workspace is due to shutdown.
func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lock *flock.Flock) notify.Condition {
return func(now time.Time) (deadline time.Time, callback func()) {
// Keep trying to regain the lock.
locked, err := lock.TryLockContext(ctx, autostopPollInterval)
if err != nil || !locked {
return time.Time{}, nil
}
ws, err := client.Workspace(ctx, workspaceID)
if err != nil {
return time.Time{}, nil
}
if ws.AutostopSchedule == "" {
return time.Time{}, nil
}
sched, err := schedule.Weekly(ws.AutostopSchedule)
if err != nil {
return time.Time{}, nil
}
deadline = sched.Next(now)
callback = func() {
ttl := deadline.Sub(now)
var title, body string
if ttl > time.Minute {
title = fmt.Sprintf(`Workspace %s stopping in %.0f mins`, ws.Name, ttl.Minutes())
body = fmt.Sprintf(
`Your Coder workspace %s is scheduled to stop at %s.`,
ws.Name,
deadline.Format(time.Kitchen),
)
} else {
title = fmt.Sprintf("Workspace %s stopping!", ws.Name)
body = fmt.Sprintf("Your Coder workspace %s is stopping any time now!", ws.Name)
}
// notify user with a native system notification (best effort)
_ = beeep.Notify(title, body, "")
}
return deadline.Truncate(time.Minute), callback
}
}

View File

@ -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
}

View File

@ -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)
}

6
go.mod
View File

@ -58,11 +58,13 @@ require (
github.com/fatih/color v1.13.0
github.com/fatih/structs v1.1.0
github.com/fullsailor/pkcs7 v0.0.0-20190404230743-d7302db945fa
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a
github.com/gliderlabs/ssh v0.3.3
github.com/go-chi/chi/v5 v5.0.7
github.com/go-chi/httprate v0.5.3
github.com/go-chi/render v1.0.1
github.com/go-playground/validator/v10 v10.11.0
github.com/gofrs/flock v0.8.1
github.com/gohugoio/hugo v0.98.0
github.com/golang-jwt/jwt v3.2.2+incompatible
github.com/golang-migrate/migrate/v4 v4.15.2
@ -159,8 +161,10 @@ require (
github.com/ghodss/yaml v1.0.0 // indirect
github.com/go-playground/locales v0.14.0 // indirect
github.com/go-playground/universal-translator v0.18.0 // indirect
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 // indirect
github.com/gobwas/glob v0.2.3 // indirect
github.com/gobwas/ws v1.1.0 // indirect
github.com/godbus/dbus/v5 v5.1.0 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/glog v1.0.0 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
@ -196,6 +200,7 @@ require (
github.com/muesli/reflow v0.3.0 // indirect
github.com/muesli/termenv v0.11.1-0.20220212125758-44cd13922739 // indirect
github.com/niklasfasching/go-org v1.6.2 // indirect
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect
github.com/opencontainers/go-digest v1.0.0 // indirect
github.com/opencontainers/image-spec v1.0.2 // indirect
github.com/opencontainers/runc v1.1.0 // indirect
@ -226,6 +231,7 @@ require (
github.com/spf13/afero v1.8.2 // indirect
github.com/spf13/cast v1.4.1 // indirect
github.com/spf13/jwalterweatherman v1.1.0 // indirect
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af // indirect
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect
github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b // indirect
github.com/tinylib/msgp v1.1.2 // indirect

13
go.sum
View File

@ -606,6 +606,8 @@ github.com/gabriel-vasile/mimetype v1.3.1/go.mod h1:fA8fi6KUiG7MgQQ+mEWotXoEOvmx
github.com/gabriel-vasile/mimetype v1.4.0/go.mod h1:fA8fi6KUiG7MgQQ+mEWotXoEOvmxRtOJlERCzSmRvr8=
github.com/garyburd/redigo v0.0.0-20150301180006-535138d7bcd7/go.mod h1:NR3MbYisc3/PwhQ00EMzDiPmrwpPxAn5GI05/YaO1SY=
github.com/garyburd/redigo v1.6.3/go.mod h1:rTb6epsqigu3kYKBnaF028A7Tf/Aw5s0cqA47doKKqw=
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a h1:fwNLHrP5Rbg/mGSXCjtPdpbqv2GucVTA/KMi8wEm6mE=
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a/go.mod h1:/WeFVhhxMOGypVKS0w8DUJxUBbHypnWkUVnW7p5c9Pw=
github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ=
github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=
@ -703,6 +705,8 @@ github.com/go-test/deep v1.0.2-0.20181118220953-042da051cf31/go.mod h1:wGDj63lr6
github.com/go-test/deep v1.0.2/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
github.com/go-test/deep v1.0.3 h1:ZrJSEWsXzPOxaZnFteGEfooLba+ju3FYIbOrS+rQd68=
github.com/go-test/deep v1.0.3/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 h1:qZNfIGkIANxGv/OqtnntR4DfOY2+BgwR60cAcu/i3SE=
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4/go.mod h1:kW3HQ4UdaAyrUCSSDR4xUzBKW6O2iA4uHhk7AtyYp10=
github.com/gobuffalo/attrs v0.0.0-20190224210810-a9411de4debd/go.mod h1:4duuawTqi2wkkpB4ePgWMaai6/Kc6WEz83bhFwpHzj0=
github.com/gobuffalo/depgen v0.0.0-20190329151759-d478694a28d3/go.mod h1:3STtPUQYuzV0gBVOY3vy6CfMm/ljR4pABfrTeHNLHUY=
github.com/gobuffalo/depgen v0.1.0/go.mod h1:+ifsuy7fhi15RWncXQQKjWS9JPkdah5sZvtHc2RXGlg=
@ -747,7 +751,11 @@ github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e/go.mod h1:bBOAhwG1umN6
github.com/godbus/dbus/v5 v5.0.3/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/godbus/dbus/v5 v5.0.6/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/gofiber/fiber/v2 v2.11.0/go.mod h1:oZTLWqYnqpMMuF922SjGbsYZsdpE1MCfh416HNdweIM=
github.com/gofrs/flock v0.8.1 h1:+gYjHKf32LDeiEEFhQaotPbLuUXjY5ZqxKgXy7n59aw=
github.com/gofrs/flock v0.8.1/go.mod h1:F1TvTiK9OcQqauNUHlbJvyl9Qa1QvF/gOUDKA14jxHU=
github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
github.com/gogo/googleapis v1.2.0/go.mod h1:Njal3psf3qN6dwBtQfUmBZh2ybovJ0tlu3o/AC7HYjU=
@ -1327,6 +1335,8 @@ github.com/neo4j/neo4j-go-driver v1.8.1-0.20200803113522-b626aa943eba/go.mod h1:
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
github.com/niklasfasching/go-org v1.6.2 h1:kQBIZlfL4oRNApJCrBgaeNBfzxWzP6XlC7/b744Polk=
github.com/niklasfasching/go-org v1.6.2/go.mod h1:wn76Xgu4/KRe43WZhsgZjxYMaloSrl3BSweGV74SwHs=
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d h1:VhgPp6v9qf9Agr/56bj7Y/xa04UccTW04VP0Qed4vnQ=
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d/go.mod h1:YUTz3bUH2ZwIWBy3CJBeOBEugqcmXREj14T+iG/4k4U=
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
github.com/oklog/run v1.0.0/go.mod h1:dlhp/R75TPv97u0XWUtDeV/lRKWPKSdTuV0TZvrmrQA=
@ -1655,6 +1665,8 @@ github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635/go.mod h1:hkRG
github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
github.com/tabbed/pqtype v0.1.1 h1:PhEcb9JZ8jr7SUjJDFjRPxny0M8fkXZrxn/a9yQfoZg=
github.com/tabbed/pqtype v0.1.1/go.mod h1:HLt2kLJPcUhODQkYn3mJkMHXVsuv3Z2n5NZEeKXL0Uk=
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af h1:6yITBqGTE2lEeTPG04SN9W+iWHCRyHqlVYILiSXziwk=
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af/go.mod h1:4F09kP5F+am0jAwlQLddpoMDM+iewkxxt6nxUQ5nq5o=
github.com/tchap/go-patricia v2.2.6+incompatible/go.mod h1:bmLyhP68RS6kStMGxByiQ23RP/odRBOTVjwp2cDyi6I=
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 h1:89CEmDvlq/F7SJEOqkIdNDGJXrQIhuIx9D2DBXjavSU=
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161/go.mod h1:wM7WEvslTq+iOEAMDLSzhVuOt5BRZ05WirO+b09GHQU=
@ -2203,6 +2215,7 @@ golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220317061510-51cd9980dadf/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220328115105-d36c6a25d886/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220502124256-b6088ccd6cba h1:AyHWHCBVlIYI5rgEM3o+1PLd0sLPcIAoaUckGQMaWtw=