From 07b7c999653d54802bf80855655f15fca6479827 Mon Sep 17 00:00:00 2001
From: David Dyke <dave@dulldave.com>
Date: Mon, 20 Jun 2016 16:49:21 +0100
Subject: [PATCH] Add timeout to health_check (#887)

* Add timeout to http get on health_check

* Add new test and up the timeout

* Tests for change to default timeout

* Only call http client once and make options more inline with current caddy directives
---
 caddyhttp/proxy/upstream.go      | 41 +++++++++++++++++++++-----
 caddyhttp/proxy/upstream_test.go | 50 ++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/caddyhttp/proxy/upstream.go b/caddyhttp/proxy/upstream.go
index 00d23e319..3e21edf63 100644
--- a/caddyhttp/proxy/upstream.go
+++ b/caddyhttp/proxy/upstream.go
@@ -31,8 +31,10 @@ type staticUpstream struct {
 	MaxFails    int32
 	MaxConns    int64
 	HealthCheck struct {
+		Client   http.Client
 		Path     string
 		Interval time.Duration
+		Timeout  time.Duration
 	}
 	WithoutPathPrefix string
 	IgnoredSubPaths   []string
@@ -99,6 +101,9 @@ func NewStaticUpstreams(c caddyfile.Dispenser) ([]Upstream, error) {
 		}
 
 		if upstream.HealthCheck.Path != "" {
+			upstream.HealthCheck.Client = http.Client{
+				Timeout: upstream.HealthCheck.Timeout,
+			}
 			go upstream.HealthCheckWorker(nil)
 		}
 		upstreams = append(upstreams, upstream)
@@ -238,14 +243,34 @@ func parseBlock(c *caddyfile.Dispenser, u *staticUpstream) error {
 			return c.ArgErr()
 		}
 		u.HealthCheck.Path = c.Val()
-		u.HealthCheck.Interval = 30 * time.Second
-		if c.NextArg() {
-			dur, err := time.ParseDuration(c.Val())
-			if err != nil {
-				return err
-			}
-			u.HealthCheck.Interval = dur
+
+		// Set defaults
+		if u.HealthCheck.Interval == 0 {
+			u.HealthCheck.Interval = 30 * time.Second
 		}
+		if u.HealthCheck.Timeout == 0 {
+			u.HealthCheck.Timeout = 60 * time.Second
+		}
+	case "health_check_interval":
+		var interval string
+		if !c.Args(&interval) {
+			return c.ArgErr()
+		}
+		dur, err := time.ParseDuration(interval)
+		if err != nil {
+			return err
+		}
+		u.HealthCheck.Interval = dur
+	case "health_check_timeout":
+		var interval string
+		if !c.Args(&interval) {
+			return c.ArgErr()
+		}
+		dur, err := time.ParseDuration(interval)
+		if err != nil {
+			return err
+		}
+		u.HealthCheck.Timeout = dur
 	case "header_upstream":
 		fallthrough
 	case "proxy_header":
@@ -289,7 +314,7 @@ func parseBlock(c *caddyfile.Dispenser, u *staticUpstream) error {
 func (u *staticUpstream) healthCheck() {
 	for _, host := range u.Hosts {
 		hostURL := host.Name + u.HealthCheck.Path
-		if r, err := http.Get(hostURL); err == nil {
+		if r, err := u.HealthCheck.Client.Get(hostURL); err == nil {
 			io.Copy(ioutil.Discard, r.Body)
 			r.Body.Close()
 			host.Unhealthy = r.StatusCode < 200 || r.StatusCode >= 400
diff --git a/caddyhttp/proxy/upstream_test.go b/caddyhttp/proxy/upstream_test.go
index 0378475f2..4fb990f6d 100644
--- a/caddyhttp/proxy/upstream_test.go
+++ b/caddyhttp/proxy/upstream_test.go
@@ -137,6 +137,56 @@ func TestAllowedPaths(t *testing.T) {
 	}
 }
 
+func TestParseBlockHealthCheck(t *testing.T) {
+	tests := []struct {
+		config   string
+		interval string
+		timeout  string
+	}{
+		// Test #1: Both options set correct time
+		{"health_check /health\n health_check_interval 10s\n health_check_timeout 20s", "10s", "20s"},
+
+		// Test #2: Health check options flipped around. Making sure health_check doesn't overwrite it
+		{"health_check_interval 10s\n health_check_timeout 20s\n health_check /health", "10s", "20s"},
+
+		// Test #3: No health_check options. So default.
+		{"health_check /health", "30s", "1m0s"},
+
+		// Test #4: Interval sets it to 15s and timeout defaults
+		{"health_check /health\n health_check_interval 15s", "15s", "1m0s"},
+
+		// Test #5: Timeout sets it to 15s and interval defaults
+		{"health_check /health\n health_check_timeout 15s", "30s", "15s"},
+
+		// Test #6: Some funky spelling to make sure it still defaults
+		{"health_check /health health_check_time 15s", "30s", "1m0s"},
+	}
+
+	for i, test := range tests {
+		u := staticUpstream{}
+		c := caddyfile.NewDispenser("Testfile", strings.NewReader(test.config))
+		for c.Next() {
+			parseBlock(&c, &u)
+		}
+		if u.HealthCheck.Interval.String() != test.interval {
+			t.Errorf(
+				"Test %d: HealthCheck interval not the same from config. Got %v. Expected: %v",
+				i+1,
+				u.HealthCheck.Interval,
+				test.interval,
+			)
+		}
+		if u.HealthCheck.Timeout.String() != test.timeout {
+			t.Errorf(
+				"Test %d: HealthCheck timeout not the same from config. Got %v. Expected: %v",
+				i+1,
+				u.HealthCheck.Timeout,
+				test.timeout,
+			)
+		}
+	}
+}
+
 func TestParseBlock(t *testing.T) {
 	tests := []struct {
 		config string