Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 4051baf

Browse files
johnstcnkylecarbs
johnstcn
authored and
kylecarbs
committed
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.
1 parent 4490e4e commit 4051baf

File tree

5 files changed

+309
-0
lines changed

5 files changed

+309
-0
lines changed

‎cli/ssh.go

+67
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,15 @@ package cli
22

33
import (
44
"context"
5+
"fmt"
56
"io"
67
"os"
8+
"path/filepath"
79
"strings"
10+
"time"
811

12+
"github.com/gen2brain/beeep"
13+
"github.com/gofrs/flock"
914
"github.com/google/uuid"
1015
"github.com/mattn/go-isatty"
1116
"github.com/spf13/cobra"
@@ -15,10 +20,15 @@ import (
1520

1621
"github.com/coder/coder/cli/cliflag"
1722
"github.com/coder/coder/cli/cliui"
23+
"github.com/coder/coder/coderd/autobuild/notify"
24+
"github.com/coder/coder/coderd/autobuild/schedule"
1825
"github.com/coder/coder/coderd/database"
1926
"github.com/coder/coder/codersdk"
2027
)
2128

29+
var autostopPollInterval = 30 * time.Second
30+
var autostopNotifyCountdown = []time.Duration{5 * time.Minute}
31+
2232
func ssh() *cobra.Command {
2333
var (
2434
stdio bool
@@ -108,6 +118,9 @@ func ssh() *cobra.Command {
108118
}
109119
defer conn.Close()
110120

121+
stopPolling := tryPollWorkspaceAutostop(cmd.Context(), client, workspace)
122+
defer stopPolling()
123+
111124
if stdio {
112125
rawSSH, err := conn.SSH()
113126
if err != nil {
@@ -179,3 +192,57 @@ func ssh() *cobra.Command {
179192

180193
return cmd
181194
}
195+
196+
// Attempt to poll workspace autostop. We write a per-workspace lockfile to
197+
// avoid spamming the user with notifications in case of multiple instances
198+
// of the CLI running simultaneously.
199+
func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace) (stop func()) {
200+
lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
201+
condition := notifyCondition(ctx, client, workspace.ID, lock)
202+
return notify.Notify(condition, autostopPollInterval, autostopNotifyCountdown...)
203+
}
204+
205+
// Notify the user if the workspace is due to shutdown.
206+
func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lock *flock.Flock) notify.Condition {
207+
return func(now time.Time) (deadline time.Time, callback func()) {
208+
// Keep trying to regain the lock.
209+
locked, err := lock.TryLockContext(ctx, autostopPollInterval)
210+
if err != nil || !locked {
211+
return time.Time{}, nil
212+
}
213+
214+
ws, err := client.Workspace(ctx, workspaceID)
215+
if err != nil {
216+
return time.Time{}, nil
217+
}
218+
219+
if ws.AutostopSchedule == "" {
220+
return time.Time{}, nil
221+
}
222+
223+
sched, err := schedule.Weekly(ws.AutostopSchedule)
224+
if err != nil {
225+
return time.Time{}, nil
226+
}
227+
228+
deadline = sched.Next(now)
229+
callback = func() {
230+
ttl := deadline.Sub(now)
231+
var title, body string
232+
if ttl > time.Minute {
233+
title = fmt.Sprintf(`Workspace %s stopping in %.0f mins`, ws.Name, ttl.Minutes())
234+
body = fmt.Sprintf(
235+
`Your Coder workspace %s is scheduled to stop at %s.`,
236+
ws.Name,
237+
deadline.Format(time.Kitchen),
238+
)
239+
} else {
240+
title = fmt.Sprintf("Workspace %s stopping!", ws.Name)
241+
body = fmt.Sprintf("Your Coder workspace %s is stopping any time now!", ws.Name)
242+
}
243+
// notify user with a native system notification (best effort)
244+
_ = beeep.Notify(title, body, "")
245+
}
246+
return deadline.Truncate(time.Minute), callback
247+
}
248+
}

‎coderd/autobuild/notify/notifier.go

+100
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
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+
lock 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+
// - It should return the deadline for the notification, as well as a
19+
// callback function to execute once the time to the deadline is
20+
// less than one of the notify attempts. If deadline is the zero
21+
// time, callback will not be executed.
22+
// - Callback is executed once for every time the difference between deadline
23+
// and the current time is less than an element of countdown.
24+
// - To enforce a minimum interval between consecutive callbacks, truncate
25+
// the returned deadline to the minimum interval.
26+
type Condition func(now time.Time) (deadline time.Time, callback func())
27+
28+
// Notify is a convenience function that initializes a new Notifier
29+
// with the given condition, interval, and countdown.
30+
// It is the responsibility of the caller to call close to stop polling.
31+
func Notify(cond Condition, interval time.Duration, countdown ...time.Duration) (close func()) {
32+
notifier := New(cond, countdown...)
33+
ticker := time.NewTicker(interval)
34+
go notifier.Poll(ticker.C)
35+
return ticker.Stop
36+
}
37+
38+
// New returns a Notifier that calls cond once every time it polls.
39+
// - Duplicate values are removed from countdown, and it is sorted in
40+
// descending order.
41+
func New(cond Condition, countdown ...time.Duration) *Notifier {
42+
// Ensure countdown is sorted in descending order and contains no duplicates.
43+
ct := unique(countdown)
44+
sort.Slice(ct, func(i, j int) bool {
45+
return ct[i] < ct[j]
46+
})
47+
48+
n := &Notifier{
49+
countdown: ct,
50+
condition: cond,
51+
notifiedAt: make(map[time.Duration]bool),
52+
}
53+
54+
return n
55+
}
56+
57+
// Poll polls once immediately, and then once for every value from ticker.
58+
// Poll exits when ticker is closed.
59+
func (n *Notifier) Poll(ticker <-chan time.Time) {
60+
// poll once immediately
61+
n.pollOnce(time.Now())
62+
for t := range ticker {
63+
n.pollOnce(t)
64+
}
65+
}
66+
67+
func (n *Notifier) pollOnce(tick time.Time) {
68+
n.lock.Lock()
69+
defer n.lock.Unlock()
70+
71+
deadline, callback := n.condition(tick)
72+
if deadline.IsZero() {
73+
return
74+
}
75+
76+
timeRemaining := deadline.Sub(tick)
77+
for _, tock := range n.countdown {
78+
if n.notifiedAt[tock] {
79+
continue
80+
}
81+
if timeRemaining > tock {
82+
continue
83+
}
84+
callback()
85+
n.notifiedAt[tock] = true
86+
return
87+
}
88+
}
89+
90+
func unique(ds []time.Duration) []time.Duration {
91+
m := make(map[time.Duration]bool)
92+
for _, d := range ds {
93+
m[d] = true
94+
}
95+
var ks []time.Duration
96+
for k := range m {
97+
ks = append(ks, k)
98+
}
99+
return ks
100+
}

‎coderd/autobuild/notify/notifier_test.go

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

‎go.mod

+6
Original file line numberDiff line numberDiff line change
@@ -58,11 +58,13 @@ 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
6465
github.com/go-chi/render v1.0.1
6566
github.com/go-playground/validator/v10 v10.11.0
67+
github.com/gofrs/flock v0.8.1
6668
github.com/gohugoio/hugo v0.98.0
6769
github.com/golang-jwt/jwt v3.2.2+incompatible
6870
github.com/golang-migrate/migrate/v4 v4.15.2
@@ -159,8 +161,10 @@ require (
159161
github.com/ghodss/yaml v1.0.0 // indirect
160162
github.com/go-playground/locales v0.14.0 // indirect
161163
github.com/go-playground/universal-translator v0.18.0 // indirect
164+
github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 // indirect
162165
github.com/gobwas/glob v0.2.3 // indirect
163166
github.com/gobwas/ws v1.1.0 // indirect
167+
github.com/godbus/dbus/v5 v5.1.0 // indirect
164168
github.com/gogo/protobuf v1.3.2 // indirect
165169
github.com/golang/glog v1.0.0 // indirect
166170
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
@@ -196,6 +200,7 @@ require (
196200
github.com/muesli/reflow v0.3.0 // indirect
197201
github.com/muesli/termenv v0.11.1-0.20220212125758-44cd13922739 // indirect
198202
github.com/niklasfasching/go-org v1.6.2 // indirect
203+
github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect
199204
github.com/opencontainers/go-digest v1.0.0 // indirect
200205
github.com/opencontainers/image-spec v1.0.2 // indirect
201206
github.com/opencontainers/runc v1.1.0 // indirect
@@ -226,6 +231,7 @@ require (
226231
github.com/spf13/afero v1.8.2 // indirect
227232
github.com/spf13/cast v1.4.1 // indirect
228233
github.com/spf13/jwalterweatherman v1.1.0 // indirect
234+
github.com/tadvi/systray v0.0.0-20190226123456-11a2b8fa57af // indirect
229235
github.com/templexxx/cpufeat v0.0.0-20180724012125-cef66df7f161 // indirect
230236
github.com/templexxx/xor v0.0.0-20191217153810-f85b25db303b // indirect
231237
github.com/tinylib/msgp v1.1.2 // indirect

0 commit comments

Comments
 (0)