diff --git a/health/probes.go b/health/probes.go new file mode 100644 index 0000000..9858a8d --- /dev/null +++ b/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")) +} diff --git a/health/probes_test.go b/health/probes_test.go new file mode 100644 index 0000000..b1f17ed --- /dev/null +++ b/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() +} diff --git a/main.go b/main.go index f22beb0..91d643d 100644 --- a/main.go +++ b/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)