Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
May 13, 2022
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions cli/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
)
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}
}
84 changes: 84 additions & 0 deletions coderd/autobuild/notify/notifier.go
Original file line number Diff line number Diff line change
@@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it make sense to move this into pollOnce? My thought is that this map never reset, and a notification can only be shown once per countdown entry (for the lifetime of the Notifier). Not sure if that's a problem or not though.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could move into Poll, but not into pollOnce.
If we did that, we could probably also remove the sync.Mutex.
I like the idea of the same Notifier instance keeping track of what it's already notified between successive calls to Poll() though.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok that makes sense. I guess it's mostly a question of how the package is to be used and what the user (might) expect to happen. I imagine this will not be a problem for ssh since when the workspace stops, so does the command, so the notifier doesn't outlive its purpose.

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
}
113 changes: 113 additions & 0 deletions coderd/autobuild/notify/notifier_test.go
Original file line number Diff line number Diff line change
@@ -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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool use of atomic, I need to remember this! ☺️

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)
}
5 changes: 5 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down