middleware/proxy: Kill goroutines on stop (#646)

* middleware/proxy: Kill goroutines on stop

Ports caddy's 59bf71c293

Excludes the proxy_test.go test part though.

Fixes #644

* Add tests

middleware/proxy/upstream.go

@@ -10,6 +10,7 @@ import (
 	"net/url"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -25,7 +26,10 @@ var (
 )
 
 type staticUpstream struct {
-	from   string
+	from string
+	stop chan struct{}  // Signals running goroutines to stop.
+	wg   sync.WaitGroup // Used to wait for running goroutines to stop.
+
 	Hosts  HostPool
 	Policy Policy
 	Spray  Policy
@@ -49,6 +53,7 @@ func NewStaticUpstreams(c *caddyfile.Dispenser) ([]Upstream, error) {
 	for c.Next() {
 		upstream := &staticUpstream{
 			from:        ".",
+			stop:        make(chan struct{}),
 			Hosts:       nil,
 			Policy:      &Random{},
 			Spray:       nil,
@@ -108,13 +113,25 @@ func NewStaticUpstreams(c *caddyfile.Dispenser) ([]Upstream, error) {
 		}
 
 		if upstream.HealthCheck.Path != "" {
-			go upstream.HealthCheckWorker(nil)
+			upstream.wg.Add(1)
+			go func() {
+				defer upstream.wg.Done()
+				upstream.HealthCheckWorker(upstream.stop)
+			}()
 		}
 		upstreams = append(upstreams, upstream)
 	}
 	return upstreams, nil
 }
 
+// Stop sends a signal to all goroutines started by this staticUpstream to exit
+// and waits for them to finish before returning.
+func (u *staticUpstream) Stop() error {
+	close(u.stop)
+	u.wg.Wait()
+	return nil
+}
+
 // RegisterPolicy adds a custom policy to the proxy.
 func RegisterPolicy(name string, policy func() Policy) {
 	supportedPolicies[name] = policy
@@ -281,9 +298,8 @@ func (u *staticUpstream) HealthCheckWorker(stop chan struct{}) {
 		case <-ticker.C:
 			u.healthCheck()
 		case <-stop:
-			// TODO: the library should provide a stop channel and global
-			// waitgroup to allow goroutines started by plugins a chance
-			// to clean themselves up.
+			ticker.Stop()
+			return
 		}
 	}
 }