From 125584c03fd8268d43301db12ed270326105442e Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 12 May 2022 13:35:22 +0000 Subject: [PATCH 1/6] 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. --- cli/ssh.go | 56 +++++++++++ coderd/autobuild/notify/notifier.go | 84 +++++++++++++++++ coderd/autobuild/notify/notifier_test.go | 113 +++++++++++++++++++++++ go.mod | 5 + go.sum | 11 +++ 5 files changed, 269 insertions(+) create mode 100644 coderd/autobuild/notify/notifier.go create mode 100644 coderd/autobuild/notify/notifier_test.go diff --git a/cli/ssh.go b/cli/ssh.go index 9dc6af13c0de2..1823cb3afd66e 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -2,10 +2,13 @@ package cli import ( "context" + "fmt" "io" "os" "strings" + "time" + "github.com/gen2brain/beeep" "github.com/google/uuid" "github.com/mattn/go-isatty" "github.com/spf13/cobra" @@ -15,6 +18,8 @@ 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" ) @@ -108,6 +113,20 @@ func ssh() *cobra.Command { } defer conn.Close() + // Notify the user if the workspace is due to shutdown. This uses + // the library gen2brain/beeep under the hood. + countdown := []time.Duration{ + 30 * time.Minute, + 10 * time.Minute, + 5 * time.Minute, + time.Minute, + } + condition := notifyWorkspaceAutostop(cmd.Context(), client, workspace.ID) + notifier := notify.New(condition, countdown...) + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + go notifier.Poll(ticker.C) + if stdio { rawSSH, err := conn.SSH() if err != nil { @@ -179,3 +198,40 @@ func ssh() *cobra.Command { return cmd } + +func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID) notify.Condition { + return func(now time.Time) (deadline time.Time, callback func()) { + 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) + } + 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..c40fc2f669215 --- /dev/null +++ b/coderd/autobuild/notify/notifier.go @@ -0,0 +1,84 @@ +package notify + +import ( + "sort" + "sync" + "time" +) + +// Notifier calls a Condition at most once for each count in countdown. +type Notifier struct { + sync.Mutex + condition Condition + notifiedAt map[time.Duration]bool + countdown []time.Duration +} + +// Condition is a function that gets executed with a certain time. +type Condition func(now time.Time) (deadline time.Time, callback func()) + +// New returns a Notifier that calls cond once every time it polls. +// - Condition is a function that returns the deadline and a callback. +// 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. +// - 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. + sort.Slice(unique(countdown), func(i, j int) bool { + return countdown[i] < countdown[j] + }) + + n := &Notifier{ + countdown: countdown, + condition: cond, + notifiedAt: make(map[time.Duration]bool), + } + + return n +} + +// Poll polls once immediately, and then once for every value from ticker. +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() + defer n.Unlock() + + deadline, callback := n.condition(tick) + if deadline.IsZero() { + return + } + + timeRemaining := deadline.Sub(tick) + for _, tock := range n.countdown { + if timeRemaining <= tock && !n.notifiedAt[tock] { + callback() + n.notifiedAt[tock] = true + } + } +} + +func unique(ds []time.Duration) []time.Duration { + m := make(map[time.Duration]bool) + for _, d := range ds { + m[d] = true + } + ks := make([]time.Duration, 0) + 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..08cac334e170b --- /dev/null +++ b/coderd/autobuild/notify/notifier_test.go @@ -0,0 +1,113 @@ +package notify_test + +import ( + "sync" + "testing" + "time" + + "github.com/coder/coder/coderd/autobuild/notify" + "github.com/stretchr/testify/require" + "go.uber.org/atomic" + "go.uber.org/goleak" +) + +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: ticks(now, 0), + ConditionDeadline: time.Time{}, + NumConditions: 1, + NumCallbacks: 0, + }, + { + Name: "no calls", + Countdown: durations(), + Ticks: ticks(now, 0), + ConditionDeadline: now, + NumConditions: 1, + NumCallbacks: 0, + }, + { + Name: "exactly one call", + Countdown: durations(time.Second), + Ticks: ticks(now, 2), + ConditionDeadline: now.Add(2 * time.Second), + NumConditions: 2, + NumCallbacks: 1, + }, + { + Name: "two calls", + Countdown: durations(4*time.Second, 2*time.Second), + Ticks: ticks(now, 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: ticks(now, 5), + ConditionDeadline: now.Add(5 * time.Second), + NumConditions: 6, + NumCallbacks: 2, + }, + } + + for _, testCase := range testCases { + testCase := testCase + t.Run(testCase.Name, func(t *testing.T) { + 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()) + }) + } +} + +func durations(ds ...time.Duration) []time.Duration { + return ds +} + +func ticks(t time.Time, n int) []time.Time { + ts := make([]time.Time, n+1) + ts = append(ts, t) + for i := 0; i < n; i++ { + ts = append(ts, t.Add(time.Duration(n)*time.Second)) + } + return ts +} + +func TestMain(m *testing.M) { + goleak.VerifyTestMain(m) +} diff --git a/go.mod b/go.mod index 73b78448a169d..edcf438c2eadf 100644 --- a/go.mod +++ b/go.mod @@ -58,6 +58,7 @@ 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 @@ -159,8 +160,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 +199,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 +230,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..2994f1468ee93 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,6 +751,8 @@ 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/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= @@ -1327,6 +1333,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 +1663,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 +2213,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= From 8313a52d345a6a67d9a4cbb494269707b77506cc Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Fri, 13 May 2022 09:34:09 +0000 Subject: [PATCH 2/6] notifier: refactor, add more unit tests --- coderd/autobuild/notify/notifier.go | 11 ++++++--- coderd/autobuild/notify/notifier_test.go | 30 +++++++++++++++--------- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go index c40fc2f669215..e62c7b1387edb 100644 --- a/coderd/autobuild/notify/notifier.go +++ b/coderd/autobuild/notify/notifier.go @@ -64,10 +64,15 @@ func (n *Notifier) pollOnce(tick time.Time) { timeRemaining := deadline.Sub(tick) for _, tock := range n.countdown { - if timeRemaining <= tock && !n.notifiedAt[tock] { - callback() - n.notifiedAt[tock] = true + if n.notifiedAt[tock] { + continue } + if timeRemaining > tock { + continue + } + callback() + n.notifiedAt[tock] = true + return } } diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go index 08cac334e170b..677cb5c4ce780 100644 --- a/coderd/autobuild/notify/notifier_test.go +++ b/coderd/autobuild/notify/notifier_test.go @@ -27,7 +27,7 @@ func TestNotifier(t *testing.T) { { Name: "zero deadline", Countdown: durations(), - Ticks: ticks(now, 0), + Ticks: fakeTicker(now, time.Second, 0), ConditionDeadline: time.Time{}, NumConditions: 1, NumCallbacks: 0, @@ -35,7 +35,7 @@ func TestNotifier(t *testing.T) { { Name: "no calls", Countdown: durations(), - Ticks: ticks(now, 0), + Ticks: fakeTicker(now, time.Second, 0), ConditionDeadline: now, NumConditions: 1, NumCallbacks: 0, @@ -43,15 +43,15 @@ func TestNotifier(t *testing.T) { { Name: "exactly one call", Countdown: durations(time.Second), - Ticks: ticks(now, 2), - ConditionDeadline: now.Add(2 * 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: ticks(now, 5), + Ticks: fakeTicker(now, time.Second, 5), ConditionDeadline: now.Add(5 * time.Second), NumConditions: 6, NumCallbacks: 2, @@ -59,11 +59,19 @@ func TestNotifier(t *testing.T) { { Name: "wrong order should not matter", Countdown: durations(2*time.Second, 4*time.Second), - Ticks: ticks(now, 5), + 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 { @@ -91,6 +99,7 @@ func TestNotifier(t *testing.T) { close(ch) wg.Wait() require.Equal(t, testCase.NumCallbacks, numCalls.Load()) + require.Equal(t, testCase.NumConditions, numConditions.Load()) }) } } @@ -99,11 +108,10 @@ func durations(ds ...time.Duration) []time.Duration { return ds } -func ticks(t time.Time, n int) []time.Time { - ts := make([]time.Time, n+1) - ts = append(ts, t) - for i := 0; i < n; i++ { - ts = append(ts, t.Add(time.Duration(n)*time.Second)) +func fakeTicker(t time.Time, d time.Duration, n int) []time.Time { + ts := make([]time.Time, 0) + for i := 1; i <= n; i++ { + ts = append(ts, t.Add(time.Duration(n)*d)) } return ts } From 67516c71e753caebe746df7fadbbc164247da7b7 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Fri, 13 May 2022 10:41:37 +0000 Subject: [PATCH 3/6] address PR comments --- cli/ssh.go | 56 +++++++++++++++++------- coderd/autobuild/notify/notifier.go | 35 ++++++++++----- coderd/autobuild/notify/notifier_test.go | 6 ++- 3 files changed, 68 insertions(+), 29 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index 1823cb3afd66e..4daf9fe2761ee 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -2,9 +2,11 @@ package cli import ( "context" + "errors" "fmt" "io" "os" + "path/filepath" "strings" "time" @@ -24,6 +26,9 @@ import ( "github.com/coder/coder/codersdk" ) +var autostopPollInterval = 30 * time.Second +var autostopNotifyCountdown = []time.Duration{5 * time.Minute, time.Minute} + func ssh() *cobra.Command { var ( stdio bool @@ -113,19 +118,8 @@ func ssh() *cobra.Command { } defer conn.Close() - // Notify the user if the workspace is due to shutdown. This uses - // the library gen2brain/beeep under the hood. - countdown := []time.Duration{ - 30 * time.Minute, - 10 * time.Minute, - 5 * time.Minute, - time.Minute, - } - condition := notifyWorkspaceAutostop(cmd.Context(), client, workspace.ID) - notifier := notify.New(condition, countdown...) - ticker := time.NewTicker(30 * time.Second) - defer ticker.Stop() - go notifier.Poll(ticker.C) + stopPolling := tryPollWorkspaceAutostop(cmd.Context(), client, workspace) + defer stopPolling() if stdio { rawSSH, err := conn.SSH() @@ -199,7 +193,36 @@ func ssh() *cobra.Command { return cmd } -func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID) notify.Condition { +// 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()) { + lockPath := filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()) + lockStat, err := os.Stat(lockPath) + if err == nil { + // Lock file already exists for this workspace. How old is it? + lockAge := lockStat.ModTime().Sub(time.Now()) + if lockAge < 3*autostopPollInterval { + // Lock file exists and is still "fresh". Do nothing. + return func() {} + } + } + if !errors.Is(err, os.ErrNotExist) { + // No permission to write to temp? Not much we can do. + return func() {} + } + lockFile, err := os.Create(lockPath) + if err != nil { + // Someone got there already? + return func() {} + } + + condition := notifyCondition(ctx, client, workspace.ID, lockFile) + 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, lockFile *os.File) notify.Condition { return func(now time.Time) (deadline time.Time, callback func()) { ws, err := client.Workspace(ctx, workspaceID) if err != nil { @@ -230,7 +253,10 @@ func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, works title = fmt.Sprintf("Workspace %s stopping!", ws.Name) body = fmt.Sprintf("Your Coder workspace %s is stopping any time now!", ws.Name) } - beeep.Notify(title, body, "") + // notify user with a native system notification (best effort) + _ = beeep.Notify(title, body, "") + // update lockFile (best effort) + _ = os.Chtimes(lockFile.Name(), now, now) } return deadline.Truncate(time.Minute), callback } diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go index e62c7b1387edb..8ba3fd150990c 100644 --- a/coderd/autobuild/notify/notifier.go +++ b/coderd/autobuild/notify/notifier.go @@ -8,18 +8,14 @@ import ( // Notifier calls a Condition at most once for each count in countdown. type Notifier struct { - sync.Mutex + lock sync.Mutex condition Condition notifiedAt map[time.Duration]bool countdown []time.Duration } // Condition is a function that gets executed with a certain time. -type Condition func(now time.Time) (deadline time.Time, callback func()) - -// New returns a Notifier that calls cond once every time it polls. -// - Condition is a function that returns the deadline and a callback. -// It should return the deadline for the notification, as well as a +// - 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. @@ -27,16 +23,30 @@ type Condition func(now time.Time) (deadline time.Time, callback func()) // 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. - sort.Slice(unique(countdown), func(i, j int) bool { - return countdown[i] < countdown[j] + ct := unique(countdown) + sort.Slice(ct, func(i, j int) bool { + return ct[i] < ct[j] }) n := &Notifier{ - countdown: countdown, + countdown: ct, condition: cond, notifiedAt: make(map[time.Duration]bool), } @@ -45,6 +55,7 @@ func New(cond Condition, countdown ...time.Duration) *Notifier { } // 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()) @@ -54,8 +65,8 @@ func (n *Notifier) Poll(ticker <-chan time.Time) { } func (n *Notifier) pollOnce(tick time.Time) { - n.Lock() - defer n.Unlock() + n.lock.Lock() + defer n.lock.Unlock() deadline, callback := n.condition(tick) if deadline.IsZero() { @@ -81,7 +92,7 @@ func unique(ds []time.Duration) []time.Duration { for _, d := range ds { m[d] = true } - ks := make([]time.Duration, 0) + var ks []time.Duration for k := range m { ks = append(ks, k) } diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go index 677cb5c4ce780..5b0b63a430c05 100644 --- a/coderd/autobuild/notify/notifier_test.go +++ b/coderd/autobuild/notify/notifier_test.go @@ -5,10 +5,11 @@ import ( "testing" "time" - "github.com/coder/coder/coderd/autobuild/notify" "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) { @@ -77,6 +78,7 @@ func TestNotifier(t *testing.T) { 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) @@ -109,7 +111,7 @@ func durations(ds ...time.Duration) []time.Duration { } func fakeTicker(t time.Time, d time.Duration, n int) []time.Time { - ts := make([]time.Time, 0) + var ts []time.Time for i := 1; i <= n; i++ { ts = append(ts, t.Add(time.Duration(n)*d)) } From 55a1de235d82f193886c96fc32af6c786a6c874e Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Fri, 13 May 2022 10:52:10 +0000 Subject: [PATCH 4/6] fixup! address PR comments --- cli/ssh.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index 4daf9fe2761ee..fe5f14e29a0a6 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -201,7 +201,7 @@ func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, work lockStat, err := os.Stat(lockPath) if err == nil { // Lock file already exists for this workspace. How old is it? - lockAge := lockStat.ModTime().Sub(time.Now()) + lockAge := time.Now().Sub(lockStat.ModTime()) if lockAge < 3*autostopPollInterval { // Lock file exists and is still "fresh". Do nothing. return func() {} @@ -224,6 +224,8 @@ func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, work // Notify the user if the workspace is due to shutdown. func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lockFile *os.File) notify.Condition { return func(now time.Time) (deadline time.Time, callback func()) { + // update lockFile (best effort) + _ = os.Chtimes(lockFile.Name(), now, now) ws, err := client.Workspace(ctx, workspaceID) if err != nil { return time.Time{}, nil @@ -255,8 +257,6 @@ func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID u } // notify user with a native system notification (best effort) _ = beeep.Notify(title, body, "") - // update lockFile (best effort) - _ = os.Chtimes(lockFile.Name(), now, now) } return deadline.Truncate(time.Minute), callback } From e830baf6e265b870bc939d16c38e7327900c7e36 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Fri, 13 May 2022 12:31:17 +0100 Subject: [PATCH 5/6] use gofrs/flock instead to handle file locking --- cli/ssh.go | 35 ++++++++++------------------------- go.mod | 1 + go.sum | 2 ++ 3 files changed, 13 insertions(+), 25 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index fe5f14e29a0a6..b42295a30adee 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -2,7 +2,6 @@ package cli import ( "context" - "errors" "fmt" "io" "os" @@ -11,6 +10,7 @@ import ( "time" "github.com/gen2brain/beeep" + "github.com/gofrs/flock" "github.com/google/uuid" "github.com/mattn/go-isatty" "github.com/spf13/cobra" @@ -197,35 +197,20 @@ func ssh() *cobra.Command { // 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()) { - lockPath := filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()) - lockStat, err := os.Stat(lockPath) - if err == nil { - // Lock file already exists for this workspace. How old is it? - lockAge := time.Now().Sub(lockStat.ModTime()) - if lockAge < 3*autostopPollInterval { - // Lock file exists and is still "fresh". Do nothing. - return func() {} - } - } - if !errors.Is(err, os.ErrNotExist) { - // No permission to write to temp? Not much we can do. - return func() {} - } - lockFile, err := os.Create(lockPath) - if err != nil { - // Someone got there already? - return func() {} - } - - condition := notifyCondition(ctx, client, workspace.ID, lockFile) + 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, lockFile *os.File) notify.Condition { +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()) { - // update lockFile (best effort) - _ = os.Chtimes(lockFile.Name(), now, now) + // 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 diff --git a/go.mod b/go.mod index edcf438c2eadf..504d97afe6984 100644 --- a/go.mod +++ b/go.mod @@ -64,6 +64,7 @@ require ( 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 diff --git a/go.sum b/go.sum index 2994f1468ee93..84f57a8fde153 100644 --- a/go.sum +++ b/go.sum @@ -754,6 +754,8 @@ github.com/godbus/dbus/v5 v5.0.6/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5x 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= From 4f39908fa03aac44288f8a022370ff3e76284471 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Fri, 13 May 2022 16:34:01 +0000 Subject: [PATCH 6/6] notify only once at most five minutes before shutdown --- cli/ssh.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/ssh.go b/cli/ssh.go index b42295a30adee..7b3bf2ae9bddd 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -27,7 +27,7 @@ import ( ) var autostopPollInterval = 30 * time.Second -var autostopNotifyCountdown = []time.Duration{5 * time.Minute, time.Minute} +var autostopNotifyCountdown = []time.Duration{5 * time.Minute} func ssh() *cobra.Command { var (