go-omni/health/health.go

// Package health provides transport-neutral liveness and readiness handlers.
package health

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/response"
)

const (
	StatusReady    = "ready"
	StatusNotReady = "not_ready"
)

// Probe is one dependency check used by readiness. Timeout is optional; a
// zero value lets the request context control the probe duration.
type Probe struct {
	Name    string
	Check   func(context.Context) error
	Timeout time.Duration
}

// CheckResult is the public result of one readiness probe.
type CheckResult struct {
	Name       string  `json:"name"`
	Status     string  `json:"status"`
	Error      string  `json:"error,omitempty"`
	DurationMs float64 `json:"durationMs"`
}

// Result is returned by Checker.Check and the readiness endpoint.
type Result struct {
	Status string        `json:"status"`
	Checks []CheckResult `json:"checks,omitempty"`
}

func (r Result) Ready() bool { return r.Status == StatusReady }

// Checker owns the probes registered by a service. It can be shared by all
// requests after startup; registration is safe before or during serving.
type Checker struct {
	mu     sync.RWMutex
	probes []Probe
}

// New creates a checker and validates its initial probes.
func New(probes ...Probe) (*Checker, error) {
	checker := &Checker{}
	for _, probe := range probes {
		if err := checker.Add(probe); err != nil {
			return nil, err
		}
	}
	return checker, nil
}

// Add registers a named readiness probe. Names must be unique and checks must
// be non-nil so configuration errors fail during startup.
func (c *Checker) Add(probe Probe) error {
	if c == nil {
		return fmt.Errorf("health: nil checker")
	}
	probe.Name = strings.TrimSpace(probe.Name)
	if probe.Name == "" {
		return fmt.Errorf("health: probe name is required")
	}
	if probe.Check == nil {
		return fmt.Errorf("health: probe %q has no check", probe.Name)
	}
	if probe.Timeout < 0 {
		return fmt.Errorf("health: probe %q timeout must not be negative", probe.Name)
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	for _, current := range c.probes {
		if current.Name == probe.Name {
			return fmt.Errorf("health: duplicate probe %q", probe.Name)
		}
	}
	c.probes = append(c.probes, probe)
	return nil
}

// Check runs every registered probe in registration order. A failed probe is
// represented in the result; the other probes still run so operators see all
// unavailable dependencies in one response.
func (c *Checker) Check(ctx context.Context) Result {
	if ctx == nil {
		ctx = context.Background()
	}
	c.mu.RLock()
	probes := append([]Probe(nil), c.probes...)
	c.mu.RUnlock()

	result := Result{Status: StatusReady, Checks: make([]CheckResult, 0, len(probes))}
	for _, probe := range probes {
		started := time.Now()
		probeContext := ctx
		cancel := func() {}
		if probe.Timeout > 0 {
			probeContext, cancel = context.WithTimeout(ctx, probe.Timeout)
		}
		err := probe.Check(probeContext)
		cancel()
		checkResult := CheckResult{
			Name:       probe.Name,
			Status:     StatusReady,
			DurationMs: time.Since(started).Seconds() * 1000,
		}
		if err != nil {
			result.Status = StatusNotReady
			checkResult.Status = StatusNotReady
			checkResult.Error = err.Error()
		}
		result.Checks = append(result.Checks, checkResult)
	}
	return result
}

// LivenessHandler reports that the process is running. It does not execute
// dependency probes and is safe to use for orchestrator liveness checks.
func (c *Checker) LivenessHandler() httpx.Handler {
	return func(ctx *httpx.Context) error {
		return ctx.OK(map[string]string{"status": "ok"})
	}
}

// ReadinessHandler reports whether all registered dependencies are ready.
// Failure uses HTTP 503 while preserving GoOmni's response envelope.
func (c *Checker) ReadinessHandler() httpx.Handler {
	return func(ctx *httpx.Context) error {
		result := c.Check(ctx.Context())
		status := http.StatusOK
		code, message := "READY", "ready"
		if !result.Ready() {
			status = http.StatusServiceUnavailable
			code, message = "NOT_READY", "service is not ready"
		}
		return ctx.JSON(status, response.Response{
			Code:      code,
			Message:   message,
			Data:      result,
			RequestID: ctx.RequestID(),
		})
	}
}