Support healthz check api (#8228)

Signed-off-by: Ndragondriver <jiquan.long@zilliz.com>

cmd/components/data_coord.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcdatacoordclient "github.com/milvus-io/milvus/internal/distributed/datacoord"
 	"github.com/milvus-io/milvus/internal/msgstream"
 )
@@ -51,3 +53,7 @@ func (s *DataCoord) Stop() error {
 	}
 	return nil
 }
+
+func (s *DataCoord) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return s.svr.GetComponentStates(ctx, request)
+}

cmd/components/data_node.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcdatanode "github.com/milvus-io/milvus/internal/distributed/datanode"
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/msgstream"
@@ -53,3 +55,7 @@ func (d *DataNode) Stop() error {
 	}
 	return nil
 }
+
+func (d *DataNode) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return d.svr.GetComponentStates(ctx, request)
+}

cmd/components/index_coord.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcindexcoord "github.com/milvus-io/milvus/internal/distributed/indexcoord"
 )
 
@@ -49,3 +51,7 @@ func (s *IndexCoord) Stop() error {
 	}
 	return nil
 }
+
+func (s *IndexCoord) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return s.svr.GetComponentStates(ctx, request)
+}

cmd/components/index_node.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcindexnode "github.com/milvus-io/milvus/internal/distributed/indexnode"
 )
 
@@ -49,3 +51,7 @@ func (n *IndexNode) Stop() error {
 	}
 	return nil
 }
+
+func (n *IndexNode) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return n.svr.GetComponentStates(ctx, request)
+}

cmd/components/proxy.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcproxy "github.com/milvus-io/milvus/internal/distributed/proxy"
 	"github.com/milvus-io/milvus/internal/msgstream"
 )
@@ -50,3 +52,7 @@ func (n *Proxy) Stop() error {
 	}
 	return nil
 }
+
+func (n *Proxy) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return n.svr.GetComponentStates(ctx, request)
+}

cmd/components/query_coord.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcquerycoord "github.com/milvus-io/milvus/internal/distributed/querycoord"
 	"github.com/milvus-io/milvus/internal/msgstream"
 )
@@ -51,3 +53,7 @@ func (qs *QueryCoord) Stop() error {
 	}
 	return nil
 }
+
+func (qs *QueryCoord) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return qs.svr.GetComponentStates(ctx, request)
+}

cmd/components/query_node.go

@@ -14,6 +14,8 @@ package components
 import (
 	"context"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	grpcquerynode "github.com/milvus-io/milvus/internal/distributed/querynode"
 	"github.com/milvus-io/milvus/internal/msgstream"
 )
@@ -52,3 +54,7 @@ func (q *QueryNode) Stop() error {
 	}
 	return nil
 }
+
+func (q *QueryNode) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return q.svr.GetComponentStates(ctx, request)
+}

cmd/components/root_coord.go

@@ -15,6 +15,8 @@ import (
 	"context"
 	"io"
 
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+
 	rc "github.com/milvus-io/milvus/internal/distributed/rootcoord"
 	"github.com/milvus-io/milvus/internal/msgstream"
 	"github.com/opentracing/opentracing-go"
@@ -55,3 +57,7 @@ func (rc *RootCoord) Stop() error {
 	}
 	return nil
 }
+
+func (rc *RootCoord) GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error) {
+	return rc.svr.GetComponentStates(ctx, request)
+}

cmd/roles/healthz_error_handler.go

+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package roles
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/util/healthz"
+	"github.com/milvus-io/milvus/internal/util/milvuserrors"
+	"go.uber.org/zap"
+)
+
+func componentsNotServingHandler(w http.ResponseWriter, r *http.Request, msg string) {
+	w.WriteHeader(http.StatusInternalServerError)
+	w.Header().Set(healthz.ContentTypeHeader, healthz.ContentTypeText)
+	_, err := fmt.Fprint(w, msg)
+	if err != nil {
+		log.Warn("failed to send response",
+			zap.Error(err))
+	}
+}
+
+func rootCoordNotServingHandler(w http.ResponseWriter, r *http.Request) {
+	componentsNotServingHandler(w, r, milvuserrors.MsgRootCoordNotServing)
+}
+
+func queryCoordNotServingHandler(w http.ResponseWriter, r *http.Request) {
+	componentsNotServingHandler(w, r, milvuserrors.MsgQueryCoordNotServing)
+}
+
+func dataCoordNotServingHandler(w http.ResponseWriter, r *http.Request) {
+	componentsNotServingHandler(w, r, milvuserrors.MsgDataCoordNotServing)
+}
+
+func indexCoordNotServingHandler(w http.ResponseWriter, r *http.Request) {
+	componentsNotServingHandler(w, r, milvuserrors.MsgIndexCoordNotServing)
+}

cmd/roles/healthz_handler.go

+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package roles
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	"github.com/milvus-io/milvus/internal/log"
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/util/healthz"
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+)
+
+func unhealthyHandler(w http.ResponseWriter, r *http.Request, reason string) {
+	w.WriteHeader(http.StatusInternalServerError)
+	w.Header().Set(healthz.ContentTypeHeader, healthz.ContentTypeText)
+	_, err := fmt.Fprint(w, reason)
+	if err != nil {
+		log.Warn("failed to send response",
+			zap.Error(err))
+	}
+}
+
+func healthyHandler(w http.ResponseWriter, r *http.Request) {
+	var err error
+
+	w.WriteHeader(http.StatusOK)
+	w.Header().Set(healthz.ContentTypeHeader, healthz.ContentTypeText)
+	_, err = fmt.Fprint(w, "OK")
+	if err != nil {
+		log.Warn("failed to send response",
+			zap.Error(err))
+	}
+}
+
+type GetComponentStatesInterface interface {
+	GetComponentStates(ctx context.Context, request *internalpb.GetComponentStatesRequest) (*internalpb.ComponentStates, error)
+}
+
+type componentsHealthzHandler struct {
+	component GetComponentStatesInterface
+}
+
+func (handler *componentsHealthzHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	states, err := handler.component.GetComponentStates(context.Background(), &internalpb.GetComponentStatesRequest{})
+
+	if err != nil {
+		unhealthyHandler(w, r, err.Error())
+		return
+	}
+
+	if states == nil {
+		unhealthyHandler(w, r, "failed to get states")
+		return
+	}
+
+	if states.Status == nil {
+		unhealthyHandler(w, r, "failed to get status")
+		return
+	}
+
+	if states.Status.ErrorCode != commonpb.ErrorCode_Success {
+		unhealthyHandler(w, r, states.Status.Reason)
+		return
+	}
+
+	healthyHandler(w, r)
+}

cmd/roles/roles.go

@@ -14,6 +14,7 @@ package roles
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"os"
 	"os/signal"
 	"path"
@@ -21,6 +22,9 @@ import (
 	"sync"
 	"syscall"
 
+	"github.com/milvus-io/milvus/internal/util/healthz"
+	"go.uber.org/zap"
+
 	"github.com/milvus-io/milvus/internal/util/metricsinfo"
 
 	"github.com/milvus-io/milvus/cmd/components"
@@ -97,6 +101,9 @@ func (mr *MilvusRoles) runRootCoord(ctx context.Context, localMsg bool) *compone
 		}
 		wg.Done()
 		_ = rc.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: rc})
+		}
 	}()
 	wg.Wait()
 
@@ -126,6 +133,9 @@ func (mr *MilvusRoles) runProxy(ctx context.Context, localMsg bool, alias string
 		}
 		wg.Done()
 		_ = pn.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: pn})
+		}
 	}()
 	wg.Wait()
 
@@ -154,6 +164,9 @@ func (mr *MilvusRoles) runQueryCoord(ctx context.Context, localMsg bool) *compon
 		}
 		wg.Done()
 		_ = qs.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: qs})
+		}
 	}()
 	wg.Wait()
 
@@ -183,6 +196,9 @@ func (mr *MilvusRoles) runQueryNode(ctx context.Context, localMsg bool, alias st
 		}
 		wg.Done()
 		_ = qn.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: qn})
+		}
 	}()
 	wg.Wait()
 
@@ -211,6 +227,9 @@ func (mr *MilvusRoles) runDataCoord(ctx context.Context, localMsg bool) *compone
 		}
 		wg.Done()
 		_ = ds.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: ds})
+		}
 	}()
 	wg.Wait()
 
@@ -240,6 +259,9 @@ func (mr *MilvusRoles) runDataNode(ctx context.Context, localMsg bool, alias str
 		}
 		wg.Done()
 		_ = dn.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: dn})
+		}
 	}()
 	wg.Wait()
 
@@ -267,6 +289,9 @@ func (mr *MilvusRoles) runIndexCoord(ctx context.Context, localMsg bool) *compon
 		}
 		wg.Done()
 		_ = is.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: is})
+		}
 	}()
 	wg.Wait()
 
@@ -295,6 +320,9 @@ func (mr *MilvusRoles) runIndexNode(ctx context.Context, localMsg bool, alias st
 		}
 		wg.Done()
 		_ = in.Run()
+		if !localMsg {
+			http.Handle(healthz.HealthzRouterPath, &componentsHealthzHandler{component: in})
+		}
 	}()
 	wg.Wait()
 
@@ -417,6 +445,39 @@ func (mr *MilvusRoles) Run(localMsg bool, alias string) {
 		}
 	}
 
+	if localMsg {
+		standaloneHealthzHandler := func(w http.ResponseWriter, r *http.Request) {
+			if rc == nil {
+				rootCoordNotServingHandler(w, r)
+				return
+			}
+			if qs == nil {
+				queryCoordNotServingHandler(w, r)
+				return
+			}
+			if ds == nil {
+				dataCoordNotServingHandler(w, r)
+				return
+			}
+			if is == nil {
+				indexCoordNotServingHandler(w, r)
+				return
+			}
+			// TODO(dragondriver): need to check node state?
+
+			w.WriteHeader(http.StatusOK)
+			w.Header().Set(healthz.ContentTypeHeader, healthz.ContentTypeText)
+			_, err := fmt.Fprint(w, "OK")
+			if err != nil {
+				log.Warn("failed to send response",
+					zap.Error(err))
+			}
+
+			// TODO(dragondriver): handle component states
+		}
+		http.HandleFunc(healthz.HealthzRouterPath, standaloneHealthzHandler)
+	}
+
 	metrics.ServeHTTP()
 
 	sc := make(chan os.Signal, 1)

internal/util/healthz/content_type.go

+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package healthz
+
+const (
+	ContentTypeHeader = "Content-Type"
+	ContentTypeText   = "text/plain"
+)

internal/util/healthz/router.go

+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+package healthz
+
+const HealthzRouterPath = "healthz"

internal/util/milvuserrors/errors.go

@@ -16,6 +16,13 @@ import (
 	"fmt"
 )
 
+const (
+	MsgRootCoordNotServing  = "root coordinator is not serving"
+	MsgQueryCoordNotServing = "query coordinator is not serving"
+	MsgDataCoordNotServing  = "data coordinator is not serving"
+	MsgIndexCoordNotServing = "index coordinator is not serving"
+)
+
 func MsgCollectionAlreadyExist(name string) string {
 	return fmt.Sprintf("Collection %s already exist", name)
 }