migrate from github

health/health.go

@@ -0,0 +1,104 @@
+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
+}