From 1910e2a246c0d8758c3c8af5891db4c342089abd Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Mon, 20 Jul 2020 18:55:39 -0400
Subject: [PATCH] checks: wait for goroutine to complete

CheckAlias already had a waitGroup, but the Add() call was happening too late, which was causing a race in tests. The add must happen before the goroutine is started.

CheckHTTP did not have a waitGroup, so I added it to match CheckAlias.

It looks like a lot of the implementation could be shared, and may not need all of channel, waitgroup and bool, but I will leave that refactor for another time.
---
 agent/checks/alias.go | 5 ++---
 agent/checks/check.go | 6 ++++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/agent/checks/alias.go b/agent/checks/alias.go
index 3a0ed36eca..cb41a586c8 100644
--- a/agent/checks/alias.go
+++ b/agent/checks/alias.go
@@ -32,8 +32,7 @@ type CheckAlias struct {
 	stop     bool
 	stopCh   chan struct{}
 	stopLock sync.Mutex
-
-	stopWg sync.WaitGroup
+	stopWg   sync.WaitGroup
 
 	structs.EnterpriseMeta
 }
@@ -55,6 +54,7 @@ func (c *CheckAlias) Start() {
 	defer c.stopLock.Unlock()
 	c.stop = false
 	c.stopCh = make(chan struct{})
+	c.stopWg.Add(1)
 	go c.run(c.stopCh)
 }
 
@@ -76,7 +76,6 @@ func (c *CheckAlias) Stop() {
 
 // run is invoked in a goroutine until Stop() is called.
 func (c *CheckAlias) run(stopCh chan struct{}) {
-	c.stopWg.Add(1)
 	defer c.stopWg.Done()
 
 	// If we have a specific node set, then use a blocking query
diff --git a/agent/checks/check.go b/agent/checks/check.go
index 1bebeffa98..a75ee912d6 100644
--- a/agent/checks/check.go
+++ b/agent/checks/check.go
@@ -348,6 +348,7 @@ type CheckHTTP struct {
 	stop       bool
 	stopCh     chan struct{}
 	stopLock   sync.Mutex
+	stopWg     sync.WaitGroup
 
 	// Set if checks are exposed through Connect proxies
 	// If set, this is the target of check()
@@ -399,6 +400,7 @@ func (c *CheckHTTP) Start() {
 
 	c.stop = false
 	c.stopCh = make(chan struct{})
+	c.stopWg.Add(1)
 	go c.run()
 }
 
@@ -410,10 +412,14 @@ func (c *CheckHTTP) Stop() {
 		c.stop = true
 		close(c.stopCh)
 	}
+
+	// Wait for the c.run() goroutine to complete before returning.
+	c.stopWg.Wait()
 }
 
 // run is invoked by a goroutine to run until Stop() is called
 func (c *CheckHTTP) run() {
+	defer c.stopWg.Done()
 	// Get the randomized initial pause time
 	initialPauseTime := lib.RandomStagger(c.Interval)
 	next := time.After(initialPauseTime)