health/go-health/readiness.go

package health

import (
	"encoding/json"
	"net/http"
	"slices"
	"time"
)

// Status of the service readiness.
type ReadinessStatus string

// String returns the string representation of the readiness status.
func (s ReadinessStatus) String() string {
	return string(s)
}

// HTTPStatus returns the corresponding HTTP status code for the readiness status.
func (s ReadinessStatus) HTTPStatus() int {
	switch s {
	case ReadinessStatusReady:
		return http.StatusOK
	case ReadinessStatusDegraded:
		return http.StatusOK // 206 could also be suitable, but let's avoid false alarms
	case ReadinessStatusFatal:
		return http.StatusServiceUnavailable
	case ReadinessStatusNotReady:
		return http.StatusOK // treat as ok to avoid false alarms
	default:
		return http.StatusOK
	}
}

const (
	// Service is ready to handle requests.
	ReadinessStatusReady ReadinessStatus = "ready"
	// Service is not ready to handle requests.
	ReadinessStatusNotReady ReadinessStatus = "not_ready"
	// Service experienced a fatal error and cannot handle requests.
	ReadinessStatusFatal ReadinessStatus = "fatal"
	// Service is degraded but can still handle requests.
	ReadinessStatusDegraded ReadinessStatus = "degraded"
)

// Returns the worst readiness status from a list of statuses.
func worstStatus(statuses []ReadinessStatus) ReadinessStatus {
	if slices.Contains(statuses, ReadinessStatusFatal) {
		return ReadinessStatusFatal
	}
	if slices.Contains(statuses, ReadinessStatusNotReady) {
		return ReadinessStatusNotReady
	}
	if slices.Contains(statuses, ReadinessStatusDegraded) {
		return ReadinessStatusDegraded
	}
	return ReadinessStatusReady
}

// Readiness check function return type.
type ReadinessReturn struct {
	// Status of the readiness check.
	Status ReadinessStatus
	// Optional message providing additional information about the readiness status.
	Message *string
}

// Function that performs a readiness check.
type ReadinessCheckFn func() ReadinessReturn

// Individual readiness check.
type ReadinessCheck struct {
	// Name of the readiness check.
	Name string
	// Function that performs the readiness check.
	Fn ReadinessCheckFn
	// Timeout for the readiness check (default: 5 seconds)
	Timeout time.Duration
}

// System-generated result from an individual readiness check.
type ReadinessCheckResult struct {
	// Name of the readiness check.
	Name string
	// Status of the readiness check.
	Status ReadinessStatus
	// Optional message providing additional information about the readiness status.
	Message *string
	// Duration of the readiness check in milliseconds.
	Duration int64
}

// System readiness status.
type Readiness struct {
	// Status of the service readiness. Aggregated as the worst status from
	// all individual readiness checks. If no checks have been performed
	// yet, the status is ReadinessStatusNotReady.
	Status ReadinessStatus
	// Start time of the readiness check in milliseconds since Unix epoch.
	Start int64
	// Duration of the readiness check in milliseconds.
	Duration int64
	// Individual readiness check results.
	Details map[string]ReadinessCheckResult
}

// Performs a full system readiness check given a list of individual checks.
func CheckReadiness(checks []ReadinessCheck) Readiness {
	start := time.Now()
	details := make(map[string]ReadinessCheckResult)
	var statuses []ReadinessStatus

	for _, check := range checks {
		checkt0 := time.Now()
		result := checkWithTimeout(check)
		duration := time.Since(checkt0).Milliseconds()

		details[check.Name] = ReadinessCheckResult{
			Name:     check.Name,
			Status:   result.Status,
			Message:  result.Message,
			Duration: duration,
		}
		statuses = append(statuses, result.Status)
	}

	return Readiness{
		Status:   worstStatus(statuses),
		Start:    start.UnixMilli(),
		Duration: time.Since(start).Milliseconds(),
		Details:  details,
	}
}

// Creates a handler function for readiness HTTP requests.
func ReadinessHandler(checks []ReadinessCheck) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		respondWithResult(w, CheckReadiness(checks))
	}
}

// Runs a readiness check with a timeout.
func checkWithTimeout(check ReadinessCheck) ReadinessReturn {
	timeout := check.Timeout
	if timeout == 0 {
		timeout = 5 * time.Second
	}

	resultCh := make(chan ReadinessReturn, 1)

	go func() {
		resultCh <- check.Fn()
	}()

	select {
	case result := <-resultCh:
		return result
	case <-time.After(timeout):
		msg := "readiness check timed out"
		return ReadinessReturn{
			Status:  ReadinessStatusFatal,
			Message: &msg,
		}
	}
}

func respondWithResult(w http.ResponseWriter, result Readiness) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(result.Status.HTTPStatus())
	json.NewEncoder(w).Encode(result)
}

// ScheduledReadiness performs periodic readiness checks and stores the latest result.
type ScheduledReadiness struct {
	checks       []ReadinessCheck
	interval     time.Duration
	currentState Readiness
	ticker       *time.Ticker
	quit         chan struct{}
}

// NewScheduledReadiness creates a new ScheduledReadiness instance.
func NewScheduledReadiness(checks []ReadinessCheck, interval time.Duration) *ScheduledReadiness {
	return &ScheduledReadiness{
		checks:   checks,
		interval: interval,
		quit:     make(chan struct{}),
	}
}

// Start begins the periodic readiness checks.
func (sr *ScheduledReadiness) Start() {
	sr.ticker = time.NewTicker(sr.interval)
	go func() {
		for {
			select {
			case <-sr.ticker.C:
				sr.currentState = CheckReadiness(sr.checks)
			case <-sr.quit:
				sr.ticker.Stop()
				return
			}
		}
	}()
}

// Stop halts the periodic readiness checks.
func (sr *ScheduledReadiness) Stop() {
	close(sr.quit)
}

// GetCurrentState returns the latest readiness state.
func (sr *ScheduledReadiness) GetCurrentState() Readiness {
	return sr.currentState
}

// Creates a handler function for scheduled readiness HTTP requests.
func (sr *ScheduledReadiness) ReadinessHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		respondWithResult(w, sr.GetCurrentState())
	}
}