support user supplied health functions in pkg/healthz

pkg/healthz/healthz.go

@@ -17,25 +17,100 @@ limitations under the License.
 package healthz
 
 import (
+	"bytes"
+	"fmt"
 	"net/http"
+	"sync"
 )
 
+var (
+	// guards names and checks
+	lock = sync.RWMutex{}
+	// used to ensure checks are performed in the order added
+	names  = []string{}
+	checks = map[string]*healthzCheck{}
+)
+
+func init() {
+	http.HandleFunc("/healthz", handleRootHealthz)
+	// add ping health check by default
+	AddHealthzFunc("ping", func(_ *http.Request) error {
+		return nil
+	})
+}
+
+// AddHealthzFunc adds a health check under the url /healhz/{name}
+func AddHealthzFunc(name string, check func(r *http.Request) error) {
+	lock.Lock()
+	defer lock.Unlock()
+	if _, found := checks[name]; !found {
+		names = append(names, name)
+	}
+	checks[name] = &healthzCheck{name, check}
+}
+
+// InstallHandler registers a handler for health checking on the path "/healthz" to mux.
+func InstallHandler(mux mux) {
+	lock.RLock()
+	defer lock.RUnlock()
+	mux.HandleFunc("/healthz", handleRootHealthz)
+	for _, check := range checks {
+		mux.HandleFunc(fmt.Sprintf("/healthz/%v", check.name), adaptCheckToHandler(check.check))
+	}
+}
+
 // mux is an interface describing the methods InstallHandler requires.
 type mux interface {
 	HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request))
 }
 
-func init() {
-	http.HandleFunc("/healthz", handleHealthz)
+type healthzCheck struct {
+	name  string
+	check func(r *http.Request) error
 }
 
-func handleHealthz(w http.ResponseWriter, r *http.Request) {
-	// TODO Support user supplied health functions too.
-	w.WriteHeader(http.StatusOK)
-	w.Write([]byte("ok"))
+func handleRootHealthz(w http.ResponseWriter, r *http.Request) {
+	lock.RLock()
+	defer lock.RUnlock()
+	failed := false
+	var verboseOut bytes.Buffer
+	for _, name := range names {
+		check, found := checks[name]
+		if !found {
+			// this should not happen
+			http.Error(w, fmt.Sprintf("Internal server error: check \"%q\" not registered", name), http.StatusInternalServerError)
+			return
+		}
+		err := check.check(r)
+		if err != nil {
+			fmt.Fprintf(&verboseOut, "[-]%v failed: %v\n", check.name, err)
+			failed = true
+		} else {
+			fmt.Fprintf(&verboseOut, "[+]%v ok\n", check.name)
+		}
+	}
+	// always be verbose on failure
+	if failed {
+		http.Error(w, fmt.Sprintf("%vhealthz check failed", verboseOut.String()), http.StatusInternalServerError)
+		return
+	}
+
+	if _, found := r.URL.Query()["verbose"]; !found {
+		fmt.Fprint(w, "ok")
+		return
+	} else {
+		verboseOut.WriteTo(w)
+		fmt.Fprint(w, "healthz check passed\n")
+	}
 }
 
-// InstallHandler registers a handler for health checking on the path "/healthz" to mux.
-func InstallHandler(mux mux) {
-	mux.HandleFunc("/healthz", handleHealthz)
+func adaptCheckToHandler(c func(r *http.Request) error) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		err := c(r)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("Internal server error: %v", err), http.StatusInternalServerError)
+		} else {
+			fmt.Fprint(w, "ok")
+		}
+	}
 }

pkg/healthz/healthz_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package healthz
 
 import (
+	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -38,3 +40,42 @@ func TestInstallHandler(t *testing.T) {
 		t.Errorf("Expected %v, got %v", "ok", w.Body.String())
 	}
 }
+
+func TestMulitipleChecks(t *testing.T) {
+	tests := []struct {
+		path             string
+		expectedResponse string
+		expectedStatus   int
+		addBadCheck      bool
+	}{
+		{"/healthz?verbose", "[+]ping ok\nhealthz check passed\n", http.StatusOK, false},
+		{"/healthz/ping", "ok", http.StatusOK, false},
+		{"/healthz", "ok", http.StatusOK, false},
+		{"/healthz?verbose", "[+]ping ok\n[-]bad failed: this will fail\nhealthz check failed\n", http.StatusInternalServerError, true},
+		{"/healthz/ping", "ok", http.StatusOK, true},
+		{"/healthz/bad", "Internal server error: this will fail\n", http.StatusInternalServerError, true},
+		{"/healthz", "[+]ping ok\n[-]bad failed: this will fail\nhealthz check failed\n", http.StatusInternalServerError, true},
+	}
+
+	for i, test := range tests {
+		mux := http.NewServeMux()
+		if test.addBadCheck {
+			AddHealthzFunc("bad", func(_ *http.Request) error {
+				return errors.New("this will fail")
+			})
+		}
+		InstallHandler(mux)
+		req, err := http.NewRequest("GET", fmt.Sprintf("http://example.com%v", test.path), nil)
+		if err != nil {
+			t.Fatalf("case[%d] Unexpected error: %v", i, err)
+		}
+		w := httptest.NewRecorder()
+		mux.ServeHTTP(w, req)
+		if w.Code != test.expectedStatus {
+			t.Errorf("case[%d] Expected: %v, got: %v", i, test.expectedStatus, w.Code)
+		}
+		if w.Body.String() != test.expectedResponse {
+			t.Errorf("case[%d] Expected:\n%v\ngot:\n%v\n", i, test.expectedResponse, w.Body.String())
+		}
+	}
+}