diff options
Diffstat (limited to 'middleware/proxy/upstream.go')
| -rw-r--r-- | middleware/proxy/upstream.go | 26 |
1 files changed, 21 insertions, 5 deletions
diff --git a/middleware/proxy/upstream.go b/middleware/proxy/upstream.go index 278b08d0f..8bc3e5306 100644 --- a/middleware/proxy/upstream.go +++ b/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 } } } |