/healthz endpoint accepts JSON now (#5558)

* /healthz endpoint accepts JSON now
* Merge refs/heads/master into json-healthz
* Merge refs/heads/master into json-healthz
* Merge refs/heads/master into json-healthz
* Merge refs/heads/master into json-healthz
* Merge refs/heads/master into json-healthz
* Merge refs/heads/master into json-healthz

WORKSPACE

@@ -1686,3 +1686,9 @@ go_repository(
     sum = "h1:KU7oHjnv3XNWfa5COkzUifxZmxp1TyI7ImMXqFxLwvQ=",
     version = "v0.2.0",
 )
+
+go_repository(
+    name = "com_github_golang_gddo",
+    commit = "3c2cc9a6329d9842b3bbdaf307a8110d740cf94c",
+    importpath = "github.com/golang/gddo",
+)

shared/prometheus/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 go_library(
     name = "go_default_library",
     srcs = [
+        "content_negotiation.go",
         "logrus_collector.go",
         "service.go",
         "simple_server.go",
@@ -11,6 +12,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//shared:go_default_library",
+        "@com_github_golang_gddo//httputil:go_default_library",
         "@com_github_prometheus_client_golang//prometheus:go_default_library",
         "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
         "@com_github_prometheus_client_golang//prometheus/promhttp:go_default_library",

shared/prometheus/content_negotiation.go

@@ -0,0 +1,59 @@
+package prometheus
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+
+	"github.com/golang/gddo/httputil"
+)
+
+const (
+	contentTypePlainText = "text/plain"
+	contentTypeJSON      = "application/json"
+)
+
+// generatedResponse is a container for response output.
+type generatedResponse struct {
+	// Err is protocol error, if any.
+	Err string `json:"error"`
+
+	// Data is response output, if any.
+	Data interface{} `json:"data"`
+}
+
+// negotiateContentType parses "Accept:" header and returns preferred content type string.
+func negotiateContentType(r *http.Request) string {
+	contentTypes := []string{
+		contentTypePlainText,
+		contentTypeJSON,
+	}
+	return httputil.NegotiateContentType(r, contentTypes, contentTypePlainText)
+}
+
+// writeResponse is content-type aware response writer.
+func writeResponse(w http.ResponseWriter, r *http.Request, response generatedResponse) error {
+	if response.Err != "" {
+		w.WriteHeader(http.StatusInternalServerError)
+	} else {
+		w.WriteHeader(http.StatusOK)
+	}
+
+	switch negotiateContentType(r) {
+	case contentTypePlainText:
+		buf, ok := response.Data.(bytes.Buffer)
+		if !ok {
+			return fmt.Errorf("unexpected data: %v", response.Data)
+		}
+		if _, err := w.Write(buf.Bytes()); err != nil {
+			return fmt.Errorf("could not write response body: %v", err)
+		}
+	case contentTypeJSON:
+		w.Header().Set("Content-Type", contentTypeJSON)
+		if err := json.NewEncoder(w).Encode(response); err != nil {
+			return err
+		}
+	}
+	return nil
+}

shared/prometheus/service.go

@@ -52,39 +52,49 @@ func NewPrometheusService(addr string, svcRegistry *shared.ServiceRegistry, addi
 	return s
 }
 
-func (s *Service) healthzHandler(w http.ResponseWriter, _ *http.Request) {
-	// Call all services in the registry.
-	// if any are not OK, write 500
-	// print the statuses of all services.
-
-	statuses := s.svcRegistry.Statuses()
-	hasError := false
-	var buf bytes.Buffer
-	for k, v := range statuses {
-		var status string
-		if v == nil {
-			status = "OK"
-		} else {
-			hasError = true
-			status = "ERROR " + v.Error()
-		}
+func (s *Service) healthzHandler(w http.ResponseWriter, r *http.Request) {
+	response := generatedResponse{}
 
-		if _, err := buf.WriteString(fmt.Sprintf("%s: %s\n", k, status)); err != nil {
-			hasError = true
+	type serviceStatus struct {
+		Name   string `json:"service"`
+		Status bool   `json:"status"`
+		Err    string `json:"error"`
+	}
+	var statuses []serviceStatus
+	for k, v := range s.svcRegistry.Statuses() {
+		s := serviceStatus{
+			Name:   fmt.Sprintf("%s", k),
+			Status: true,
+		}
+		if v != nil {
+			s.Status = false
+			s.Err = v.Error()
 		}
+		statuses = append(statuses, s)
 	}
+	response.Data = statuses
+
+	// Handle plain text content.
+	if contentType := negotiateContentType(r); contentType == contentTypePlainText {
+		var buf bytes.Buffer
+		for _, s := range statuses {
+			var status string
+			if s.Status {
+				status = "OK"
+			} else {
+				status = "ERROR " + s.Err
+			}
 
-	// Write status header
-	if hasError {
-		w.WriteHeader(http.StatusInternalServerError)
-		log.WithField("statuses", buf.String()).Warn("Node is unhealthy!")
-	} else {
-		w.WriteHeader(http.StatusOK)
+			if _, err := buf.WriteString(fmt.Sprintf("%s: %s\n", s.Name, status)); err != nil {
+				response.Err = err.Error()
+				break
+			}
+		}
+		response.Data = buf
 	}
 
-	// Write http body
-	if _, err := w.Write(buf.Bytes()); err != nil {
-		log.Errorf("Could not write healthz body %v", err)
+	if err := writeResponse(w, r, response); err != nil {
+		log.Errorf("Error writing response: %v", err)
 	}
 }
 

shared/prometheus/service_test.go

@@ -2,15 +2,22 @@ package prometheus
 
 import (
 	"errors"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/prysmaticlabs/prysm/shared"
+	"github.com/sirupsen/logrus"
 )
 
+func init() {
+	logrus.SetLevel(logrus.DebugLevel)
+	logrus.SetOutput(ioutil.Discard)
+}
+
 func TestLifecycle(t *testing.T) {
 	prometheusService := NewPrometheusService(":2112", nil)
 	prometheusService.Start()
@@ -87,8 +94,8 @@ func TestHealthz(t *testing.T) {
 	rr = httptest.NewRecorder()
 	handler.ServeHTTP(rr, req)
 
-	if status := rr.Code; status != http.StatusInternalServerError {
-		t.Errorf("expected error status but got %v", rr.Code)
+	if status := rr.Code; status != http.StatusOK {
+		t.Errorf("expected OK status but got %v", rr.Code)
 	}
 
 	body = rr.Body.String()
@@ -109,3 +116,74 @@ func TestStatus(t *testing.T) {
 		t.Errorf("Wanted: %v, got: %v", s.failStatus, s.Status())
 	}
 }
+
+func TestContentNegotiation(t *testing.T) {
+	t.Run("/healthz all services are ok", func(t *testing.T) {
+		registry := shared.NewServiceRegistry()
+		m := &mockService{}
+		if err := registry.RegisterService(m); err != nil {
+			t.Fatalf("failed to registry service %v", err)
+		}
+		s := NewPrometheusService("", registry)
+
+		req, err := http.NewRequest("GET", "/healthz", nil /* body */)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		handler := http.HandlerFunc(s.healthzHandler)
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		body := rr.Body.String()
+		if !strings.Contains(body, "*prometheus.mockService: OK") {
+			t.Errorf("Expected body to contain mockService status, but got %q", body)
+		}
+
+		// Request response as JSON.
+		req.Header.Add("Accept", "application/json, */*;q=0.5")
+		rr = httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		body = rr.Body.String()
+		expectedJSON := "{\"error\":\"\",\"data\":[{\"service\":\"*prometheus.mockService\",\"status\":true,\"error\":\"\"}]}"
+		if !strings.Contains(body, expectedJSON) {
+			t.Errorf("Unexpected data, want: %q got %q", expectedJSON, body)
+		}
+	})
+
+	t.Run("/healthz failed service", func(t *testing.T) {
+		registry := shared.NewServiceRegistry()
+		m := &mockService{}
+		m.status = errors.New("something is wrong")
+		if err := registry.RegisterService(m); err != nil {
+			t.Fatalf("failed to registry service %v", err)
+		}
+		s := NewPrometheusService("", registry)
+
+		req, err := http.NewRequest("GET", "/healthz", nil /* body */)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		handler := http.HandlerFunc(s.healthzHandler)
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		body := rr.Body.String()
+		if !strings.Contains(body, "*prometheus.mockService: ERROR something is wrong") {
+			t.Errorf("Expected body to contain mockService status, but got %q", body)
+		}
+
+		// Request response as JSON.
+		req.Header.Add("Accept", "application/json, */*;q=0.5")
+		rr = httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+
+		body = rr.Body.String()
+		expectedJSON := "{\"error\":\"\",\"data\":[{\"service\":\"*prometheus.mockService\",\"status\":false,\"error\":\"something is wrong\"}]}"
+		if !strings.Contains(body, expectedJSON) {
+			t.Errorf("Unexpected data, want: %q got %q", expectedJSON, body)
+		}
+	})
+}