Skip to content

Commit 125584c

Browse files
johnstcnjohnstcn
committed
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.
1 parent 64e408c commit 125584c

File tree

5 files changed

+269
-0
lines changed

5 files changed

+269
-0
lines changed

‎cli/ssh.go

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,13 @@ package cli
22

33
import (
44
"context"
5+
"fmt"
56
"io"
67
"os"
78
"strings"
9+
"time"
810

11+
"github.com/gen2brain/beeep"
912
"github.com/google/uuid"
1013
"github.com/mattn/go-isatty"
1114
"github.com/spf13/cobra"
@@ -15,6 +18,8 @@ import (
1518

1619
"github.com/coder/coder/cli/cliflag"
1720
"github.com/coder/coder/cli/cliui"
21+
"github.com/coder/coder/coderd/autobuild/notify"
22+
"github.com/coder/coder/coderd/autobuild/schedule"
1823
"github.com/coder/coder/coderd/database"
1924
"github.com/coder/coder/codersdk"
2025
)
@@ -108,6 +113,20 @@ func ssh() *cobra.Command {
108113
}
109114
defer conn.Close()
110115

116+
// Notify the user if the workspace is due to shutdown. This uses
117+
// the library gen2brain/beeep under the hood.
118+
countdown := []time.Duration{
119+
30 * time.Minute,
120+
10 * time.Minute,
121+
5 * time.Minute,
122+
time.Minute,
123+
}
124+
condition := notifyWorkspaceAutostop(cmd.Context(), client, workspace.ID)
125+
notifier := notify.New(condition, countdown...)
126+
ticker := time.NewTicker(30 * time.Second)
127+
defer ticker.Stop()
128+
go notifier.Poll(ticker.C)
129+
111130
if stdio {
112131
rawSSH, err := conn.SSH()
113132
if err != nil {
@@ -179,3 +198,40 @@ func ssh() *cobra.Command {
179198

180199
return cmd
181200
}
201+
202+
func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID) notify.Condition {
203+
return func(now time.Time) (deadline time.Time, callback func()) {
204+
ws, err := client.Workspace(ctx, workspaceID)
205+
if err != nil {
206+
return time.Time{}, nil
207+
}
208+
209+
if ws.AutostopSchedule == "" {
210+
return time.Time{}, nil
211+
}
212+
213+
sched, err := schedule.Weekly(ws.AutostopSchedule)
214+
if err != nil {
215+
return time.Time{}, nil
216+
}
217+
218+
deadline = sched.Next(now)
219+
callback = func() {
220+
ttl := deadline.Sub(now)
221+
var title, body string
222+
if ttl > time.Minute {
223+
title = fmt.Sprintf(`Workspace %s stopping in %.0f mins`, ws.Name, ttl.Minutes())
224+
body = fmt.Sprintf(
225+
`Your Coder workspace %s is scheduled to stop at %s.`,
226+
ws.Name,
227+
deadline.Format(time.Kitchen),
228+
)
229+
} else {
230+
title = fmt.Sprintf("Workspace %s stopping!", ws.Name)
231+
body = fmt.Sprintf("Your Coder workspace %s is stopping any time now!", ws.Name)
232+
}
233+
beeep.Notify(title, body, "")
234+
}
235+
return deadline.Truncate(time.Minute), callback
236+
}
237+
}

‎coderd/autobuild/notify/notifier.go

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package notify
2+
3+
import (
4+
"sort"
5+
"sync"
6+
"time"
7+
)
8+
9+
// Notifier calls a Condition at most once for each count in countdown.
10+
type Notifier struct {
11+
sync.Mutex
12+
condition Condition
13+
notifiedAt map[time.Duration]bool
14+
countdown []time.Duration
15+
}
16+
17+
// Condition is a function that gets executed with a certain time.
18+
type Condition func(now time.Time) (deadline time.Time, callback func())
19+
20+
// New returns a Notifier that calls cond once every time it polls.
21+
// - Condition is a function that returns the deadline and a callback.
22+
// It should return the deadline for the notification, as well as a
23+
// callback function to execute once the time to the deadline is
24+
// less than one of the notify attempts. If deadline is the zero
25+
// time, callback will not be executed.
26+
// - Callback is executed once for every time the difference between deadline
27+
// and the current time is less than an element of countdown.
28+
// - To enforce a minimum interval between consecutive callbacks, truncate
29+
// the returned deadline to the minimum interval.
30+
// - Duplicate values are removed from countdown, and it is sorted in
31+
// descending order.
32+
func New(cond Condition, countdown ...time.Duration) *Notifier {
33+
// Ensure countdown is sorted in descending order and contains no duplicates.
34+
sort.Slice(unique(countdown), func(i, j int) bool {
35+
return countdown[i] < countdown[j]
36+
})
37+
38+
n := &Notifier{
39+
countdown: countdown,
40+
condition: cond,
41+
notifiedAt: make(map[time.Duration]bool),
42+
}
43+
44+
return n
45+
}
46+
47+
// Poll polls once immediately, and then once for every value from ticker.
48+
func (n *Notifier) Poll(ticker <-chan time.Time) {
49+
// poll once immediately
50+
n.pollOnce(time.Now())
51+
for t := range ticker {
52+
n.pollOnce(t)
53+
}
54+
}
55+
56+
func (n *Notifier) pollOnce(tick time.Time) {
57+
n.Lock()
58+
defer n.Unlock()
59+
60+
deadline, callback := n.condition(tick)
61+
if deadline.IsZero() {
62+
return
63+
}
64+
65+
timeRemaining := deadline.Sub(tick)
66+
for _, tock := range n.countdown {
67+
if timeRemaining <= tock && !n.notifiedAt[tock] {
68+
callback()
69+
n.notifiedAt[tock] = true
70+
}
71+
}
72+
}
73+
74+
func unique(ds []time.Duration) []time.Duration {
75+
m := make(map[time.Duration]bool)
76+
for _, d := range ds {
77+
m[d] = true
78+
}
79+
ks := make([]time.Duration, 0)
80+
for k := range m {
81+
ks = append(ks, k)
82+
}
83+
return ks
84+
}

‎coderd/autobuild/notify/notifier_test.go

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package notify_test
2+
3+
import (
4+
"sync"
5+
"testing"
6+
"time"
7+
8+
"github.com/coder/coder/coderd/autobuild/notify"
9+
"github.com/stretchr/testify/require"
10+
"go.uber.org/atomic"
11+
"go.uber.org/goleak"
12+
)
13+
14+
func TestNotifier(t *testing.T) {
15+
t.Parallel()
16+
17+
now := time.Now()
18+
19+
testCases := []struct {
20+
Name string
21+
Countdown []time.Duration
22+
Ticks []time.Time
23+
ConditionDeadline time.Time
24+
NumConditions int64
25+
NumCallbacks int64
26+
}{
27+
{
28+
Name: "zero deadline",
29+
Countdown: durations(),
30+
Ticks: ticks(now, 0),
31+
ConditionDeadline: time.Time{},
32+
NumConditions: 1,
33+
NumCallbacks: 0,
34+
},
35+
{
36+
Name: "no calls",
37+
Countdown: durations(),
38+
Ticks: ticks(now, 0),
39+
ConditionDeadline: now,
40+
NumConditions: 1,
41+
NumCallbacks: 0,
42+
},
43+
{
44+
Name: "exactly one call",
45+
Countdown: durations(time.Second),
46+
Ticks: ticks(now, 2),
47+
ConditionDeadline: now.Add(2 * time.Second),
48+
NumConditions: 2,
49+
NumCallbacks: 1,
50+
},
51+
{
52+
Name: "two calls",
53+
Countdown: durations(4*time.Second, 2*time.Second),
54+
Ticks: ticks(now, 5),
55+
ConditionDeadline: now.Add(5 * time.Second),
56+
NumConditions: 6,
57+
NumCallbacks: 2,
58+
},
59+
{
60+
Name: "wrong order should not matter",
61+
Countdown: durations(2*time.Second, 4*time.Second),
62+
Ticks: ticks(now, 5),
63+
ConditionDeadline: now.Add(5 * time.Second),
64+
NumConditions: 6,
65+
NumCallbacks: 2,
66+
},
67+
}
68+
69+
for _, testCase := range testCases {
70+
testCase := testCase
71+
t.Run(testCase.Name, func(t *testing.T) {
72+
ch := make(chan time.Time)
73+
numConditions := atomic.NewInt64(0)
74+
numCalls := atomic.NewInt64(0)
75+
cond := func(time.Time) (time.Time, func()) {
76+
numConditions.Inc()
77+
return testCase.ConditionDeadline, func() {
78+
numCalls.Inc()
79+
}
80+
}
81+
var wg sync.WaitGroup
82+
go func() {
83+
n := notify.New(cond, testCase.Countdown...)
84+
n.Poll(ch)
85+
wg.Done()
86+
}()
87+
wg.Add(1)
88+
for _, tick := range testCase.Ticks {
89+
ch <- tick
90+
}
91+
close(ch)
92+
wg.Wait()
93+
require.Equal(t, testCase.NumCallbacks, numCalls.Load())
94+
})
95+
}
96+
}
97+
98+
func durations(ds ...time.Duration) []time.Duration {
99+
return ds
100+
}
101+
102+
func ticks(t time.Time, n int) []time.Time {
103+
ts := make([]time.Time, n+1)
104+
ts = append(ts, t)
105+
for i := 0; i < n; i++ {
106+
ts = append(ts, t.Add(time.Duration(n)*time.Second))
107+
}
108+
return ts
109+
}
110+
111+
func TestMain(m *testing.M) {
112+
goleak.VerifyTestMain(m)
113+
}

‎go.mod

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ require (
5858
github.com/fatih/color v1.13.0
5959
github.com/fatih/structs v1.1.0
6060
github.com/fullsailor/pkcs7 v0.0.0-20190404230743-d7302db945fa
61+
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a
6162
github.com/gliderlabs/ssh v0.3.3
6263
github.com/go-chi/chi/v5 v5.0.7
6364
github.com/go-chi/httprate v0.5.3
@@ -159,8 +160,10 @@ require (
159160
github.com/ghodss/yaml v1.0.0 // indirect
160161
github.com/go-playground/locales v0.14.0 // indirect
161162
github.com/go-playground/universal-translator v0.18.0 // indirect
163+
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 // indirect
162164
github.com/gobwas/glob v0.2.3 // indirect
163165
github.com/gobwas/ws v1.1.0 // indirect
166+
github.com/godbus/dbus/v5 v5.1.0 // indirect
164167
github.com/gogo/protobuf v1.3.2 // indirect
165168
github.com/golang/glog v1.0.0 // indirect
166169
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
@@ -196,6 +199,7 @@ require (
196199
github.com/muesli/reflow v0.3.0 // indirect
197200
github.com/muesli/termenv v0.11.1-0.20220212125758-44cd13922739 // indirect
198201
github.com/niklasfasching/go-org v1.6.2 // indirect
202+
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect
199203
github.com/opencontainers/go-digest v1.0.0 // indirect
200204
github.com/opencontainers/image-spec v1.0.2 // indirect
201205
github.com/opencontainers/runc v1.1.0 // indirect
@@ -226,6 +230,7 @@ require (
226230
github.com/spf13/afero v1.8.2 // indirect
227231
github.com/spf13/cast v1.4.1 // indirect
228232
github.com/spf13/jwalterweatherman v1.1.0 // indirect
233+
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af // indirect
229234
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect
230235
github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b // indirect
231236
github.com/tinylib/msgp v1.1.2 // indirect

‎go.sum

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,6 +606,8 @@ github.com/gabriel-vasile/mimetype v1.3.1/go.mod h1:fA8fi6KUiG7MgQQ+mEWotXoEOvmx
606606
github.com/gabriel-vasile/mimetype v1.4.0/go.mod h1:fA8fi6KUiG7MgQQ+mEWotXoEOvmxRtOJlERCzSmRvr8=
607607
github.com/garyburd/redigo v0.0.0-20150301180006-535138d7bcd7/go.mod h1:NR3MbYisc3/PwhQ00EMzDiPmrwpPxAn5GI05/YaO1SY=
608608
github.com/garyburd/redigo v1.6.3/go.mod h1:rTb6epsqigu3kYKBnaF028A7Tf/Aw5s0cqA47doKKqw=
609+
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a h1:fwNLHrP5Rbg/mGSXCjtPdpbqv2GucVTA/KMi8wEm6mE=
610+
github.com/gen2brain/beeep v0.0.0-20220402123239-6a3042f4b71a/go.mod h1:/WeFVhhxMOGypVKS0w8DUJxUBbHypnWkUVnW7p5c9Pw=
609611
github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ=
610612
github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
611613
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
703705
github.com/go-test/deep v1.0.2/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
704706
github.com/go-test/deep v1.0.3 h1:ZrJSEWsXzPOxaZnFteGEfooLba+ju3FYIbOrS+rQd68=
705707
github.com/go-test/deep v1.0.3/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
708+
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 h1:qZNfIGkIANxGv/OqtnntR4DfOY2+BgwR60cAcu/i3SE=
709+
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4/go.mod h1:kW3HQ4UdaAyrUCSSDR4xUzBKW6O2iA4uHhk7AtyYp10=
706710
github.com/gobuffalo/attrs v0.0.0-20190224210810-a9411de4debd/go.mod h1:4duuawTqi2wkkpB4ePgWMaai6/Kc6WEz83bhFwpHzj0=
707711
github.com/gobuffalo/depgen v0.0.0-20190329151759-d478694a28d3/go.mod h1:3STtPUQYuzV0gBVOY3vy6CfMm/ljR4pABfrTeHNLHUY=
708712
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
747751
github.com/godbus/dbus/v5 v5.0.3/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
748752
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
749753
github.com/godbus/dbus/v5 v5.0.6/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
754+
github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
755+
github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
750756
github.com/gofiber/fiber/v2 v2.11.0/go.mod h1:oZTLWqYnqpMMuF922SjGbsYZsdpE1MCfh416HNdweIM=
751757
github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
752758
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:
13271333
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
13281334
github.com/niklasfasching/go-org v1.6.2 h1:kQBIZlfL4oRNApJCrBgaeNBfzxWzP6XlC7/b744Polk=
13291335
github.com/niklasfasching/go-org v1.6.2/go.mod h1:wn76Xgu4/KRe43WZhsgZjxYMaloSrl3BSweGV74SwHs=
1336+
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d h1:VhgPp6v9qf9Agr/56bj7Y/xa04UccTW04VP0Qed4vnQ=
1337+
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d/go.mod h1:YUTz3bUH2ZwIWBy3CJBeOBEugqcmXREj14T+iG/4k4U=
13301338
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
13311339
github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
13321340
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
16551663
github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
16561664
github.com/tabbed/pqtype v0.1.1 h1:PhEcb9JZ8jr7SUjJDFjRPxny0M8fkXZrxn/a9yQfoZg=
16571665
github.com/tabbed/pqtype v0.1.1/go.mod h1:HLt2kLJPcUhODQkYn3mJkMHXVsuv3Z2n5NZEeKXL0Uk=
1666+
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af h1:6yITBqGTE2lEeTPG04SN9W+iWHCRyHqlVYILiSXziwk=
1667+
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af/go.mod h1:4F09kP5F+am0jAwlQLddpoMDM+iewkxxt6nxUQ5nq5o=
16581668
github.com/tchap/go-patricia v2.2.6+incompatible/go.mod h1:bmLyhP68RS6kStMGxByiQ23RP/odRBOTVjwp2cDyi6I=
16591669
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 h1:89CEmDvlq/F7SJEOqkIdNDGJXrQIhuIx9D2DBXjavSU=
16601670
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
22032213
golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
22042214
golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
22052215
golang.org/x/sys v0.0.0-20220317061510-51cd9980dadf/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
2216+
golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
22062217
golang.org/x/sys v0.0.0-20220328115105-d36c6a25d886/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
22072218
golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
22082219
golang.org/x/sys v0.0.0-20220502124256-b6088ccd6cba h1:AyHWHCBVlIYI5rgEM3o+1PLd0sLPcIAoaUckGQMaWtw=

0 commit comments

Comments
 (0)