Cleanup goroutines on health.Stop (#1325)

api/health/health.go

@@ -28,7 +28,12 @@ type Health interface {
 	Registerer
 	Reporter
 
+	// Start running periodic health checks at the specified frequency.
+	// Repeated calls to Start will be no-ops.
 	Start(ctx context.Context, freq time.Duration)
+
+	// Stop running periodic health checks. Stop should only be called after
+	// Start. Once Stop returns, no more health checks will be executed.
 	Stop()
 }
 

api/health/worker.go

@@ -31,6 +31,7 @@ type worker struct {
 
 	startOnce sync.Once
 	closeOnce sync.Once
+	wg        sync.WaitGroup
 	closer    chan struct{}
 }
 
@@ -126,9 +127,13 @@ func (w *worker) Results(tags ...string) (map[string]Result, bool) {
 func (w *worker) Start(ctx context.Context, freq time.Duration) {
 	w.startOnce.Do(func() {
 		detachedCtx := utils.Detach(ctx)
+		w.wg.Add(1)
 		go func() {
 			ticker := time.NewTicker(freq)
-			defer ticker.Stop()
+			defer func() {
+				ticker.Stop()
+				w.wg.Done()
+			}()
 
 			w.runChecks(detachedCtx)
 			for {
@@ -146,6 +151,7 @@ func (w *worker) Start(ctx context.Context, freq time.Duration) {
 func (w *worker) Stop() {
 	w.closeOnce.Do(func() {
 		close(w.closer)
+		w.wg.Wait()
 	})
 }