105 lines
2.5 KiB
Go
105 lines
2.5 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
yoreseedocpb "github.com/XingfenD/yoresee_doc/collab-go/pkg/gen/yoresee_doc/v1"
|
|
)
|
|
|
|
type HealthChecker struct {
|
|
systemService yoreseedocpb.SystemServiceClient
|
|
draining atomic.Bool
|
|
}
|
|
|
|
type probeResponse struct {
|
|
Status string `json:"status"`
|
|
Detail string `json:"detail,omitempty"`
|
|
Backend string `json:"backend,omitempty"`
|
|
}
|
|
|
|
func NewHealthChecker(systemService yoreseedocpb.SystemServiceClient) *HealthChecker {
|
|
return &HealthChecker{
|
|
systemService: systemService,
|
|
}
|
|
}
|
|
|
|
func (h *HealthChecker) SetDraining(v bool) {
|
|
h.draining.Store(v)
|
|
}
|
|
|
|
func (h *HealthChecker) RegisterProbeRoutes(mux *http.ServeMux) {
|
|
mux.HandleFunc("/health", h.Health)
|
|
mux.HandleFunc("/readyz", h.Readiness)
|
|
mux.HandleFunc("/livez", h.Liveness)
|
|
}
|
|
|
|
func (h *HealthChecker) Liveness(w http.ResponseWriter, _ *http.Request) {
|
|
writeProbeResponse(w, http.StatusOK, "ok", "", "")
|
|
}
|
|
|
|
func (h *HealthChecker) Readiness(w http.ResponseWriter, _ *http.Request) {
|
|
if err := h.readinessErr(); err != nil {
|
|
writeProbeResponse(w, http.StatusServiceUnavailable, "not_ready", err.Error(), probeBackendStatusFromErr(err))
|
|
return
|
|
}
|
|
writeProbeResponse(w, http.StatusOK, "ok", "", "ok")
|
|
}
|
|
|
|
func (h *HealthChecker) Health(w http.ResponseWriter, _ *http.Request) {
|
|
detail := ""
|
|
backend := "ok"
|
|
if err := h.readinessErr(); err != nil {
|
|
detail = err.Error()
|
|
backend = probeBackendStatusFromErr(err)
|
|
}
|
|
writeProbeResponse(w, http.StatusOK, "ok", detail, backend)
|
|
}
|
|
|
|
func (h *HealthChecker) readinessErr() error {
|
|
if h.draining.Load() {
|
|
return fmt.Errorf("server is draining")
|
|
}
|
|
return h.checkBackend()
|
|
}
|
|
|
|
func writeProbeResponse(w http.ResponseWriter, statusCode int, status, detail, backend string) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(statusCode)
|
|
_ = json.NewEncoder(w).Encode(probeResponse{
|
|
Status: status,
|
|
Detail: detail,
|
|
Backend: backend,
|
|
})
|
|
}
|
|
|
|
func probeBackendStatusFromErr(err error) string {
|
|
if err == nil {
|
|
return "ok"
|
|
}
|
|
if err.Error() == "server is draining" {
|
|
return "unknown"
|
|
}
|
|
return "unhealthy"
|
|
}
|
|
|
|
func (h *HealthChecker) checkBackend() error {
|
|
if h.systemService == nil {
|
|
return nil
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
|
|
if _, err := h.systemService.Health(ctx, &yoreseedocpb.HealthRequest{}); err != nil {
|
|
log.Printf("backend health check failed: %v", err)
|
|
return err
|
|
}
|
|
return nil
|
|
}
|