You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
110 lines
2.1 KiB
110 lines
2.1 KiB
12 months ago
|
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)
|
||
|
}
|