Skip to content

Commit 67516c7

Browse files
johnstcnjohnstcn
committed
address PR comments
1 parent 8313a52 commit 67516c7

File tree

3 files changed

+68
-29
lines changed

3 files changed

+68
-29
lines changed

‎cli/ssh.go

Lines changed: 41 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@ package cli
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"io"
78
"os"
9+
"path/filepath"
810
"strings"
911
"time"
1012

@@ -24,6 +26,9 @@ import (
2426
"github.com/coder/coder/codersdk"
2527
)
2628

29+
var autostopPollInterval = 30 * time.Second
30+
var autostopNotifyCountdown = []time.Duration{5 * time.Minute, time.Minute}
31+
2732
func ssh() *cobra.Command {
2833
var (
2934
stdio bool
@@ -113,19 +118,8 @@ func ssh() *cobra.Command {
113118
}
114119
defer conn.Close()
115120

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)
121+
stopPolling := tryPollWorkspaceAutostop(cmd.Context(), client, workspace)
122+
defer stopPolling()
129123

130124
if stdio {
131125
rawSSH, err := conn.SSH()
@@ -199,7 +193,36 @@ func ssh() *cobra.Command {
199193
return cmd
200194
}
201195

202-
func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID) notify.Condition {
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+
lockPath := filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String())
201+
lockStat, err := os.Stat(lockPath)
202+
if err == nil {
203+
// Lock file already exists for this workspace. How old is it?
204+
lockAge := lockStat.ModTime().Sub(time.Now())
205+
if lockAge < 3*autostopPollInterval {
206+
// Lock file exists and is still "fresh". Do nothing.
207+
return func() {}
208+
}
209+
}
210+
if !errors.Is(err, os.ErrNotExist) {
211+
// No permission to write to temp? Not much we can do.
212+
return func() {}
213+
}
214+
lockFile, err := os.Create(lockPath)
215+
if err != nil {
216+
// Someone got there already?
217+
return func() {}
218+
}
219+
220+
condition := notifyCondition(ctx, client, workspace.ID, lockFile)
221+
return notify.Notify(condition, autostopPollInterval, autostopNotifyCountdown...)
222+
}
223+
224+
// Notify the user if the workspace is due to shutdown.
225+
func notifyCondition(ctx context.Context, client *codersdk.Client, workspaceID uuid.UUID, lockFile *os.File) notify.Condition {
203226
return func(now time.Time) (deadline time.Time, callback func()) {
204227
ws, err := client.Workspace(ctx, workspaceID)
205228
if err != nil {
@@ -230,7 +253,10 @@ func notifyWorkspaceAutostop(ctx context.Context, client *codersdk.Client, works
230253
title = fmt.Sprintf("Workspace %s stopping!", ws.Name)
231254
body = fmt.Sprintf("Your Coder workspace %s is stopping any time now!", ws.Name)
232255
}
233-
beeep.Notify(title, body, "")
256+
// notify user with a native system notification (best effort)
257+
_ = beeep.Notify(title, body, "")
258+
// update lockFile (best effort)
259+
_ = os.Chtimes(lockFile.Name(), now, now)
234260
}
235261
return deadline.Truncate(time.Minute), callback
236262
}

‎coderd/autobuild/notify/notifier.go

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,35 +8,45 @@ import (
88

99
// Notifier calls a Condition at most once for each count in countdown.
1010
type Notifier struct {
11-
sync.Mutex
11+
lock sync.Mutex
1212
condition Condition
1313
notifiedAt map[time.Duration]bool
1414
countdown []time.Duration
1515
}
1616

1717
// 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
18+
// - It should return the deadline for the notification, as well as a
2319
// callback function to execute once the time to the deadline is
2420
// less than one of the notify attempts. If deadline is the zero
2521
// time, callback will not be executed.
2622
// - Callback is executed once for every time the difference between deadline
2723
// and the current time is less than an element of countdown.
2824
// - To enforce a minimum interval between consecutive callbacks, truncate
2925
// 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.
3039
// - Duplicate values are removed from countdown, and it is sorted in
3140
// descending order.
3241
func New(cond Condition, countdown ...time.Duration) *Notifier {
3342
// 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]
43+
ct := unique(countdown)
44+
sort.Slice(ct, func(i, j int) bool {
45+
return ct[i] < ct[j]
3646
})
3747

3848
n := &Notifier{
39-
countdown: countdown,
49+
countdown: ct,
4050
condition: cond,
4151
notifiedAt: make(map[time.Duration]bool),
4252
}
@@ -45,6 +55,7 @@ func New(cond Condition, countdown ...time.Duration) *Notifier {
4555
}
4656

4757
// Poll polls once immediately, and then once for every value from ticker.
58+
// Poll exits when ticker is closed.
4859
func (n *Notifier) Poll(ticker <-chan time.Time) {
4960
// poll once immediately
5061
n.pollOnce(time.Now())
@@ -54,8 +65,8 @@ func (n *Notifier) Poll(ticker <-chan time.Time) {
5465
}
5566

5667
func (n *Notifier) pollOnce(tick time.Time) {
57-
n.Lock()
58-
defer n.Unlock()
68+
n.lock.Lock()
69+
defer n.lock.Unlock()
5970

6071
deadline, callback := n.condition(tick)
6172
if deadline.IsZero() {
@@ -81,7 +92,7 @@ func unique(ds []time.Duration) []time.Duration {
8192
for _, d := range ds {
8293
m[d] = true
8394
}
84-
ks := make([]time.Duration, 0)
95+
var ks []time.Duration
8596
for k := range m {
8697
ks = append(ks, k)
8798
}

‎coderd/autobuild/notify/notifier_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,11 @@ import (
55
"testing"
66
"time"
77

8-
"github.com/coder/coder/coderd/autobuild/notify"
98
"github.com/stretchr/testify/require"
109
"go.uber.org/atomic"
1110
"go.uber.org/goleak"
11+
12+
"github.com/coder/coder/coderd/autobuild/notify"
1213
)
1314

1415
func TestNotifier(t *testing.T) {
@@ -77,6 +78,7 @@ func TestNotifier(t *testing.T) {
7778
for _, testCase := range testCases {
7879
testCase := testCase
7980
t.Run(testCase.Name, func(t *testing.T) {
81+
t.Parallel()
8082
ch := make(chan time.Time)
8183
numConditions := atomic.NewInt64(0)
8284
numCalls := atomic.NewInt64(0)
@@ -109,7 +111,7 @@ func durations(ds ...time.Duration) []time.Duration {
109111
}
110112

111113
func fakeTicker(t time.Time, d time.Duration, n int) []time.Time {
112-
ts := make([]time.Time, 0)
114+
var ts []time.Time
113115
for i := 1; i <= n; i++ {
114116
ts = append(ts, t.Add(time.Duration(n)*d))
115117
}

0 commit comments

Comments
 (0)