feat(server): add HA capacity metrics, scoped locks, and graceful drain

[EmilienM]

Summary

Multi-replica gateways need capacity signals, bounded watch-query load, and a paced shutdown path. This PR adds per-replica metrics, scoped mutation locks, batched watch polling, and graceful supervisor-session draining, plus an optional Helm HPA.

The PR was initially drafted by Opus 5.5 xhigh and reviewed by GPT-6-Astra xhigh through independent adversarial code and architecture reviews.

Related Issue

Part of #3528. The changes address its seven acceptance criteria, but leave the broader goal of balancing established sessions after scale-out open. Draining controls when sessions leave; connect-time placement or a rebalancer must decide where they land. #3661 covers placement.

The issue has no acceptance or agent-workflow labels. Implementation and this review were directly requested; labels are unchanged.

Changes

• Drain PostgreSQL gateways that advertise a peer endpoint: close supervisor admission, report 503 draining on /readyz and /health, allow 3 seconds for endpoint propagation, then pace session closes over at most 12 seconds. /healthz stays healthy. Supervisor ownership cleanup has a separate 10-second timeout; compute-driver cleanup and final trace export can add time. The Helm termination grace default rises from 5 to 30 seconds.
• Reset supervisor reconnect backoff after an accepted session and add jitter. Queue SessionAccepted before exposing a session to new relay requests. A draining owner fails relay requests promptly once it no longer holds a usable session.
• Export bounded-label metrics for held sessions, drain state, pending relays and capacity, relay rejection/expiry/claim latency, peer RPC outcomes and latency, mutation-lock contention, and watch-poller cost. New latency series are histograms; existing duration summaries stay compatible.
• Replace the fleet-wide mutation lock with hierarchical intention locks. Global policy/settings and platform profiles exclude all mutation scopes; workspace provider/profile writes exclude sandbox writes in that workspace; unrelated sandbox writes can proceed concurrently. Lifecycle paths retain local locking and cross-replica compare-and-swap. The legacy global key preserves exclusion during mixed-version rollouts.
• Use a separate, reusable PostgreSQL lock pool capped at four connections. Lock acquisition has a deadline, cancellation cleanup, and retryable UNAVAILABLE errors with reason MUTATION_LOCK_TIMEOUT. The complete connection return, including SQLx's final ping, has a five-second timeout so a stalled connection cannot permanently consume a pool slot. Provider refresh validates the credential expiry before staging anything and removes staged credential handles when lock acquisition fails. Each replica can open up to 14 database connections: 10 data and 4 lock connections.
• Batch watched sandbox version reads in groups of up to 1,000 IDs, replacing one full-record query per watched sandbox. Startup endpoint-status reconciliation also uses bounded concurrency and per-sandbox guards.
• Add an optional autoscaling/v2 HPA with CPU/memory targets, custom metrics, and conservative scale-down. When enabled, the workload omits spec.replicas; chart validation uses maxReplicas. Certgen hook pods no longer match the gateway workload selector.
• Add PostgreSQL, contention, relay saturation, large-watch-set, and rollout tests. Update the HA guide, metrics reference, architecture overview, API error reference, and related skills. Before terminating an orphaned database session, the troubleshooting skill requires conclusive ownership and confirmation that its gateway can no longer write. No protobuf changes or new gateway TOML settings.

Operational limits:

• Draining moves sessions off a stopping replica; it does not balance the fleet or move existing sessions on scale-out.
• A moving sandbox can briefly report Provisioning, causing new exec/SSH/forward requests to fail with FAILED_PRECONDITION. On Kubernetes, since (Fix) ha sandbox resilience with k8s #3644, a reconnect that takes more than a few seconds can also show the Ready condition as DependenciesNotReady while the workload keeps running. Streams through a stopping pod disconnect when it exits. Older supervisors retain their previous reconnect backoff until their sandboxes are recreated.
• A lone PostgreSQL replica still drains, adding downtime to a StatefulSet restart. helm upgrade --reuse-values retains an older termination grace period unless explicitly overridden.
• Enabling HPA on an existing release can briefly reset the workload to one replica while HPA takes ownership; the chart README describes the migration. HPA runtime behavior still needs validation with a metrics server/adapter.

Testing

Rebased on main at d7f92119, which includes #3644. Writes with explicit resource versions still surface conflicts. Internal writes using version 0, including provider refresh and lifecycle helpers, can retry up to five attempts while their guards remain held. Verification after the rebase and the review follow-ups:

• mise run rust:lint: clean.
• Server unit/integration tests: 1,939 passed. openshell-driver-kubernetes: 271 passed. Both suites and the 17 PostgreSQL tests were rerun during the repeat adversarial review.
• mise run test:rust:postgres: 17 passed against disposable PostgreSQL 17, including the stalled-connection recovery regression.
• mise run helm:test: 204 gateway tests and 6 workspace tests passed; chart ownership check passed.
• mise run docs, docs:build:strict, markdown:lint, and shellcheck -x tasks/scripts/run-postgres-tests.sh passed.
• Kubernetes HA e2e on a local kind cluster (rootless Podman, two replicas, Envoy, external PostgreSQL 17, images built from the runtime changes in this commit): mise run e2e:kubernetes:ha-rebalancing passed all tests: the new supervisor_sessions_redistribute_across_gateway_pod_rolls (kubectl rollout restart with four sandboxes: draining=1 observed on a terminating pod, recorded termination intervals at most 28 seconds, all sandboxes Ready again, per-pod supervisor_sessions summing to the Ready count, exec through each new pod) in 155 seconds, sandbox_exec_rebalances_across_gateway_scale_and_rollout in 219 seconds, and sandbox_file_sync_survives_gateway_pod_rolls in 27 seconds. Timing uses polling observations and does not establish that every pod avoided SIGKILL.
• Both medium review findings reproduced before their fixes: a stalled pool return retained an advisory lock; a timed-out provider refresh left two credential handles where one should remain. The corrected credential test also verifies that the original credential stays usable.
• Repeat review corrected an unsafe backend-termination procedure. Fresh scenario checks cover a live gateway with idle lock holders, a confirmed orphan from a stopped gateway, and ambiguous sessions behind a shared proxy. Both reviewers found no remaining critical/high/medium findings after the correction; mise run pre-commit passed again. Runtime code was unchanged in this pass, so the Kubernetes HA e2e was not repeated.
• mise run ci stops in two unchanged openshell-binary-identity tests (resolves_current_process_from_live_executable and resolver_cache_is_owned_and_only_explicit_clones_share_it) because this environment denies /proc/<pid>/exe access. This is not a full CI pass.

The HA CI lane requires the test:e2e-kubernetes label; PostgreSQL-specific tests currently have no CI job.

• mise run pre-commit passes.
• Unit tests added/updated.
• E2E tests added/updated.

Checklist

• Follows Conventional Commits.
• Single commit, signed off for DCO.
• Architecture docs updated.
• Published docs and related skills updated.

Merge coordination

With #3661 (consistent-hash placement), whichever PR lands second:

• Membership during drain: feat(server): place supervisor sessions by consistent hash #3661's membership worker deregisters on shutdown_rx, which this PR sends only after the drain (up to 15 seconds). Until then, peers keep redirecting drained supervisors to the draining replica, which rejects them with UNAVAILABLE. Pass draining_rx (already in scope in run_server) to spawn_membership_worker instead. Peers can still hold the old ring for up to one 10-second refresh, longer than the 3-second propagation delay, so when a redirected attempt fails without SessionAccepted, the supervisor's fallback hello should keep redirected = true (redirected = redirected && !accepted in the error arm) instead of being redirected back.
• TLS on redirected connections: Deployment peer endpoints use pod IPs, while the chart's gateway certificate covers Service DNS names. feat(server): place supervisor sessions by consistent hash #3661's supervisor redirect must preserve the gateway TLS server name when dialing a pod IP, and needs a TLS-enabled integration test. The existing OPENSHELL_PEER_TLS_SERVER_NAME configures gateway-to-gateway calls; it does not configure the supervisor's client. A successful compile or plaintext test merge does not validate this path.
• Supervisor run_session_loop: keep this PR's single accepted read before the match, drop feat(server): place supervisor sessions by consistent hash #3661's per-arm ready_tx.send_replace(false), keep feat(server): place supervisor sessions by consistent hash #3661's redirect arm and target reset, and use ReconnectBackoff::next_delay in the error arm.
• Removed store API: this PR removes Store::list_by_type. GatewayMemberIndex::live_members should call list_by_type_after(MEMBER_OBJECT_TYPE, None, MEMBER_LIST_LIMIT), which runs the same query without an offset. The textual merge is clean, but the build fails without this change.
• Session setup: keep feat(server): place supervisor sessions by consistent hash #3661's redirect check before track_session; sessions registered after it use register_awaiting_accept and mark_accepts_relays from this PR. The new SupervisorHello test literals use ..Default::default(), so feat(server): place supervisor sessions by consistent hash #3661's redirected field needs no test changes.

#3384 may overlap in Helm values, helpers, and documentation; the new chart keys here are top-level (autoscaling.*, podLifecycle.*).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[EmilienM]

cc @drew

[mrunalp]

/ok to test 14df55d

[github-actions]

Label test:e2e applied, but pull-request/3710 does not exist yet. A maintainer needs to comment /ok to test 14df55d7941a7c7e3f4ca6200fd16f15c6d79e46 to mirror this PR. Once the mirror exists, re-apply the label or re-run Branch E2E Checks from the Actions tab.

[sjenning]

/ok to test 14df55d

[FrostGod]

just a heads up, there is slight design change to optimize the current gateway HA workflow.

#3661
please do ensure this PR takes the new changes into account. most likely it still works
and will review this PR myself, and please feel free to review posted PR.

[EmilienM]

@FrostGod Thanks for the heads-up! I rebased on main (with #3644) and did a local test merge with #3661. They work together fine. There are a few small things for whoever lands second, and I listed them under "Merge coordination" in the description. The main one: #3661's membership worker should stop on draining_rx rather than shutdown_rx, so peers stop redirecting supervisors to a pod that's draining. Let me know if your newer changes affect any of that.

[EmilienM]

I ran another round of adversarial code and architecture reviews with GPT-6-Astra xhigh on this iteration and addressed the findings. The PR is now ready for review.