1package http
  2
  3import (
  4	"context"
  5	"net/http"
  6	"strings"
  7	"sync"
  8	"time"
  9)
 10
 11const readinessCheckTimeout = 2 * time.Second
 12
 13type readinessCheck struct {
 14	name  string
 15	check func(context.Context) error
 16}
 17
 18type Readiness struct {
 19	mu       sync.Mutex
 20	checks   []readinessCheck
 21	started  bool
 22	stopping bool
 23}
 24
 25func NewReadiness() *Readiness { return &Readiness{} }
 26
 27func (r *Readiness) Register(name string, check func(context.Context) error) {
 28	r.mu.Lock()
 29	defer r.mu.Unlock()
 30	r.checks = append(r.checks, readinessCheck{name: name, check: check})
 31}
 32
 33func (r *Readiness) MarkStarted() {
 34	r.mu.Lock()
 35	defer r.mu.Unlock()
 36	r.started = true
 37}
 38
 39func (r *Readiness) MarkStopping() {
 40	r.mu.Lock()
 41	defer r.mu.Unlock()
 42	r.stopping = true
 43}
 44
 45func (r *Readiness) Report(ctx context.Context) (bool, string) {
 46	r.mu.Lock()
 47	started, stopping := r.started, r.stopping
 48	checks := append([]readinessCheck(nil), r.checks...)
 49	r.mu.Unlock()
 50	switch {
 51	case stopping:
 52		return false, "shutting down\n"
 53	case !started:
 54		return false, "not started\n"
 55	}
 56	ctx, cancel := context.WithTimeout(ctx, readinessCheckTimeout)
 57	defer cancel()
 58	var failing strings.Builder
 59	for _, c := range checks {
 60		if err := c.check(ctx); err != nil {
 61			failing.WriteString(c.name + "\n")
 62		}
 63	}
 64	if failing.Len() > 0 {
 65		return false, failing.String()
 66	}
 67	return true, "ok\n"
 68}
 69
 70func OpsHandler(rd *Readiness) http.Handler {
 71	return OpsHandlerWithMetrics(rd, nil)
 72}
 73
 74func OpsHandlerWithMetrics(rd *Readiness, metrics http.Handler) http.Handler {
 75	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 76		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 77		w.Header().Set("Cache-Control", "no-store")
 78		if r.Method != http.MethodGet && r.Method != http.MethodHead {
 79			http.NotFound(w, r)
 80			return
 81		}
 82		switch r.URL.Path {
 83		case "/healthz":
 84			w.Write([]byte("ok\n"))
 85		case "/readyz":
 86			ready, body := rd.Report(r.Context())
 87			if !ready {
 88				w.WriteHeader(http.StatusServiceUnavailable)
 89			}
 90			w.Write([]byte(body))
 91		case "/metrics":
 92			if metrics == nil {
 93				http.NotFound(w, r)
 94				return
 95			}
 96			metrics.ServeHTTP(w, r)
 97		default:
 98			http.NotFound(w, r)
 99		}
100	})
101}