diff --git a/go/vt/vttablet/tabletserver/tabletserver.go b/go/vt/vttablet/tabletserver/tabletserver.go index 582ce8aaaee..9959847708c 100644 --- a/go/vt/vttablet/tabletserver/tabletserver.go +++ b/go/vt/vttablet/tabletserver/tabletserver.go @@ -1534,19 +1534,21 @@ var okMessage = []byte("ok\n") // Health check // Returns ok if we are in the desired serving state func (tsv *TabletServer) registerHealthzHealthHandler() { - tsv.exporter.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { - if err := acl.CheckAccessHTTP(r, acl.MONITORING); err != nil { - acl.SendError(w, err) - return - } - if tsv.sm.wantState == StateServing && !tsv.sm.IsServing() { - http.Error(w, "500 internal server error: vttablet is not serving", http.StatusInternalServerError) - return - } - w.Header().Set("Content-Length", fmt.Sprintf("%v", len(okMessage))) - w.WriteHeader(http.StatusOK) - w.Write(okMessage) - }) + tsv.exporter.HandleFunc("/healthz", tsv.healthzHandler) +} + +func (tsv *TabletServer) healthzHandler(w http.ResponseWriter, r *http.Request) { + if err := acl.CheckAccessHTTP(r, acl.MONITORING); err != nil { + acl.SendError(w, err) + return + } + if (tsv.sm.wantState == StateServing || tsv.sm.wantState == StateNotConnected) && !tsv.sm.IsServing() { + http.Error(w, "500 internal server error: vttablet is not serving", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Length", fmt.Sprintf("%v", len(okMessage))) + w.WriteHeader(http.StatusOK) + w.Write(okMessage) } // Query service health check diff --git a/go/vt/vttablet/tabletserver/tabletserver_test.go b/go/vt/vttablet/tabletserver/tabletserver_test.go index bb2a1dfa9fb..a09f9adcd7a 100644 --- a/go/vt/vttablet/tabletserver/tabletserver_test.go +++ b/go/vt/vttablet/tabletserver/tabletserver_test.go @@ -22,6 +22,8 @@ import ( "io" "io/ioutil" "math/rand" + "net/http" + "net/http/httptest" "os" "reflect" "strings" @@ -56,6 +58,62 @@ import ( vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" ) +func TestTabletServerHealthz(t *testing.T) { + db, tsv := setupTabletServerTest(t, "") + defer tsv.StopService() + defer db.Close() + + req, err := http.NewRequest("GET", "/healthz", nil) + if err != nil { + t.Fatal(err) + } + + rr := httptest.NewRecorder() + handler := http.HandlerFunc(tsv.healthzHandler) + handler.ServeHTTP(rr, req) + + expectedCode := http.StatusOK + if status := rr.Code; status != expectedCode { + t.Errorf("handler returned wrong status code: got %v want %v", + status, expectedCode) + } + + expected := "ok\n" + if rr.Body.String() != expected { + t.Errorf("handler returned unexpected body: got %v want %v", + rr.Body.String(), expected) + } +} + +func TestTabletServerHealthzNotConnected(t *testing.T) { + db, tsv := setupTabletServerTest(t, "") + defer tsv.StopService() + defer db.Close() + + tsv.sm.SetServingType(topodatapb.TabletType_MASTER, time.Time{}, StateNotConnected, "test disconnected") + + req, err := http.NewRequest("GET", "/healthz", nil) + if err != nil { + t.Fatal(err) + } + + rr := httptest.NewRecorder() + handler := http.HandlerFunc(tsv.healthzHandler) + handler.ServeHTTP(rr, req) + + expectedCode := http.StatusInternalServerError + if status := rr.Code; status != expectedCode { + t.Errorf("handler returned wrong status code: got %v want %v", + status, expectedCode) + } + + expected := "500 internal server error: vttablet is not serving\n" + if rr.Body.String() != expected { + t.Errorf("handler returned unexpected body: got %v want %v", + rr.Body.String(), expected) + } +} + func TestBeginOnReplica(t *testing.T) { db, tsv := setupTabletServerTest(t, "") defer tsv.StopService()