diff --git a/cli/ssh.go b/cli/ssh.go index 9dc6af13c0de2..7b3bf2ae9bddd 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -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 + } +} diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go new file mode 100644 index 0000000000000..8ba3fd150990c --- /dev/null +++ b/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 +} diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go new file mode 100644 index 0000000000000..5b0b63a430c05 --- /dev/null +++ b/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) +} diff --git a/go.mod b/go.mod index 73b78448a169d..504d97afe6984 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 99ea1d9935b8f..84f57a8fde153 100644 --- a/go.sum +++ b/go.sum @@ -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=