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

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 all commits
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
67 changes: 67 additions & 0 deletions cli/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,15 @@ package cli

import (
"context"
"fmt"
"io"
"os"
"path/filepath"
"strings"
"time"

"github.com/gen2brain/beeep"
"github.com/gofrs/flock"
"github.com/google/uuid"
"github.com/mattn/go-isatty"
"github.com/spf13/cobra"
Expand All @@ -15,10 +20,15 @@ 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"
)

var autostopPollInterval = 30 * time.Second
var autostopNotifyCountdown = []time.Duration{5 * time.Minute}

func ssh() *cobra.Command {
var (
stdio bool
Expand Down Expand Up @@ -108,6 +118,9 @@ func ssh() *cobra.Command {
}
defer conn.Close()

stopPolling := tryPollWorkspaceAutostop(cmd.Context(), client, workspace)
defer stopPolling()

if stdio {
rawSSH, err := conn.SSH()
if err != nil {
Expand Down Expand Up @@ -179,3 +192,57 @@ func ssh() *cobra.Command {

return cmd
}

// 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()) {
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, lock *flock.Flock) notify.Condition {
return func(now time.Time) (deadline time.Time, callback func()) {
// Keep trying to regain the lock.
locked, err := lock.TryLockContext(ctx, autostopPollInterval)
if err != nil || !locked {
return time.Time{}, nil
}
Comment on lines +196 to +212
Copy link
Member

Choose a reason for hiding this comment

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

I wonder if it'd be best to place as much of this as possible in the notify package. eg. for coder port-forward we'd probably want to do the exact same type of thing.

I could imagine for the frontend too, we would send a notification via WebSocket or something!

Copy link
Member Author

Choose a reason for hiding this comment

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

That could be a good idea when we have a better idea of what different use-cases crop up.
I fear I've already over-engineered this enough as it is!


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)
}
// notify user with a native system notification (best effort)
_ = beeep.Notify(title, body, "")
}
return deadline.Truncate(time.Minute), callback
}
}
100 changes: 100 additions & 0 deletions coderd/autobuild/notify/notifier.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
package notify

import (
"sort"
"sync"
"time"
)

// Notifier calls a Condition at most once for each count in countdown.
type Notifier struct {
lock 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.
// - 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.
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.
ct := unique(countdown)
sort.Slice(ct, func(i, j int) bool {
return ct[i] < ct[j]
})

n := &Notifier{
countdown: ct,
condition: cond,
notifiedAt: make(map[time.Duration]bool),
}

return n
}

// 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())
for t := range ticker {
n.pollOnce(t)
}
}

func (n *Notifier) pollOnce(tick time.Time) {
n.lock.Lock()
defer n.lock.Unlock()

deadline, callback := n.condition(tick)
if deadline.IsZero() {
return
}

timeRemaining := deadline.Sub(tick)
for _, tock := range n.countdown {
if n.notifiedAt[tock] {
continue
}
if timeRemaining > tock {
continue
}
callback()
n.notifiedAt[tock] = true
return
}
}

func unique(ds []time.Duration) []time.Duration {
m := make(map[time.Duration]bool)
for _, d := range ds {
m[d] = true
}
var ks []time.Duration
for k := range m {
ks = append(ks, k)
}
return ks
}
123 changes: 123 additions & 0 deletions coderd/autobuild/notify/notifier_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
package notify_test

import (
"sync"
"testing"
"time"

"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) {
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: fakeTicker(now, time.Second, 0),
ConditionDeadline: time.Time{},
NumConditions: 1,
NumCallbacks: 0,
},
{
Name: "no calls",
Countdown: durations(),
Ticks: fakeTicker(now, time.Second, 0),
ConditionDeadline: now,
NumConditions: 1,
NumCallbacks: 0,
},
{
Name: "exactly one call",
Countdown: durations(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: fakeTicker(now, time.Second, 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: 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 {
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)
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())
require.Equal(t, testCase.NumConditions, numConditions.Load())
})
}
}

func durations(ds ...time.Duration) []time.Duration {
return ds
}

func fakeTicker(t time.Time, d time.Duration, n int) []time.Time {
var ts []time.Time
for i := 1; i <= n; i++ {
ts = append(ts, t.Add(time.Duration(n)*d))
}
return ts
}

func TestMain(m *testing.M) {
goleak.VerifyTestMain(m)
}
6 changes: 6 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -58,11 +58,13 @@ 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
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
Expand Down Expand Up @@ -159,8 +161,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 +200,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 +231,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
Loading