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())
+		}
+	}
+}

pkg/kubelet/server.go

@@ -33,6 +33,7 @@ import (
 
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/latest"
+	"github.com/GoogleCloudPlatform/kubernetes/pkg/healthz"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/httplog"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/types"
@@ -108,7 +109,9 @@ func NewServer(host HostInterface, enableDebuggingHandlers bool) Server {
 
 // InstallDefaultHandlers registers the default set of supported HTTP request patterns with the mux.
 func (s *Server) InstallDefaultHandlers() {
-	s.mux.HandleFunc("/healthz", s.handleHealthz)
+	healthz.AddHealthzFunc("docker", s.dockerHealthCheck)
+	healthz.AddHealthzFunc("hostname", s.hostnameHealthCheck)
+	healthz.InstallHandler(s.mux)
 	s.mux.HandleFunc("/podInfo", s.handlePodInfoOld)
 	s.mux.HandleFunc("/api/v1beta1/podInfo", s.handlePodInfoVersioned)
 	s.mux.HandleFunc("/boundPods", s.handleBoundPods)
@@ -151,38 +154,35 @@ func isValidDockerVersion(ver []uint) (bool, string) {
 	return true, ""
 }
 
-// handleHealthz handles /healthz request and checks Docker version
-func (s *Server) handleHealthz(w http.ResponseWriter, req *http.Request) {
+func (s *Server) dockerHealthCheck(req *http.Request) error {
 	versions, err := s.host.GetDockerVersion()
 	if err != nil {
-		s.error(w, errors.New("unknown Docker version"))
-		return
+		return errors.New("unknown Docker version")
 	}
 	valid, version := isValidDockerVersion(versions)
 	if !valid {
-		s.error(w, errors.New("Docker version is too old ("+version+")"))
-		return
+		return fmt.Errorf("Docker version is too old (%v)", version)
 	}
+	return nil
+}
 
+func (s *Server) hostnameHealthCheck(req *http.Request) error {
 	masterHostname, _, err := net.SplitHostPort(req.Host)
 	if err != nil {
 		if !strings.Contains(req.Host, ":") {
 			masterHostname = req.Host
 		} else {
-			msg := fmt.Sprintf("Could not parse hostname from http request: %v", err)
-			s.error(w, errors.New(msg))
-			return
+			return fmt.Errorf("Could not parse hostname from http request: %v", err)
 		}
 	}
 
 	// Check that the hostname known by the master matches the hostname
 	// the kubelet knows
 	hostname := s.host.GetHostname()
 	if masterHostname != hostname && masterHostname != "127.0.0.1" && masterHostname != "localhost" {
-		s.error(w, errors.New("Kubelet hostname \""+hostname+"\" does not match the hostname expected by the master \""+masterHostname+"\""))
-		return
+		return fmt.Errorf("Kubelet hostname \"%v\" does not match the hostname expected by the master \"%v\"", hostname, masterHostname)
 	}
-	w.Write([]byte("ok"))
+	return nil
 }
 
 // handleContainerLogs handles containerLogs request against the Kubelet