feat(health): serve healthz and readyz

health/probes.go

@@ -0,0 +1,86 @@
+// Package health answers the liveness and readiness probes a deployment gates
+// on.
+package health
+
+import (
+	"net/http"
+	"sync/atomic"
+)
+
+// Probes reports whether a run still responds, and whether it has finished
+// starting.
+//
+// The two answer different questions, and conflating them restarts a pod that
+// is working. Liveness asks whether the process is still there. Readiness asks
+// whether it has finished a startup that funds accounts, deploys contracts and
+// prewarms, which takes minutes against a cold chain. A liveness probe that
+// waited for all of that would kill the run before its first transaction, and
+// kill the next attempt at the same point.
+//
+// So /healthz answers as soon as the server binds and never consults the
+// startup sequence. Only /readyz gates on it.
+type Probes struct {
+	state atomic.Pointer[state]
+}
+
+// state is one consistent answer. The phase sits beside the flag in a single
+// stored value, so a reader cannot pair a stale phase with a fresh flag.
+type state struct {
+	ready bool
+	phase string
+}
+
+// New returns probes reporting the given phase, not yet ready.
+func New(phase string) *Probes {
+	probes := &Probes{}
+	probes.state.Store(&state{phase: phase})
+	return probes
+}
+
+// Enter records the phase a run is working through. It does not make the run
+// ready. The phase is what /readyz reports while it is still refusing, so an
+// operator watching a ten-minute startup reads the step rather than a bare 503.
+func (p *Probes) Enter(phase string) {
+	p.state.Store(&state{phase: phase})
+}
+
+// Ready marks the run started and serving.
+func (p *Probes) Ready() {
+	p.state.Store(&state{ready: true, phase: "running"})
+}
+
+// NotReady takes a run out of service without reporting it dead. A shutting-down
+// run answers /healthz until its server stops, so the kubelet lets it finish
+// rather than killing it as unresponsive.
+func (p *Probes) NotReady(phase string) {
+	p.state.Store(&state{phase: phase})
+}
+
+// Register mounts both endpoints on mux.
+func (p *Probes) Register(mux *http.ServeMux) {
+	mux.HandleFunc("/healthz", p.serveLive)
+	mux.HandleFunc("/readyz", p.serveReady)
+}
+
+// serveLive answers for as long as the server runs. It reads no state on
+// purpose: see the type's documentation.
+func (p *Probes) serveLive(w http.ResponseWriter, _ *http.Request) {
+	writeText(w, http.StatusOK, "ok")
+}
+
+func (p *Probes) serveReady(w http.ResponseWriter, _ *http.Request) {
+	current := p.state.Load()
+	if !current.ready {
+		writeText(w, http.StatusServiceUnavailable, current.phase)
+		return
+	}
+	writeText(w, http.StatusOK, current.phase)
+}
+
+func writeText(w http.ResponseWriter, status int, body string) {
+	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+	w.WriteHeader(status)
+	// The probe reads the status line. A short body is for an operator running
+	// curl, so a write that fails changes nothing worth reporting.
+	_, _ = w.Write([]byte(body + "\n"))
+}

health/probes_test.go

@@ -0,0 +1,156 @@
+package health_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/sei-protocol/sei-load/health"
+)
+
+// get returns the status and body one endpoint answers with.
+func get(t *testing.T, probes *health.Probes, path string) (int, string) {
+	t.Helper()
+	mux := http.NewServeMux()
+	probes.Register(mux)
+	recorder := httptest.NewRecorder()
+	mux.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, path, nil))
+	return recorder.Code, strings.TrimSpace(recorder.Body.String())
+}
+
+// TestLivenessDoesNotWaitForStartup guards the failure that makes these probes
+// worse than none. Funding, deployment and prewarm take minutes against a cold
+// chain. A liveness probe that reported the run dead for that whole window would
+// restart the pod before it sent a transaction, and restart the next attempt at
+// the same point, so the run never happens and the cause looks like a crash.
+func TestLivenessDoesNotWaitForStartup(t *testing.T) {
+	probes := health.New("starting")
+
+	for _, phase := range []string{"starting", "deploying contracts", "funding accounts", "prewarming accounts"} {
+		probes.Enter(phase)
+		status, _ := get(t, probes, "/healthz")
+		require.Equal(t, http.StatusOK, status,
+			"liveness failed during %q, so the kubelet restarts the pod mid-startup and the run never reaches its first transaction", phase)
+	}
+}
+
+// TestReadinessWaitsForStartup is the other half: a run that has not deployed or
+// funded anything must not be reported as serving.
+func TestReadinessWaitsForStartup(t *testing.T) {
+	probes := health.New("starting")
+
+	status, _ := get(t, probes, "/readyz")
+	require.Equal(t, http.StatusServiceUnavailable, status,
+		"readiness passed before the dispatcher started, so the probe cannot tell a run that is working from one still funding")
+
+	probes.Ready()
+	status, _ = get(t, probes, "/readyz")
+	require.Equal(t, http.StatusOK, status,
+		"readiness still fails after the run started, so a startupProbe would exhaust its budget and kill a healthy run")
+}
+
+// TestReadinessNamesThePhaseItIsWaitingOn keeps the body useful. A ten-minute
+// startup that answers only "503" tells an operator nothing about which step is
+// slow.
+func TestReadinessNamesThePhaseItIsWaitingOn(t *testing.T) {
+	probes := health.New("starting")
+	probes.Enter("funding accounts")
+
+	status, body := get(t, probes, "/readyz")
+	require.Equal(t, http.StatusServiceUnavailable, status)
+	require.Equal(t, "funding accounts", body,
+		"the body does not name the phase, so a slow startup reports no more than a bare failure")
+}
+
+// TestShutdownLeavesServiceWithoutReportingDeath covers the window where a run
+// holds the pod open for the post-summary scrape. Readiness must drop so nothing
+// routes to it. Liveness must hold, or the kubelet reads the deliberate hold as
+// a hang and kills the process before the scrape lands.
+func TestShutdownLeavesServiceWithoutReportingDeath(t *testing.T) {
+	probes := health.New("starting")
+	probes.Ready()
+	probes.NotReady("shutting down")
+
+	ready, body := get(t, probes, "/readyz")
+	require.Equal(t, http.StatusServiceUnavailable, ready,
+		"a shutting-down run still reports ready, so traffic routes to a process that is leaving")
+	require.Equal(t, "shutting down", body)
+
+	live, _ := get(t, probes, "/healthz")
+	require.Equal(t, http.StatusOK, live,
+		"liveness failed during shutdown, so the kubelet kills the run before its final metrics are scraped")
+}
+
+// TestAReadyStatusNeverCarriesAStartupPhase pins the reason the flag and the
+// phase are stored as one value rather than as two atomics. Stored separately,
+// a writer setting the flag and then the phase leaves a window where a reader
+// sees the run serving while the body still names the step it was on. The
+// status and the body would then disagree about the same instant, and an
+// operator reading the body would act on a phase the run had left.
+func TestAReadyStatusNeverCarriesAStartupPhase(t *testing.T) {
+	probes := health.New("starting")
+	phases := []string{"deploying contracts", "funding accounts", "prewarming accounts"}
+
+	var group sync.WaitGroup
+	group.Add(2)
+	stop := make(chan struct{})
+	go func() {
+		defer group.Done()
+		defer close(stop)
+		for i := 0; i < 20_000; i++ {
+			probes.Enter(phases[i%len(phases)])
+			probes.Ready()
+			probes.NotReady("shutting down")
+		}
+	}()
+	go func() {
+		defer group.Done()
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+			}
+			if status, body := get(t, probes, "/readyz"); status == http.StatusOK {
+				require.Equal(t, "running", body,
+					"a serving status carried the phase %q, so the flag and the phase are not stored as one value", body)
+			}
+		}
+	}()
+	group.Wait()
+}
+
+// TestConcurrentPhasesAndProbesDoNotRace runs the real pairing: the run goroutine
+// advances phases while the probe goroutine reads. Under -race this fails if the
+// phase and the flag are ever stored separately.
+func TestConcurrentPhasesAndProbesDoNotRace(t *testing.T) {
+	probes := health.New("starting")
+	phases := []string{"deploying contracts", "funding accounts", "prewarming accounts"}
+
+	var group sync.WaitGroup
+	group.Add(2)
+	go func() {
+		defer group.Done()
+		for i := 0; i < 200; i++ {
+			probes.Enter(phases[i%len(phases)])
+		}
+		probes.Ready()
+	}()
+	go func() {
+		defer group.Done()
+		for i := 0; i < 200; i++ {
+			status, body := get(t, probes, "/readyz")
+			// Never a ready status paired with a startup phase: the two are stored
+			// as one value, so a reader cannot see half of an update.
+			if status == http.StatusOK {
+				require.Equal(t, "running", body,
+					"a ready status carried a startup phase, so the flag and the phase are not stored together")
+			}
+		}
+	}()
+	group.Wait()
+}

main.go

@@ -24,6 +24,7 @@ import (
 	"github.com/sei-protocol/sei-load/config"
 	"github.com/sei-protocol/sei-load/funder"
 	"github.com/sei-protocol/sei-load/generator"
+	"github.com/sei-protocol/sei-load/health"
 	"github.com/sei-protocol/sei-load/observability"
 	"github.com/sei-protocol/sei-load/sender"
 	"github.com/sei-protocol/sei-load/stats"
@@ -171,6 +172,10 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 	listenAddr := cmd.Flag("metricsListenAddr").Value.String()
 	log.Printf("serving metrics at %s/metrics", listenAddr)
 
+	// Built before the server so /readyz answers from the first scrape rather
+	// than from whenever the run reaches its first phase.
+	probes := health.New("starting")
+
 	obsShutdown, err := observability.Setup(ctx, observability.Config{
 		RunScope:     observability.RunScopeFromEnv(),
 		OTLPEndpoint: os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"),
@@ -189,6 +194,7 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 	// EnableOpenMetrics is load-bearing: the default promhttp.Handler() strips
 	// exemplars regardless of the scraper's Accept header.
 	mux := http.NewServeMux()
+	probes.Register(mux)
 	mux.Handle("/metrics", promhttp.HandlerFor(
 		prometheus.DefaultGatherer,
 		promhttp.HandlerOpts{EnableOpenMetrics: true},
@@ -237,6 +243,7 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 		}
 
 		// Create the generator from the config struct
+		probes.Enter("deploying contracts")
 		gen, err := generator.NewGenerator(ctx, rng, cfg, deployer)
 		if err != nil {
 			return fmt.Errorf("failed to create generator: %w", err)
@@ -328,6 +335,7 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 				for _, a := range gen.Accounts() {
 					addrs = append(addrs, a.Address)
 				}
+				probes.Enter("funding accounts")
 				if err := funder.FundAccounts(ctx, cfg, deployer, addrs); err != nil {
 					return fmt.Errorf("failed to fund accounts: %w", err)
 				}
@@ -342,6 +350,7 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 
 		// Set up prewarming if enabled
 		if cfg.Settings.Prewarm {
+			probes.Enter("prewarming accounts")
 			log.Printf("🔥 Creating prewarm generator...")
 			if err := gen.Prewarm(ctx, rng, cfg, snd); err != nil {
 				return fmt.Errorf("gen.Prewarm(): %w", err)
@@ -357,6 +366,16 @@ func runLoadTest(ctx context.Context, cmd *cobra.Command) error {
 		s.SpawnBgNamed("generator", func() error { return gen.Run(ctx, rng, snd) })
 
 		log.Printf("✅ Started dispatcher")
+		// Everything a run needs is up: contracts deployed, accounts funded and
+		// prewarmed, sender and dispatcher running.
+		probes.Ready()
+		// Deferred because every path out of this run leaves service, not only
+		// the signal below. A duration deadline and a failed background worker
+		// both return early, and the run then logs its summary and holds the pod
+		// open for the scrape window — the whole time readiness is meant to
+		// cover. /healthz keeps answering through it, so the kubelet does not
+		// read that hold as a hang.
+		defer probes.NotReady("shutting down")
 
 		// Set up signal handling for graceful shutdown
 		sigChan := make(chan os.Signal, 1)