package health import ( "fmt" "github.com/zeromicro/go-zero/core/syncx" "net/http" "strings" "sync" ) var defaultHealthManager = newComboHealthManager() type ( Probe interface { MarkReady() MarkNotReady() IsReady() bool Name() string } healthManager struct { ready syncx.AtomicBool name string } comboHealthManager struct { mu sync.Mutex probes []Probe } ) func AddProbe(probe Probe) { defaultHealthManager.addProbe(probe) } func CreateHttpHandler() http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { if defaultHealthManager.IsReady() { _, _ = w.Write([]byte("OK")) } else { http.Error(w, "Service Unavailable\n"+defaultHealthManager.verboseInfo(), http.StatusServiceUnavailable) } } } func NewHealthManager(name string) Probe { return &healthManager{ name: name, } } func (h *healthManager) MarkReady() { h.ready.Set(true) } func (h *healthManager) MarkNotReady() { h.ready.Set(false) } func (h *healthManager) IsReady() bool { return h.ready.True() } func (h *healthManager) Name() string { return h.name } func newComboHealthManager() *comboHealthManager { return &comboHealthManager{} } func (p *comboHealthManager) MarkReady() { p.mu.Lock() defer p.mu.Unlock() for _, probe := range p.probes { probe.MarkReady() } } func (p *comboHealthManager) MarkNotReady() { p.mu.Lock() defer p.mu.Unlock() for _, probe := range p.probes { probe.MarkNotReady() } } func (p *comboHealthManager) IsReady() bool { p.mu.Lock() defer p.mu.Unlock() for _, probe := range p.probes { if !probe.IsReady() { return false } } return true } func (p *comboHealthManager) verboseInfo() string { p.mu.Lock() defer p.mu.Unlock() var info strings.Builder for _, probe := range p.probes { if probe.IsReady() { info.WriteString(fmt.Sprintf("%s is ready\n", probe.Name())) } else { info.WriteString(fmt.Sprintf("%s is not ready\n", probe.Name())) } } return info.String() } func (p *comboHealthManager) addProbe(probe Probe) { p.mu.Lock() defer p.mu.Unlock() p.probes = append(p.probes, probe) }