From e753cbb04ba2ecb0a6599940f7ad646af599b4fd Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Tue, 21 Apr 2020 23:58:53 +0300 Subject: [PATCH] /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 | 6 ++ shared/prometheus/BUILD.bazel | 2 + shared/prometheus/content_negotiation.go | 59 +++++++++++++++++ shared/prometheus/service.go | 64 ++++++++++-------- shared/prometheus/service_test.go | 82 +++++++++++++++++++++++- 5 files changed, 184 insertions(+), 29 deletions(-) create mode 100644 shared/prometheus/content_negotiation.go diff --git a/WORKSPACE b/WORKSPACE index eede235d7e5f..92922883a02e 100644 --- a/WORKSPACE +++ b/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", +) diff --git a/shared/prometheus/BUILD.bazel b/shared/prometheus/BUILD.bazel index 716fd57b8539..cf95d2ba8865 100644 --- a/shared/prometheus/BUILD.bazel +++ b/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", diff --git a/shared/prometheus/content_negotiation.go b/shared/prometheus/content_negotiation.go new file mode 100644 index 000000000000..74d765ce279a --- /dev/null +++ b/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 +} diff --git a/shared/prometheus/service.go b/shared/prometheus/service.go index 70e6fe27a62f..71e139bc9450 100644 --- a/shared/prometheus/service.go +++ b/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) } } diff --git a/shared/prometheus/service_test.go b/shared/prometheus/service_test.go index cc9f020ca14e..cd0fe5d34188 100644 --- a/shared/prometheus/service_test.go +++ b/shared/prometheus/service_test.go @@ -2,6 +2,7 @@ package prometheus import ( "errors" + "io/ioutil" "net/http" "net/http/httptest" "strings" @@ -9,8 +10,14 @@ import ( "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) + } + }) +}