mw/health: poll other middleware (#976)

This add the infrastructure to let other middleware report their health
status back to the health middleware. A health.Healther interface is
introduced and a middleware needs to implement that. A middleware
that supports healthchecks is statically configured.

Every second each supported middleware is queried and the global health
state is updated.

Actual tests have been disabled as no other middleware implements this
at the moment.

6 files changed, 147 insertions, 21 deletions

diff --git a/middleware/health/README.md b/middleware/health/README.md
index aee23ceae..195a480c0 100644
--- a/middleware/health/README.md
+++ b/middleware/health/README.md
@@ -1,7 +1,6 @@
 # health
 
-This module enables a simple health check endpoint.
-By default it will listen on port 8080.
+This module enables a simple health check endpoint. By default it will listen on port 8080.
 
 ## Syntax
 
@@ -9,13 +8,16 @@ By default it will listen on port 8080.
 health [ADDRESS]
 ~~~
 
-Optionally takes an address; the default is `:8080`. The health path is fixed to `/health`. It
-will just return "OK" when CoreDNS is healthy, which currently mean: it is up and running.
-
-This middleware only needs to be enabled once.
+Optionally takes an address; the default is `:8080`. The health path is fixed to `/health`. The
+health endpoint returns a 200 response code and the word "OK" when CoreDNS is healthy. It returns
+a 503. *health* periodically (1s) polls middleware that exports health information. If any of the
+middleware signals that it is unhealthy, the server will go unhealthy too. Each middleware that
+supports health checks has a section "Health" in their README.
 
 ## Examples
 
+Run another health endpoint on http://localhost:8091.
+
 ~~~
 health localhost:8091
 ~~~
diff --git a/middleware/health/health.go b/middleware/health/health.go
index 08aa847d2..5e5c4c32f 100644
--- a/middleware/health/health.go
+++ b/middleware/health/health.go
@@ -16,6 +16,11 @@ type health struct {
 
 	ln  net.Listener
 	mux *http.ServeMux
+
+	// A slice of Healthers that the health middleware will poll every second for their health status.
+	h []Healther
+	sync.RWMutex
+	ok bool // ok is the global boolean indicating an all healthy middleware stack
 }
 
 func (h *health) Startup() error {
@@ -35,7 +40,12 @@ func (h *health) Startup() error {
 		h.mux = http.NewServeMux()
 
 		h.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
-			io.WriteString(w, ok)
+			if h.Ok() {
+				w.WriteHeader(http.StatusOK)
+				io.WriteString(w, ok)
+				return
+			}
+			w.WriteHeader(http.StatusServiceUnavailable)
 		})
 
 		go func() {
diff --git a/middleware/health/health_test.go b/middleware/health/health_test.go
index de95eb103..ea22ad8b6 100644
--- a/middleware/health/health_test.go
+++ b/middleware/health/health_test.go
@@ -1,16 +1,11 @@
 package health
 
-import (
-	"fmt"
-	"io/ioutil"
-	"net/http"
-	"testing"
-)
-
+// TODO(miek): enable again if middleware gets health check.
+/*
 func TestHealth(t *testing.T) {
-	// We use a random port instead of a fixed port like 8080 that may have been
-	// occupied by some other process.
 	h := health{Addr: ":0"}
+	h.h = append(h.h, &erratic.Erratic{})
+
 	if err := h.Startup(); err != nil {
 		t.Fatalf("Unable to startup the health server: %v", err)
 	}
@@ -19,11 +14,23 @@ func TestHealth(t *testing.T) {
 	// Reconstruct the http address based on the port allocated by operating system.
 	address := fmt.Sprintf("http://%s%s", h.ln.Addr().String(), path)
 
+	// Norhing set should be unhealthy
 	response, err := http.Get(address)
 	if err != nil {
 		t.Fatalf("Unable to query %s: %v", address, err)
 	}
-	defer response.Body.Close()
+	if response.StatusCode != 503 {
+		t.Errorf("Invalid status code: expecting '503', got '%d'", response.StatusCode)
+	}
+	response.Body.Close()
+
+	// Make healthy
+	h.Poll()
+
+	response, err = http.Get(address)
+	if err != nil {
+		t.Fatalf("Unable to query %s: %v", address, err)
+	}
 	if response.StatusCode != 200 {
 		t.Errorf("Invalid status code: expecting '200', got '%d'", response.StatusCode)
 	}
@@ -31,7 +38,10 @@ func TestHealth(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Unable to get response body from %s: %v", address, err)
 	}
-	if string(content) != "OK" {
+	response.Body.Close()
+
+	if string(content) != ok {
 		t.Errorf("Invalid response body: expecting 'OK', got '%s'", string(content))
 	}
 }
+*/
diff --git a/middleware/health/healther.go b/middleware/health/healther.go
new file mode 100644
index 000000000..1b794e03b
--- /dev/null
+++ b/middleware/health/healther.go
@@ -0,0 +1,42 @@
+package health
+
+// Healther interface needs to be implemented by each middleware willing to
+// provide healthhceck information to the health middleware. As a second step
+// the middleware needs to registered against the health middleware, by addding
+// it to healthers map. Note this method should return quickly, i.e. just
+// checking a boolean status, as it is called every second from the health
+// middleware.
+type Healther interface {
+	// Health returns a boolean indicating the health status of a middleware.
+	// False indicates unhealthy.
+	Health() bool
+}
+
+// Ok returns the global health status of all middleware configured in this server.
+func (h *health) Ok() bool {
+	h.RLock()
+	defer h.RUnlock()
+	return h.ok
+}
+
+// SetOk sets the global health status of all middleware configured in this server.
+func (h *health) SetOk(ok bool) {
+	h.Lock()
+	defer h.Unlock()
+	h.ok = ok
+}
+
+// poll polls all healthers and sets the global state.
+func (h *health) poll() {
+	for _, m := range h.h {
+		if !m.Health() {
+			h.SetOk(false)
+			return
+		}
+	}
+	h.SetOk(true)
+}
+
+// Middleware that implements the Healther interface.
+// TODO(miek): none yet.
+var healthers = map[string]bool{}
diff --git a/middleware/health/setup.go b/middleware/health/setup.go
index f698d9603..c54a6ea53 100644
--- a/middleware/health/setup.go
+++ b/middleware/health/setup.go
@@ -1,6 +1,10 @@
 package health
 
 import (
+	"net"
+	"time"
+
+	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/middleware"
 
 	"github.com/mholt/caddy"
@@ -20,12 +24,32 @@ func setup(c *caddy.Controller) error {
 	}
 
 	h := &health{Addr: addr}
+
+	c.OnStartup(func() error {
+		for he := range healthers {
+			m := dnsserver.GetConfig(c).Handler(he)
+			if x, ok := m.(Healther); ok {
+				h.h = append(h.h, x)
+			}
+		}
+		return nil
+	})
+
+	c.OnStartup(func() error {
+		h.poll()
+		go func() {
+			for {
+				<-time.After(1 * time.Second)
+				h.poll()
+			}
+		}()
+		return nil
+	})
+
 	c.OnStartup(h.Startup)
 	c.OnShutdown(h.Shutdown)
 
-	// Don't do AddMiddleware, as health is not *really* a middleware just a separate
-	// webserver running.
-
+	// Don't do AddMiddleware, as health is not *really* a middleware just a separate webserver running.
 	return nil
 }
 
@@ -38,6 +62,9 @@ func healthParse(c *caddy.Controller) (string, error) {
 		case 0:
 		case 1:
 			addr = args[0]
+			if _, _, e := net.SplitHostPort(addr); e != nil {
+				return "", e
+			}
 		default:
 			return "", c.ArgErr()
 		}
diff --git a/middleware/health/setup_test.go b/middleware/health/setup_test.go
new file mode 100644
index 000000000..87f4fc5fd
--- /dev/null
+++ b/middleware/health/setup_test.go
@@ -0,0 +1,35 @@
+package health
+
+import (
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+func TestSetupHealth(t *testing.T) {
+	tests := []struct {
+		input     string
+		shouldErr bool
+	}{
+		{`health`, false},
+		{`health localhost:1234`, false},
+		{`health bla:a`, false},
+		{`health bla`, true},
+		{`health bla bla`, true},
+	}
+
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		_, err := healthParse(c)
+
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %d: Expected error but found %s for input %s", i, err, test.input)
+		}
+
+		if err != nil {
+			if !test.shouldErr {
+				t.Errorf("Test %d: Expected no error but found one for input %s. Error was: %v", i, test.input, err)
+			}
+		}
+	}
+}