From 3f9f183b2403fc472821b6c8f86d640b98915818 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 27 Jul 2022 19:24:06 +0300
Subject: [PATCH 1/3] fix: Close notifier Poll goroutine on stop

Fix towards #3221.
---
 coderd/autobuild/notify/notifier.go | 33 +++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go
index 8ba3fd150990c..e459c4baebb16 100644
--- a/coderd/autobuild/notify/notifier.go
+++ b/coderd/autobuild/notify/notifier.go
@@ -1,6 +1,7 @@
 package notify
 
 import (
+	"context"
 	"sort"
 	"sync"
 	"time"
@@ -8,6 +9,10 @@ import (
 
 // 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,26 @@ 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 := <-ticker:
+			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()

From 117ed09ae57ecd7676e04077983bb374a01dceb3 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 27 Jul 2022 20:00:45 +0300
Subject: [PATCH 2/3] fix: Handle closed ticker in notifier Poll

---
 coderd/autobuild/notify/notifier.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go
index e459c4baebb16..f7e1058aee3b9 100644
--- a/coderd/autobuild/notify/notifier.go
+++ b/coderd/autobuild/notify/notifier.go
@@ -77,7 +77,10 @@ func (n *Notifier) Poll(ticker <-chan time.Time) {
 		select {
 		case <-n.ctx.Done():
 			return
-		case t := <-ticker:
+		case t, ok := <-ticker:
+			if !ok {
+				return
+			}
 			n.pollOnce(t)
 		}
 	}

From 6c6b949d7a1e9c57e741355b4ab1d370799e54f9 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 27 Jul 2022 20:01:50 +0300
Subject: [PATCH 3/3] fix: Clean up Notifier resources via Close

---
 coderd/autobuild/notify/notifier_test.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go
index 5b0b63a430c05..fb547defd5b36 100644
--- 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 {