From 1a1c1bb813ad4ce156fd4145686acc78f9114e2e Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@prysmaticlabs.com>
Date: Tue, 28 Jul 2020 17:31:21 -0700
Subject: [PATCH] Return 503 on healthz with error (#6754)

* Return 503 on healthz with error
* fix other test
* Update shared/prometheus/service_test.go

Co-authored-by: Victor Farazdagi <simple.square@gmail.com>
---
 shared/prometheus/service.go      | 8 ++++++++
 shared/prometheus/service_test.go | 7 +++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/shared/prometheus/service.go b/shared/prometheus/service.go
index 711e0b9fccb..7301d9611ae 100644
--- a/shared/prometheus/service.go
+++ b/shared/prometheus/service.go
@@ -61,6 +61,7 @@ func (s *Service) healthzHandler(w http.ResponseWriter, r *http.Request) {
 		Status bool   `json:"status"`
 		Err    string `json:"error"`
 	}
+	var hasError bool
 	var statuses []serviceStatus
 	for k, v := range s.svcRegistry.Statuses() {
 		s := serviceStatus{
@@ -70,11 +71,18 @@ func (s *Service) healthzHandler(w http.ResponseWriter, r *http.Request) {
 		if v != nil {
 			s.Status = false
 			s.Err = v.Error()
+			if s.Err != "" {
+				hasError = true
+			}
 		}
 		statuses = append(statuses, s)
 	}
 	response.Data = statuses
 
+	if hasError {
+		w.WriteHeader(http.StatusServiceUnavailable)
+	}
+
 	// Handle plain text content.
 	if contentType := negotiateContentType(r); contentType == contentTypePlainText {
 		var buf bytes.Buffer
diff --git a/shared/prometheus/service_test.go b/shared/prometheus/service_test.go
index cd0fe5d3418..086a7c0f790 100644
--- a/shared/prometheus/service_test.go
+++ b/shared/prometheus/service_test.go
@@ -94,8 +94,8 @@ func TestHealthz(t *testing.T) {
 	rr = httptest.NewRecorder()
 	handler.ServeHTTP(rr, req)
 
-	if status := rr.Code; status != http.StatusOK {
-		t.Errorf("expected OK status but got %v", rr.Code)
+	if status := rr.Code; status != http.StatusServiceUnavailable {
+		t.Errorf("expected StatusServiceUnavailable status but got %v", rr.Code)
 	}
 
 	body = rr.Body.String()
@@ -185,5 +185,8 @@ func TestContentNegotiation(t *testing.T) {
 		if !strings.Contains(body, expectedJSON) {
 			t.Errorf("Unexpected data, want: %q got %q", expectedJSON, body)
 		}
+		if rr.Code < 500 {
+			t.Errorf("Expected a server error response code, but got %d", rr.Code)
+		}
 	})
 }