126 lines
3.3 KiB
Go
126 lines
3.3 KiB
Go
package grpchealth
|
||
|
||
import (
|
||
"context"
|
||
"time"
|
||
|
||
"google.golang.org/grpc/codes"
|
||
healthgrpc "google.golang.org/grpc/health/grpc_health_v1"
|
||
healthpb "google.golang.org/grpc/health/grpc_health_v1"
|
||
"google.golang.org/grpc/status"
|
||
)
|
||
|
||
const (
|
||
// StatusOK 表示单个健康检查项通过。
|
||
StatusOK = "ok"
|
||
// StatusFail 表示单个健康检查项失败。
|
||
StatusFail = "fail"
|
||
)
|
||
|
||
// Check 描述一项 readiness 子检查；gRPC health 只返回总状态，子项用于服务内聚合。
|
||
type Check struct {
|
||
Name string
|
||
Status string
|
||
Message string
|
||
}
|
||
|
||
// Checker 把服务自己的 ready 语义适配到标准 gRPC health protocol。
|
||
type Checker interface {
|
||
ServiceName() string
|
||
Ready(ctx context.Context) []Check
|
||
}
|
||
|
||
// Server 实现 grpc.health.v1.Health，并把 SERVING 严格映射到 Checker.Ready。
|
||
type Server struct {
|
||
healthgrpc.UnimplementedHealthServer
|
||
|
||
checker Checker
|
||
}
|
||
|
||
// NewServer 创建标准 gRPC health server。
|
||
func NewServer(checker Checker) *Server {
|
||
return &Server{checker: checker}
|
||
}
|
||
|
||
// Check 返回当前服务是否 ready；未知服务名按标准返回 NotFound。
|
||
func (s *Server) Check(ctx context.Context, req *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
|
||
if !s.supports(req.GetService()) {
|
||
return nil, status.Error(codes.NotFound, "unknown service")
|
||
}
|
||
|
||
return &healthpb.HealthCheckResponse{Status: s.status(ctx)}, nil
|
||
}
|
||
|
||
// Watch 定时推送 ready 状态变化，兼容支持 gRPC health watch 的服务发现系统。
|
||
func (s *Server) Watch(req *healthpb.HealthCheckRequest, stream healthgrpc.Health_WatchServer) error {
|
||
if !s.supports(req.GetService()) {
|
||
if err := stream.Send(&healthpb.HealthCheckResponse{Status: healthpb.HealthCheckResponse_SERVICE_UNKNOWN}); err != nil {
|
||
return status.Error(codes.Canceled, "health watch stream ended")
|
||
}
|
||
<-stream.Context().Done()
|
||
return status.Error(codes.Canceled, "health watch stream ended")
|
||
}
|
||
|
||
ticker := time.NewTicker(2 * time.Second)
|
||
defer ticker.Stop()
|
||
|
||
var last healthpb.HealthCheckResponse_ServingStatus = -1
|
||
for {
|
||
current := s.status(stream.Context())
|
||
if current != last {
|
||
if err := stream.Send(&healthpb.HealthCheckResponse{Status: current}); err != nil {
|
||
return status.Error(codes.Canceled, "health watch stream ended")
|
||
}
|
||
last = current
|
||
}
|
||
|
||
select {
|
||
case <-ticker.C:
|
||
case <-stream.Context().Done():
|
||
return status.Error(codes.Canceled, "health watch stream ended")
|
||
}
|
||
}
|
||
}
|
||
|
||
// Healthy 聚合子检查，任何失败项都会让服务从 ready 池摘除。
|
||
func Healthy(checks []Check) bool {
|
||
for _, check := range checks {
|
||
if check.Status != StatusOK {
|
||
return false
|
||
}
|
||
}
|
||
|
||
return true
|
||
}
|
||
|
||
// OK 创建通过项。
|
||
func OK(name string) Check {
|
||
return Check{Name: name, Status: StatusOK}
|
||
}
|
||
|
||
// Fail 创建失败项。
|
||
func Fail(name string, message string) Check {
|
||
return Check{Name: name, Status: StatusFail, Message: message}
|
||
}
|
||
|
||
// Bool 根据布尔条件创建检查项。
|
||
func Bool(name string, ok bool, message string) Check {
|
||
if ok {
|
||
return OK(name)
|
||
}
|
||
|
||
return Fail(name, message)
|
||
}
|
||
|
||
func (s *Server) supports(service string) bool {
|
||
return service == "" || service == s.checker.ServiceName()
|
||
}
|
||
|
||
func (s *Server) status(ctx context.Context) healthpb.HealthCheckResponse_ServingStatus {
|
||
if Healthy(s.checker.Ready(ctx)) {
|
||
return healthpb.HealthCheckResponse_SERVING
|
||
}
|
||
|
||
return healthpb.HealthCheckResponse_NOT_SERVING
|
||
}
|