Day 106 — Health Checks & Readiness Probes

Key Concepts

Kubernetes probe types:

Rule: Liveness checks only detect catastrophic failure (deadlock, panic). Don't check upstream dependencies in liveness — a slow database should not cause your pod to restart.

Exercise — Health Check Framework

package health

import (
    "context"
    "encoding/json"
    "fmt"
    "net/http"
    "sync"
    "time"
)

// CheckFunc is a function that checks if a subsystem is healthy.
// It must respect context cancellation and return within the deadline.
type CheckFunc func(ctx context.Context) error

type checkEntry struct {
    name    string
    fn      CheckFunc
    timeout time.Duration
    // For readiness: is this required for traffic?
    required bool
}

// Registry holds all registered health checks
type Registry struct {
    mu     sync.RWMutex
    checks []checkEntry
}

func NewRegistry() *Registry {
    return &Registry{}
}

// Register adds a check. Use required=true for readiness-gating checks.
func (r *Registry) Register(name string, fn CheckFunc, timeout time.Duration, required bool) {
    r.mu.Lock()
    defer r.mu.Unlock()
    r.checks = append(r.checks, checkEntry{name: name, fn: fn, timeout: timeout, required: required})
}

type CheckResult struct {
    Status  string `json:"status"` // "ok" or "fail"
    Error   string `json:"error,omitempty"`
    Latency string `json:"latency_ms"`
}

type HealthResponse struct {
    Status string                 `json:"status"` // "ok" or "degraded" or "unhealthy"
    Checks map[string]CheckResult `json:"checks"`
    Took   string                 `json:"took_ms"`
}

func (r *Registry) run(ctx context.Context) HealthResponse {
    r.mu.RLock()
    checks := make([]checkEntry, len(r.checks))
    copy(checks, r.checks)
    r.mu.RUnlock()

    start := time.Now()
    results := make(map[string]CheckResult, len(checks))
    var mu sync.Mutex
    var wg sync.WaitGroup

    for _, c := range checks {
        wg.Add(1)
        go func(c checkEntry) {
            defer wg.Done()
            checkCtx, cancel := context.WithTimeout(ctx, c.timeout)
            defer cancel()

            checkStart := time.Now()
            err := c.fn(checkCtx)
            latency := time.Since(checkStart)

            result := CheckResult{
                Latency: fmt.Sprintf("%.1f", float64(latency.Microseconds())/1000.0),
            }
            if err != nil {
                result.Status = "fail"
                result.Error = err.Error()
            } else {
                result.Status = "ok"
            }

            mu.Lock()
            results[c.name] = result
            mu.Unlock()
        }(c)
    }
    wg.Wait()

    // Aggregate status
    status := "ok"
    for _, c := range checks {
        res := results[c.name]
        if res.Status == "fail" && c.required {
            status = "unhealthy"
            break
        } else if res.Status == "fail" {
            if status == "ok" {
                status = "degraded"
            }
        }
    }

    return HealthResponse{
        Status: status,
        Checks: results,
        Took:   fmt.Sprintf("%.1f", float64(time.Since(start).Microseconds())/1000.0),
    }
}

// LivezHandler handles GET /healthz (liveness — is process alive and not deadlocked?)
// Only runs lightweight checks — never checks upstream dependencies.
func (r *Registry) LivezHandler(w http.ResponseWriter, req *http.Request) {
    // Liveness is simple: if this handler runs, the process is alive.
    // Run only non-required checks with very short timeouts.
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

// ReadyzHandler handles GET /readyz (readiness — ready for traffic?)
// Runs all required checks. Fails if any required check fails.
func (r *Registry) ReadyzHandler(w http.ResponseWriter, req *http.Request) {
    ctx, cancel := context.WithTimeout(req.Context(), 5*time.Second)
    defer cancel()

    resp := r.run(ctx)
    w.Header().Set("Content-Type", "application/json")
    if resp.Status == "unhealthy" {
        w.WriteHeader(http.StatusServiceUnavailable)
    }
    json.NewEncoder(w).Encode(resp)
}

// StartupHandler handles GET /startup — allows slow initialization.
// After startup completes, this always returns 200.
type StartupHandler struct {
    ready chan struct{}
    once  sync.Once
}

func NewStartupHandler() *StartupHandler {
    return &StartupHandler{ready: make(chan struct{})}
}

func (sh *StartupHandler) MarkReady() {
    sh.once.Do(func() { close(sh.ready) })
}

func (sh *StartupHandler) Handler(w http.ResponseWriter, r *http.Request) {
    select {
    case <-sh.ready:
        w.WriteHeader(http.StatusOK)
        w.Write([]byte(`{"status":"ready"}`))
    default:
        w.WriteHeader(http.StatusServiceUnavailable)
        w.Write([]byte(`{"status":"starting"}`))
    }
}

// Example checks for common subsystems
func NewTCPCheck(addr string) CheckFunc {
    return func(ctx context.Context) error {
        d := net.Dialer{}
        conn, err := d.DialContext(ctx, "tcp", addr)
        if err != nil {
            return fmt.Errorf("tcp connect to %s: %w", addr, err)
        }
        conn.Close()
        return nil
    }
}

func NewHTTPCheck(url string) CheckFunc {
    client := &http.Client{Timeout: 2 * time.Second}
    return func(ctx context.Context) error {
        req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
        if err != nil {
            return err
        }
        resp, err := client.Do(req)
        if err != nil {
            return fmt.Errorf("http get %s: %w", url, err)
        }
        resp.Body.Close()
        if resp.StatusCode >= 500 {
            return fmt.Errorf("upstream returned %d", resp.StatusCode)
        }
        return nil
    }
}

func NewGoroutineLeakCheck(baseline int, maxMultiplier float64) CheckFunc {
    return func(ctx context.Context) error {
        current := runtime.NumGoroutine()
        limit := int(float64(baseline) * maxMultiplier)
        if current > limit {
            return fmt.Errorf("goroutine count %d exceeds %dx baseline (%d)", current, int(maxMultiplier), baseline)
        }
        return nil
    }
}