kit/internal/health/health.go

package health

import (
	"fmt"
	"git.diulo.com/mogfee/kit/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)
}