85 lines
3.0 KiB
Go
85 lines
3.0 KiB
Go
// Package healthcheck 提供 im-service 的 HTTP 和 gRPC 健康检查统一状态模型。
|
||
package healthcheck
|
||
|
||
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"
|
||
)
|
||
|
||
// GRPCServer 实现标准 grpc.health.v1.Health，并把 SERVING 映射到 ready。
|
||
type GRPCServer struct {
|
||
// UnimplementedHealthServer 保持 gRPC health proto 的向前兼容行为。
|
||
healthgrpc.UnimplementedHealthServer
|
||
|
||
// state 是健康状态来源，gRPC health 不维护独立状态。
|
||
state *State
|
||
}
|
||
|
||
// NewGRPCServer 创建 gRPC health server。
|
||
func NewGRPCServer(state *State) *GRPCServer {
|
||
return &GRPCServer{state: state}
|
||
}
|
||
|
||
// Check 按标准 gRPC health protocol 返回服务 ready 状态。
|
||
func (s *GRPCServer) Check(ctx context.Context, req *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
|
||
if !s.supports(req.GetService()) {
|
||
// 标准 health 协议要求未知 service 返回 NotFound。
|
||
return nil, status.Error(codes.NotFound, "unknown service")
|
||
}
|
||
|
||
return &healthpb.HealthCheckResponse{Status: s.status(ctx)}, nil
|
||
}
|
||
|
||
// Watch 周期性推送 ready 状态变化，支持依赖健康检查原生 watch 的服务发现系统。
|
||
func (s *GRPCServer) Watch(req *healthpb.HealthCheckRequest, stream healthgrpc.Health_WatchServer) error {
|
||
if !s.supports(req.GetService()) {
|
||
// 未知服务先发送 SERVICE_UNKNOWN，再等待客户端结束 stream。
|
||
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 {
|
||
// 只有状态变化时推送，减少健康 watch 的无意义网络噪音。
|
||
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():
|
||
// 客户端断开 watch 时按 canceled 返回，符合标准 health watch 语义。
|
||
return status.Error(codes.Canceled, "health watch stream ended")
|
||
}
|
||
}
|
||
}
|
||
|
||
func (s *GRPCServer) supports(service string) bool {
|
||
// 空 service 表示检查整个进程，具名 service 只接受 im-service。
|
||
return service == "" || service == s.state.ServiceName()
|
||
}
|
||
|
||
func (s *GRPCServer) status(ctx context.Context) healthpb.HealthCheckResponse_ServingStatus {
|
||
if Healthy(s.state.Ready(ctx)) {
|
||
// gRPC health 的 SERVING 对应 HTTP ready 成功，而不是单纯进程 live。
|
||
return healthpb.HealthCheckResponse_SERVING
|
||
}
|
||
|
||
return healthpb.HealthCheckResponse_NOT_SERVING
|
||
}
|