go-omni/health/health.go

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

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

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

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"`
	cause  error
}

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.cause = errors.Join(result.cause, fmt.Errorf("%s: %w", probe.Name, err))
		}
		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(options ...HandlerOption) httpx.Handler {
	settings := handlerSettings(options)
	return func(ctx *httpx.Context) error {
		data := map[string]string{"status": "ok"}
		if settings.raw {
			return ctx.JSON(http.StatusOK, data)
		}
		return ctx.RespondStatus(http.StatusOK, data, nil)
	}
}

// ReadinessHandler reports whether all registered dependencies are ready.
// Failure always uses HTTP 503. App response configuration is used by default;
// WithRawResponse returns just health data. Internal probe errors remain in the
// recorded failure for logs, not in the HTTP response.
func (c *Checker) ReadinessHandler(options ...HandlerOption) httpx.Handler {
	settings := handlerSettings(options)
	return func(ctx *httpx.Context) error {
		var result Result
		if settings.shutdown != nil && settings.shutdown.Err() != nil {
			result = Result{Status: StatusNotReady, cause: context.Cause(settings.shutdown)}
		} else {
			result = c.Check(ctx.Context())
			// A stop may arrive while the probes are running.
			if settings.shutdown != nil && settings.shutdown.Err() != nil {
				result.Status = StatusNotReady
				result.cause = errors.Join(result.cause, context.Cause(settings.shutdown))
			}
		}
		for i := range result.Checks {
			result.Checks[i].Error = ""
		}
		status := http.StatusOK
		var failure error
		if !result.Ready() {
			status = http.StatusServiceUnavailable
			failure = apierrors.New(status, "NOT_READY", "service is not ready").WithCause(result.cause)
		}
		if settings.raw {
			if failure != nil {
				resolved := ctx.ResolveError(failure)
				resolved.Status = status
				ctx.RecordFailure(resolved)
			}
			return ctx.JSON(status, result)
		}
		return ctx.RespondStatus(status, result, failure)
	}
}

type handlerConfig struct {
	raw      bool
	shutdown context.Context
}

type HandlerOption func(*handlerConfig)

// WithRawResponse bypasses the App envelope for this probe's health data.
func WithRawResponse() HandlerOption { return func(c *handlerConfig) { c.raw = true } }

// WithShutdownContext marks readiness false as soon as service shutdown begins.
// It never changes liveness or closes any project resource.
func WithShutdownContext(ctx context.Context) HandlerOption {
	return func(c *handlerConfig) { c.shutdown = ctx }
}

func handlerSettings(options []HandlerOption) handlerConfig {
	var cfg handlerConfig
	for _, option := range options {
		if option != nil {
			option(&cfg)
		}
	}
	return cfg
}