84 lines
2.2 KiB
Go
84 lines
2.2 KiB
Go
package grpchealth
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
const defaultDependencyBudget = 300 * time.Millisecond
|
|
|
|
// Dependency 是 readiness 中必须可用的外部依赖。
|
|
type Dependency struct {
|
|
Name string
|
|
Budget time.Duration
|
|
Check func(ctx context.Context) error
|
|
}
|
|
|
|
// ServingChecker 覆盖无外部硬依赖的占位或 stub gRPC 服务 readiness。
|
|
type ServingChecker struct {
|
|
serviceName string
|
|
dependencies []Dependency
|
|
serving atomic.Bool
|
|
draining atomic.Bool
|
|
}
|
|
|
|
// NewServingChecker 创建只依赖 gRPC listener 和 draining 状态的 checker。
|
|
func NewServingChecker(serviceName string, dependencies ...Dependency) *ServingChecker {
|
|
return &ServingChecker{serviceName: serviceName, dependencies: dependencies}
|
|
}
|
|
|
|
// ServiceName 返回标准 gRPC health 使用的服务边界名。
|
|
func (c *ServingChecker) ServiceName() string {
|
|
return c.serviceName
|
|
}
|
|
|
|
// MarkServing 标记 gRPC listener 已进入 Serve 循环。
|
|
func (c *ServingChecker) MarkServing() {
|
|
c.serving.Store(true)
|
|
}
|
|
|
|
// MarkStopped 标记 gRPC listener 已退出 Serve 循环。
|
|
func (c *ServingChecker) MarkStopped() {
|
|
c.serving.Store(false)
|
|
}
|
|
|
|
// MarkDraining 让 ready 立即失败，服务下线时不再进入可用池。
|
|
func (c *ServingChecker) MarkDraining() {
|
|
c.draining.Store(true)
|
|
}
|
|
|
|
// Ready 返回 listener、draining 和外部硬依赖的聚合 readiness。
|
|
func (c *ServingChecker) Ready(ctx context.Context) []Check {
|
|
checks := []Check{
|
|
Bool("grpc_listener", c.serving.Load(), "gRPC listener is not serving"),
|
|
Bool("draining", !c.draining.Load(), "node is draining"),
|
|
}
|
|
for _, dependency := range c.dependencies {
|
|
checks = append(checks, dependency.Ready(ctx))
|
|
}
|
|
return checks
|
|
}
|
|
|
|
// Ready executes a dependency check with a short bounded budget.
|
|
func (d Dependency) Ready(ctx context.Context) Check {
|
|
name := strings.TrimSpace(d.Name)
|
|
if name == "" {
|
|
name = "dependency"
|
|
}
|
|
if d.Check == nil {
|
|
return Fail(name, "dependency check is not configured")
|
|
}
|
|
budget := d.Budget
|
|
if budget <= 0 {
|
|
budget = defaultDependencyBudget
|
|
}
|
|
checkCtx, cancel := context.WithTimeout(ctx, budget)
|
|
defer cancel()
|
|
if err := d.Check(checkCtx); err != nil {
|
|
return Fail(name, err.Error())
|
|
}
|
|
return OK(name)
|
|
}
|