fix: Close notifier Poll goroutine on stop (#3252)

mafredri · web-flow · commit a0320f455a8b · 2022-07-27T20:26:13.000+03:00
Fix towards #3221.
diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go
@@ -1,13 +1,18 @@
 package notify
 
 import (
+	"context"
 	"sort"
 	"sync"
 	"time"
 )
 
 // Notifier calls a Condition at most once for each count in countdown.
 type Notifier struct {
+	ctx      context.Context
+	cancel   context.CancelFunc
+	pollDone chan struct{}
+
 	lock       sync.Mutex
 	condition  Condition
 	notifiedAt map[time.Duration]bool
@@ -28,11 +33,14 @@ 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()) {
+func Notify(cond Condition, interval time.Duration, countdown ...time.Duration) (closeFunc func()) {
 	notifier := New(cond, countdown...)
 	ticker := time.NewTicker(interval)
 	go notifier.Poll(ticker.C)
-	return ticker.Stop
+	return func() {
+		ticker.Stop()
+		_ = notifier.Close()
+	}
 }
 
 // New returns a Notifier that calls cond once every time it polls.
@@ -45,7 +53,11 @@ func New(cond Condition, countdown ...time.Duration) *Notifier {
 		return ct[i] < ct[j]
 	})
 
+	ctx, cancel := context.WithCancel(context.Background())
 	n := &Notifier{
+		ctx:        ctx,
+		cancel:     cancel,
+		pollDone:   make(chan struct{}),
 		countdown:  ct,
 		condition:  cond,
 		notifiedAt: make(map[time.Duration]bool),
@@ -57,13 +69,29 @@ func New(cond Condition, countdown ...time.Duration) *Notifier {
 // 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) {
+	defer close(n.pollDone)
+
 	// poll once immediately
 	n.pollOnce(time.Now())
-	for t := range ticker {
-		n.pollOnce(t)
+	for {
+		select {
+		case <-n.ctx.Done():
+			return
+		case t, ok := <-ticker:
+			if !ok {
+				return
+			}
+			n.pollOnce(t)
+		}
 	}
 }
 
+func (n *Notifier) Close() error {
+	n.cancel()
+	<-n.pollDone
+	return nil
+}
+
 func (n *Notifier) pollOnce(tick time.Time) {
 	n.lock.Lock()
 	defer n.lock.Unlock()
diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go
@@ -90,9 +90,10 @@ func TestNotifier(t *testing.T) {
 			}
 			var wg sync.WaitGroup
 			go func() {
+				defer wg.Done()
 				n := notify.New(cond, testCase.Countdown...)
+				defer n.Close()
 				n.Poll(ch)
-				wg.Done()
 			}()
 			wg.Add(1)
 			for _, tick := range testCase.Ticks {
