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

Multi-replica gateways route correctly, but they exposed no capacity
signals, dropped every supervisor session at once on shutdown,
serialized all cross-object mutations fleet-wide on one PostgreSQL
advisory lock, and polled one full sandbox record per watched sandbox
per second on every replica.

Graceful drain: on SIGTERM the gateway closes supervisor admission and
reports "draining" on /readyz and /health (/healthz stays 200).
Gateways with PostgreSQL and a peer endpoint (every Helm install on an
external database) keep the listener open for a 3-second propagation
delay, then close their sessions at most 100 ms apart within 12
seconds; each session keeps serving relays and heartbeats until its
slot. The existing shutdown then runs with up to 10 seconds of
cleanup, 25 seconds in total. A draining owner fails peer relays for
sessions it no longer holds immediately, and disconnect cleanup skips
the endpoint-status lock when a replacement session exists.
Supervisors reset their reconnect backoff after an accepted session
and jitter retries, and a new session receives RelayOpen only after
SessionAccepted is queued.

Capacity metrics: a gateway_metrics module exports per-replica held
supervisor sessions, draining state, pending relays against the 256
and 32 caps, relay rejections, expiries and claim latency, peer
request rate, outcome (with the owner's code before the remap to
UNAVAILABLE) and latency, mutation-lock waits and timeouts, and
watch-poller cost. Labels are bounded, series are zero-initialized,
new latency metrics are histograms, and existing duration summaries
are unchanged.

Scoped mutation locks: hierarchical intention locks on a global key,
a key per workspace, and a key per sandbox replace the single
fleet-wide key. Global settings and platform profiles hold the global
key exclusively; provider and workspace-profile writes hold their
workspace key exclusively; every sandbox mutation, admin or
supervisor, holds only its sandbox key exclusively. Keys are taken in
a process-local RwLock table, then as PostgreSQL advisory locks in
ascending order on one connection from a dedicated 4-connection lock
pool, with one 10-second deadline. Timeouts return UNAVAILABLE with
reason MUTATION_LOCK_TIMEOUT instead of INTERNAL. The global key keeps
its legacy value, so mixed-version rollouts stay mutually exclusive.
DeleteProvider now takes its workspace guard, the settings mutex is
removed, and startup endpoint-status reconciliation takes one sandbox
guard at a time instead of holding the global key for the whole scan.

Watch polling: Store::get_resource_versions reads only id and
resource_version (PostgreSQL id = ANY($2), SQLite IN, 1000 ids per
statement), and the cross-replica poller makes one batched call per
tick with unchanged notification semantics.

Helm: the default termination grace period rises from 5 to 30 seconds,
and an optional autoscaling/v2 HorizontalPodAutoscaler (disabled by
default) omits spec.replicas from the workload when enabled. Chart
validation covers the effective maximum replicas and resource requests
or limits for utilization targets, and the certgen hook pods no longer
match the gateway selector.

Tests and docs: mise run test:rust:postgres runs the ignored
PostgreSQL-backed tests (batched lookups, cross-store lock exclusion,
cancellation, pool bounds, a drain-rate envelope); a Kubernetes HA e2e
test verifies that supervisor sessions move off gateway pods during a
rollout; unit and integration tests cover relay saturation, 5000
watched sandboxes, drain pacing, and lock scopes. The HA guide, a new
gateway metrics reference, the architecture notes, the API errors
reference, and the cluster debugging skills describe the drain
lifecycle, capacity signals, autoscaling, and PostgreSQL connection
sizing.

Part of #3528

Signed-off-by: Emilien Macchi <emacchi@redhat.com>

Author: EmilienM

.agents/skills/helm-dev-environment/SKILL.md

@@ -134,6 +134,10 @@ namespace with a `uri` key. For local manual testing, either create your own
 PostgreSQL Secret or use the e2e PostgreSQL fixture manifest in
 `e2e/kubernetes/postgres-fixture.yaml`.
 
+Gateway pods in the `high-availability` profile use external PostgreSQL, so
+they drain their supervisor sessions when deleted or rolled and can take up to
+30 seconds to terminate.
+
 For the `high-availability` profile, return to the repository root and apply the
 GatewayClass and BackendTrafficPolicy manifest after Skaffold has installed
 Envoy Gateway:
@@ -271,16 +275,24 @@ The kube e2e wrapper creates only one port-forward, to `svc/openshell`; it no
 longer forwards the unauthenticated health listener or runs a `/readyz` e2e
 target. `/readyz` remains covered by server unit/integration tests.
 
-Use `mise run e2e:kubernetes:ha-rebalancing` for full-suite HA coverage. The
-task creates an external PostgreSQL fixture, installs Envoy Gateway, applies
+Use `mise run e2e:kubernetes:ha-rebalancing` for HA coverage. The task creates
+an external PostgreSQL fixture, installs Envoy Gateway, applies
 `deploy/kube/manifests/envoy-gateway-openshell.yaml`, enables the chart
-`GRPCRoute`, and runs the full Kubernetes e2e suite, including
-`kubernetes_ha_rebalancing`. That coverage validates sandbox create/watch and
+`GRPCRoute`, and runs the CLI conformance profile and the
+`kubernetes_ha_rebalancing` tests. That coverage validates sandbox create/watch and
 exec through the Envoy proxy while gateway replicas scale up, scale down, and
 rotate. It also keeps a long-running sandbox alive and runs upload/download
 operations while gateway pods roll, so file sync exercises the same relay retry
 path as interactive sessions.
 
+`supervisor_sessions_redistribute_across_gateway_pod_rolls` restarts the
+gateway Deployment, reads `openshell_server_draining` from the terminating
+pods, and checks that the new pods' `openshell_server_supervisor_sessions` add
+up to the Ready sandbox count. It scrapes each pod through
+`kubectl get --raw /api/v1/namespaces/<ns>/pods/<pod>:9090/proxy/metrics`.
+Gateway pods take up to 30 seconds to terminate because each drains its
+sessions, and the HA tests are allowed ten minutes each.
+
 If you reuse an existing Skaffold cluster for the full kube suite, make sure the
 chart has `server.hostGatewayIP` set so sandbox pods can resolve
 `host.openshell.internal` back to the test host. The e2e wrapper detects this on
@@ -446,6 +458,7 @@ for dependencies still declared in `Chart.yaml`.
 | `deploy/helm/openshell/ci/values-cert-manager.yaml` | cert-manager PKI overlay (opt-in; disables pkiInitJob) |
 | `deploy/helm/openshell/ci/values-gateway.yaml` | Envoy Gateway GRPCRoute + Gateway overlay |
 | `deploy/helm/openshell/ci/values-high-availability.yaml` | HA test overlay (`replicaCount: 2` with external PostgreSQL Secret) |
+| `deploy/helm/openshell/ci/values-autoscaling.yaml` | Render-only overlay for the optional gateway HorizontalPodAutoscaler (helm lint and helm-unittest) |
 | `deploy/helm/openshell/ci/values-keycloak.yaml` | Keycloak OIDC overlay |
 | `deploy/helm/openshell/ci/values-spire.yaml` | SPIFFE/SPIRE provider token grant overlay |
 | `deploy/helm/openshell/ci/values-spire-stack.yaml` | SPIRE hardened chart values for local dev |

.config/nextest.toml

@@ -26,6 +26,10 @@ kubernetes-ha = { max-threads = 1 }
 [[profile.e2e-kubernetes.overrides]]
 filter = "test(/gateway_(scale_and_rollout|pod_rolls)$/)"
 test-group = "kubernetes-ha"
+# Each rolled gateway pod drains its supervisor sessions for up to the 30s
+# termination grace period, so these tests get ten minutes, matching
+# HA_SYNC_TIMEOUT in kubernetes_ha_rebalancing.rs.
+slow-timeout = { period = "60s", terminate-after = 10 }
 
 # Relative to the profile store dir (`e2e/rust/target/nextest/e2e-kubernetes/`).
 [profile.e2e-kubernetes.junit]

TESTING.md

@@ -48,6 +48,23 @@ mise run test:rust     # cargo test --workspace
 
 Rust validation checks tracked Cargo lockfiles; run `mise run rust:lockfiles:check` to check them directly. If one is stale, refresh it with Cargo using its adjacent manifest, review the diff, and commit the update.
 
+### PostgreSQL-backed tests
+
+Tests that need a real PostgreSQL server, such as advisory-lock concurrency
+across two stores, are `#[ignore]`d and named `postgres_*`. Run them with:
+
+```shell
+mise run test:rust:postgres
+```
+
+The task starts a disposable PostgreSQL container with Docker or Podman
+(set `CONTAINER_ENGINE` to choose), runs the tests one at a time, and removes
+the container. Each test works in its own temporary schema. To use your own
+disposable database, set `OPENSHELL_TEST_POSTGRES_URL`. Never point it at a
+database that a running gateway uses: the tests take fleet-wide advisory locks.
+CI does not run these tests; the Kubernetes HA e2e suite covers PostgreSQL end
+to end.
+
 ### Native Windows validation
 
 Use `mise run --skip-tools pre-commit` with the existing Rust/MSVC toolchain.
@@ -403,6 +420,7 @@ Available task variants:
 |---|---|
 | `e2e:kubernetes` | Default Rust e2e against Helm-deployed gateway |
 | `e2e:kubernetes:db` | All database backend scenarios (SQLite + external PostgreSQL) |
+| `e2e:kubernetes:ha-rebalancing` | Two gateway replicas behind Envoy with external PostgreSQL: scale, pod deletion, rollout drain and session redistribution, and file sync during pod rolls |
 | `e2e:kubernetes:sidecar` | Supervisor sidecar topology overlay |
 | `e2e:kubernetes:credential-drivers` | Kubernetes Secrets and Vault credential storage |
 | `e2e:kubernetes:workspace-managed` | Managed workspace mode (auto-created namespaces) |

architecture/gateway.md

@@ -415,10 +415,44 @@ validates the current supervisor session and keeps the in-memory evidence; a
 non-owner never accepts evidence from a stale local session or projects a
 remote session as disconnected.
 
-Nothing redistributes established sessions, so after a rolling restart the last
-surviving replica holds most sessions and a new replica serves none until
-sandboxes reconnect. That skew decays only as sandboxes churn. Client traffic
-stays correct throughout because a non-owner relays to the owner.
+Planned shutdown moves ownership instead of waiting for churn. On SIGTERM a
+gateway closes supervisor admission and reports `draining` on `/readyz`. A
+gateway that uses an external PostgreSQL database and advertises a peer
+endpoint (every Helm chart install on an external database: Deployment pods
+advertise their pod IP and StatefulSet pods their stable DNS name) then keeps
+its listener open for a 3-second propagation delay, so endpoint removal
+reaches kube-proxy and ingress, and closes its supervisor sessions at most 100 ms apart within 12
+seconds. Peers keep following the durable owner record, which changes only
+when a session is closed and reconnects elsewhere. Until its turn, a session keeps serving relays and
+heartbeats. Each supervisor redials the gateway Service, lands on a ready
+replica, and publishes ownership there. Supervisors reset their reconnect
+backoff after an accepted session and jitter failed retries, so a drained
+session normally reconnects within about a second. The old replica demotes the
+sandbox to `Provisioning` when it closes the session, so for that second new
+exec, SSH, and forward requests fail the readiness check with
+`FAILED_PRECONDITION`; requests already routed wait up to 15 seconds for the
+supervisor. A new session receives relay requests only after the gateway has
+sent `SessionAccepted`, and a draining owner answers peer relays for sessions
+it no longer holds with `UNAVAILABLE` at once so the requester re-reads
+ownership. Up to 120 sessions close one every 100 ms; above that the close rate
+is the session count divided by 12 seconds. Each reconnect takes one or two
+short mutation locks on the receiving replicas, so the mutation-lock pool
+bounds how many sessions a replica can drain without queuing. Other gateways
+skip the drain. A lone replica, such as a single-replica StatefulSet, has no
+peer to receive its sessions, so its drain adds up to 15 seconds to shutdown;
+a StatefulSet creates the replacement only after the old pod exits, so the
+restart outage grows by the same amount. Supervisors keep backing off while no
+replica is ready, so some reconnect up to one maximum backoff interval (30
+seconds) after the replacement is ready. Skipping the drain needs a view of
+live replicas, which the gateway does not have yet.
+
+The drain decides when sessions leave, not where they land. Reconnects go to
+whichever replicas are ready at that moment, so a two-replica surge rollout
+leaves the first replacement pod with most sessions (about 62/38), and pod
+deletion, StatefulSet updates, and scale-out leave new replicas nearly empty
+until sandboxes churn. Connect-time placement and rebalancing are future work.
+Client traffic stays correct throughout because a non-owner relays to the
+owner.
 
 File upload and download use tar-over-SSH through the same relay path. A gateway
 pod termination drops the active SSH proxy byte stream, so the CLI retries the
@@ -434,18 +468,54 @@ also trust the chart CA, present the chart-generated client certificate for
 mTLS, and verify the stable gateway Service DNS name even when connecting to a
 Deployment pod IP.
 
-`WatchSandbox` uses the local update bus for same-replica writes. On
-multi-replica backends one shared poller per gateway observes resource-version
-changes made by other replicas and feeds that bus for all local watchers,
-avoiding a database poll per client stream. SQLite deployments do not run the
-poller because they are single-replica and the local bus already sees every
-write.
-
-Mutations whose invariants span sandbox, provider-profile, policy, or provider
-records take a process-local mutex and a shared PostgreSQL advisory lock. The
-database session remains dedicated to the request and closes when the guard is
-dropped, which releases the lock on normal completion, cancellation, or error.
-SQLite deployments use only the local mutex because they are single-replica.
+Each replica exports per-replica capacity signals on its metrics listener:
+held supervisor sessions, draining state, pending relays against the fixed
+limits of 256 per replica and 32 per sandbox, relay rejections and expiries,
+outbound peer RPC outcomes and latency, mutation-lock waits and timeouts, and
+watch-poller cost. Session and pending-relay gauges are held by the registry
+entries themselves, so every removal path keeps them exact, and peer request
+metrics record the owner's status before relay failures are reported to
+clients as `UNAVAILABLE`. Labels are bounded; no metric carries a sandbox,
+channel, endpoint, or replica identifier, because the scrape target already
+identifies the replica.
+
+`WatchSandbox` uses the local update bus for same-replica writes. Every
+PostgreSQL-backed gateway, even with a single replica, runs one shared poller
+that observes writes made by other replicas and feeds that bus for all local
+watchers. Every second it reads the id and `resource_version` of each sandbox
+that has a local watcher in one store call, with one statement per 1000 ids
+and no payload reads, so database load follows the number of distinct watched
+sandboxes on the replica rather than the number of client streams. It
+notifies watchers when it first observes a sandbox, when the version changes,
+and when the row disappears. A failed read keeps the last known versions and
+retries on the next tick. SQLite deployments do not run the poller because
+they are single-replica and the local bus already sees every write.
+
+Mutations whose invariants span sandbox, provider, provider-profile, policy, or
+settings records take a hierarchical mutation guard. It names a global key, one
+key per workspace, and one key per sandbox, each held shared or exclusive.
+Global policy and settings updates and platform-scope profile changes hold the
+global key exclusively. Provider and workspace-scoped profile mutations hold
+the global key shared and their workspace key exclusively. Sandbox-scoped
+mutations, including every supervisor report, hold the global and workspace
+keys shared and their sandbox key exclusively. Unrelated sandboxes proceed
+concurrently, and a provider change still excludes every sandbox mutation in
+its workspace. Each replica takes the keys in a process-local lock table first,
+then, on PostgreSQL, as session-level advisory locks in ascending key order on
+one connection from a dedicated four-connection lock pool. Returning that
+connection runs `pg_advisory_unlock_all()`, and a cancelled acquisition closes
+its session. Acquisition is bounded at 10 seconds, PostgreSQL enforces the
+remaining deadline on every lock wait, and a timeout fails with `UNAVAILABLE`.
+The global key is the legacy cross-object key, so a replica from an earlier
+release, which holds it exclusively for every mutation, still excludes new
+replicas during a rolling upgrade. Lifecycle, driver-watch, and reconcile paths
+take only process-local keys, the global key shared and their sandbox key
+exclusively, and rely on compare-and-swap across replicas. Provisioning-deadline
+reconciliation also holds its workspace key shared, because it re-derives
+configuration from provider and profile records. SQLite deployments use only
+the local table. Startup endpoint-status reconciliation takes each candidate's
+full sandbox-scoped guard, up to four at a time, and never holds a guard across
+the whole scan.
 
 ## API Surface
 
@@ -936,13 +1006,13 @@ coverage:
 | Provider | `MustCreate` | `update_message_cas` | `list_messages` |
 | ProviderProfile | `MustCreate` | `MatchResourceVersion` | `list_messages` |
 | SandboxPolicy | scoped versioning | scoped versioning | scoped query |
-| Settings | `Mutex`-guarded | `Mutex`-guarded | single-row |
+| Settings | mutation-guarded | mutation-guarded | single-row |
 
-Global settings updates use a Tokio `Mutex` to serialize multi-step
-validation within a single gateway process, with CAS on the underlying
-persistence write as defense in depth. In an HA deployment with multiple
-gateways, the Mutex alone would be insufficient. Sandbox-scoped settings
-rely entirely on CAS without a Mutex.
+Global settings and policy updates hold the global mutation key exclusively,
+and sandbox-scoped settings updates hold their sandbox key with the global key
+shared. The precedence check between a sandbox setting and a globally managed
+key therefore cannot interleave with a global change on any replica. Settings
+writes also use CAS as defense in depth.
 
 The `resource_version` is surfaced to clients through `ObjectMeta` in proto
 responses. Provider profiles are the exception: custom profile get/list/export
@@ -952,13 +1022,14 @@ requests also carry an explicit target profile ID; the payload ID must match the
 target so an edited export cannot overwrite a different profile. Database
 migrations backfill existing rows with version 1.
 
-Provider profile imports, updates, and deletes hold the sandbox synchronization
-guard while checking attached-sandbox dynamic token grant ambiguity or in-use
-state and writing the profile record. Sandbox creation with initial providers and
-sandbox provider attach/detach use the same guard, so gateway replicas cannot
-interleave a profile mutation with a sandbox provider-set mutation that would
-leave an ambiguous final dynamic-token state or a deleted custom profile that is
-still referenced by a sandbox.
+Provider profile imports, updates, and deletes hold their workspace mutation
+key exclusively (the global key for platform-scope profiles) while checking
+attached-sandbox dynamic token grant ambiguity or in-use state and writing the
+profile record. Sandbox creation with initial providers and sandbox provider
+attach/detach hold the same workspace key shared plus their sandbox key, so
+gateway replicas cannot interleave a profile mutation with a sandbox
+provider-set mutation that would leave an ambiguous final dynamic-token state or
+a deleted custom profile that is still referenced by a sandbox.
 
 Policy and runtime settings are delivered together through the effective sandbox
 config path. A gateway-global policy can override sandbox-scoped policy. The
@@ -1081,14 +1152,21 @@ The same relay pattern backs interactive SSH, command execution, file sync, and
 local service forwarding. The gateway tracks live sessions in memory and
 persists session records so tokens can expire or be revoked.
 
-Graceful gateway shutdown closes supervisor-session admission before stopping
-local compute. It then signals the remaining control sessions to exit and waits
-up to ten seconds for their cleanup, including conditional deletion of persisted
-ownership. Pending connection setup and sessions already removed from the live
-registry remain tracked until cleanup finishes. This lets a replacement
-supervisor claim ownership immediately after restart without deleting a newer
-replica's claim. An incomplete drain is reported as a shutdown error. Closing
-these control sessions does not stop Kubernetes-owned workloads.
+Graceful gateway shutdown first closes supervisor-session admission and reports
+`draining` from readiness on every gateway. A gateway that uses an external
+PostgreSQL database and advertises a peer endpoint then drains its sessions
+with the listener still open (see HA Supervisor Ownership): 3 seconds of
+propagation delay, then paced closes within 12 seconds. Other gateways skip
+the drain. The gateway then stops its listener and local compute, signals the
+remaining control sessions to exit, and waits up to ten seconds for their
+cleanup, including conditional deletion of persisted ownership. Pending
+connection setup and sessions already removed from the live registry remain
+tracked until cleanup finishes. This lets a replacement supervisor claim
+ownership immediately after restart without deleting a newer replica's claim.
+An incomplete cleanup is reported as a shutdown error. The worst case is 25
+seconds, inside the chart's 30-second termination grace period; the final OTLP
+trace flush runs after that and is not counted. Closing these control sessions
+does not stop Kubernetes-owned workloads.
 
 Relay liveness has two backstops so a reset supervisor session cannot leave a
 request parked forever. The gateway runs server-side HTTP/2 keepalive on
@@ -1268,11 +1346,11 @@ and that span continues incoming W3C trace context when present or starts a new
 trace otherwise. It is named for the RPC and carries the request ID that also
 appears in the gateway's logs — the identifier that lets an operator pivot
 between a trace and its log lines. Store and compute-driver spans become
-children of the request span. Reconciliation, provider refresh, and
-driver-watch loops create their own operation spans because they have no
-inbound request to provide a parent. gRPC status is recorded when response
-trailers arrive. Gateway spans carry resource attributes for the gateway
-identity and configured compute driver.
+children of the request span. Reconciliation, provider refresh, the
+cross-replica watch poller, and driver-watch loops create their own operation
+spans because they have no inbound request to provide a parent. gRPC status is
+recorded when response trailers arrive. Gateway spans carry resource attributes
+for the gateway identity and configured compute driver.
 
 The gateway forwards OTLP configuration, its configured gateway name, and W3C
 trace context to managed external drivers. Built-in drivers use dedicated