From 42a0efe983a0ea0b9282347304c33c04688a9a47 Mon Sep 17 00:00:00 2001 From: Emilien Macchi Date: Fri, 25 Sep 2026 06:50:43 -0400 Subject: [PATCH] feat(server): add HA capacity metrics, scoped locks, and graceful drain Multi-replica gateways need capacity signals, paced supervisor handoff, and bounded shared-database contention for production scaling. On SIGTERM every gateway reports 503 "draining" on /readyz and /health. Peer-routed PostgreSQL gateways then drain their supervisor sessions before stopping their listener. Supervisors reset their reconnect backoff after an accepted session and jitter retries, and a new session receives relays only after SessionAccepted is queued. Expose per-replica session, relay, peer-RPC, mutation-lock, and watch-poller metrics with bounded labels. Replace the fleet-wide mutation guard with global/workspace/sandbox intention locks while preserving the legacy global key for rolling upgrades. Use a dedicated four-connection PostgreSQL lock pool with bounded acquisition and return. Lock timeouts fail with UNAVAILABLE, reason MUTATION_LOCK_TIMEOUT, and a one-second retry delay. Provider refresh validates the credential expiry before staging and removes staged credentials when lock acquisition fails. Reconcile startup endpoint status per sandbox and batch watched resource-version reads in groups of 1,000. Raise the Helm termination grace default to 30 seconds and add optional autoscaling/v2 HPA support. Give certgen hook pods their own labels so gateway selectors and the HPA no longer match them. Document placement, rollout, connection-pool, and HPA limitations. Add a test:rust:postgres task with PostgreSQL concurrency and stalled-connection regressions, credential cleanup coverage, and Kubernetes rollout tests. Part of #3528 Signed-off-by: Emilien Macchi --- .agents/skills/helm-dev-environment/SKILL.md | 21 +- .config/nextest.toml | 4 + TESTING.md | 18 + architecture/gateway.md | 166 +- architecture/sandbox.md | 9 +- crates/openshell-server/src/compute/mod.rs | 215 +- .../src/compute/mutation_guard.rs | 1208 +++++++++++ .../src/compute/provisioning_deadline.rs | 19 +- .../openshell-server/src/gateway_metrics.rs | 723 +++++++ crates/openshell-server/src/grpc/mod.rs | 36 +- crates/openshell-server/src/grpc/policy.rs | 108 +- .../src/grpc/policy/endpoint_status.rs | 265 ++- .../src/grpc/policy/endpoint_status_tests.rs | 465 ++++- crates/openshell-server/src/grpc/provider.rs | 305 ++- .../src/grpc/provider_readiness_tests.rs | 2 +- crates/openshell-server/src/grpc/sandbox.rs | 137 +- crates/openshell-server/src/http.rs | 126 +- crates/openshell-server/src/lib.rs | 116 +- .../openshell-server/src/persistence/mod.rs | 108 +- .../src/persistence/mutation_lock.rs | 267 +++ .../src/persistence/mutation_lock_pg_tests.rs | 872 ++++++++ .../src/persistence/postgres.rs | 310 ++- .../src/persistence/sqlite.rs | 64 +- .../src/persistence/test_postgres.rs | 77 + .../openshell-server/src/persistence/tests.rs | 124 ++ .../openshell-server/src/provider_refresh.rs | 109 +- crates/openshell-server/src/sandbox_watch.rs | 610 +++++- .../src/supervisor_session.rs | 1764 +++++++++++++++-- .../tests/supervisor_relay_integration.rs | 153 ++ .../src/supervisor_session.rs | 101 +- deploy/helm/openshell/README.md | 39 +- deploy/helm/openshell/README.md.gotmpl | 30 + .../helm/openshell/ci/values-autoscaling.yaml | 33 + .../ci/values-high-availability.yaml | 5 + .../helm/openshell/templates/_autoscaling.tpl | 64 + deploy/helm/openshell/templates/_helpers.tpl | 24 +- deploy/helm/openshell/templates/certgen.yaml | 4 +- .../helm/openshell/templates/deployment.yaml | 2 + deploy/helm/openshell/templates/hpa.yaml | 42 + .../helm/openshell/templates/statefulset.yaml | 2 + .../openshell/tests/autoscaling_test.yaml | 317 +++ deploy/helm/openshell/tests/certgen_test.yaml | 35 + .../openshell/tests/gateway_config_test.yaml | 6 + deploy/helm/openshell/values.yaml | 51 +- docs/kubernetes/high-availability.mdx | 373 +++- docs/kubernetes/setup.mdx | 4 + docs/observability/gateway-metrics.mdx | 254 +++ docs/sdk/api-errors.mdx | 1 + e2e/rust/tests/kubernetes_ha_rebalancing.rs | 731 ++++++- skills/debug-openshell-cluster/SKILL.md | 147 +- tasks/scripts/run-postgres-tests.sh | 82 + tasks/test.toml | 5 + 52 files changed, 10084 insertions(+), 669 deletions(-) create mode 100644 crates/openshell-server/src/compute/mutation_guard.rs create mode 100644 crates/openshell-server/src/gateway_metrics.rs create mode 100644 crates/openshell-server/src/persistence/mutation_lock.rs create mode 100644 crates/openshell-server/src/persistence/mutation_lock_pg_tests.rs create mode 100644 crates/openshell-server/src/persistence/test_postgres.rs create mode 100644 deploy/helm/openshell/ci/values-autoscaling.yaml create mode 100644 deploy/helm/openshell/templates/_autoscaling.tpl create mode 100644 deploy/helm/openshell/templates/hpa.yaml create mode 100644 deploy/helm/openshell/tests/autoscaling_test.yaml create mode 100644 docs/observability/gateway-metrics.mdx create mode 100755 tasks/scripts/run-postgres-tests.sh diff --git a/.agents/skills/helm-dev-environment/SKILL.md b/.agents/skills/helm-dev-environment/SKILL.md index cb3023f476..266b1f7c45 100644 --- a/.agents/skills/helm-dev-environment/SKILL.md +++ b/.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//pods/: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 | diff --git a/.config/nextest.toml b/.config/nextest.toml index 19b4c21769..ec7bd72738 100644 --- a/.config/nextest.toml +++ b/.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] diff --git a/TESTING.md b/TESTING.md index 7a19ab73c0..91877c9fd0 100644 --- a/TESTING.md +++ b/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) | diff --git a/architecture/gateway.md b/architecture/gateway.md index 45cb5efd70..658f337d05 100644 --- a/architecture/gateway.md +++ b/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,56 @@ 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()`. The complete return, including its +health check, is bounded at five seconds; a stalled connection is discarded +and its pool slot released. 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 +1008,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 +1024,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 +1154,23 @@ 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 drain schedule and +session cleanup budgets total 25 seconds, leaving margin in the chart's +30-second termination grace period for the built-in Kubernetes driver. +Compute-driver cleanup and the final OTLP trace flush are outside those +budgets, so they are not an overall process-exit deadline. 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 +1350,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 diff --git a/architecture/sandbox.md b/architecture/sandbox.md index b1169ca64e..3924748f22 100644 --- a/architecture/sandbox.md +++ b/architecture/sandbox.md @@ -794,9 +794,12 @@ record and restartable storage survive cleanup, including for ephemeral creates. Explicit start is blocked while cleanup is pending, then creates a fresh attempt using the latest configuration. Configuration edits alone never restart an expired sandbox. Legacy provisioning records receive one persisted rollout -window. Cross-object configuration serialization uses the gateway's existing -single-writer guard; enabling concurrent configuration writers still requires -the database-backed invariant work tracked by #1255. +window. Configuration writers take the hierarchical mutation guard described in +`architecture/gateway.md`: sandbox-scoped writes hold their sandbox key, +provider and workspace-profile writes hold their workspace key, and global +policy and settings hold the global key. The deadline scan holds the same keys +process-locally (global and workspace shared, sandbox exclusive) and relies on +compare-and-swap across replicas. Docker startup health remains unready during policy quarantine. A failed probe does not terminate a live provisioning supervisor; the gateway deadline owns diff --git a/crates/openshell-server/src/compute/mod.rs b/crates/openshell-server/src/compute/mod.rs index 4a38361e1b..4f39f35dda 100644 --- a/crates/openshell-server/src/compute/mod.rs +++ b/crates/openshell-server/src/compute/mod.rs @@ -5,9 +5,13 @@ pub mod driver_config; pub mod lease; +mod mutation_guard; pub mod provisioning_deadline; pub mod rootfs_tar; +pub use mutation_guard::MutationScope; +use mutation_guard::{LocalMutationGuard, LocalMutationLocks, MutationGuard}; + use crate::grpc::policy::SANDBOX_SETTINGS_OBJECT_TYPE; use crate::otel_tracing::TraceContextInterceptor; use crate::persistence::{ @@ -196,10 +200,21 @@ impl LifecycleGateRegistry { async fn lock_for(&self, sandbox_id: &str) -> SandboxLifecycleGuard { let gate = self.gate_for(sandbox_id); SandboxLifecycleGuard { + sandbox_id: sandbox_id.to_string(), _guard: gate.lock_owned().await, } } + /// Take the gate only when nobody holds it. Never waits, so a caller may + /// use it while holding the sandbox's local mutation lock. + fn try_lock_for(&self, sandbox_id: &str) -> Option { + let guard = self.gate_for(sandbox_id).try_lock_owned().ok()?; + Some(SandboxLifecycleGuard { + sandbox_id: sandbox_id.to_string(), + _guard: guard, + }) + } + fn gate_for(&self, sandbox_id: &str) -> Arc> { let mut gates = self .gates @@ -227,11 +242,12 @@ impl LifecycleGateRegistry { /// Proof that the current operation holds its sandbox-ID lifecycle gate. /// -/// Lifecycle code must acquire this guard before taking `ComputeRuntime::sync_lock`. -/// Passing it to `lock_global_for_lifecycle` makes that ordering visible at -/// every global-lock acquisition in a lifecycle path. +/// Lifecycle code must acquire this guard before taking the sandbox's local +/// mutation lock (`lock_sandbox_for_lifecycle`). Passing it there makes that +/// ordering visible at every mutation-lock acquisition in a lifecycle path. #[derive(Debug)] pub struct SandboxLifecycleGuard { + sandbox_id: String, _guard: tokio::sync::OwnedMutexGuard<()>, } @@ -636,7 +652,7 @@ pub struct ComputeRuntime { sandbox_watch_bus: SandboxWatchBus, tracing_log_bus: TracingLogBus, supervisor_sessions: Arc, - sync_lock: Arc>, + mutation_locks: Arc, lifecycle_gates: Arc, replica_id: String, /// Gateway-issued staging slots for rootfs tar archives. Shared across @@ -645,13 +661,6 @@ pub struct ComputeRuntime { rootfs_tar_staging: Arc, } -pub struct SandboxSyncGuard { - // Drop the database guard before the local mutex so another local waiter - // cannot race ahead while this replica still owns the cluster-wide lock. - _distributed: crate::persistence::DistributedMutationGuard, - _local: tokio::sync::OwnedMutexGuard<()>, -} - impl fmt::Debug for ComputeRuntime { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ComputeRuntime").finish_non_exhaustive() @@ -734,48 +743,13 @@ impl ComputeRuntime { sandbox_watch_bus, tracing_log_bus, supervisor_sessions, - sync_lock: Arc::new(Mutex::new(())), + mutation_locks: Arc::new(LocalMutationLocks::new()), lifecycle_gates: Arc::new(LifecycleGateRegistry::default()), replica_id: lease::replica_id(), rootfs_tar_staging, }) } - /// Serializes sandbox/provider-profile invariant checks and object writes - /// across gateway replicas. - /// - /// The local mutex preserves lock ordering within one process. `PostgreSQL` - /// deployments also hold a session-level advisory lock for the duration. - pub(crate) async fn sandbox_sync_guard( - &self, - ) -> crate::persistence::PersistenceResult { - let local = self.sync_lock.clone().lock_owned().await; - let distributed = self.store.acquire_distributed_mutation_guard().await?; - Ok(SandboxSyncGuard { - _distributed: distributed, - _local: local, - }) - } - - pub(crate) async fn sandbox_create_guards( - &self, - sandbox_id: &str, - ) -> crate::persistence::PersistenceResult<(SandboxLifecycleGuard, SandboxSyncGuard)> { - let lifecycle_guard = self.lifecycle_gates.lock_for(sandbox_id).await; - let global_guard = self.sandbox_sync_guard().await?; - Ok((lifecycle_guard, global_guard)) - } - - /// Acquires the process-wide lock for code that already holds the - /// sandbox-ID lifecycle gate. The guard parameter documents and enforces - /// that callers acquire locks in lifecycle-gate -> global-lock order. - async fn lock_global_for_lifecycle( - &self, - _lifecycle_guard: &SandboxLifecycleGuard, - ) -> tokio::sync::OwnedMutexGuard<()> { - self.sync_lock.clone().lock_owned().await - } - #[cfg(test)] pub(crate) fn lifecycle_gate_entry_count(&self) -> usize { self.lifecycle_gates.entry_count() @@ -990,8 +964,8 @@ impl ComputeRuntime { launch_authentication: Option>, await_main_process_attachment: bool, ) -> Result { - let (lifecycle_guard, global_guard) = self - .sandbox_create_guards(sandbox.object_id()) + let (lifecycle_guard, mutation_guard) = self + .sandbox_create_guards(sandbox.object_workspace(), sandbox.object_id()) .await .map_err(|error| { crate::grpc::persistence_error_to_status(error, "acquire sandbox mutation lock") @@ -1002,7 +976,7 @@ impl ComputeRuntime { launch_authentication, await_main_process_attachment, lifecycle_guard, - global_guard, + mutation_guard, )) .await } @@ -1015,7 +989,7 @@ impl ComputeRuntime { launch_authentication: Option>, await_main_process_attachment: bool, lifecycle_guard: SandboxLifecycleGuard, - global_guard: SandboxSyncGuard, + mutation_guard: MutationGuard, ) -> Result { self.validate_caller_driver_config( sandbox @@ -1080,7 +1054,7 @@ impl ComputeRuntime { if let Some(metadata) = sandbox.metadata.as_mut() { metadata.resource_version = result.resource_version; } - drop(global_guard); + drop(mutation_guard); if let Some(token) = sandbox_token && let Some(spec) = driver_sandbox.spec.as_mut() @@ -1113,7 +1087,7 @@ impl ComputeRuntime { if let Some(staged) = staged.as_mut() { staged.disarm(); } - let global_guard = self.lock_global_for_lifecycle(&lifecycle_guard).await; + let sandbox_guard = self.lock_sandbox_for_lifecycle(&lifecycle_guard).await; if self.supports_sandbox_authentication() && runtime_identity.is_empty() { let status = Status::internal("compute driver did not return a runtime identity"); @@ -1122,7 +1096,7 @@ impl ComputeRuntime { &sandbox_id, sandbox.object_name(), lifecycle_guard, - global_guard, + sandbox_guard, status, ) .await); @@ -1148,7 +1122,7 @@ impl ComputeRuntime { &sandbox_id, sandbox.object_name(), lifecycle_guard, - global_guard, + sandbox_guard, status, ) .await); @@ -1194,7 +1168,7 @@ impl ComputeRuntime { sandbox_id: &str, sandbox_name: &str, lifecycle_guard: SandboxLifecycleGuard, - global_guard: tokio::sync::OwnedMutexGuard<()>, + sandbox_guard: LocalMutationGuard, original: Status, ) -> Status { let transition = match self @@ -1212,7 +1186,7 @@ impl ComputeRuntime { ); } Err(error) => { - drop(global_guard); + drop(sandbox_guard); let delete_result = self .delete_backend_after_failed_create(sandbox_id, sandbox_name) .await; @@ -1236,7 +1210,7 @@ impl ComputeRuntime { }; self.sandbox_index.update_from_sandbox(&transition.deleting); self.sandbox_watch_bus.notify(sandbox_id); - drop(global_guard); + drop(sandbox_guard); let delete_result = self .delete_backend_after_failed_create(sandbox_id, sandbox_name) @@ -1319,7 +1293,7 @@ impl ComputeRuntime { let sandbox_id = candidate.object_id().to_string(); let sandbox_name = candidate.object_name().to_string(); let lifecycle_guard = self.lifecycle_gates.lock_for(&sandbox_id).await; - let global_guard = self.lock_global_for_lifecycle(&lifecycle_guard).await; + let sandbox_guard = self.lock_sandbox_for_lifecycle(&lifecycle_guard).await; let current = self .store .get_message::(&sandbox_id) @@ -1365,7 +1339,7 @@ impl ComputeRuntime { self.sandbox_watch_bus.notify(&sandbox_id); (previous, stopping) }; - drop(global_guard); + drop(sandbox_guard); // Once the durable transition is committed, request cancellation must // not cancel the driver operation and strand the sandbox in @@ -1424,7 +1398,7 @@ impl ComputeRuntime { match result { Ok(_) => { - let _global_guard = self.lock_global_for_lifecycle(&lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(&lifecycle_guard).await; let latest = self .store .get_message::(&sandbox_id) @@ -1501,7 +1475,7 @@ impl ComputeRuntime { let sandbox_id = candidate.object_id().to_string(); let sandbox_name = candidate.object_name().to_string(); let lifecycle_guard = self.lifecycle_gates.lock_for(&sandbox_id).await; - let global_guard = self.lock_global_for_lifecycle(&lifecycle_guard).await; + let sandbox_guard = self.lock_sandbox_for_lifecycle(&lifecycle_guard).await; let mut current = self .store .get_message::(&sandbox_id) @@ -1635,7 +1609,7 @@ impl ComputeRuntime { } } }; - drop(global_guard); + drop(sandbox_guard); // The durable `Starting` transition commits the operation. Let an // owned worker finish it even if the initiating RPC is canceled. @@ -1779,7 +1753,7 @@ impl ComputeRuntime { .compensate_successful_start(&lifecycle_guard, &starting, &previous, status) .await); } - let global_guard = self.lock_global_for_lifecycle(&lifecycle_guard).await; + let sandbox_guard = self.lock_sandbox_for_lifecycle(&lifecycle_guard).await; let latest = if self.supports_sandbox_authentication() { let driver_name = self.configured_driver_name().to_string(); let persisted = self @@ -1798,7 +1772,7 @@ impl ComputeRuntime { match persisted { Ok(sandbox) => sandbox, Err(error) => { - drop(global_guard); + drop(sandbox_guard); let status = Status::internal(format!( "persist compute runtime identity failed: {error}" )); @@ -1942,7 +1916,7 @@ impl ComputeRuntime { ); } - let _global_guard = self.lock_global_for_lifecycle(lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(lifecycle_guard).await; if self.restore_lifecycle_snapshot(starting, previous).await { original } else { @@ -1960,8 +1934,9 @@ impl ComputeRuntime { /// state before deciding whether the pre-operation snapshot is still true. /// /// A transport error can arrive after the runtime applied stop or start. - /// The driver lookup deliberately runs without the process-wide lock; the - /// exact transition resource version then fences the recovery write. + /// The driver lookup deliberately runs without the sandbox's local + /// mutation lock; the exact transition resource version then fences the + /// recovery write. async fn recover_failed_lifecycle( &self, lifecycle_guard: &SandboxLifecycleGuard, @@ -1977,7 +1952,7 @@ impl ComputeRuntime { ) .await .unwrap_or_else(|_| Err("compute lifecycle reconciliation timed out".into())); - let _global_guard = self.lock_global_for_lifecycle(lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(lifecycle_guard).await; match observed { Ok(Some(snapshot)) if snapshot.id == sandbox_id && snapshot.status.is_some() => { @@ -2172,7 +2147,7 @@ impl ComputeRuntime { target: SandboxDeleteTarget, ) -> Result { let delete_guard = self.lifecycle_gates.lock_for(&target.sandbox_id).await; - let global_guard = self.lock_global_for_lifecycle(&delete_guard).await; + let sandbox_guard = self.lock_sandbox_for_lifecycle(&delete_guard).await; // There is no await between acquiring the initial guards and spawning // the worker. From this commitment point onward, request cancellation @@ -2183,7 +2158,7 @@ impl ComputeRuntime { let request_span = tracing::Span::current(); tokio::spawn( async move { - Box::pin(runtime.delete_sandbox_inner(target, delete_guard, global_guard)).await + Box::pin(runtime.delete_sandbox_inner(target, delete_guard, sandbox_guard)).await } .instrument(request_span), ) @@ -2199,7 +2174,7 @@ impl ComputeRuntime { &self, target: SandboxDeleteTarget, delete_guard: SandboxLifecycleGuard, - guard: tokio::sync::OwnedMutexGuard<()>, + guard: LocalMutationGuard, ) -> Result { let current = self .store @@ -2377,7 +2352,7 @@ impl ComputeRuntime { delete_guard: &SandboxLifecycleGuard, sandbox_id: &str, ) -> bool { - let _guard = self.lock_global_for_lifecycle(delete_guard).await; + let _guard = self.lock_sandbox_for_lifecycle(delete_guard).await; for attempt in 1..=DELETE_PHASE_CAS_RETRY_LIMIT { let record = match self.store.get(Sandbox::object_type(), sandbox_id).await { Ok(Some(record)) => record, @@ -2435,10 +2410,10 @@ impl ComputeRuntime { } /// Removes the sandbox by stable ID only when the expected resource - /// version still owns the row. The caller holds `sync_lock`; a successful - /// delete also removes sandbox-owned records, while successful or - /// already-completed removal clears this replica's index and watch/log - /// buses. + /// version still owns the row. The caller holds this sandbox's local + /// mutation lock; a successful delete also removes sandbox-owned records, + /// while successful or already-completed removal clears this replica's + /// index and watch/log buses. async fn remove_sandbox_record_if_version_locked( &self, sandbox_id: &str, @@ -2502,10 +2477,11 @@ impl ComputeRuntime { /// Resolves an ambiguous driver delete error without overwriting newer /// gateway state. /// - /// The external lookup runs without `sync_lock`. Recovery then uses the - /// exact `Deleting` resource version to apply one of three outcomes: - /// reconcile an observed backend snapshot, remove a confirmed-absent - /// backend, or restore the pre-delete snapshot when lookup is inconclusive. + /// The external lookup runs without the sandbox's local mutation lock. + /// Recovery then uses the exact `Deleting` resource version to apply one of + /// three outcomes: reconcile an observed backend snapshot, remove a + /// confirmed-absent backend, or restore the pre-delete snapshot when lookup + /// is inconclusive. async fn recover_failed_delete( &self, delete_guard: &SandboxLifecycleGuard, @@ -2515,9 +2491,9 @@ impl ComputeRuntime { let sandbox_name = transition.deleting.object_name(); let deleting_resource_version = sandbox_resource_version(&transition.deleting); - // The driver lookup is deliberately outside the process-wide guard. + // The driver lookup is deliberately outside the local mutation lock. let observed = self.get_driver_sandbox(sandbox_id, sandbox_name).await; - let _guard = self.lock_global_for_lifecycle(delete_guard).await; + let _guard = self.lock_sandbox_for_lifecycle(delete_guard).await; match observed { Ok(Some(snapshot)) if snapshot.id == sandbox_id && snapshot.status.is_some() => { @@ -2661,7 +2637,8 @@ impl ComputeRuntime { } } - /// Handles a recovery CAS conflict while the caller holds `sync_lock`. + /// Handles a recovery CAS conflict while the caller holds the sandbox's + /// local mutation lock. /// Another replica may have removed the durable row during the external /// driver lookup; in that case this replica still needs local cleanup. async fn handle_delete_recovery_conflict( @@ -3316,7 +3293,7 @@ impl ComputeRuntime { } async fn mark_sandbox_error(&self, sandbox: &Sandbox, reason: &str, message: &str) { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox.object_id()).await; let sandbox_id = sandbox.object_id().to_string(); let reason = reason.to_string(); let message = message.to_string(); @@ -3359,7 +3336,7 @@ impl ComputeRuntime { /// `Provisioning` with a `Resumed` Ready condition. Returns `true` if the /// store update succeeded. async fn clear_recoverable_error(&self, sandbox: &Sandbox) -> bool { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox.object_id()).await; let sandbox_id = sandbox.object_id().to_string(); match self .store @@ -3711,7 +3688,7 @@ impl ComputeRuntime { } async fn apply_sandbox_update(&self, mut incoming: DriverSandbox) -> Result<(), String> { - let guard = self.sync_lock.lock().await; + let guard = self.lock_sandbox_local(&incoming.id).await; let mut existing = self .store .get(Sandbox::object_type(), &incoming.id) @@ -3732,9 +3709,9 @@ impl ComputeRuntime { // durable phase to Starting. In particular, an old-generation Ready // event followed by its terminal event can otherwise promote and then // stop the new generation before the replacement supervisor connects. - // Release the global watch lock, wait for that lifecycle operation, - // and then reread both the driver and store before applying an - // authoritative observation. Taking the per-sandbox gate only for + // Release this sandbox's local mutation lock, wait for that lifecycle + // operation, and then reread both the driver and store before applying + // an authoritative observation. Taking the per-sandbox gate only for // this ambiguous phase avoids delaying unrelated watch events behind // slow lifecycle operations. let existing_name = existing_sandbox.as_ref().map_or_else( @@ -3744,7 +3721,7 @@ impl ComputeRuntime { drop(guard); let _lifecycle_guard = self.lifecycle_gates.lock_for(&incoming.id).await; let observed = self.get_driver_sandbox(&incoming.id, &existing_name).await; - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(&incoming.id).await; existing = self .store .get(Sandbox::object_type(), &incoming.id) @@ -3866,7 +3843,7 @@ impl ComputeRuntime { sandbox_id: &str, terminal_delivery_finalized: bool, ) -> Result<(), String> { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox_id).await; // A replacement session may already belong to another gateway. Do not // let cleanup from this replica overwrite the replacement's Ready state. @@ -3946,7 +3923,7 @@ impl ComputeRuntime { instance_id: Option<&str>, terminal_delivery_finalized: bool, ) -> Result<(), String> { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox_id).await; let existing = self .store .get_message::(sandbox_id) @@ -4090,7 +4067,7 @@ impl ComputeRuntime { instance_id: &str, exit_code: i32, ) -> Result<(), String> { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox_id).await; let Some(existing) = self .store .get_message::(sandbox_id) @@ -4167,7 +4144,7 @@ impl ComputeRuntime { sandbox_id: &str, instance_id: &str, ) -> Result<(), String> { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox_id).await; let Some(sandbox) = self .store .get_message::(sandbox_id) @@ -4230,7 +4207,7 @@ impl ComputeRuntime { } async fn apply_deleted(&self, sandbox_id: &str) -> Result<(), String> { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(sandbox_id).await; self.apply_deleted_locked(sandbox_id).await } @@ -4318,9 +4295,9 @@ impl ComputeRuntime { /// The gate check is synchronous, but the actual `DeleteSandbox` RPC is /// always deferred to a background task, never awaited inline: both /// call sites run while holding a broader lock (the watch loop's - /// sequential event processing; the prune sweep's gateway-wide - /// `sync_lock`), and a slow or stuck driver call must never block that - /// wider scope. The gate itself is held for the background call's + /// sequential event processing; the prune sweep's local mutation lock + /// for this sandbox), and a slow or stuck driver call must never block + /// that wider scope. The gate itself is held for the background call's /// duration, so this still can't race a concurrent request-side /// operation — only the potentially-slow RPC is backgrounded. fn spawn_driver_sandbox_cleanup(&self, sandbox_id: &str, sandbox_name: &str) { @@ -4490,7 +4467,7 @@ impl ComputeRuntime { delete_guard: &SandboxLifecycleGuard, sandbox_id: &str, ) -> Result { - let _guard = self.lock_global_for_lifecycle(delete_guard).await; + let _guard = self.lock_sandbox_for_lifecycle(delete_guard).await; let record = self .store .get(Sandbox::object_type(), sandbox_id) @@ -4521,7 +4498,7 @@ impl ComputeRuntime { sweep_started_at_ms: i64, ) -> Result<(), String> { let expected_resource_version = { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(&snapshot.id).await; let Some(existing) = self .store .get(Sandbox::object_type(), &snapshot.id) @@ -4544,7 +4521,7 @@ impl ComputeRuntime { return Ok(()); }; - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(&snapshot.id).await; let Some(existing) = self .store .get(Sandbox::object_type(), &snapshot.id) @@ -4570,7 +4547,7 @@ impl ComputeRuntime { grace_ms: i64, ) -> Result<(), String> { let (sandbox_id, sandbox_name, expected_resource_version, age_ms) = { - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(&record.id).await; let Some(current_record) = self .store .get(Sandbox::object_type(), &record.id) @@ -4602,7 +4579,7 @@ impl ComputeRuntime { let current = self.get_driver_sandbox(&sandbox_id, &sandbox_name).await?; - let _guard = self.sync_lock.lock().await; + let _guard = self.lock_sandbox_local(&sandbox_id).await; let Some(current_record) = self .store .get(Sandbox::object_type(), &sandbox_id) @@ -4677,10 +4654,10 @@ impl ComputeRuntime { ); // The driver's own snapshot never reported this sandbox, so no // request-side DeleteSandbox call is coming for it either — release - // driver-owned resources in the background. This function holds - // `sync_lock` (the gateway-wide state guard) through the rest of its - // body, so the driver call must not be awaited here: doing so would - // block every other sandbox operation gateway-wide on a single, + // driver-owned resources in the background. This function holds this + // sandbox's local mutation lock through the rest of its body, so the + // driver call must not be awaited here: doing so would block every + // operation on this sandbox and any global mutation on a single, // potentially slow or stuck driver RPC. self.spawn_driver_sandbox_cleanup(&sandbox_id, &sandbox_name); self.apply_deleted_if_version_locked(&sandbox, expected_resource_version) @@ -6197,7 +6174,7 @@ pub fn new_test_runtime_with_driver( sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), - sync_lock: Arc::new(Mutex::new(())), + mutation_locks: Arc::new(LocalMutationLocks::new()), lifecycle_gates: Arc::new(LifecycleGateRegistry::default()), replica_id: "test-replica".to_string(), rootfs_tar_staging: Arc::new(rootfs_tar::RootfsTarStagingRegistry::disabled()), @@ -7244,7 +7221,7 @@ mod tests { sandbox_watch_bus: SandboxWatchBus::new(), tracing_log_bus: TracingLogBus::new(), supervisor_sessions: Arc::new(SupervisorSessionRegistry::new()), - sync_lock: Arc::new(Mutex::new(())), + mutation_locks: Arc::new(LocalMutationLocks::new()), lifecycle_gates: Arc::new(LifecycleGateRegistry::default()), replica_id: "test-replica".to_string(), rootfs_tar_staging: Arc::new(rootfs_tar::RootfsTarStagingRegistry::disabled()), @@ -11055,8 +11032,8 @@ mod tests { ); // The driver call is backgrounded (see `spawn_driver_sandbox_cleanup`) - // so the prune sweep never awaits it while holding the gateway-wide - // sync_lock; wait for it to actually land before asserting on it. + // so the prune sweep never awaits it while holding the sandbox's local + // mutation lock; wait for it to actually land before asserting on it. tokio::time::timeout(Duration::from_secs(1), driver.delete_started.notified()) .await .expect("background driver cleanup did not run"); @@ -11069,7 +11046,7 @@ mod tests { #[tokio::test] async fn prune_sweep_does_not_block_on_a_stuck_driver_delete_call() { // Regression test: the prune sweep's driver cleanup must not be - // awaited while holding `sync_lock` (the gateway-wide state guard). + // awaited while holding the sandbox's local mutation lock. // Block the driver's delete call indefinitely and confirm the sweep // itself still completes promptly and removes the store record. let driver = ControlledDriver::new(); @@ -13923,7 +13900,7 @@ mod tests { let runtime = test_runtime(driver.clone()).await; let sandbox = seed_provisioning_attempt(&runtime).await; let gate = runtime.lifecycle_gates.lock_for("sb-ttl").await; - let global = runtime.lock_global_for_lifecycle(&gate).await; + let sandbox_guard = runtime.lock_sandbox_for_lifecycle(&gate).await; assert!( runtime .claim_provisioning_timeout(&sandbox, 299_999) @@ -13936,7 +13913,7 @@ mod tests { .await .unwrap() .unwrap(); - drop(global); + drop(sandbox_guard); assert_eq!(expired.phase(), i32::from(SandboxPhase::Error)); assert!( ready_condition(&expired) @@ -14030,13 +14007,13 @@ mod tests { let runtime = test_runtime(driver.clone()).await; let sandbox = seed_provisioning_attempt(&runtime).await; let gate = runtime.lifecycle_gates.lock_for("sb-ttl").await; - let global = runtime.lock_global_for_lifecycle(&gate).await; + let sandbox_guard = runtime.lock_sandbox_for_lifecycle(&gate).await; let expired = runtime .claim_provisioning_timeout(&sandbox, 300_000) .await .unwrap() .unwrap(); - drop(global); + drop(sandbox_guard); runtime .reclaim_provisioning_timeout(&expired, &gate) .await @@ -14284,13 +14261,13 @@ mod tests { let runtime = test_runtime(driver.clone()).await; let sandbox = seed_provisioning_attempt(&runtime).await; let gate = runtime.lifecycle_gates.lock_for("sb-ttl").await; - let global = runtime.lock_global_for_lifecycle(&gate).await; + let sandbox_guard = runtime.lock_sandbox_for_lifecycle(&gate).await; let expired = runtime .claim_provisioning_timeout(&sandbox, 300_000) .await .unwrap() .unwrap(); - drop(global); + drop(sandbox_guard); runtime .reclaim_provisioning_timeout(&expired, &gate) .await diff --git a/crates/openshell-server/src/compute/mutation_guard.rs b/crates/openshell-server/src/compute/mutation_guard.rs new file mode 100644 index 0000000000..b0bfbbe395 --- /dev/null +++ b/crates/openshell-server/src/compute/mutation_guard.rs @@ -0,0 +1,1208 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Hierarchical mutation guards: a process-local keyed lock table plus, on +//! `PostgreSQL`, the matching advisory locks on the dedicated lock pool. +//! +//! See [`crate::persistence::mutation_lock`] for the key hierarchy and the +//! ordering rules. In short: lifecycle gate first, then local keys in +//! ascending order, then `PostgreSQL` keys in ascending order on one +//! connection. A task never acquires a guard while it holds one. + +use super::{ComputeRuntime, SandboxLifecycleGuard}; +use crate::gateway_metrics::{self, LockScope}; +use crate::grpc::workspace::DEFAULT_WORKSPACE_NAME; +use crate::persistence::mutation_lock::MUTATION_LOCK_TIMEOUT; +use crate::persistence::{ + DistributedMutationGuard, LockMode, MutationLockKey, MutationLockSet, PersistenceError, + PersistenceResult, +}; +use openshell_core::ObjectWorkspace; +use openshell_core::proto::Sandbox; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex as StdMutex, Weak}; +use std::time::Duration; +use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock}; +use tracing::warn; + +/// What a guarded mutation reads and writes, which selects its lock set. +/// +/// Every mutation whose invariant spans several persisted objects must take +/// the narrowest scope that still conflicts with every writer of the objects +/// it validates against. Global settings and policy writers and +/// platform-scope profile writers hold the global key exclusively; provider +/// and workspace-scoped profile writers hold their workspace key exclusively; +/// sandbox-scoped settings and policy writers hold only their sandbox key +/// exclusively. A new cross-object writer needs a scope from the same table. +#[derive(Clone, Copy, Debug)] +pub enum MutationScope<'a> { + /// Global policy/settings and platform-scope profiles. Excludes every + /// other scope fleet-wide. + Global, + /// Provider and workspace-scoped profile mutations. `""` (platform) + /// behaves as `Global`. + Workspace(&'a str), + /// Any mutation of one sandbox's records, admin or supervisor. + Sandbox { + workspace: &'a str, + sandbox_id: &'a str, + }, +} + +impl<'a> MutationScope<'a> { + pub(crate) const fn sandbox(workspace: &'a str, sandbox_id: &'a str) -> Self { + Self::Sandbox { + workspace, + sandbox_id, + } + } + + /// Keys and modes of this scope: + /// + /// - `Global` and `Workspace("")`: X(global). + /// - `Workspace(ws)`: S(global) X(workspace). + /// - `Sandbox`: S(global) S(workspace) X(sandbox). A legacy sandbox with + /// an empty workspace locks the default workspace, where its providers + /// resolve. + pub(crate) fn lock_set(&self) -> MutationLockSet { + let mut set = MutationLockSet::default(); + match *self { + Self::Global => set.insert(MutationLockKey::Global, LockMode::Exclusive), + Self::Workspace("") => { + set.insert(MutationLockKey::Global, LockMode::Exclusive); + } + Self::Workspace(workspace) => { + set.insert(MutationLockKey::Global, LockMode::Shared); + set.insert(MutationLockKey::Workspace(workspace), LockMode::Exclusive); + } + Self::Sandbox { + workspace, + sandbox_id, + } => { + set.insert(MutationLockKey::Global, LockMode::Shared); + set.insert( + MutationLockKey::Workspace(sandbox_workspace_key(workspace)), + LockMode::Shared, + ); + set.insert(MutationLockKey::Sandbox(sandbox_id), LockMode::Exclusive); + } + } + set + } + + /// The metric label of this scope. + pub(crate) const fn lock_scope(&self) -> LockScope { + match self { + Self::Global => LockScope::Global, + Self::Workspace(workspace) if workspace.is_empty() => LockScope::Global, + Self::Workspace(_) => LockScope::Workspace, + Self::Sandbox { .. } => LockScope::Sandbox, + } + } +} + +/// Workspace key of a sandbox scope. Legacy sandboxes may carry an empty +/// workspace; their providers resolve in the default workspace. +fn sandbox_workspace_key(workspace: &str) -> &str { + if workspace.is_empty() { + DEFAULT_WORKSPACE_NAME + } else { + workspace + } +} + +/// Process-local table of mutation lock keys. +/// +/// Entries are weak, so a key disappears once no guard holds it and no +/// acquisition waits on it. Tokio's `RwLock` is fair and write-preferring: a +/// queued exclusive request blocks later shared requests on the same key. +#[derive(Debug)] +pub struct LocalMutationLocks { + entries: StdMutex>>>, + timeout_ms: AtomicU64, +} + +impl LocalMutationLocks { + pub(crate) fn new() -> Self { + Self { + entries: StdMutex::new(HashMap::new()), + timeout_ms: AtomicU64::new(duration_millis(MUTATION_LOCK_TIMEOUT)), + } + } + + fn lock_for(&self, key: i64) -> Arc> { + let mut entries = self + .entries + .lock() + .expect("mutation lock registry lock poisoned"); + entries.retain(|_, lock| lock.strong_count() > 0); + + if let Some(lock) = entries.get(&key).and_then(Weak::upgrade) { + return lock; + } + + let lock = Arc::new(RwLock::new(())); + entries.insert(key, Arc::downgrade(&lock)); + lock + } + + /// Acquire `set` in ascending key order. Dropping the future releases the + /// keys already taken and leaves no queue entry behind. + async fn acquire(&self, set: &MutationLockSet) -> LocalMutationGuard { + let mut guards = Vec::new(); + for (key, mode) in set.iter() { + let lock = self.lock_for(key); + guards.push(match mode { + LockMode::Shared => LocalKeyGuard::Shared { + _guard: lock.read_owned().await, + }, + LockMode::Exclusive => LocalKeyGuard::Exclusive { + _guard: lock.write_owned().await, + }, + }); + } + LocalMutationGuard { _guards: guards } + } + + fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.load(Ordering::Relaxed)) + } + + #[cfg(test)] + pub(crate) fn set_timeout_for_tests(&self, timeout: Duration) { + self.timeout_ms + .store(duration_millis(timeout), Ordering::Relaxed); + } + + /// Entries in the table, including released keys that `lock_for` has not + /// pruned yet. + #[cfg(test)] + pub(crate) fn entry_count(&self) -> usize { + self.entries + .lock() + .expect("mutation lock registry lock poisoned") + .len() + } + + /// Keys still held by a guard or awaited by an acquisition. Does not + /// prune, so it cannot hide a leak in `lock_for`. + #[cfg(test)] + pub(crate) fn live_entry_count(&self) -> usize { + self.entries + .lock() + .expect("mutation lock registry lock poisoned") + .values() + .filter(|lock| lock.strong_count() > 0) + .count() + } +} + +fn duration_millis(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} + +enum LocalKeyGuard { + Shared { _guard: OwnedRwLockReadGuard<()> }, + Exclusive { _guard: OwnedRwLockWriteGuard<()> }, +} + +/// Process-local keys of one mutation or lifecycle operation. +#[must_use = "dropping the guard releases the local mutation locks"] +pub struct LocalMutationGuard { + _guards: Vec, +} + +/// Local and, on `PostgreSQL`, distributed keys of one guarded mutation. +#[must_use = "dropping the guard releases the mutation locks"] +pub struct MutationGuard { + // Field order is drop order: the database guard goes first so its + // connection returns to the lock pool (and is unlocked) as early as + // possible. + _distributed: DistributedMutationGuard, + _local: LocalMutationGuard, +} + +impl ComputeRuntime { + /// Serialize a cross-object mutation against every conflicting mutation + /// on this and, on `PostgreSQL`, every other replica. + /// + /// One `MUTATION_LOCK_TIMEOUT` deadline covers the local keys, the + /// lock-pool connection, and the advisory locks. Missing it fails with + /// [`PersistenceError::LockTimeout`], after which nothing was written. + pub(crate) async fn mutation_guard( + &self, + scope: MutationScope<'_>, + ) -> PersistenceResult { + let started = tokio::time::Instant::now(); + let deadline = started + self.mutation_locks.timeout(); + let result = self + .acquire_mutation_guard(&scope.lock_set(), deadline) + .await; + match &result { + Ok(_) => gateway_metrics::record_lock_wait(scope.lock_scope(), started.elapsed()), + Err(PersistenceError::LockTimeout(detail)) => { + gateway_metrics::record_lock_timeout(scope.lock_scope()); + warn!( + scope = scope.lock_scope().label(), + waited_ms = duration_millis(started.elapsed()), + detail = %detail, + "mutation lock acquisition timed out" + ); + } + Err(_) => {} + } + result + } + + async fn acquire_mutation_guard( + &self, + set: &MutationLockSet, + deadline: tokio::time::Instant, + ) -> PersistenceResult { + let local = tokio::time::timeout_at(deadline, self.mutation_locks.acquire(set)) + .await + .map_err(|_| { + PersistenceError::LockTimeout("waiting for a local mutation lock".into()) + })?; + let distributed = self + .store + .acquire_distributed_mutation_guard(set, deadline) + .await?; + Ok(MutationGuard { + _distributed: distributed, + _local: local, + }) + } + + /// Sandbox-scoped guard for paths that know only the sandbox id, such as + /// supervisor reports. + /// + /// A sandbox's workspace never changes, so one read before locking + /// derives the key set. Callers must re-read the sandbox after locking and + /// never validate against this read. Returns `Ok(None)` when the sandbox + /// does not exist. + pub(crate) async fn sandbox_mutation_guard_by_id( + &self, + sandbox_id: &str, + ) -> PersistenceResult> { + let Some(sandbox) = self.store.get_message::(sandbox_id).await? else { + return Ok(None); + }; + self.mutation_guard(MutationScope::sandbox( + sandbox.object_workspace(), + sandbox_id, + )) + .await + .map(Some) + } + + /// Lifecycle gate, then the sandbox-scoped mutation guard, for a new + /// sandbox. + pub(crate) async fn sandbox_create_guards( + &self, + workspace: &str, + sandbox_id: &str, + ) -> PersistenceResult<(SandboxLifecycleGuard, MutationGuard)> { + let lifecycle_guard = self.lifecycle_gates.lock_for(sandbox_id).await; + let mutation_guard = self + .mutation_guard(MutationScope::sandbox(workspace, sandbox_id)) + .await?; + Ok((lifecycle_guard, mutation_guard)) + } + + /// Local S(global) X(sandbox) for code that already holds the sandbox's + /// lifecycle gate. The guard parameter documents and enforces the + /// lifecycle-gate -> mutation-lock order. + pub(super) async fn lock_sandbox_for_lifecycle( + &self, + lifecycle_guard: &SandboxLifecycleGuard, + ) -> LocalMutationGuard { + self.lock_sandbox_local(&lifecycle_guard.sandbox_id).await + } + + /// Local S(global) X(sandbox) for lifecycle, driver-watch, and reconcile + /// paths. They write only this sandbox and its owned records and rely on + /// compare-and-swap across replicas, so they take no database lock and + /// never exclude another sandbox or a provider writer. + pub(super) async fn lock_sandbox_local(&self, sandbox_id: &str) -> LocalMutationGuard { + self.mutation_locks + .acquire(&MutationLockSet::sandbox_lifecycle(sandbox_id)) + .await + } + + /// Local S(global) S(workspace) X(sandbox), for provisioning-deadline + /// reconciliation, which re-derives configuration from provider and + /// profile records and must not interleave with their local writers. + pub(super) async fn lock_sandbox_local_in_workspace( + &self, + workspace: &str, + sandbox_id: &str, + ) -> LocalMutationGuard { + self.mutation_locks + .acquire(&MutationScope::sandbox(workspace, sandbox_id).lock_set()) + .await + } + + /// Shorten the mutation lock deadline of this runtime and its clones. + #[cfg(test)] + pub(crate) fn set_mutation_lock_timeout_for_tests(&self, timeout: Duration) { + self.mutation_locks.set_timeout_for_tests(timeout); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::gateway_metrics::MetricsCapture; + use crate::persistence::Store; + use crate::persistence::mutation_lock::GLOBAL_MUTATION_LOCK_KEY; + use crate::persistence::test_postgres::TestSchema; + use openshell_core::GetResourceVersion; + use openshell_core::proto::SandboxPhase; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + use std::sync::atomic::{AtomicBool, AtomicIsize}; + use tokio::task::JoinHandle; + use uuid::Uuid; + + const BLOCKED_FOR: Duration = Duration::from_millis(100); + const PROCEEDS_WITHIN: Duration = Duration::from_secs(5); + + async fn test_runtime() -> ComputeRuntime { + let store = Arc::new( + Store::connect("sqlite::memory:?cache=shared") + .await + .expect("in-memory store"), + ); + super::super::new_test_runtime_for_driver(store, "test").await + } + + fn spawn_guard( + runtime: &ComputeRuntime, + scope: MutationScope<'static>, + ) -> JoinHandle { + let runtime = runtime.clone(); + tokio::spawn(async move { + runtime + .mutation_guard(scope) + .await + .expect("mutation guard acquired") + }) + } + + fn spawn_local(runtime: &ComputeRuntime, sandbox_id: &'static str) -> JoinHandle<()> { + let runtime = runtime.clone(); + tokio::spawn(async move { + drop(runtime.lock_sandbox_local(sandbox_id).await); + }) + } + + async fn assert_blocked(handle: &mut JoinHandle, what: &str) { + assert!( + tokio::time::timeout(BLOCKED_FOR, handle).await.is_err(), + "{what} should wait" + ); + } + + async fn assert_proceeds(handle: JoinHandle, what: &str) -> T { + tokio::time::timeout(PROCEEDS_WITHIN, handle) + .await + .unwrap_or_else(|_| panic!("{what} should proceed")) + .expect("guard task") + } + + fn keys(entries: &[(MutationLockKey<'_>, LockMode)]) -> Vec<(i64, LockMode)> { + let mut keys: Vec<_> = entries + .iter() + .map(|(key, mode)| (key.advisory_key(), *mode)) + .collect(); + keys.sort_unstable(); + keys + } + + #[test] + fn scope_lock_sets_follow_the_hierarchy() { + use LockMode::{Exclusive, Shared}; + use MutationLockKey::{Global, Sandbox as SandboxKey, Workspace}; + + let cases = [ + (MutationScope::Global, keys(&[(Global, Exclusive)])), + (MutationScope::Workspace(""), keys(&[(Global, Exclusive)])), + ( + MutationScope::Workspace("team-a"), + keys(&[(Global, Shared), (Workspace("team-a"), Exclusive)]), + ), + ( + MutationScope::sandbox("team-a", "sb-1"), + keys(&[ + (Global, Shared), + (Workspace("team-a"), Shared), + (SandboxKey("sb-1"), Exclusive), + ]), + ), + ( + MutationScope::sandbox("", "sb-1"), + keys(&[ + (Global, Shared), + (Workspace("default"), Shared), + (SandboxKey("sb-1"), Exclusive), + ]), + ), + ]; + for (scope, expected) in cases { + assert_eq!( + scope.lock_set().iter().collect::>(), + expected, + "{scope:?}" + ); + } + assert_eq!( + MutationLockSet::sandbox_lifecycle("sb-1") + .iter() + .collect::>(), + keys(&[(Global, Shared), (SandboxKey("sb-1"), Exclusive)]) + ); + assert!( + MutationScope::Global + .lock_set() + .iter() + .eq([(GLOBAL_MUTATION_LOCK_KEY, Exclusive)]) + ); + } + + #[test] + fn scope_labels_map_platform_workspace_to_global() { + assert_eq!(MutationScope::Global.lock_scope(), LockScope::Global); + assert_eq!(MutationScope::Workspace("").lock_scope(), LockScope::Global); + assert_eq!( + MutationScope::Workspace("team-a").lock_scope(), + LockScope::Workspace + ); + assert_eq!( + MutationScope::sandbox("", "sb-1").lock_scope(), + LockScope::Sandbox + ); + } + + #[tokio::test] + async fn sandbox_guards_for_different_sandboxes_proceed_concurrently() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::sandbox("w1", "a")) + .await + .unwrap(); + + let other = spawn_guard(&runtime, MutationScope::sandbox("w1", "b")); + drop(assert_proceeds(other, "a different sandbox in the same workspace").await); + drop(held); + } + + #[tokio::test] + async fn same_sandbox_guard_waits_until_release() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::sandbox("w1", "a")) + .await + .unwrap(); + + let mut same = spawn_guard(&runtime, MutationScope::sandbox("w1", "a")); + assert_blocked(&mut same, "the same sandbox").await; + drop(held); + drop(assert_proceeds(same, "the same sandbox after release").await); + } + + #[tokio::test] + async fn workspace_guard_blocks_sandbox_scope_in_that_workspace_only() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::Workspace("w1")) + .await + .unwrap(); + + let mut same_workspace = spawn_guard(&runtime, MutationScope::sandbox("w1", "a")); + assert_blocked(&mut same_workspace, "a sandbox in the held workspace").await; + let other_sandbox = spawn_guard(&runtime, MutationScope::sandbox("w2", "b")); + drop(assert_proceeds(other_sandbox, "a sandbox in another workspace").await); + let other_workspace = spawn_guard(&runtime, MutationScope::Workspace("w2")); + drop(assert_proceeds(other_workspace, "another workspace").await); + + drop(held); + drop(assert_proceeds(same_workspace, "the sandbox after release").await); + } + + #[tokio::test] + async fn global_guard_blocks_every_scope_and_lifecycle_lock() { + let runtime = test_runtime().await; + let held = runtime.mutation_guard(MutationScope::Global).await.unwrap(); + + let mut waiting_guards = vec![ + spawn_guard(&runtime, MutationScope::Global), + spawn_guard(&runtime, MutationScope::Workspace("")), + spawn_guard(&runtime, MutationScope::Workspace("w1")), + spawn_guard(&runtime, MutationScope::sandbox("w1", "a")), + ]; + for waiting in &mut waiting_guards { + assert_blocked(waiting, "a guard behind the global guard").await; + } + let mut lifecycle = spawn_local(&runtime, "b"); + assert_blocked(&mut lifecycle, "a lifecycle lock behind the global guard").await; + let reconcile_runtime = runtime.clone(); + let mut reconcile = tokio::spawn(async move { + drop( + reconcile_runtime + .lock_sandbox_local_in_workspace("w1", "c") + .await, + ); + }); + assert_blocked(&mut reconcile, "a reconcile lock behind the global guard").await; + + drop(held); + for waiting in waiting_guards { + drop(assert_proceeds(waiting, "a guard after release").await); + } + assert_proceeds(lifecycle, "the lifecycle lock after release").await; + assert_proceeds(reconcile, "the reconcile lock after release").await; + } + + #[tokio::test] + async fn queued_global_guard_is_not_starved() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .unwrap(); + + let mut global = spawn_guard(&runtime, MutationScope::Global); + assert_blocked(&mut global, "the global guard behind a sandbox guard").await; + let mut later = spawn_guard(&runtime, MutationScope::sandbox("w", "b")); + assert_blocked(&mut later, "a sandbox guard queued behind the global guard").await; + + drop(held); + let global = assert_proceeds(global, "the queued global guard").await; + assert_blocked(&mut later, "a sandbox guard while the global guard holds").await; + drop(global); + drop(assert_proceeds(later, "the later sandbox guard").await); + } + + #[tokio::test] + async fn lifecycle_lock_excludes_same_sandbox_only() { + let runtime = test_runtime().await; + let held = runtime.lock_sandbox_local("a").await; + + let mut same = spawn_guard(&runtime, MutationScope::sandbox("w", "a")); + assert_blocked(&mut same, "the sandbox held by a lifecycle lock").await; + let other = spawn_guard(&runtime, MutationScope::sandbox("w", "b")); + drop(assert_proceeds(other, "another sandbox").await); + let provider = spawn_guard(&runtime, MutationScope::Workspace("w")); + drop(assert_proceeds(provider, "a provider writer").await); + + drop(held); + drop(assert_proceeds(same, "the sandbox after release").await); + } + + #[tokio::test] + async fn legacy_empty_workspace_sandbox_conflicts_with_default_workspace_writer() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::Workspace("default")) + .await + .unwrap(); + + let mut legacy = spawn_guard(&runtime, MutationScope::sandbox("", "a")); + assert_blocked( + &mut legacy, + "a legacy sandbox behind a default-workspace writer", + ) + .await; + drop(held); + drop(assert_proceeds(legacy, "the legacy sandbox after release").await); + } + + #[tokio::test] + async fn local_registry_drops_released_entries() { + let runtime = test_runtime().await; + let sandbox = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .unwrap(); + let lifecycle = runtime.lock_sandbox_local("b").await; + assert_eq!(runtime.mutation_locks.entry_count(), 4); + + drop(sandbox); + drop(lifecycle); + assert_eq!(runtime.mutation_locks.live_entry_count(), 0); + + // The next acquisition prunes the released keys, so the table holds + // only the new guard's global and sandbox keys. + let fresh = runtime.lock_sandbox_local("c").await; + assert_eq!(runtime.mutation_locks.entry_count(), 2); + drop(fresh); + } + + #[tokio::test] + async fn cancelled_acquisition_leaves_no_queue_entry() { + let runtime = test_runtime().await; + let held = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .unwrap(); + let mut waiter = spawn_guard(&runtime, MutationScope::sandbox("w", "a")); + assert_blocked(&mut waiter, "the waiter").await; + waiter.abort(); + let Err(error) = waiter.await else { + panic!("the aborted waiter should not acquire the guard"); + }; + assert!(error.is_cancelled()); + drop(held); + + let next = spawn_guard(&runtime, MutationScope::sandbox("w", "a")); + drop(assert_proceeds(next, "a new acquisition after the cancelled one").await); + assert_eq!(runtime.mutation_locks.live_entry_count(), 0); + } + + #[tokio::test] + async fn local_timeout_returns_lock_timeout_and_unavailable() { + let runtime = test_runtime().await; + runtime.set_mutation_lock_timeout_for_tests(Duration::from_millis(50)); + let held = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .unwrap(); + + let Err(error) = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + else { + panic!("the second acquisition should time out"); + }; + assert!( + matches!(error, PersistenceError::LockTimeout(_)), + "{error:?}" + ); + let status = crate::grpc::persistence_error_to_status(error, "op"); + assert_eq!(status.code(), tonic::Code::Unavailable); + let details = openshell_core::rpc_error::decode_details(&status).expect("error details"); + assert_eq!( + details.error_info().expect("error info").reason, + "MUTATION_LOCK_TIMEOUT" + ); + drop(held); + assert_eq!(runtime.mutation_locks.live_entry_count(), 0); + } + + #[tokio::test] + async fn lock_metrics_record_wait_and_timeout() { + const WAITS: &str = "openshell_server_mutation_lock_wait_seconds_count{scope=\"sandbox\"}"; + const TIMEOUTS: &str = "openshell_server_mutation_lock_timeouts_total{scope=\"sandbox\"}"; + let metrics = MetricsCapture::install(); + let runtime = test_runtime().await; + runtime.set_mutation_lock_timeout_for_tests(Duration::from_millis(50)); + + let held = runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .unwrap(); + assert_eq!(metrics.value(WAITS), Some(1)); + assert_eq!(metrics.value(TIMEOUTS), None); + + assert!( + runtime + .mutation_guard(MutationScope::sandbox("w", "a")) + .await + .is_err() + ); + assert_eq!(metrics.value(TIMEOUTS), Some(1)); + assert_eq!(metrics.value(WAITS), Some(1)); + + drop(runtime.lock_sandbox_local("b").await); + assert_eq!(metrics.value(WAITS), Some(1)); + drop(held); + } + + /// Workspaces and sandboxes the random scope mix draws from. + const MIX_WORKSPACES: usize = 3; + const MIX_SANDBOXES: usize = 8; + + /// Names of the random scope mix's workspaces and sandboxes. + struct MixNames { + workspaces: [String; MIX_WORKSPACES], + sandboxes: [String; MIX_SANDBOXES], + } + + impl MixNames { + fn new(prefix: &str) -> Self { + Self { + workspaces: std::array::from_fn(|index| format!("{prefix}w{index}")), + sandboxes: std::array::from_fn(|index| format!("{prefix}s{index}")), + } + } + } + + /// One step of a random-mix task, as indices into [`MixNames`]. + #[derive(Clone, Copy)] + enum MixOp { + Global, + Workspace(usize), + Sandbox(usize, usize), + Lifecycle(usize), + GatedLifecycle(usize), + } + + impl MixOp { + fn random(rng: &mut StdRng) -> Self { + match rng.random_range(0..4) { + 0 => Self::Global, + 1 => Self::Workspace(rng.random_range(0..MIX_WORKSPACES)), + 2 => Self::Sandbox( + rng.random_range(0..MIX_WORKSPACES), + rng.random_range(0..MIX_SANDBOXES), + ), + _ => { + let sandbox = rng.random_range(0..MIX_SANDBOXES); + if rng.random_bool(0.5) { + Self::GatedLifecycle(sandbox) + } else { + Self::Lifecycle(sandbox) + } + } + } + } + + /// Mutation guards also take `PostgreSQL` advisory locks, so they + /// exclude conflicting guards on every replica. Lifecycle locks are + /// process-local. + const fn is_distributed(self) -> bool { + matches!(self, Self::Global | Self::Workspace(_) | Self::Sandbox(..)) + } + } + + /// Holders of each key, changed only while the matching guard is held, so + /// lost exclusion panics instead of passing silently. A counter is `-1` + /// under an exclusive holder and otherwise counts shared holders. A + /// sandbox flag marks the holder of that sandbox key, which does not + /// depend on the workspace. + #[derive(Default)] + struct Occupancy { + global: AtomicIsize, + workspaces: [AtomicIsize; MIX_WORKSPACES], + sandboxes: [AtomicBool; MIX_SANDBOXES], + } + + impl Occupancy { + fn share(counter: &AtomicIsize) { + assert!( + counter.fetch_add(1, Ordering::SeqCst) >= 0, + "shared holder entered under an exclusive holder" + ); + } + + fn exclude(counter: &AtomicIsize) { + assert!( + counter + .compare_exchange(0, -1, Ordering::SeqCst, Ordering::SeqCst) + .is_ok(), + "exclusive holder entered while the key was occupied" + ); + } + + fn claim(sandbox: &AtomicBool) { + assert!( + !sandbox.swap(true, Ordering::SeqCst), + "sandbox entered twice" + ); + } + + fn enter(&self, op: MixOp) { + match op { + MixOp::Global => Self::exclude(&self.global), + MixOp::Workspace(workspace) => { + Self::share(&self.global); + Self::exclude(&self.workspaces[workspace]); + } + MixOp::Sandbox(workspace, sandbox) => { + Self::share(&self.global); + Self::share(&self.workspaces[workspace]); + Self::claim(&self.sandboxes[sandbox]); + } + MixOp::Lifecycle(sandbox) | MixOp::GatedLifecycle(sandbox) => { + Self::share(&self.global); + Self::claim(&self.sandboxes[sandbox]); + } + } + } + + fn leave(&self, op: MixOp) { + match op { + MixOp::Global => self.global.store(0, Ordering::SeqCst), + MixOp::Workspace(workspace) => { + self.workspaces[workspace].store(0, Ordering::SeqCst); + self.global.fetch_sub(1, Ordering::SeqCst); + } + MixOp::Sandbox(workspace, sandbox) => { + self.sandboxes[sandbox].store(false, Ordering::SeqCst); + self.workspaces[workspace].fetch_sub(1, Ordering::SeqCst); + self.global.fetch_sub(1, Ordering::SeqCst); + } + MixOp::Lifecycle(sandbox) | MixOp::GatedLifecycle(sandbox) => { + self.sandboxes[sandbox].store(false, Ordering::SeqCst); + self.global.fetch_sub(1, Ordering::SeqCst); + } + } + } + + fn assert_empty(&self) { + assert_eq!(self.global.load(Ordering::SeqCst), 0); + assert!( + self.workspaces + .iter() + .all(|workspace| workspace.load(Ordering::SeqCst) == 0) + ); + assert!( + self.sandboxes + .iter() + .all(|sandbox| !sandbox.load(Ordering::SeqCst)) + ); + } + } + + /// Enter `op`'s keys, keep them for `hold`, then leave them in reverse + /// order. `local` is the occupancy of the replica that runs `op`; mutation + /// guards also enter `fleet`. The caller holds `op`'s guards throughout. + async fn occupy(local: &Occupancy, fleet: &Occupancy, op: MixOp, hold: Duration) { + local.enter(op); + if op.is_distributed() { + fleet.enter(op); + } + tokio::time::sleep(hold).await; + if op.is_distributed() { + fleet.leave(op); + } + local.leave(op); + } + + /// Run `tasks` tasks of `iterations` random operations each, spread + /// round-robin over `replicas`, holding each operation's locks for 0-2 ms. + /// Every operation must acquire its locks and finish `within`. + async fn run_random_scope_mix( + replicas: &[ComputeRuntime], + prefix: &str, + tasks: usize, + iterations: usize, + within: Duration, + ) { + let names = Arc::new(MixNames::new(prefix)); + let fleet = Arc::new(Occupancy::default()); + let locals: Vec> = replicas.iter().map(|_| Arc::default()).collect(); + let mut rng = StdRng::seed_from_u64(3528); + let mut handles = Vec::new(); + for task in 0..tasks { + let plan: Vec<(MixOp, u64)> = (0..iterations) + .map(|_| { + let op = MixOp::random(&mut rng); + (op, rng.random_range(0..=2)) + }) + .collect(); + let replica = task % replicas.len(); + let runtime = replicas[replica].clone(); + let names = Arc::clone(&names); + let fleet = Arc::clone(&fleet); + let local = Arc::clone(&locals[replica]); + handles.push(tokio::spawn(async move { + for (op, hold_ms) in plan { + let hold = Duration::from_millis(hold_ms); + match op { + MixOp::Global => { + let _guard = runtime + .mutation_guard(MutationScope::Global) + .await + .expect("guard"); + occupy(&local, &fleet, op, hold).await; + } + MixOp::Workspace(workspace) => { + let scope = MutationScope::Workspace(&names.workspaces[workspace]); + let _guard = runtime.mutation_guard(scope).await.expect("guard"); + occupy(&local, &fleet, op, hold).await; + } + MixOp::Sandbox(workspace, sandbox) => { + let scope = MutationScope::sandbox( + &names.workspaces[workspace], + &names.sandboxes[sandbox], + ); + let _guard = runtime.mutation_guard(scope).await.expect("guard"); + occupy(&local, &fleet, op, hold).await; + } + MixOp::Lifecycle(sandbox) => { + let _guard = + runtime.lock_sandbox_local(&names.sandboxes[sandbox]).await; + occupy(&local, &fleet, op, hold).await; + } + MixOp::GatedLifecycle(sandbox) => { + let gate = runtime + .lifecycle_gates + .lock_for(&names.sandboxes[sandbox]) + .await; + let _guard = runtime.lock_sandbox_for_lifecycle(&gate).await; + occupy(&local, &fleet, op, hold).await; + } + } + } + })); + } + + tokio::time::timeout(within, async { + for handle in handles { + handle.await.expect("stress task"); + } + }) + .await + .expect("random scope mix finished without a deadlock"); + for runtime in replicas { + assert_eq!(runtime.mutation_locks.live_entry_count(), 0); + } + fleet.assert_empty(); + for local in &locals { + local.assert_empty(); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn random_scope_mix_never_deadlocks() { + let runtime = test_runtime().await; + run_random_scope_mix(&[runtime], "", 64, 50, Duration::from_secs(20)).await; + } + + #[tokio::test] + async fn unrelated_supervisor_state_update_does_not_wait_for_sandbox_guard() { + let runtime = test_runtime().await; + let mut sandbox = Sandbox { + metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { + id: "mutation-guard-unrelated-a".to_string(), + name: "mutation-guard-unrelated-a".to_string(), + workspace: "default".to_string(), + ..Default::default() + }), + ..Default::default() + }; + sandbox.set_phase(SandboxPhase::Provisioning as i32); + runtime.store.put_message(&sandbox).await.unwrap(); + let held = runtime + .mutation_guard(MutationScope::sandbox( + "default", + "mutation-guard-unrelated-b", + )) + .await + .unwrap(); + + tokio::time::timeout( + PROCEEDS_WITHIN, + runtime.supervisor_session_connected("mutation-guard-unrelated-a", "i"), + ) + .await + .expect("an unrelated supervisor update should not wait") + .expect("supervisor session connected"); + let stored = runtime + .store + .get_message::("mutation-guard-unrelated-a") + .await + .unwrap() + .unwrap(); + assert_eq!(stored.phase(), SandboxPhase::Ready as i32); + drop(held); + } + + /// A runtime on its own store connected to `schema`, like one gateway + /// replica. + async fn postgres_runtime(schema: &TestSchema) -> ComputeRuntime { + let store = Arc::new(schema.connect_store().await); + super::super::new_test_runtime_for_driver(store, "test").await + } + + fn stored_sandbox(sandbox_id: &str, workspace: &str) -> Sandbox { + Sandbox { + metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { + id: sandbox_id.to_string(), + name: sandbox_id.to_string(), + workspace: workspace.to_string(), + ..Default::default() + }), + ..Default::default() + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_compute_guards_random_scope_mix_never_deadlocks() { + let schema = TestSchema::create("mix").await; + let replicas = [ + postgres_runtime(&schema).await, + postgres_runtime(&schema).await, + ]; + // Advisory locks are database-wide, so the mix uses fresh names. + let prefix = format!("{}-", Uuid::new_v4().simple()); + + run_random_scope_mix(&replicas, &prefix, 32, 20, Duration::from_mins(1)).await; + + for replica in &replicas { + replica.store.close().await; + } + schema.drop_schema().await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_compute_guard_by_id_uses_the_sandbox_workspace() { + let schema = TestSchema::create("guard").await; + let replica_a = postgres_runtime(&schema).await; + let replica_b = postgres_runtime(&schema).await; + let workspace = format!("ws-{}", Uuid::new_v4()); + let sandbox_id = format!("sb-{}", Uuid::new_v4()); + replica_a + .store + .put_message(&stored_sandbox(&sandbox_id, &workspace)) + .await + .expect("seed the sandbox"); + + // Only the database locks connect the two replicas. + let held = replica_b + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .expect("workspace guard on replica B"); + let mut by_id = { + let replica_a = replica_a.clone(); + let sandbox_id = sandbox_id.clone(); + tokio::spawn(async move { replica_a.sandbox_mutation_guard_by_id(&sandbox_id).await }) + }; + assert!( + tokio::time::timeout(Duration::from_millis(200), &mut by_id) + .await + .is_err(), + "the by-id guard should wait for the sandbox's workspace" + ); + drop(held); + let guard = tokio::time::timeout(PROCEEDS_WITHIN, by_id) + .await + .expect("the by-id guard should proceed after release") + .expect("guard task") + .expect("by-id guard"); + assert!(guard.is_some(), "the seeded sandbox exists"); + drop(guard); + + assert!( + replica_a + .sandbox_mutation_guard_by_id(&format!("sb-{}", Uuid::new_v4())) + .await + .expect("by-id guard for an unknown sandbox") + .is_none() + ); + + replica_a.store.close().await; + replica_b.store.close().await; + schema.drop_schema().await; + } + + /// Measures the drain envelope: the drain's fastest pacing (one session + /// every 12 ms, as for 1000 sessions in its 12 s close window), all into + /// one receiving replica with the production lock pool. Run it with + /// `--no-capture` to see the wait percentiles. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_lock_pool_absorbs_the_fastest_drain_reconnect_rate() { + /// Sessions that move to the receiving replica. + const RECONNECTS: usize = 500; + /// The drain's close interval for 1000 sessions in its 12 s window. + const RECONNECT_INTERVAL: Duration = Duration::from_millis(12); + /// Guarded operations per reconnect on the receiver: the pre-ack + /// endpoint-status reset and one endpoint report. + const GUARDED_OPS_PER_RECONNECT: usize = 2; + /// Extra time under the guard, so each critical section takes about + /// 20 ms, like one against a managed database. + const CRITICAL_SECTION_PADDING: Duration = Duration::from_millis(15); + + let schema = TestSchema::create("envelope").await; + // The production lock pool, as on a real receiving replica. + let receiver = postgres_runtime(&schema).await; + let workspace = format!("ws-{}", Uuid::new_v4()); + let mut sandbox_ids = Vec::with_capacity(RECONNECTS); + for _ in 0..RECONNECTS { + let sandbox_id = format!("sb-{}", Uuid::new_v4()); + receiver + .store + .put_message(&stored_sandbox(&sandbox_id, &workspace)) + .await + .expect("seed a sandbox"); + sandbox_ids.push(sandbox_id); + } + + let started = tokio::time::Instant::now(); + let reconnects: Vec<_> = sandbox_ids + .into_iter() + .enumerate() + .map(|(index, sandbox_id)| { + let receiver = receiver.clone(); + let arrives = started + + RECONNECT_INTERVAL * u32::try_from(index).expect("reconnect index fits u32"); + tokio::spawn(async move { + tokio::time::sleep_until(arrives).await; + let mut waits = Vec::with_capacity(GUARDED_OPS_PER_RECONNECT); + for op in 0..GUARDED_OPS_PER_RECONNECT { + let called = tokio::time::Instant::now(); + let guard = receiver + .sandbox_mutation_guard_by_id(&sandbox_id) + .await? + .expect("seeded sandbox"); + waits.push(called.elapsed()); + let sandbox = receiver + .store + .get_message::(&sandbox_id) + .await? + .expect("seeded sandbox"); + receiver + .store + .update_message_cas::( + &sandbox_id, + sandbox.get_resource_version(), + |sandbox| { + sandbox + .metadata + .as_mut() + .expect("sandbox metadata") + .labels + .insert("envelope-op".to_string(), op.to_string()); + }, + ) + .await?; + tokio::time::sleep(CRITICAL_SECTION_PADDING).await; + drop(guard); + } + Ok::<_, PersistenceError>(waits) + }) + }) + .collect(); + let mut waits = Vec::with_capacity(RECONNECTS * GUARDED_OPS_PER_RECONNECT); + for reconnect in reconnects { + match reconnect.await.expect("reconnect task") { + Ok(reconnect_waits) => waits.extend(reconnect_waits), + Err(error) => panic!("a guarded reconnect operation failed: {error:?}"), + } + } + + waits.sort_unstable(); + let percentile = |percent: usize| waits[(waits.len() * percent).div_ceil(100) - 1]; + let (p50, p99) = (percentile(50), percentile(99)); + let max = waits[waits.len() - 1]; + eprintln!( + "drain envelope: {} guarded ops from {RECONNECTS} reconnects {RECONNECT_INTERVAL:?} \ + apart into one receiver: lock wait p50 {p50:?}, p99 {p99:?}, max {max:?}", + waits.len() + ); + // Waits must stay far from the lock timeout, where requests fail. + assert!( + p99 * 5 < MUTATION_LOCK_TIMEOUT, + "p99 lock wait {p99:?} is too close to the {MUTATION_LOCK_TIMEOUT:?} timeout" + ); + + receiver.store.close().await; + schema.drop_schema().await; + } +} diff --git a/crates/openshell-server/src/compute/provisioning_deadline.rs b/crates/openshell-server/src/compute/provisioning_deadline.rs index 54d1f83682..673f86cef1 100644 --- a/crates/openshell-server/src/compute/provisioning_deadline.rs +++ b/crates/openshell-server/src/compute/provisioning_deadline.rs @@ -339,7 +339,7 @@ impl super::ComputeRuntime { pub(super) async fn reconcile_provisioning_deadlines(&self, now_ms: i64) -> Result<(), String> { use crate::persistence::{ObjectListQuery, ObjectType}; use openshell_core::{ - ObjectId, + ObjectId, ObjectWorkspace, proto::{Sandbox, SandboxPhase}, }; use prost::Message; @@ -361,7 +361,9 @@ impl super::ComputeRuntime { // Expiration can fence Starting while its driver RPC owns the // lifecycle gate. Cleanup waits for that gate; Error never waits // for compute I/O, matching the existing driver-observation fence. - let global = self.sync_lock.clone().lock_owned().await; + let local_guard = self + .lock_sandbox_local_in_workspace(candidate.object_workspace(), &record.id) + .await; let Some(mut current) = self .store .get_message::(&record.id) @@ -387,7 +389,7 @@ impl super::ComputeRuntime { if let Some(expired) = self.claim_provisioning_timeout(¤t, now_ms).await? { current = expired; } - drop(global); + drop(local_guard); if timed_out(¤t) && current .status @@ -402,10 +404,9 @@ impl super::ComputeRuntime { .is_none_or(|t| t <= now_ms) }) { - let Ok(guard) = self.lifecycle_gates.gate_for(&record.id).try_lock_owned() else { + let Some(gate) = self.lifecycle_gates.try_lock_for(&record.id) else { continue; }; - let gate = super::SandboxLifecycleGuard { _guard: guard }; let runtime = self.clone(); tokio::spawn(async move { if let Err(error) = runtime.reclaim_provisioning_timeout(¤t, &gate).await @@ -419,7 +420,7 @@ impl super::ComputeRuntime { } /// Claim expiration durably before touching the backend. The caller owns the - /// global configuration guard; CAS fences concurrent lifecycle operations. + /// sandbox's local mutation lock; CAS fences concurrent lifecycle operations. /// The separate cleanup step also requires the per-sandbox lifecycle gate. pub(crate) async fn claim_provisioning_timeout( &self, @@ -505,7 +506,7 @@ impl super::ComputeRuntime { use openshell_core::proto::compute::v1::StopSandboxRequest; use openshell_core::{ObjectId, ObjectName}; let current = { - let _global_guard = self.lock_global_for_lifecycle(lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(lifecycle_guard).await; self.store .get_message::(expired.object_id()) .await @@ -554,7 +555,7 @@ impl super::ComputeRuntime { // Cross-replica cleanup claim. A replacement leader waits longer than // the bounded driver call before retrying an interrupted reclamation. { - let _global_guard = self.lock_global_for_lifecycle(lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(lifecycle_guard).await; self.store .update_message_cas::( expired.object_id(), @@ -595,7 +596,7 @@ impl super::ComputeRuntime { .await; let reclaimed = matches!(&result, Ok(Ok(_))) || matches!(&result, Ok(Err(error)) if error.code() == tonic::Code::NotFound); - let _global_guard = self.lock_global_for_lifecycle(lifecycle_guard).await; + let _sandbox_guard = self.lock_sandbox_for_lifecycle(lifecycle_guard).await; let Some(current) = self .store .get_message::(&sandbox_id) diff --git a/crates/openshell-server/src/gateway_metrics.rs b/crates/openshell-server/src/gateway_metrics.rs new file mode 100644 index 0000000000..58103c00ec --- /dev/null +++ b/crates/openshell-server/src/gateway_metrics.rs @@ -0,0 +1,723 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Gateway capacity and HA metrics. +//! +//! The metric names are an operator-facing contract, documented in +//! docs/observability/gateway-metrics.mdx. Labels are bounded enums only: never a sandbox, +//! channel, endpoint, token, or replica id. The scrape target already identifies the replica. +//! +//! Metric handles bind to whichever recorder is current when a macro runs. `run_server` +//! installs the Prometheus recorder after it builds `ServerState`, so never cache a handle in a +//! static or in state built before [`install_global_recorder`]. [`GaugeSlot`] acquires its +//! handle when the tracked object is created and releases it through the same handle, so a +//! slot can never drive a series negative. + +use std::time::{Duration, Instant}; + +use metrics::{ + Gauge, Unit, counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram, +}; +use metrics_exporter_prometheus::{BuildError, Matcher, PrometheusBuilder, PrometheusHandle}; +use tonic::{Code, Status}; + +// Gauges +pub const SUPERVISOR_SESSIONS: &str = "openshell_server_supervisor_sessions"; +pub const DRAINING: &str = "openshell_server_draining"; +pub const RELAY_PENDING: &str = "openshell_server_relay_pending"; +pub const RELAY_PENDING_CAPACITY: &str = "openshell_server_relay_pending_capacity"; +pub const RELAY_PENDING_PER_SANDBOX_CAPACITY: &str = + "openshell_server_relay_pending_per_sandbox_capacity"; +pub const SANDBOX_WATCH_POLLED_SANDBOXES: &str = "openshell_server_sandbox_watch_polled_sandboxes"; +// Counters +pub const RELAY_REJECTED_TOTAL: &str = "openshell_server_relay_rejected_total"; +pub const RELAY_EXPIRED_TOTAL: &str = "openshell_server_relay_expired_total"; +pub const PEER_REQUESTS_TOTAL: &str = "openshell_server_peer_requests_total"; +pub const MUTATION_LOCK_TIMEOUTS_TOTAL: &str = "openshell_server_mutation_lock_timeouts_total"; +pub const SANDBOX_WATCH_POLL_ERRORS_TOTAL: &str = + "openshell_server_sandbox_watch_poll_errors_total"; +// Histograms (explicit buckets, see BUCKETED_HISTOGRAMS) +pub const RELAY_CLAIM_DURATION_SECONDS: &str = "openshell_server_relay_claim_duration_seconds"; +pub const PEER_REQUEST_DURATION_SECONDS: &str = "openshell_server_peer_request_duration_seconds"; +pub const MUTATION_LOCK_WAIT_SECONDS: &str = "openshell_server_mutation_lock_wait_seconds"; +pub const SANDBOX_WATCH_POLL_DURATION_SECONDS: &str = + "openshell_server_sandbox_watch_poll_duration_seconds"; + +const LABEL_REASON: &str = "reason"; +const LABEL_RPC: &str = "rpc"; +const LABEL_OUTCOME: &str = "outcome"; +const LABEL_CODE: &str = "code"; +const LABEL_SCOPE: &str = "scope"; + +/// Buckets for the new latency histograms, 1 ms to 15 s. The top buckets cover the 10 s relay +/// claim and lock timeouts and the 15 s routed-relay wait. +const LATENCY_BUCKETS_SECONDS: [f64; 14] = [ + 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, +]; + +/// Only these names render as Prometheus histograms. Every existing `*_duration_seconds` metric +/// keeps its summary format, so current dashboards are unaffected. +const BUCKETED_HISTOGRAMS: [&str; 4] = [ + RELAY_CLAIM_DURATION_SECONDS, + PEER_REQUEST_DURATION_SECONDS, + MUTATION_LOCK_WAIT_SECONDS, + SANDBOX_WATCH_POLL_DURATION_SECONDS, +]; + +/// Which pending-relay cap rejected an open. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RelayRejection { + GlobalCapacity, + SandboxCapacity, +} + +impl RelayRejection { + pub const ALL: [Self; 2] = [Self::GlobalCapacity, Self::SandboxCapacity]; + + pub const fn label(self) -> &'static str { + match self { + Self::GlobalCapacity => "global_capacity", + Self::SandboxCapacity => "sandbox_capacity", + } + } +} + +/// Outbound peer RPC. Labels match the existing `method` label of +/// `openshell_server_grpc_requests_total` on the owning replica. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PeerRpc { + Relay, + ReportProviderReadiness, + ReportEndpointStatus, + GetSandboxProviderStatus, +} + +impl PeerRpc { + pub const ALL: [Self; 4] = [ + Self::Relay, + Self::ReportProviderReadiness, + Self::ReportEndpointStatus, + Self::GetSandboxProviderStatus, + ]; + + pub const fn label(self) -> &'static str { + match self { + Self::Relay => "PeerRelay", + Self::ReportProviderReadiness => "PeerReportProviderReadiness", + Self::ReportEndpointStatus => "PeerReportEndpointStatus", + Self::GetSandboxProviderStatus => "PeerGetSandboxProviderStatus", + } + } +} + +/// Mutation lock scope kind. The platform scope ("" workspace) maps to `Global`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LockScope { + Global, + Workspace, + Sandbox, +} + +impl LockScope { + pub const ALL: [Self; 3] = [Self::Global, Self::Workspace, Self::Sandbox]; + + pub const fn label(self) -> &'static str { + match self { + Self::Global => "global", + Self::Workspace => "workspace", + Self::Sandbox => "sandbox", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PeerOutcome { + Ok, + ClientError, + RpcError, +} + +impl PeerOutcome { + const fn label(self) -> &'static str { + match self { + Self::Ok => "ok", + Self::ClientError => "client_error", + Self::RpcError => "rpc_error", + } + } +} + +/// Snake-case gRPC status name for the `code` label. Exhaustive on purpose, so a new tonic +/// variant fails to compile instead of producing an unbounded label. +const fn grpc_code_label(code: Code) -> &'static str { + match code { + Code::Ok => "ok", + Code::Cancelled => "cancelled", + Code::Unknown => "unknown", + Code::InvalidArgument => "invalid_argument", + Code::DeadlineExceeded => "deadline_exceeded", + Code::NotFound => "not_found", + Code::AlreadyExists => "already_exists", + Code::PermissionDenied => "permission_denied", + Code::ResourceExhausted => "resource_exhausted", + Code::FailedPrecondition => "failed_precondition", + Code::Aborted => "aborted", + Code::OutOfRange => "out_of_range", + Code::Unimplemented => "unimplemented", + Code::Internal => "internal", + Code::Unavailable => "unavailable", + Code::DataLoss => "data_loss", + Code::Unauthenticated => "unauthenticated", + } +} + +/// Relay caps published as `openshell_server_relay_pending_capacity` and +/// `openshell_server_relay_pending_per_sandbox_capacity`. The caller passes the values that +/// enforce the caps, so this module does not depend on the relay registry. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RelayCapacity { + /// Pending relays allowed on one replica. + pub global: usize, + /// Pending relays allowed for one sandbox on one replica. + pub per_sandbox: usize, +} + +/// Apply the bucket overrides. Tests build local recorders from the same builder. +pub fn configure_exporter(builder: PrometheusBuilder) -> Result { + BUCKETED_HISTOGRAMS + .iter() + .try_fold(builder, |builder, name| { + builder.set_buckets_for_metric( + Matcher::Full((*name).to_string()), + &LATENCY_BUCKETS_SECONDS, + ) + }) +} + +/// Install the process-wide recorder, then describe and zero-initialize the catalog. Call +/// once, from `run_server`. +pub fn install_global_recorder(relay: RelayCapacity) -> Result { + let handle = configure_exporter(PrometheusBuilder::new())?.install_recorder()?; + describe_and_initialize(relay); + Ok(handle) +} + +/// Emit HELP metadata, create every fixed-label series at 0, and publish the relay caps. An +/// idle replica then exports 0 instead of "no data", which HPA and `rate()` need. +pub fn describe_and_initialize(relay: RelayCapacity) { + describe_gauge!( + SUPERVISOR_SESSIONS, + Unit::Count, + "Supervisor control sessions registered on this gateway replica." + ); + describe_gauge!( + DRAINING, + "1 while this gateway replica is draining supervisor sessions before shutdown, otherwise 0." + ); + describe_gauge!( + RELAY_PENDING, + Unit::Count, + "Relay channels on this replica waiting for the supervisor to connect back, including channels opened for peer replicas." + ); + describe_gauge!( + RELAY_PENDING_CAPACITY, + Unit::Count, + "Maximum pending relay channels on one gateway replica." + ); + describe_gauge!( + RELAY_PENDING_PER_SANDBOX_CAPACITY, + Unit::Count, + "Maximum pending relay channels for one sandbox on one gateway replica." + ); + describe_gauge!( + SANDBOX_WATCH_POLLED_SANDBOXES, + Unit::Count, + "Sandboxes with a local WatchSandbox follower checked on the last watch-poller tick." + ); + describe_counter!( + RELAY_REJECTED_TOTAL, + Unit::Count, + "Relay opens rejected because a pending relay cap was reached." + ); + describe_counter!( + RELAY_EXPIRED_TOTAL, + Unit::Count, + "Pending relay channels dropped because the supervisor did not connect back in time." + ); + describe_counter!( + PEER_REQUESTS_TOTAL, + Unit::Count, + "Requests this replica sent to the replica that owns a sandbox's supervisor session." + ); + describe_counter!( + MUTATION_LOCK_TIMEOUTS_TOTAL, + Unit::Count, + "Mutation lock acquisitions that timed out." + ); + describe_counter!( + SANDBOX_WATCH_POLL_ERRORS_TOTAL, + Unit::Count, + "Watch-poller ticks whose version lookup failed." + ); + describe_histogram!( + RELAY_CLAIM_DURATION_SECONDS, + Unit::Seconds, + "Time from opening a relay channel to the supervisor claiming it." + ); + describe_histogram!( + PEER_REQUEST_DURATION_SECONDS, + Unit::Seconds, + "Latency of requests to the owning replica. For PeerRelay, until the owner's supervisor claimed the relay." + ); + describe_histogram!( + MUTATION_LOCK_WAIT_SECONDS, + Unit::Seconds, + "Time spent acquiring the mutation lock for a scope." + ); + describe_histogram!( + SANDBOX_WATCH_POLL_DURATION_SECONDS, + Unit::Seconds, + "Duration of the batched version lookup for one watch-poller tick." + ); + + // `increment(0)` registers a series without overwriting a value recorded earlier. + gauge!(SUPERVISOR_SESSIONS).increment(0.0); + gauge!(RELAY_PENDING).increment(0.0); + gauge!(DRAINING).increment(0.0); + gauge!(SANDBOX_WATCH_POLLED_SANDBOXES).increment(0.0); + gauge!(RELAY_PENDING_CAPACITY).set(count_as_f64(relay.global)); + gauge!(RELAY_PENDING_PER_SANDBOX_CAPACITY).set(count_as_f64(relay.per_sandbox)); + for reason in RelayRejection::ALL { + counter!(RELAY_REJECTED_TOTAL, LABEL_REASON => reason.label()).increment(0); + } + counter!(RELAY_EXPIRED_TOTAL).increment(0); + for scope in LockScope::ALL { + counter!(MUTATION_LOCK_TIMEOUTS_TOTAL, LABEL_SCOPE => scope.label()).increment(0); + } + counter!(SANDBOX_WATCH_POLL_ERRORS_TOTAL).increment(0); + for rpc in PeerRpc::ALL { + counter!( + PEER_REQUESTS_TOTAL, + LABEL_RPC => rpc.label(), + LABEL_OUTCOME => PeerOutcome::Ok.label(), + LABEL_CODE => grpc_code_label(Code::Ok) + ) + .increment(0); + } +} + +/// Counts in this module stay far below 2^53, so the conversion is exact. +#[allow(clippy::cast_precision_loss)] +pub fn count_as_f64(count: usize) -> f64 { + count as f64 +} + +/// One unit of an exact gauge, held for as long as the tracked object lives. Dropping it +/// decrements through the same handle it incremented, so every removal path is counted exactly +/// once, including paths added in the future. +#[must_use = "dropping a GaugeSlot immediately releases it"] +pub struct GaugeSlot(Gauge); + +impl GaugeSlot { + /// Share of `openshell_server_supervisor_sessions`. + pub fn supervisor_session() -> Self { + Self::acquire(SUPERVISOR_SESSIONS) + } + + /// Share of `openshell_server_relay_pending`. + pub fn relay_pending() -> Self { + Self::acquire(RELAY_PENDING) + } + + fn acquire(name: &'static str) -> Self { + let gauge = gauge!(name); + gauge.increment(1.0); + Self(gauge) + } +} + +impl Drop for GaugeSlot { + fn drop(&mut self) { + self.0.decrement(1.0); + } +} + +/// Mark this replica as draining (`true`) for the rest of the process lifetime. +pub fn set_draining(draining: bool) { + gauge!(DRAINING).set(if draining { 1.0 } else { 0.0 }); +} + +pub fn record_relay_rejected(reason: RelayRejection) { + counter!(RELAY_REJECTED_TOTAL, LABEL_REASON => reason.label()).increment(1); +} + +/// `count` pending relays were dropped unclaimed (late claim or reaper). +pub fn record_relay_expired(count: usize) { + if count > 0 { + counter!(RELAY_EXPIRED_TOTAL).increment(count as u64); + } +} + +pub fn record_relay_claimed(waited: Duration) { + histogram!(RELAY_CLAIM_DURATION_SECONDS).record(waited); +} + +/// Time to acquire every key of one mutation guard (local registry plus Postgres), recorded +/// on success only. +pub fn record_lock_wait(scope: LockScope, waited: Duration) { + histogram!(MUTATION_LOCK_WAIT_SECONDS, LABEL_SCOPE => scope.label()).record(waited); +} + +/// A guard acquisition that timed out (local wait or Postgres `lock_timeout` / SQLSTATE 55P03). +/// The caller maps it to `Status::unavailable`. +pub fn record_lock_timeout(scope: LockScope) { + counter!(MUTATION_LOCK_TIMEOUTS_TOTAL, LABEL_SCOPE => scope.label()).increment(1); +} + +/// Times one outbound peer request and records it exactly once. Dropping an unfinished timer +/// (the caller's future was cancelled) records `client_error` / `cancelled`. +#[must_use = "finish the timer with client_error() or finish()"] +pub struct PeerRequestTimer { + rpc: PeerRpc, + started: Instant, + recorded: bool, +} + +impl PeerRequestTimer { + pub fn start(rpc: PeerRpc) -> Self { + Self { + rpc, + started: Instant::now(), + recorded: false, + } + } + + /// The request failed before it reached the peer: token, channel, headers, or stream setup. + pub fn client_error(&mut self, status: &Status) { + self.record(PeerOutcome::ClientError, status.code()); + } + + /// Record the raw tonic result of the RPC itself. Call this BEFORE any remap to + /// `Unavailable`, so the owner's code (for example `resource_exhausted`) is kept. + pub fn finish(&mut self, result: &Result) { + match result { + Ok(_) => self.record(PeerOutcome::Ok, Code::Ok), + Err(status) => self.record(PeerOutcome::RpcError, status.code()), + } + } + + fn record(&mut self, outcome: PeerOutcome, code: Code) { + if self.recorded { + return; + } + self.recorded = true; + counter!( + PEER_REQUESTS_TOTAL, + LABEL_RPC => self.rpc.label(), + LABEL_OUTCOME => outcome.label(), + LABEL_CODE => grpc_code_label(code) + ) + .increment(1); + histogram!( + PEER_REQUEST_DURATION_SECONDS, + LABEL_RPC => self.rpc.label(), + LABEL_OUTCOME => outcome.label() + ) + .record(self.started.elapsed()); + } +} + +impl Drop for PeerRequestTimer { + fn drop(&mut self) { + self.record(PeerOutcome::ClientError, Code::Cancelled); + } +} + +/// Captures metrics recorded on the current thread through a configured Prometheus recorder. +/// +/// Works in `#[test]` and in the default current-thread `#[tokio::test]`, where tasks spawned on +/// the runtime share the thread. It does not work in `multi_thread` tests or inside +/// `spawn_blocking`. Never pass it into an `async fn` helper: it is `!Send`, and clippy +/// `future_not_send` (nursery) would fire. +#[cfg(test)] +pub struct MetricsCapture { + handle: PrometheusHandle, + _guard: metrics::LocalRecorderGuard<'static>, +} + +#[cfg(test)] +impl MetricsCapture { + pub fn install() -> Self { + // Leaked (test only, one small allocation per test) so the guard can borrow it for 'static. + let recorder: &'static metrics_exporter_prometheus::PrometheusRecorder = + Box::leak(Box::new( + configure_exporter(PrometheusBuilder::new()) + .expect("valid exporter config") + .build_recorder(), + )); + let handle = recorder.handle(); + let guard = metrics::set_default_local_recorder(recorder); + Self { + handle, + _guard: guard, + } + } + + pub fn render(&self) -> String { + self.handle.render() + } + + /// Integer value of one exact series, such as `name` or `name{a="b"}`. Parses as i64 to + /// avoid clippy `float_cmp` and so a negative gauge is visible. `None` if the series is absent. + pub fn value(&self, series: &str) -> Option { + series_value(&self.handle, series) + } + + /// Reads one series like [`Self::value`], from code that cannot hold `self`, such as a + /// waker that runs while the value is being produced. + pub fn value_reader(&self, series: &'static str) -> Box Option + Send + Sync> { + let handle = self.handle.clone(); + Box::new(move || series_value(&handle, series)) + } +} + +#[cfg(test)] +fn series_value(handle: &PrometheusHandle, series: &str) -> Option { + handle + .render() + .lines() + .find_map(|line| line.strip_prefix(series)?.strip_prefix(' ')?.parse().ok()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + #[test] + fn describe_and_initialize_exports_capacity_and_zero_series() { + let metrics = MetricsCapture::install(); + describe_and_initialize(RelayCapacity { + global: 256, + per_sandbox: 32, + }); + + for (series, expected) in [ + ("openshell_server_relay_pending_capacity", 256), + ("openshell_server_relay_pending_per_sandbox_capacity", 32), + ("openshell_server_supervisor_sessions", 0), + ("openshell_server_relay_pending", 0), + ("openshell_server_draining", 0), + ("openshell_server_sandbox_watch_polled_sandboxes", 0), + ( + "openshell_server_relay_rejected_total{reason=\"global_capacity\"}", + 0, + ), + ( + "openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}", + 0, + ), + ("openshell_server_relay_expired_total", 0), + ( + "openshell_server_mutation_lock_timeouts_total{scope=\"global\"}", + 0, + ), + ( + "openshell_server_mutation_lock_timeouts_total{scope=\"workspace\"}", + 0, + ), + ( + "openshell_server_mutation_lock_timeouts_total{scope=\"sandbox\"}", + 0, + ), + ("openshell_server_sandbox_watch_poll_errors_total", 0), + ] { + assert_eq!(metrics.value(series), Some(expected), "{series}"); + } + for rpc in [ + "PeerRelay", + "PeerReportProviderReadiness", + "PeerReportEndpointStatus", + "PeerGetSandboxProviderStatus", + ] { + let series = format!( + "openshell_server_peer_requests_total{{rpc=\"{rpc}\",outcome=\"ok\",code=\"ok\"}}" + ); + assert_eq!(metrics.value(&series), Some(0), "{series}"); + } + assert!( + metrics + .render() + .contains("# HELP openshell_server_supervisor_sessions ") + ); + } + + #[test] + fn configured_exporter_buckets_only_new_latency_histograms() { + let metrics = MetricsCapture::install(); + let sample = Duration::from_millis(3); + histogram!(RELAY_CLAIM_DURATION_SECONDS).record(sample); + histogram!(PEER_REQUEST_DURATION_SECONDS, LABEL_RPC => "PeerRelay", LABEL_OUTCOME => "ok") + .record(sample); + histogram!(MUTATION_LOCK_WAIT_SECONDS, LABEL_SCOPE => "sandbox").record(sample); + histogram!(SANDBOX_WATCH_POLL_DURATION_SECONDS).record(sample); + histogram!( + "openshell_server_grpc_request_duration_seconds", + "method" => "ListSandboxes", + "code" => "0" + ) + .record(sample); + histogram!( + "openshell_server_http_request_duration_seconds", + "path" => "/healthz", + "status" => "200" + ) + .record(sample); + histogram!( + "openshell_server_readiness_database_probe_duration_seconds", + "outcome" => "success" + ) + .record(sample); + histogram!("openshell_gateway_interceptor_latency_seconds").record(sample); + + let rendered = metrics.render(); + for name in BUCKETED_HISTOGRAMS { + assert!( + rendered.contains(&format!("# TYPE {name} histogram")), + "{name} should render as a histogram" + ); + } + for name in [ + "openshell_server_grpc_request_duration_seconds", + "openshell_server_http_request_duration_seconds", + "openshell_server_readiness_database_probe_duration_seconds", + "openshell_gateway_interceptor_latency_seconds", + ] { + assert!( + rendered.contains(&format!("# TYPE {name} summary")), + "{name} should keep its summary format" + ); + } + assert!( + rendered.contains("openshell_server_relay_claim_duration_seconds_bucket{le=\"0.001\"}") + ); + assert!( + rendered.contains("openshell_server_relay_claim_duration_seconds_bucket{le=\"15\"}") + ); + } + + #[test] + fn gauge_slot_counts_until_dropped() { + let metrics = MetricsCapture::install(); + let first = GaugeSlot::relay_pending(); + let second = GaugeSlot::relay_pending(); + assert_eq!(metrics.value(RELAY_PENDING), Some(2)); + drop(first); + assert_eq!(metrics.value(RELAY_PENDING), Some(1)); + drop(second); + assert_eq!(metrics.value(RELAY_PENDING), Some(0)); + + let session = GaugeSlot::supervisor_session(); + assert_eq!(metrics.value(SUPERVISOR_SESSIONS), Some(1)); + drop(session); + assert_eq!(metrics.value(SUPERVISOR_SESSIONS), Some(0)); + } + + #[test] + fn gauge_slot_acquired_before_recorder_never_goes_negative() { + // No capture is installed yet, so this slot binds to the no-op recorder. + let slot = GaugeSlot::relay_pending(); + let metrics = MetricsCapture::install(); + drop(slot); + assert_eq!(metrics.value(RELAY_PENDING), None); + } + + #[test] + fn peer_request_timer_records_outcome_code_and_latency() { + let metrics = MetricsCapture::install(); + + let mut relay = PeerRequestTimer::start(PeerRpc::Relay); + relay.finish(&Err::<(), _>(Status::resource_exhausted("x"))); + drop(relay); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerRelay\",outcome=\"rpc_error\",code=\"resource_exhausted\"}" + ), + Some(1) + ); + assert_eq!( + metrics.value( + "openshell_server_peer_request_duration_seconds_count{rpc=\"PeerRelay\",outcome=\"rpc_error\"}" + ), + Some(1) + ); + + let mut endpoint = PeerRequestTimer::start(PeerRpc::ReportEndpointStatus); + endpoint.client_error(&Status::unavailable("x")); + drop(endpoint); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerReportEndpointStatus\",outcome=\"client_error\",code=\"unavailable\"}" + ), + Some(1) + ); + + let mut provider_status = PeerRequestTimer::start(PeerRpc::GetSandboxProviderStatus); + provider_status.finish(&Ok::<(), Status>(())); + drop(provider_status); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerGetSandboxProviderStatus\",outcome=\"ok\",code=\"ok\"}" + ), + Some(1) + ); + } + + #[test] + fn peer_request_timer_records_cancelled_when_dropped_unfinished() { + let metrics = MetricsCapture::install(); + drop(PeerRequestTimer::start(PeerRpc::ReportProviderReadiness)); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerReportProviderReadiness\",outcome=\"client_error\",code=\"cancelled\"}" + ), + Some(1) + ); + } + + #[test] + fn peer_request_timer_records_once() { + let metrics = MetricsCapture::install(); + let mut timer = PeerRequestTimer::start(PeerRpc::Relay); + timer.finish(&Ok::<(), Status>(())); + timer.client_error(&Status::unavailable("x")); + drop(timer); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerRelay\",outcome=\"ok\",code=\"ok\"}" + ), + Some(1) + ); + assert!(!metrics.render().contains("outcome=\"client_error\"")); + } + + #[test] + fn grpc_code_labels_are_distinct_snake_case() { + let labels: HashSet<&str> = (0..=16) + .map(|code| grpc_code_label(Code::from(code))) + .collect(); + for label in &labels { + assert!( + label.chars().all(|c| c.is_ascii_lowercase() || c == '_'), + "{label} is not snake_case" + ); + } + assert_eq!(labels.len(), 17); + assert_eq!(grpc_code_label(Code::DeadlineExceeded), "deadline_exceeded"); + assert_eq!( + grpc_code_label(Code::ResourceExhausted), + "resource_exhausted" + ); + assert_eq!(grpc_code_label(Code::Ok), "ok"); + } +} diff --git a/crates/openshell-server/src/grpc/mod.rs b/crates/openshell-server/src/grpc/mod.rs index 5ac58307a4..50f893092b 100644 --- a/crates/openshell-server/src/grpc/mod.rs +++ b/crates/openshell-server/src/grpc/mod.rs @@ -78,8 +78,9 @@ use crate::ServerState; /// Map a `PersistenceError` to an appropriate gRPC `Status`. /// /// CAS conflicts (optimistic concurrency failures) are mapped to `ABORTED` -/// to signal that the client should retry with fresh data. Other persistence -/// errors are mapped to `INTERNAL`. +/// to signal that the client should retry with fresh data. Mutation lock +/// timeouts are mapped to `UNAVAILABLE` with a retry delay, because nothing +/// was written. Other persistence errors are mapped to `INTERNAL`. pub fn persistence_error_to_status( err: crate::persistence::PersistenceError, operation: &str, @@ -97,6 +98,11 @@ pub fn persistence_error_to_status( ), current_resource_version, ), + PersistenceError::LockTimeout(_) => openshell_core::rpc_error::unavailable( + "MUTATION_LOCK_TIMEOUT", + format!("{operation} timed out waiting for a concurrent mutation; retry the request"), + std::time::Duration::from_secs(1), + ), other => Status::internal(format!("{operation} failed: {other}")), } } @@ -1159,6 +1165,32 @@ mod tests { assert!(gpu.count_selection_supported); } + #[test] + fn persistence_error_to_status_maps_mutation_lock_timeout_to_unavailable() { + let status = persistence_error_to_status( + crate::persistence::PersistenceError::LockTimeout( + "waiting for a PostgreSQL advisory lock".into(), + ), + "acquire provider mutation lock", + ); + + assert_eq!(status.code(), tonic::Code::Unavailable); + assert_eq!( + status.message(), + "acquire provider mutation lock timed out waiting for a concurrent mutation; \ + retry the request" + ); + let details = openshell_core::rpc_error::decode_details(&status).expect("error details"); + assert_eq!( + details.error_info().expect("error info").reason, + "MUTATION_LOCK_TIMEOUT" + ); + assert_eq!( + details.retry_info().expect("retry info").retry_delay, + Some(std::time::Duration::from_secs(1)) + ); + } + #[test] fn public_resource_capabilities_preserves_absence() { let absent: Option = None; diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index 74fac49de8..d2a8435377 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -25,6 +25,7 @@ pub use endpoint_status::{ use crate::ServerState; use crate::auth::principal::Principal; use crate::auth::workspace_authz::{MinWorkspaceRole, require_platform_admin}; +use crate::compute::MutationScope; use crate::pagination::Pagination; use crate::persistence::{ DraftChunkRecord, ObjectId, ObjectListQuery, ObjectName, ObjectType, ObjectWorkspace, @@ -3602,10 +3603,13 @@ async fn handle_update_config_inner( "annotations are only supported for sandbox-scoped updates", )); } - let _settings_guard = state.settings_mutex.lock().await; - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire policy mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire policy mutation lock") + })?; if has_merge_ops { return Err(Status::invalid_argument( @@ -3739,7 +3743,8 @@ async fn handle_update_config_inner( } // Deleting global policy changes the report's effective configuration. - // Keep settings -> sandbox lock order for all global policy mutations. + // The global guard taken at the top of this branch excludes every + // sandbox-scoped settings, policy, and report mutation. let mut global_settings = load_global_settings(state.store.as_ref()).await?; let provider_composition_was_enabled = provider_policy_composition_enabled_in(&global_settings)?; @@ -3793,10 +3798,13 @@ async fn handle_update_config_inner( let mut response_annotations = sandbox_metadata_annotations(&sandbox); if has_setting { - let _settings_guard = state.settings_mutex.lock().await; - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire policy mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::sandbox(&workspace, &sandbox_id)) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire policy mutation lock") + })?; if key == POLICY_SETTING_KEY { return Err(Status::invalid_argument( @@ -3891,9 +3899,13 @@ async fn handle_update_config_inner( )); } - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire policy mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::sandbox(&workspace, &sandbox_id)) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire policy mutation lock") + })?; if has_merge_ops { let global_settings = load_global_settings(state.store.as_ref()).await?; if global_settings.settings.contains_key(POLICY_SETTING_KEY) { @@ -4424,9 +4436,16 @@ pub(super) async fn handle_report_sandbox_configuration( if reported == ConfigurationAdmissionState::Unspecified { return Err(Status::invalid_argument("admission state is required")); } - let _guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire configuration admission lock") - })?; + let Some(_mutation_guard) = state + .compute + .sandbox_mutation_guard_by_id(&sandbox_id) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire configuration admission lock") + })? + else { + return Err(Status::not_found("sandbox not found")); + }; let mut sandbox = state .store .get_message::(&sandbox_id) @@ -4626,9 +4645,16 @@ pub(super) async fn handle_report_policy_status( .supersede_older_policies(&req.sandbox_id, version) .await; - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire policy mutation lock") - })?; + let Some(_mutation_guard) = state + .compute + .sandbox_mutation_guard_by_id(&req.sandbox_id) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire policy mutation lock") + })? + else { + return Err(Status::not_found("sandbox not found")); + }; let sandbox = state .store .get_message::(&req.sandbox_id) @@ -19915,6 +19941,52 @@ mod tests { ); } + #[tokio::test] + async fn sandbox_setting_update_does_not_wait_for_unrelated_sandbox_guard() { + let state = test_server_state().await; + state + .store + .put_message(&test_sandbox( + "sb-setting-target", + "setting-target", + ProtoSandboxPolicy::default(), + Vec::new(), + )) + .await + .unwrap(); + let unrelated_guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "sb-setting-unrelated")) + .await + .unwrap(); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + handle_update_config( + &state, + authed_request(UpdateConfigRequest { + sandbox: "setting-target".to_string(), + workspace_scope: Some(openshell_core::proto::workspace_selector( + "default".to_string(), + )), + setting_key: "ocsf_json_enabled".to_string(), + setting_value: Some(SettingValue { + value: Some(setting_value::Value::BoolValue(true)), + }), + ..Default::default() + }), + ), + ) + .await + .expect("a sandbox setting update should not wait for an unrelated sandbox") + .expect("sandbox setting update succeeds"); + let settings = load_sandbox_settings(state.store.as_ref(), "default", "setting-target") + .await + .unwrap(); + assert!(settings.settings.contains_key("ocsf_json_enabled")); + drop(unrelated_guard); + } + #[tokio::test] async fn update_config_global_policy_rejects_reserved_provider_key() { let state = test_server_state().await; diff --git a/crates/openshell-server/src/grpc/policy/endpoint_status.rs b/crates/openshell-server/src/grpc/policy/endpoint_status.rs index 462099198b..7576251b19 100644 --- a/crates/openshell-server/src/grpc/policy/endpoint_status.rs +++ b/crates/openshell-server/src/grpc/policy/endpoint_status.rs @@ -10,11 +10,15 @@ use super::{ deterministic_policy_hash, load_global_settings, policy_static_credential_endpoint_bindings, }; use crate::ServerState; -use crate::persistence::{ObjectId, ObjectWorkspace}; +use crate::compute::MutationScope; +use crate::persistence::{ + ObjectCursor, ObjectId, ObjectListQuery, ObjectWorkspace, PersistenceError, +}; use crate::policy_store::PolicyStoreExt; use crate::provider_profile_sources::EffectiveProviderProfileCatalog; use crate::supervisor_owner::{OWNER_TTL, SupervisorOwnerIndex}; use crate::supervisor_session::EndpointReportCursor; +use futures::TryStreamExt; use openshell_core::GetResourceVersion; use openshell_core::endpoint_status::initial_endpoint_status; use openshell_core::mcp::is_mcp_protocol; @@ -30,6 +34,11 @@ use tonic::{Request, Response, Status}; use tracing::warn; const ENDPOINT_STARTUP_RECONCILIATION_PAGE_SIZE: u32 = 100; +/// Sandboxes reconciled at once, matching the four-connection mutation lock +/// pool. +const ENDPOINT_STARTUP_RECONCILIATION_CONCURRENCY: usize = 4; +/// Guarded attempts per sandbox before a concurrent write fails startup. +const ENDPOINT_STARTUP_RECONCILIATION_ATTEMPTS: usize = 5; const ENDPOINT_DISCONNECT_RETRY_INITIAL_BACKOFF: std::time::Duration = std::time::Duration::from_millis(100); const ENDPOINT_DISCONNECT_RETRY_MAX_BACKOFF: std::time::Duration = @@ -111,9 +120,19 @@ async fn handle_report_endpoint_status_inner( // Session validation, configuration derivation, and persistence share the // sandbox mutation boundary. A newly registered supervisor can therefore // invalidate its predecessor before any stale report reaches the CAS. - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::super::persistence_error_to_status(error, "acquire endpoint status mutation lock") - })?; + let Some(_mutation_guard) = state + .compute + .sandbox_mutation_guard_by_id(&req.sandbox_id) + .await + .map_err(|error| { + super::super::persistence_error_to_status( + error, + "acquire endpoint status mutation lock", + ) + })? + else { + return Err(Status::not_found("sandbox not found")); + }; if !state .supervisor_sessions .is_endpoint_status_authority(&req.sandbox_id, &req.supervisor_session_id) @@ -362,9 +381,19 @@ pub async fn reset_endpoint_status_for_supervisor_session( sandbox_id: &str, supervisor_session_id: &str, ) -> Result<(), Status> { - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::super::persistence_error_to_status(error, "acquire endpoint status mutation lock") - })?; + let Some(_mutation_guard) = state + .compute + .sandbox_mutation_guard_by_id(sandbox_id) + .await + .map_err(|error| { + super::super::persistence_error_to_status( + error, + "acquire endpoint status mutation lock", + ) + })? + else { + return Err(Status::not_found("sandbox not found")); + }; if !state .supervisor_sessions .is_current_session(sandbox_id, supervisor_session_id) @@ -404,22 +433,33 @@ pub async fn reset_endpoint_status_for_supervisor_session( /// Reset endpoint observations after the active supervisor stream disconnects. /// /// A concurrently registered replacement owns its own pre-acknowledgement -/// reset, so this path leaves that session's cursor alone. +/// reset, so this path leaves that session's cursor alone. A replacement that +/// already exists is detected before the mutation guard, so the reset does not +/// wait behind other mutations of the sandbox. pub async fn reset_endpoint_status_after_supervisor_disconnect( state: &Arc, sandbox_id: &str, ) -> Result<(), Status> { - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::super::persistence_error_to_status(error, "acquire endpoint status mutation lock") - })?; - if state - .supervisor_sessions - .current_session_id(sandbox_id) - .is_some() - { + if disconnect_reset_is_superseded(state, sandbox_id).await? { return Ok(()); } - if has_fresh_shared_owner(state, sandbox_id).await? { + let Some(_mutation_guard) = state + .compute + .sandbox_mutation_guard_by_id(sandbox_id) + .await + .map_err(|error| { + super::super::persistence_error_to_status( + error, + "acquire endpoint status mutation lock", + ) + })? + else { + return Err(Status::not_found("sandbox not found")); + }; + // Re-check under the guard: a replacement's pre-acknowledgement reset runs + // under the same sandbox key, and resetting after it would wipe the + // replacement's fresh evidence. + if disconnect_reset_is_superseded(state, sandbox_id).await? { return Ok(()); } let sandbox = state @@ -484,68 +524,153 @@ pub async fn retry_endpoint_status_after_supervisor_disconnect( /// Supervisor sessions are intentionally process-local. This reconciliation /// runs before gateway listeners are bound. A fresh shared owner preserves its /// evidence; records without one are reset so stale success is never served. +/// +/// Each sandbox is reset under its own mutation guard, so other replicas keep +/// mutating unrelated sandboxes during the scan. Keyset paging keeps the scan +/// stable when sandboxes are deleted mid-scan. pub async fn invalidate_endpoint_status_on_startup(state: &Arc) -> Result<(), Status> { - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::super::persistence_error_to_status( - error, - "acquire endpoint status startup reconciliation lock", - ) - })?; - let mut offset = 0; + let mut cursor: Option = None; loop { - let sandboxes = state + let page = state .store - .list_all_messages::(ENDPOINT_STARTUP_RECONCILIATION_PAGE_SIZE, offset) + .list_message_page::( + ObjectListQuery::AllWorkspaces, + cursor.as_ref(), + ENDPOINT_STARTUP_RECONCILIATION_PAGE_SIZE, + ) .await .map_err(|error| { Status::internal(format!( "list sandboxes for tool server endpoint-status startup reconciliation failed: {error}" )) })?; - if sandboxes.is_empty() { + // Each future holds at most its own sandbox guard and never waits on + // another, so running them in one task cannot deadlock. + futures::stream::iter( + page.messages + .iter() + .filter(|sandbox| has_endpoint_status(sandbox)) + .map(Ok::<_, Status>), + ) + .try_for_each_concurrent(ENDPOINT_STARTUP_RECONCILIATION_CONCURRENCY, |candidate| { + invalidate_sandbox_endpoint_status_on_startup(state, candidate) + }) + .await?; + let Some(next_cursor) = page.next_cursor else { return Ok(()); - } + }; + cursor = Some(next_cursor); + } +} - for sandbox in &sandboxes { - let has_endpoint_status = sandbox - .status - .as_ref() - .is_some_and(|status| !status.endpoint_statuses.is_empty()); - if !has_endpoint_status { - continue; - } - let sandbox_id = sandbox.object_id(); - if has_fresh_shared_owner(state, sandbox_id).await? { - continue; +async fn invalidate_sandbox_endpoint_status_on_startup( + state: &Arc, + candidate: &Sandbox, +) -> Result<(), Status> { + // A live owner keeps its evidence, so it costs no guard. + if has_fresh_shared_owner(state, candidate.object_id()).await? { + return Ok(()); + } + retry_startup_reconciliation(|| invalidate_sandbox_endpoint_status_once(state, candidate)).await +} + +/// Outcome of one guarded startup-reconciliation attempt. +enum StartupAttempt { + Done, + /// The write hit a concurrent change, a row deleted after the re-read, or + /// another database error; re-read and try again. + Retry(PersistenceError), +} + +/// Run `attempt` until it is done, at most +/// `ENDPOINT_STARTUP_RECONCILIATION_ATTEMPTS` times. An error from `attempt` +/// is returned without a retry. +async fn retry_startup_reconciliation(mut attempt: F) -> Result<(), Status> +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempts = 1; + loop { + match attempt().await? { + StartupAttempt::Done => return Ok(()), + StartupAttempt::Retry(error) + if attempts >= ENDPOINT_STARTUP_RECONCILIATION_ATTEMPTS => + { + return Err(super::super::persistence_error_to_status( + error, + "invalidate tool server endpoint status during gateway startup", + )); } - let expected_resource_version = sandbox.get_resource_version(); - let updated = state - .store - .update_message_cas::( - sandbox_id, - expected_resource_version, - invalidate_endpoint_status_without_session, - ) - .await - .map_err(|error| { - super::super::persistence_error_to_status( - error, - "invalidate tool server endpoint status during gateway startup", - ) - })?; - state.sandbox_index.update_from_sandbox(&updated); + StartupAttempt::Retry(_) => attempts += 1, } + } +} - let page_len = sandboxes.len() as u32; - if page_len < ENDPOINT_STARTUP_RECONCILIATION_PAGE_SIZE { - return Ok(()); - } - offset = offset.checked_add(page_len).ok_or_else(|| { - Status::internal( - "sandbox pagination overflow during tool server endpoint-status reconciliation", +async fn invalidate_sandbox_endpoint_status_once( + state: &Arc, + candidate: &Sandbox, +) -> Result { + let sandbox_id = candidate.object_id(); + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::sandbox( + candidate.object_workspace(), + sandbox_id, + )) + .await + .map_err(|error| { + super::super::persistence_error_to_status( + error, + "acquire endpoint status startup reconciliation lock", ) })?; + // A supervisor may have connected to another replica while this waited. + if has_fresh_shared_owner(state, sandbox_id).await? { + return Ok(StartupAttempt::Done); } + let Some(current) = state + .store + .get_message::(sandbox_id) + .await + .map_err(|error| Status::internal(format!("fetch sandbox failed: {error}")))? + else { + return Ok(StartupAttempt::Done); + }; + if !has_endpoint_status(¤t) { + return Ok(StartupAttempt::Done); + } + match state + .store + .update_message_cas::( + sandbox_id, + current.get_resource_version(), + invalidate_endpoint_status_without_session, + ) + .await + { + Ok(updated) => { + state.sandbox_index.update_from_sandbox(&updated); + Ok(StartupAttempt::Done) + } + // Lifecycle writers on other replicas take no distributed guard, and a + // delete between the re-read and the write surfaces as a database + // error. Retry with a fresh read after this guard drops. + Err(error @ (PersistenceError::Conflict { .. } | PersistenceError::Database(_))) => { + Ok(StartupAttempt::Retry(error)) + } + Err(error) => Err(super::super::persistence_error_to_status( + error, + "invalidate tool server endpoint status during gateway startup", + )), + } +} + +fn has_endpoint_status(sandbox: &Sandbox) -> bool { + sandbox + .status + .as_ref() + .is_some_and(|status| !status.endpoint_statuses.is_empty()) } async fn has_fresh_shared_owner( @@ -559,6 +684,22 @@ async fn has_fresh_shared_owner( .map_err(|error| Status::unavailable(format!("resolve supervisor owner failed: {error}"))) } +/// True when a replacement supervisor session (local, or a fresh owner on a +/// peer) now owns endpoint observation, so a disconnect must not reset. +async fn disconnect_reset_is_superseded( + state: &Arc, + sandbox_id: &str, +) -> Result { + if state + .supervisor_sessions + .current_session_id(sandbox_id) + .is_some() + { + return Ok(true); + } + has_fresh_shared_owner(state, sandbox_id).await +} + fn invalidate_endpoint_status_without_session(sandbox: &mut Sandbox) { let Some(status) = sandbox.status.as_mut() else { return; diff --git a/crates/openshell-server/src/grpc/policy/endpoint_status_tests.rs b/crates/openshell-server/src/grpc/policy/endpoint_status_tests.rs index 4c951863a0..800d4cf362 100644 --- a/crates/openshell-server/src/grpc/policy/endpoint_status_tests.rs +++ b/crates/openshell-server/src/grpc/policy/endpoint_status_tests.rs @@ -408,7 +408,11 @@ async fn global_policy_update_waits_for_endpoint_report_guard() { let before = load_global_settings(state.store.as_ref()) .await .expect("read settings before update"); - let guard = state.compute.sandbox_sync_guard().await.unwrap(); + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "endpoint-report-guard")) + .await + .unwrap(); let mut pending = Box::pin(handle_update_config(&state, authed_request(update))); // Poll the actual writer while an endpoint report owns the mutation @@ -442,6 +446,39 @@ async fn global_policy_update_waits_for_endpoint_report_guard() { } } +#[tokio::test] +async fn report_endpoint_status_does_not_wait_for_unrelated_sandbox_guard() { + let sandbox_id = "endpoint-unrelated-guard"; + let (state, mut report) = sandbox_with_accepted_endpoint_result(sandbox_id, true).await; + let unrelated_guard = state + .compute + .mutation_guard(MutationScope::sandbox( + "default", + "endpoint-unrelated-other", + )) + .await + .expect("hold an unrelated sandbox guard"); + + report.report_sequence = 2; + report.observations[0].result = EndpointResult::TransportFailed as i32; + tokio::time::timeout( + std::time::Duration::from_secs(5), + handle_report_endpoint_status(&state, with_sandbox(Request::new(report), sandbox_id)), + ) + .await + .expect("an endpoint report must not wait for an unrelated sandbox mutation") + .expect("accept endpoint result"); + let status = stored_sandbox(&state, sandbox_id) + .await + .status + .expect("sandbox status"); + assert_eq!( + status.endpoint_statuses[0].last_result, + EndpointResult::TransportFailed as i32 + ); + drop(unrelated_guard); +} + #[test] fn expected_endpoint_statuses_canonicalize_identity_and_distinguish_paths() { let policy = ProtoSandboxPolicy { @@ -836,6 +873,432 @@ async fn startup_reconciliation_invalidates_status_from_previous_sessions() { assert_eq!(status.conditions, vec![ready_condition()]); } +/// Store a sandbox carrying endpoint evidence from an earlier gateway process +/// and return the status startup reconciliation must leave behind. +async fn seed_stale_endpoint_status(state: &ServerState, sandbox_id: &str) -> EndpointStatus { + let mut sandbox = test_sandbox( + sandbox_id, + sandbox_id, + mcp_policy_with_versions(&["2025-11-25"]), + Vec::new(), + ); + let initial = test_initial_endpoint_status(sandbox_id, "api.example.com", "/mcp"); + sandbox.status = Some(SandboxStatus { + endpoint_statuses: vec![EndpointStatus { + last_result: EndpointResult::HttpResponseReceived as i32, + last_reported_time: Some(timestamp("2026-09-05T01:01:00.000Z")), + ..initial.clone() + }], + conditions: vec![ready_condition()], + ..Default::default() + }); + state + .store + .put_message(&sandbox) + .await + .expect("store prior session status"); + initial +} + +fn spawn_startup_reconciliation( + state: &Arc, +) -> tokio::task::JoinHandle> { + let state = state.clone(); + tokio::spawn(async move { invalidate_endpoint_status_on_startup(&state).await }) +} + +#[tokio::test] +async fn startup_reconciliation_does_not_hold_a_fleet_guard() { + let state = test_server_state().await; + let sandbox_id = "endpoint-startup-fleet-target"; + let initial = seed_stale_endpoint_status(&state, sandbox_id).await; + let unrelated_guard = state + .compute + .mutation_guard(MutationScope::sandbox( + "default", + "endpoint-startup-fleet-unrelated", + )) + .await + .expect("hold an unrelated sandbox guard"); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + invalidate_endpoint_status_on_startup(&state), + ) + .await + .expect("startup reconciliation must not wait for an unrelated sandbox mutation") + .expect("startup reconciliation"); + let status = stored_sandbox(&state, sandbox_id) + .await + .status + .expect("status remains present"); + assert_eq!(status.endpoint_statuses, vec![initial]); + drop(unrelated_guard); +} + +#[tokio::test] +async fn startup_reconciliation_skips_sandbox_deleted_while_waiting() { + let state = test_server_state().await; + let sandbox_id = "endpoint-startup-deleted"; + seed_stale_endpoint_status(&state, sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", sandbox_id)) + .await + .expect("hold the target sandbox guard"); + let mut reconciliation = spawn_startup_reconciliation(&state); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut reconciliation) + .await + .is_err(), + "startup reconciliation must wait for the target sandbox guard" + ); + + // The sandbox was listed before this delete, so the reconciliation that + // wakes up must treat the missing row as done rather than fail startup. + assert!( + state + .store + .delete("sandbox", sandbox_id) + .await + .expect("delete sandbox") + ); + drop(guard); + tokio::time::timeout(std::time::Duration::from_secs(5), reconciliation) + .await + .expect("startup reconciliation finishes after the guard is released") + .expect("startup reconciliation task") + .expect("a sandbox deleted while reconciliation waited is skipped"); +} + +#[tokio::test] +async fn startup_reconciliation_rereads_under_guard() { + let state = test_server_state().await; + let sandbox_id = "endpoint-startup-reread"; + let initial = seed_stale_endpoint_status(&state, sandbox_id).await; + let listed_version = stored_sandbox(&state, sandbox_id) + .await + .get_resource_version(); + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", sandbox_id)) + .await + .expect("hold the target sandbox guard"); + let mut reconciliation = spawn_startup_reconciliation(&state); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut reconciliation) + .await + .is_err(), + "startup reconciliation must wait for the target sandbox guard" + ); + + // Lifecycle writers on other replicas take no distributed guard, so the + // listed version can go stale while reconciliation waits. + let bumped = state + .store + .update_message_cas::(sandbox_id, listed_version, |sandbox| { + sandbox + .metadata + .as_mut() + .expect("sandbox metadata") + .labels + .insert("concurrent-write".to_string(), "true".to_string()); + }) + .await + .expect("concurrent unrelated write"); + assert!(bumped.get_resource_version() > listed_version); + drop(guard); + tokio::time::timeout(std::time::Duration::from_secs(5), reconciliation) + .await + .expect("startup reconciliation finishes after the guard is released") + .expect("startup reconciliation task") + .expect("startup reconciliation"); + + let sandbox = stored_sandbox(&state, sandbox_id).await; + assert!(sandbox.get_resource_version() > bumped.get_resource_version()); + assert_eq!( + sandbox + .metadata + .as_ref() + .expect("sandbox metadata") + .labels + .get("concurrent-write") + .map(String::as_str), + Some("true") + ); + assert_eq!( + sandbox + .status + .expect("status remains present") + .endpoint_statuses, + vec![initial] + ); +} + +/// Assert that the evidence `seed_stale_endpoint_status` stored survived. +async fn assert_endpoint_evidence_kept(state: &ServerState, sandbox_id: &str) { + let status = stored_sandbox(state, sandbox_id) + .await + .status + .expect("status remains present"); + assert_eq!( + status.endpoint_statuses[0].last_result, + EndpointResult::HttpResponseReceived as i32 + ); + assert!(status.endpoint_statuses[0].last_reported_time.is_some()); +} + +#[tokio::test] +async fn startup_reconciliation_keeps_evidence_of_a_live_owner() { + let state = test_server_state().await; + let sandbox_id = "endpoint-startup-live-owner"; + seed_stale_endpoint_status(&state, sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", sandbox_id)) + .await + .expect("hold the target sandbox guard"); + let mut reconciliation = spawn_startup_reconciliation(&state); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut reconciliation) + .await + .is_err(), + "startup reconciliation must wait for the target sandbox guard" + ); + + // The supervisor connects to a peer after the unguarded owner check, so + // only the re-check under the guard can see it. + SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .publish( + sandbox_id, + "session", + "supervisor", + 1, + "peer-replica", + "https://peer", + ) + .await + .expect("publish a live owner on a peer"); + drop(guard); + tokio::time::timeout(std::time::Duration::from_secs(5), reconciliation) + .await + .expect("startup reconciliation finishes after the guard is released") + .expect("startup reconciliation task") + .expect("startup reconciliation"); + assert_endpoint_evidence_kept(&state, sandbox_id).await; + + // With the owner already live, the check before the guard skips the + // sandbox without waiting for its mutation. + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", sandbox_id)) + .await + .expect("hold the target sandbox guard"); + tokio::time::timeout( + std::time::Duration::from_secs(5), + invalidate_endpoint_status_on_startup(&state), + ) + .await + .expect("a live owner is skipped without waiting for the sandbox guard") + .expect("startup reconciliation"); + drop(guard); + assert_endpoint_evidence_kept(&state, sandbox_id).await; +} + +#[tokio::test] +async fn startup_reconciliation_retries_a_conflict_then_succeeds() { + let mut calls = 0; + retry_startup_reconciliation(|| { + calls += 1; + let call = calls; + async move { + Ok(if call == 1 { + StartupAttempt::Retry(PersistenceError::Conflict { + current_resource_version: Some(2), + }) + } else { + StartupAttempt::Done + }) + } + }) + .await + .expect("a conflict is retried"); + assert_eq!(calls, 2); +} + +#[tokio::test] +async fn startup_reconciliation_stops_after_the_attempt_limit() { + let mut calls = 0; + let error = retry_startup_reconciliation(|| { + calls += 1; + async { + Ok(StartupAttempt::Retry(PersistenceError::Database( + "object sb not found".to_string(), + ))) + } + }) + .await + .expect_err("retries are bounded"); + assert_eq!(error.code(), Code::Internal); + assert_eq!(calls, ENDPOINT_STARTUP_RECONCILIATION_ATTEMPTS); + + let mut calls = 0; + let error = retry_startup_reconciliation(|| { + calls += 1; + async { Err(Status::unavailable("resolve supervisor owner failed")) } + }) + .await + .expect_err("an attempt error is returned"); + assert_eq!(error.code(), Code::Unavailable); + assert_eq!(calls, 1, "an attempt error is not retried"); +} + +/// Store accepted endpoint evidence, then end the supervisor session that +/// reported it, as the session task does before its disconnect reset. +async fn disconnected_sandbox_with_endpoint_result(sandbox_id: &str) -> Arc { + let (state, _report) = sandbox_with_accepted_endpoint_result(sandbox_id, true).await; + assert!( + state + .supervisor_sessions + .remove_if_current(sandbox_id, "session-a") + .is_some() + ); + assert_eq!( + stored_sandbox(&state, sandbox_id) + .await + .status + .expect("sandbox status") + .endpoint_statuses[0] + .last_result, + EndpointResult::HttpResponseReceived as i32 + ); + state +} + +#[tokio::test] +async fn disconnect_reset_skips_guard_when_local_session_replaced() { + let sandbox_id = "endpoint-disconnect-local-replacement"; + let state = disconnected_sandbox_with_endpoint_result(sandbox_id).await; + register_session(&state, sandbox_id, "session-b"); + let before = stored_sandbox(&state, sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .expect("hold a guard that excludes every sandbox guard"); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + reset_endpoint_status_after_supervisor_disconnect(&state, sandbox_id), + ) + .await + .expect("a disconnect with a local replacement must not wait for the guard") + .expect("disconnect reset"); + assert_eq!(stored_sandbox(&state, sandbox_id).await, before); + drop(guard); +} + +#[tokio::test] +async fn disconnect_reset_skips_guard_when_peer_owns_session() { + let sandbox_id = "endpoint-disconnect-peer-owner"; + let state = disconnected_sandbox_with_endpoint_result(sandbox_id).await; + SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .publish( + sandbox_id, + "peer-s", + "inst", + 1, + "peer-replica", + "https://peer:8080", + ) + .await + .expect("publish a live owner on a peer"); + let before = stored_sandbox(&state, sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .expect("hold a guard that excludes every sandbox guard"); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + reset_endpoint_status_after_supervisor_disconnect(&state, sandbox_id), + ) + .await + .expect("a disconnect with a live peer owner must not wait for the guard") + .expect("disconnect reset"); + assert_eq!(stored_sandbox(&state, sandbox_id).await, before); + drop(guard); +} + +#[tokio::test] +async fn disconnect_reset_waits_for_guard_without_replacement() { + let sandbox_id = "endpoint-disconnect-no-replacement"; + let state = disconnected_sandbox_with_endpoint_result(sandbox_id).await; + let before = stored_sandbox(&state, sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .expect("hold a guard that excludes every sandbox guard"); + let mut pending = Box::pin(reset_endpoint_status_after_supervisor_disconnect( + &state, sandbox_id, + )); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), pending.as_mut()) + .await + .is_err(), + "a disconnect without a replacement must wait for the guard" + ); + assert_eq!(stored_sandbox(&state, sandbox_id).await, before); + + drop(guard); + tokio::time::timeout(std::time::Duration::from_secs(5), pending) + .await + .expect("disconnect reset finishes after the guard is released") + .expect("disconnect reset"); + let status = stored_sandbox(&state, sandbox_id) + .await + .status + .expect("status remains present"); + assert_eq!( + status.endpoint_statuses[0].last_result, + EndpointResult::NoObservedExchange as i32 + ); + assert!(status.endpoint_statuses[0].last_reported_time.is_none()); +} + +#[tokio::test] +async fn disconnect_reset_rechecks_for_replacement_under_guard() { + let sandbox_id = "endpoint-disconnect-late-replacement"; + let state = disconnected_sandbox_with_endpoint_result(sandbox_id).await; + let guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .expect("hold a guard that excludes every sandbox guard"); + let mut pending = Box::pin(reset_endpoint_status_after_supervisor_disconnect( + &state, sandbox_id, + )); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), pending.as_mut()) + .await + .is_err(), + "a disconnect without a replacement must wait for the guard" + ); + + // The replacement registers after the unguarded check, so only the + // re-check under the guard can keep its evidence. + register_session(&state, sandbox_id, "session-b"); + let before = stored_sandbox(&state, sandbox_id).await; + drop(guard); + tokio::time::timeout(std::time::Duration::from_secs(5), pending) + .await + .expect("disconnect reset finishes after the guard is released") + .expect("disconnect reset"); + assert_eq!(stored_sandbox(&state, sandbox_id).await, before); +} + #[tokio::test] async fn report_endpoint_status_is_session_bound_and_retry_idempotent() { let state = test_server_state().await; diff --git a/crates/openshell-server/src/grpc/provider.rs b/crates/openshell-server/src/grpc/provider.rs index 7cbc2749fa..8a5d652703 100644 --- a/crates/openshell-server/src/grpc/provider.rs +++ b/crates/openshell-server/src/grpc/provider.rs @@ -5,6 +5,7 @@ #![allow(clippy::result_large_err)] // gRPC handlers return Result, Status> +use crate::compute::MutationScope; #[cfg(test)] use crate::credentials::RefreshMaterialScope; use crate::pagination::Pagination; @@ -2548,9 +2549,13 @@ pub(super) async fn handle_create_provider( )); } let provider_type = provider.r#type.clone(); - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire provider mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire provider mutation lock") + })?; let catalog = state .provider_profile_sources .snapshot_catalog(state.store.as_ref(), &workspace) @@ -2777,10 +2782,11 @@ pub(super) async fn handle_import_provider_profiles( .ensure_active()?; let (profiles, mut diagnostics) = profiles_from_import_items(&request.profiles); add_empty_profile_set_diagnostic(&profiles, &mut diagnostics); - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire provider mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire provider mutation lock"))?; let catalog = state .provider_profile_sources .snapshot_catalog(state.store.as_ref(), &workspace) @@ -2872,10 +2878,11 @@ pub(super) async fn handle_update_provider_profiles( let (profiles, mut diagnostics) = profiles_from_import_items(&items); add_empty_profile_set_diagnostic(&profiles, &mut diagnostics); let target_id = normalize_profile_id_request(&request.id)?; - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire provider mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire provider mutation lock"))?; let catalog = state .provider_profile_sources .snapshot_catalog(state.store.as_ref(), &workspace) @@ -3035,10 +3042,11 @@ pub(super) async fn handle_delete_provider_profile( .name; let id = req.id; let id = normalize_profile_id_request(&id)?; - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire provider mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire provider mutation lock"))?; let catalog = state .provider_profile_sources .snapshot_catalog(state.store.as_ref(), &workspace) @@ -3841,10 +3849,15 @@ pub(super) async fn handle_update_provider( .name; // Provider material contributes to the route-report configuration epoch. // Serialize its mutation with route-status validation so a report derived - // from the prior revision cannot commit after this update. - let _sandbox_sync_guard = state.compute.sandbox_sync_guard().await.map_err(|error| { - super::persistence_error_to_status(error, "acquire provider mutation lock") - })?; + // from the prior revision cannot commit after this update. The workspace + // key excludes every sandbox mutation in this workspace. + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|error| { + super::persistence_error_to_status(error, "acquire provider mutation lock") + })?; let Some(mut provider) = req.provider else { emit_provider_lifecycle( "custom", @@ -4599,11 +4612,14 @@ pub(super) async fn handle_configure_provider_refresh( // persist further down are otherwise separate steps: two concurrent // configures of providers attached to the same sandbox could each pass // validation before either persisted and both reserve the same key (CWE-362). - // This is the same guard sandbox create/attach and profile changes take. - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire provider mutation lock") - })?; + // Every provider attached to one sandbox lives in this workspace, so holding + // the workspace key exclusively also excludes sandbox create and attach, + // which hold it shared. + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire provider mutation lock"))?; let provider = state .store @@ -5047,6 +5063,14 @@ pub(super) async fn handle_delete_provider( let workspace = super::workspace::resolve_workspace(state.store.as_ref(), &authz.workspace) .await? .name; + // A sandbox create or attach in this workspace holds the workspace key + // shared, so no sandbox can start referencing the provider between the + // attached-sandbox check and the delete. + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::Workspace(&workspace)) + .await + .map_err(|e| super::persistence_error_to_status(e, "acquire provider mutation lock"))?; let name = req.name; let provider_profile = provider_profile_for_name(state.store.as_ref(), &workspace, &name).await; let result = delete_provider_record_with_credentials( @@ -5563,9 +5587,13 @@ mod tests { } #[tokio::test] - async fn import_provider_profile_waits_for_sandbox_sync_guard() { + async fn import_provider_profile_waits_for_sandbox_mutation_in_workspace() { let state = test_server_state().await; - let guard = state.compute.sandbox_sync_guard().await.unwrap(); + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "profile-import-guard")) + .await + .unwrap(); let task_state = state.clone(); let task = tokio::spawn(async move { handle_import_provider_profiles( @@ -5587,7 +5615,7 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(20)).await; assert!( !task.is_finished(), - "profile import should wait for sandbox sync guard" + "profile import should wait for a sandbox mutation in its workspace" ); drop(guard); @@ -8738,7 +8766,7 @@ mod tests { } #[tokio::test] - async fn delete_provider_profile_waits_for_sandbox_sync_guard() { + async fn delete_provider_profile_waits_for_sandbox_mutation_in_workspace() { let state = test_server_state().await; state .store @@ -8746,7 +8774,11 @@ mod tests { .await .unwrap(); - let guard = state.compute.sandbox_sync_guard().await.unwrap(); + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "profile-delete-guard")) + .await + .unwrap(); let task_state = state.clone(); let task = tokio::spawn(async move { handle_delete_provider_profile( @@ -8766,7 +8798,7 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(20)).await; assert!( !task.is_finished(), - "profile delete should wait for sandbox sync guard" + "profile delete should wait for a sandbox mutation in its workspace" ); drop(guard); @@ -8801,7 +8833,11 @@ mod tests { .await .unwrap(); - let guard = state.compute.sandbox_sync_guard().await.unwrap(); + let guard = state + .compute + .mutation_guard(MutationScope::Workspace("default")) + .await + .unwrap(); let task_state = state.clone(); let task = tokio::spawn(async move { let mut provider = provider_with_values("guarded-provider", "guarded-create"); @@ -8839,6 +8875,213 @@ mod tests { ); } + fn default_workspace_selector() -> openshell_core::proto::WorkspaceSelector { + openshell_core::proto::workspace_selector("default".to_string()) + } + + async fn create_openai_provider(state: &Arc, name: &str) -> Provider { + let provider = provider_with_credential_value(name, "openai", "OPENAI_API_KEY", "sk-test"); + handle_create_provider( + state, + authed_request(CreateProviderRequest { + request_id: String::new(), + provider: Some(provider), + workspace_scope: Some(default_workspace_selector()), + }), + ) + .await + .expect("create provider") + .into_inner() + .provider + .expect("created provider") + } + + fn provider_config_update(current: &Provider) -> Request { + let mut provider = current.clone(); + provider.credential_handles.clear(); + provider + .config + .insert("NEW_CONFIG".to_string(), "new-value".to_string()); + authed_request(UpdateProviderRequest { + request_id: String::new(), + provider: Some(provider), + credential_expiration_times: HashMap::new(), + clear_credential_expiration_keys: Vec::new(), + workspace_scope: Some(default_workspace_selector()), + }) + } + + fn delete_provider_request(name: &str) -> Request { + authed_request(DeleteProviderRequest { + request_id: String::new(), + allow_missing: false, + name: name.to_string(), + workspace_scope: Some(default_workspace_selector()), + }) + } + + fn sandbox_in_default_workspace(id: &str, providers: Vec) -> Sandbox { + Sandbox { + metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta { + id: id.to_string(), + name: id.to_string(), + workspace: "default".to_string(), + ..Default::default() + }), + spec: Some(SandboxSpec { + providers, + ..Default::default() + }), + ..Default::default() + } + } + + #[tokio::test] + async fn delete_provider_waits_for_sandbox_mutation_in_workspace() { + let state = test_server_state().await; + create_openai_provider(&state, "guarded-delete-provider").await; + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "x")) + .await + .unwrap(); + + let task_state = state.clone(); + let mut delete = tokio::spawn(async move { + handle_delete_provider( + &task_state, + delete_provider_request("guarded-delete-provider"), + ) + .await + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut delete) + .await + .is_err(), + "provider delete should wait for a sandbox mutation in its workspace" + ); + drop(guard); + + let response = tokio::time::timeout(std::time::Duration::from_secs(5), delete) + .await + .expect("delete should finish after guard release") + .expect("join delete task") + .expect("delete should succeed") + .into_inner(); + assert_eq!( + response.outcome(), + openshell_core::proto::DeletionOutcome::Completed + ); + } + + #[tokio::test] + async fn delete_provider_rejects_provider_attached_while_waiting() { + let state = test_server_state().await; + create_openai_provider(&state, "raced-provider").await; + let sandbox = sandbox_in_default_workspace("raced-sandbox", Vec::new()); + state.store.put_message(&sandbox).await.unwrap(); + // An attach to this sandbox holds its sandbox scope while it writes + // the provider into the spec. + let attach_guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", sandbox.object_id())) + .await + .unwrap(); + + let task_state = state.clone(); + let mut delete = tokio::spawn(async move { + handle_delete_provider(&task_state, delete_provider_request("raced-provider")).await + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut delete) + .await + .is_err(), + "provider delete should wait for the in-flight attach" + ); + state + .store + .update_message_cas::(sandbox.object_id(), 0, |sandbox| { + sandbox + .spec + .get_or_insert_with(Default::default) + .providers + .push("raced-provider".to_string()); + }) + .await + .unwrap(); + drop(attach_guard); + + let error = tokio::time::timeout(std::time::Duration::from_secs(5), delete) + .await + .expect("delete should finish after the attach") + .expect("join delete task") + .expect_err("a provider attached while the delete waited must not be deleted"); + assert_eq!(error.code(), Code::FailedPrecondition); + assert!(error.message().contains("attached to sandbox"), "{error}"); + assert!( + state + .store + .get_message_by_name::("default", "raced-provider") + .await + .unwrap() + .is_some() + ); + } + + #[tokio::test] + async fn update_provider_waits_for_sandbox_mutation_in_same_workspace() { + let state = test_server_state().await; + let current = create_openai_provider(&state, "guarded-update-provider").await; + let guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", "x")) + .await + .unwrap(); + + let task_state = state.clone(); + let request = provider_config_update(¤t); + let mut update = + tokio::spawn(async move { handle_update_provider(&task_state, request).await }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut update) + .await + .is_err(), + "provider update should wait for a sandbox mutation in its workspace" + ); + drop(guard); + + let response = tokio::time::timeout(std::time::Duration::from_secs(5), update) + .await + .expect("update should finish after guard release") + .expect("join update task") + .expect("update should succeed") + .into_inner(); + assert!(response.provider.unwrap().config.contains_key("NEW_CONFIG")); + } + + #[tokio::test] + async fn update_provider_does_not_wait_for_sandbox_mutation_in_other_workspace() { + let state = test_server_state().await; + let current = create_openai_provider(&state, "unguarded-update-provider").await; + // The "team-a" workspace row is not needed to hold its keys. + let other_workspace_guard = state + .compute + .mutation_guard(MutationScope::sandbox("team-a", "x")) + .await + .unwrap(); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(5), + handle_update_provider(&state, provider_config_update(¤t)), + ) + .await + .expect("provider update should not wait for another workspace") + .expect("update should succeed") + .into_inner(); + assert!(response.provider.unwrap().config.contains_key("NEW_CONFIG")); + drop(other_workspace_guard); + } + #[tokio::test] async fn provider_crud_round_trip_and_semantics() { let store = test_store().await; diff --git a/crates/openshell-server/src/grpc/provider_readiness_tests.rs b/crates/openshell-server/src/grpc/provider_readiness_tests.rs index d83bb80ef4..4e82b3bba0 100644 --- a/crates/openshell-server/src/grpc/provider_readiness_tests.rs +++ b/crates/openshell-server/src/grpc/provider_readiness_tests.rs @@ -824,7 +824,7 @@ async fn attach_waiting_for_update_captures_published_revision_and_becomes_ready .provider .unwrap(); - // The credential driver's gate holds UpdateProvider inside the shared + // The credential driver's gate holds UpdateProvider inside the workspace // mutation guard while the attach request reaches that same guard. let (store_hit, release_store) = state.credentials.gate_next_store(); let update_state = Arc::clone(&state); diff --git a/crates/openshell-server/src/grpc/sandbox.rs b/crates/openshell-server/src/grpc/sandbox.rs index d8263e2eb4..1f44d3d624 100644 --- a/crates/openshell-server/src/grpc/sandbox.rs +++ b/crates/openshell-server/src/grpc/sandbox.rs @@ -14,6 +14,7 @@ use crate::auth::workspace_authz::{ AuthorizedWorkspaceScope, MinWorkspaceRole, authorize_list_workspace_selector, authorize_sandbox_workspace, authorize_workspace, }; +use crate::compute::MutationScope; use crate::pagination::Pagination; use crate::persistence::{ ObjectLabels, ObjectListQuery, ObjectType, WriteCondition, generate_name, @@ -488,9 +489,9 @@ async fn handle_create_sandbox_inner( } else { request.name.clone() }; - let (sandbox_lifecycle_guard, sandbox_sync_guard) = state + let (sandbox_lifecycle_guard, mutation_guard) = state .compute - .sandbox_create_guards(&id) + .sandbox_create_guards(&workspace, &id) .await .map_err(|err| super::persistence_error_to_status(err, "acquire sandbox mutation lock"))?; @@ -649,7 +650,7 @@ async fn handle_create_sandbox_inner( launch_authentication, await_main_process_attachment, sandbox_lifecycle_guard, - sandbox_sync_guard, + mutation_guard, )) .await?; @@ -1297,10 +1298,14 @@ pub(super) async fn handle_attach_sandbox_provider( if let Some(probe) = attach_wait_probe { probe.notify_one(); } - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire sandbox mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::sandbox( + sandbox.object_workspace(), + sandbox.object_id(), + )) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire sandbox mutation lock"))?; let provider_record = get_provider_record(state.store.as_ref(), &workspace, &request.provider) .await .map_err(|err| { @@ -1470,10 +1475,14 @@ pub(super) async fn handle_detach_sandbox_provider( ))); } - let _sandbox_sync_guard = - state.compute.sandbox_sync_guard().await.map_err(|err| { - super::persistence_error_to_status(err, "acquire sandbox mutation lock") - })?; + let _mutation_guard = state + .compute + .mutation_guard(MutationScope::sandbox( + sandbox.object_workspace(), + sandbox.object_id(), + )) + .await + .map_err(|err| super::persistence_error_to_status(err, "acquire sandbox mutation lock"))?; let sandbox_name = sandbox.object_name().to_string(); let sandbox_id = sandbox .metadata @@ -5401,8 +5410,13 @@ mod tests { state.store.put_message(&original).await.unwrap(); // Hold the global guard so the handler can resolve the original ID and - // acquire its delete gate, but cannot yet revalidate or mutate it. - let global_guard = state.compute.sandbox_sync_guard().await.unwrap(); + // acquire its delete gate, but cannot yet take the sandbox's local + // lifecycle lock (shared global key) to revalidate or mutate it. + let global_guard = state + .compute + .mutation_guard(MutationScope::Global) + .await + .unwrap(); let delete_state = state.clone(); let delete = tokio::spawn(async move { handle_delete_sandbox_inner( @@ -5583,6 +5597,91 @@ mod tests { assert_eq!(providers, vec!["work-github"]); } + fn attach_request(sandbox: &str, provider: &str) -> Request { + authed_request(AttachSandboxProviderRequest { + request_id: String::new(), + sandbox: sandbox.to_string(), + workspace_scope: Some(openshell_core::proto::workspace_selector( + "default".to_string(), + )), + provider: provider.to_string(), + expected_resource_version: 0, + }) + } + + #[tokio::test] + async fn attach_provider_does_not_wait_for_unrelated_sandbox_guard() { + let state = test_server_state().await; + state + .store + .put_message(&test_provider("work-github", "github")) + .await + .unwrap(); + state + .store + .put_message(&test_sandbox("work", Vec::new())) + .await + .unwrap(); + let unrelated = test_sandbox("unrelated", Vec::new()); + state.store.put_message(&unrelated).await.unwrap(); + let unrelated_guard = state + .compute + .mutation_guard(MutationScope::sandbox("default", unrelated.object_id())) + .await + .unwrap(); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(5), + handle_attach_sandbox_provider(&state, attach_request("work", "work-github")), + ) + .await + .expect("attach should not wait for an unrelated sandbox mutation") + .expect("attach should succeed") + .into_inner(); + assert!(response.attached); + drop(unrelated_guard); + } + + #[tokio::test] + async fn attach_provider_waits_for_workspace_guard() { + let state = test_server_state().await; + state + .store + .put_message(&test_provider("work-github", "github")) + .await + .unwrap(); + state + .store + .put_message(&test_sandbox("work", Vec::new())) + .await + .unwrap(); + let workspace_guard = state + .compute + .mutation_guard(MutationScope::Workspace("default")) + .await + .unwrap(); + + let task_state = state.clone(); + let mut attach = tokio::spawn(async move { + handle_attach_sandbox_provider(&task_state, attach_request("work", "work-github")).await + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut attach) + .await + .is_err(), + "attach should wait for a provider writer in its workspace" + ); + drop(workspace_guard); + + let response = tokio::time::timeout(std::time::Duration::from_secs(5), attach) + .await + .expect("attach should finish after the workspace guard is released") + .expect("join attach task") + .expect("attach should succeed") + .into_inner(); + assert!(response.attached); + } + #[tokio::test] async fn detach_sandbox_provider_is_idempotent_and_removes_all_matches() { let state = test_server_state().await; @@ -6997,7 +7096,7 @@ mod tests { } #[tokio::test] - async fn create_sandbox_with_providers_waits_for_sandbox_sync_guard() { + async fn create_sandbox_with_providers_waits_for_workspace_mutation_guard() { let state = test_server_state().await; state .store @@ -7005,7 +7104,13 @@ mod tests { .await .unwrap(); - let guard = state.compute.sandbox_sync_guard().await.unwrap(); + // A provider writer in the workspace excludes creates that validate + // against its providers. + let guard = state + .compute + .mutation_guard(MutationScope::Workspace("default")) + .await + .unwrap(); let task_state = state.clone(); let task = tokio::spawn(async move { handle_create_sandbox( @@ -7033,7 +7138,7 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(20)).await; assert!( !task.is_finished(), - "sandbox create with initial providers should wait for sandbox sync guard" + "sandbox create with initial providers should wait for the workspace mutation guard" ); drop(guard); diff --git a/crates/openshell-server/src/http.rs b/crates/openshell-server/src/http.rs index 63b59edf04..4793f06943 100644 --- a/crates/openshell-server/src/http.rs +++ b/crates/openshell-server/src/http.rs @@ -12,7 +12,8 @@ //! `503 Service Unavailable` when the latest background check failed. //! Handler latency is sub-millisecond: the database is never pinged from //! inside the request path, so the response cannot race the kubelet's -//! probe timeout. +//! probe timeout. Once gateway shutdown begins it returns `503` with +//! status `draining`, whatever the database state. //! - `/health` — Alias of `/readyz` for external monitors //! that conventionally probe `/health`. @@ -34,6 +35,7 @@ use crate::readiness::{DatabaseHealthMonitor, HealthError, HealthState}; const STATUS_HEALTHY: &str = "healthy"; const STATUS_UNHEALTHY: &str = "unhealthy"; +const STATUS_DRAINING: &str = "draining"; const DATABASE_INITIALIZING_ERROR: &str = "readiness monitor still initializing"; const DATABASE_UNAVAILABLE_ERROR: &str = "database unavailable"; const DATABASE_TIMEOUT_ERROR: &str = "database health check timed out"; @@ -41,6 +43,9 @@ const DATABASE_TIMEOUT_ERROR: &str = "database health check timed out"; #[derive(Clone)] struct HealthRouterState { health: watch::Receiver, + /// Flips to `true` at gateway shutdown. Readiness then reports `503` + /// regardless of database health; liveness is unaffected. + draining: watch::Receiver, } /// Per-dependency check entry exposed under `checks` in the JSON payload. @@ -67,7 +72,9 @@ pub struct HealthChecks { /// Readiness response payload. #[derive(Debug, Serialize)] pub struct HealthResponse { - /// Overall status: `"healthy"` if every dependency is healthy. + /// Overall status: `"healthy"`, `"unhealthy"`, or `"draining"`. It is + /// `"healthy"` only when every dependency is healthy and the gateway is + /// not shutting down. pub status: &'static str, /// Service version. @@ -82,33 +89,31 @@ async fn healthz() -> impl IntoResponse { StatusCode::OK } -/// Kubernetes readiness probe — reflects the cached background DB state. +/// Kubernetes readiness probe — reflects the cached background DB state and +/// the gateway draining signal. async fn readyz(State(state): State>) -> impl IntoResponse { - render_response(&state.health.borrow()) + render_response(&state.health.borrow(), *state.draining.borrow()) } /// Convenience alias of [`readyz`] for monitors that probe `/health`. async fn health(State(state): State>) -> impl IntoResponse { - render_response(&state.health.borrow()) + render_response(&state.health.borrow(), *state.draining.borrow()) } -fn render_response(state: &HealthState) -> (StatusCode, Json) { +fn render_response(state: &HealthState, draining: bool) -> (StatusCode, Json) { let database = render_database(state); - let healthy = state.is_healthy(); + let (status, code) = if draining { + (STATUS_DRAINING, StatusCode::SERVICE_UNAVAILABLE) + } else if state.is_healthy() { + (STATUS_HEALTHY, StatusCode::OK) + } else { + (STATUS_UNHEALTHY, StatusCode::SERVICE_UNAVAILABLE) + }; let response = HealthResponse { - status: if healthy { - STATUS_HEALTHY - } else { - STATUS_UNHEALTHY - }, + status, version: openshell_core::VERSION, checks: HealthChecks { database }, }; - let code = if healthy { - StatusCode::OK - } else { - StatusCode::SERVICE_UNAVAILABLE - }; (code, Json(response)) } @@ -152,12 +157,35 @@ pub fn health_router(store: Arc) -> Router { health_router_from_receiver(monitor.subscribe()) } +/// Like [`health_router`], with readiness forced to `draining` once +/// `draining` is `true`. +pub fn health_router_with_drain(store: Arc, draining: watch::Receiver) -> Router { + let monitor = DatabaseHealthMonitor::spawn(store); + health_router_from_parts(monitor.subscribe(), draining) +} + /// Build the health router from an existing monitor receiver. /// /// Crate-internal: used by [`health_router`] and by tests that drive the /// `HealthState` directly without spinning up the polling task. pub fn health_router_from_receiver(receiver: watch::Receiver) -> Router { - let state = Arc::new(HealthRouterState { health: receiver }); + health_router_from_parts(receiver, not_draining()) +} + +/// A draining signal that never flips. The sender is dropped on purpose; +/// `borrow()` keeps returning `false`. +fn not_draining() -> watch::Receiver { + watch::channel(false).1 +} + +fn health_router_from_parts( + health_state: watch::Receiver, + draining: watch::Receiver, +) -> Router { + let state = Arc::new(HealthRouterState { + health: health_state, + draining, + }); Router::new() .route("/health", get(health)) @@ -471,6 +499,18 @@ mod readiness_tests { health_router_from_receiver(rx) } + /// Like [`router_with_state`], with the draining signal supplied by the + /// caller so a test can hold the sender and flip it. + fn router_with_state_and_drain(state: HealthState, draining: watch::Receiver) -> Router { + let (_tx, rx) = watch::channel(state); + health_router_from_parts(rx, draining) + } + + /// A draining signal that is already `true`. + fn draining_signal() -> watch::Receiver { + watch::channel(true).1 + } + #[tokio::test] async fn healthz_is_minimal_and_does_not_touch_the_database() { // Liveness must succeed even when the database is unreachable — @@ -563,4 +603,54 @@ mod readiness_tests { "timeout state has no completed-call latency" ); } + + #[tokio::test] + async fn readyz_reports_draining_even_when_database_is_healthy() { + let router = + router_with_state_and_drain(HealthState::Healthy { latency_ms: 1 }, draining_signal()); + let (status, body) = get(router, "/readyz").await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["status"], "draining"); + assert_eq!(body["checks"]["database"]["status"], "healthy"); + } + + #[tokio::test] + async fn health_alias_reports_draining() { + // Draining takes precedence over an unhealthy database, while the + // database check still reports its own state. + let router = router_with_state_and_drain( + HealthState::Unhealthy(HealthError::Timeout), + draining_signal(), + ); + let (status, body) = get(router, "/health").await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["status"], "draining"); + assert_eq!(body["checks"]["database"]["status"], "unhealthy"); + } + + #[tokio::test] + async fn healthz_stays_ok_while_draining() { + // A draining gateway is still a live process; only readiness changes. + let router = + router_with_state_and_drain(HealthState::Healthy { latency_ms: 1 }, draining_signal()); + let (status, body) = get(router, "/healthz").await; + assert_eq!(status, StatusCode::OK); + assert!(body.is_null(), "healthz must return an empty body"); + } + + #[tokio::test] + async fn readyz_flips_to_draining_when_signal_changes() { + let (draining_tx, draining_rx) = watch::channel(false); + let router = + router_with_state_and_drain(HealthState::Healthy { latency_ms: 1 }, draining_rx); + + let (status, body) = get(router.clone(), "/readyz").await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body["status"], "healthy"); + + draining_tx.send_replace(true); + let (status, body) = get(router, "/readyz").await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["status"], "draining"); + } } diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index a51e17a9ad..4bd4b60f18 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -22,6 +22,7 @@ mod config_update_operation; mod credentials; mod defaults; mod gateway_listener; +mod gateway_metrics; mod grpc; mod http; mod middleware; @@ -51,7 +52,6 @@ mod tracing_setup; mod watch_cursor; mod ws_tunnel; -use metrics_exporter_prometheus::PrometheusBuilder; use openshell_core::net::set_tcp_nodelay_best_effort; use openshell_core::telemetry::TelemetryComputeDriver; use openshell_core::{Config, Error, ObjectLabels, Result}; @@ -285,12 +285,6 @@ pub struct ServerState { /// Active SSH tunnel connection counts per sandbox id. pub ssh_connections_by_sandbox: Mutex>, - /// Serializes settings mutations (global and sandbox) to prevent - /// read-modify-write races. Held for the duration of any setting - /// set/delete operation, including the precedence check on sandbox - /// mutations that reads global state. - pub settings_mutex: tokio::sync::Mutex<()>, - /// Registry of active supervisor sessions and pending relay channels. /// /// Stored as `Arc` so compiled compute drivers can be constructed before @@ -429,7 +423,6 @@ impl ServerState { telemetry: telemetry::TelemetryState::new(), ssh_connections_by_token: Mutex::new(HashMap::new()), ssh_connections_by_sandbox: Mutex::new(HashMap::new()), - settings_mutex: tokio::sync::Mutex::new(()), supervisor_sessions, gateway_shutting_down: AtomicBool::new(false), replica_id, @@ -521,6 +514,7 @@ pub(crate) async fn run_server( legacy_compute_driver_env_seen: _, } = startup; let (shutdown_tx, shutdown_rx) = watch::channel(false); + let (draining_tx, draining_rx) = watch::channel(false); auth::descriptor_authz::init() .map_err(|error| Error::config(format!("invalid gRPC authorization metadata: {error}")))?; @@ -877,10 +871,10 @@ pub(crate) async fn run_server( )) })?; info!(address = %health_bind_address, "Health server listening"); - // `health_router` returns immediately; the listener serves + // `health_router_with_drain` returns immediately; the listener serves // `Initializing → 503` until the background monitor publishes the // first real probe outcome, so the endpoint is always responsive. - let router = health_router(store.clone()); + let router = http::health_router_with_drain(store.clone(), draining_rx.clone()); tokio::spawn(async move { if let Err(e) = axum::serve(health_listener, router.into_make_service()).await { error!("Health server error: {e}"); @@ -892,9 +886,9 @@ pub(crate) async fn run_server( // Bind the Prometheus metrics endpoint on a dedicated port when configured. if let Some(metrics_bind_address) = config.metrics_bind_address { - let prometheus_handle = PrometheusBuilder::new() - .install_recorder() - .map_err(|e| Error::config(format!("failed to install metrics recorder: {e}")))?; + let prometheus_handle = + gateway_metrics::install_global_recorder(supervisor_session::RELAY_CAPACITY) + .map_err(|e| Error::config(format!("failed to install metrics recorder: {e}")))?; let metrics_listener = TcpListener::bind(metrics_bind_address).await.map_err(|e| { Error::transport(format!( "failed to bind metrics port {metrics_bind_address}: {e}", @@ -998,6 +992,8 @@ pub(crate) async fn run_server( info!("Shutdown signal received; stopping gateway"); state.gateway_shutting_down.store(true, Ordering::Release); state.supervisor_sessions.close_admission(); + draining_tx.send_replace(true); + drain_supervisor_sessions(&state.supervisor_sessions, peer_routing_expected).await; let _ = shutdown_tx.send(true); if let Err(err) = listener_task.await { @@ -1009,7 +1005,7 @@ pub(crate) async fn run_server( // record. Drain it even when compute cleanup failed before exiting Tokio. let session_cleanup = state .supervisor_sessions - .shutdown(Duration::from_secs(10)) + .shutdown(supervisor_session::SESSION_CLEANUP_TIMEOUT) .await; if let Err(err) = &session_cleanup { warn!(error = %err, "Gateway supervisor session cleanup incomplete"); @@ -1199,6 +1195,35 @@ fn spawn_gateway_connection( } } +/// Close supervisor sessions on a paced schedule while the gateway listener +/// keeps accepting, so each supervisor reconnects to another replica, peers +/// can still reach sessions that have not moved yet, and lagged supervisor +/// dials get `UNAVAILABLE` quickly. Only gateways with peer routing drain: +/// elsewhere no peer can reach these sessions, so shutdown continues at once. +async fn drain_supervisor_sessions( + sessions: &supervisor_session::SupervisorSessionRegistry, + peer_routing_expected: bool, +) { + if !peer_routing_expected { + return; + } + gateway_metrics::set_draining(true); + let started = tokio::time::Instant::now(); + info!( + sessions = sessions.session_count(), + propagation_delay_ms = supervisor_session::DRAIN_PROPAGATION_DELAY.as_millis(), + max_close_window_ms = supervisor_session::DRAIN_CLOSE_WINDOW.as_millis(), + "Draining supervisor sessions before stopping the gateway listener" + ); + let summary = supervisor_session::drain_for_shutdown(sessions).await; + info!( + planned = summary.planned, + signaled = summary.signaled, + elapsed_ms = started.elapsed().as_millis(), + "Supervisor session drain finished; stopping gateway listener" + ); +} + async fn shutdown_signal() { #[cfg(unix)] { @@ -1893,9 +1918,11 @@ mod tests { BoundGatewayListener, ConfiguredComputeDriver, ConnectionProtocol, ExtensionKind, MultiplexService, ServerState, TlsAcceptor, allow_plaintext_service_http, bind_gateway_listener, classify_initial_bytes, configured_compute_driver, - extension_token_ttl, is_benign_tls_handshake_failure, mint_gateway_extension_credential, - serve_gateway_listener, validate_peer_endpoint_scheme, + drain_supervisor_sessions, extension_token_ttl, is_benign_tls_handshake_failure, + mint_gateway_extension_credential, serve_gateway_listener, validate_peer_endpoint_scheme, }; + use crate::gateway_metrics::{MetricsCapture, describe_and_initialize}; + use crate::supervisor_session::{LocalSessionRoute, RELAY_CAPACITY}; use openshell_core::{ Config, proto::{HealthRequest, open_shell_client::OpenShellClient}, @@ -1910,7 +1937,8 @@ mod tests { use tempfile::{TempDir, tempdir}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; - use tokio::sync::watch; + use tokio::sync::oneshot::error::TryRecvError; + use tokio::sync::{mpsc, oneshot, watch}; use crate::tls_test_utils::generate_test_certs_with_ca; @@ -1948,6 +1976,60 @@ mod tests { validate_peer_endpoint_scheme(&config, "http://10.0.0.1:8080").unwrap(); } + /// A registry with one Ready session for `sbx`, and that session's + /// shutdown receiver. + fn registry_with_one_session() -> ( + crate::supervisor_session::SupervisorSessionRegistry, + oneshot::Receiver<()>, + ) { + let registry = crate::supervisor_session::SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel(1); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + registry.register("sbx".into(), "s1".into(), tx, shutdown_tx); + (registry, shutdown_rx) + } + + #[tokio::test] + async fn drain_supervisor_sessions_skips_gateways_without_peer_routing() { + let metrics = MetricsCapture::install(); + describe_and_initialize(RELAY_CAPACITY); + let (registry, mut shutdown_rx) = registry_with_one_session(); + + tokio::time::timeout( + Duration::from_millis(500), + drain_supervisor_sessions(®istry, false), + ) + .await + .expect("a gateway without peer routing must not wait for a drain"); + + assert!(matches!(shutdown_rx.try_recv(), Err(TryRecvError::Empty))); + assert!(registry.has_session("sbx")); + assert_eq!(metrics.value("openshell_server_draining"), Some(0)); + } + + #[tokio::test] + async fn drain_supervisor_sessions_serves_through_propagation_delay_when_peer_routed() { + let metrics = MetricsCapture::install(); + describe_and_initialize(RELAY_CAPACITY); + let (registry, mut shutdown_rx) = registry_with_one_session(); + + let drain = drain_supervisor_sessions(®istry, true); + tokio::pin!(drain); + tokio::time::timeout(Duration::from_millis(250), &mut drain) + .await + .expect_err("a peer-routed gateway holds its sessions for the propagation delay"); + + assert_eq!(metrics.value("openshell_server_draining"), Some(1)); + assert!( + matches!(shutdown_rx.try_recv(), Err(TryRecvError::Empty)), + "sessions keep serving until the propagation delay ends" + ); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Ready + ); + } + static DETECTION_PROBE_ORDER: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); diff --git a/crates/openshell-server/src/persistence/mod.rs b/crates/openshell-server/src/persistence/mod.rs index c3a6649252..4048575fa3 100644 --- a/crates/openshell-server/src/persistence/mod.rs +++ b/crates/openshell-server/src/persistence/mod.rs @@ -4,6 +4,7 @@ //! Persistence layer for `OpenShell` Server. mod legacy_time_wire; +pub mod mutation_lock; mod postgres; mod sqlite; @@ -17,6 +18,7 @@ use rand::Rng; use std::collections::HashMap; use thiserror::Error; +pub use mutation_lock::{LockMode, MutationLockKey, MutationLockSet}; pub use postgres::PostgresStore; pub use sqlite::SqliteStore; @@ -33,6 +35,13 @@ pub type PersistenceResult = Result; /// requests into independently retryable, bounded write statements. pub const DELETE_MANY_BATCH_SIZE: usize = 128; +/// Maximum number of object ids read by one resource-version lookup statement. +/// +/// `PostgreSQL` binds each batch as a single `TEXT[]` parameter. `SQLite` +/// expands it into an `IN` list, so a batch plus the object-type parameter must +/// stay below `SQLite`'s bind-variable limit (32766 in the bundled library). +pub const RESOURCE_VERSION_BATCH_SIZE: usize = 1000; + /// Persistence-layer error type. #[derive(Debug, Error, Clone)] pub enum PersistenceError { @@ -58,6 +67,10 @@ pub enum PersistenceError { Conflict { current_resource_version: Option, }, + /// A mutation lock could not be acquired before its deadline. Nothing was + /// written; the operation is safe to retry. + #[error("mutation lock timeout: {0}")] + LockTimeout(String), } impl PersistenceError { @@ -202,6 +215,21 @@ pub struct DistributedMutationGuard { _postgres: Option, } +#[cfg(test)] +impl DistributedMutationGuard { + /// Backend process id of the `PostgreSQL` session holding the locks, or + /// `None` on `SQLite`. + pub(crate) async fn postgres_backend_pid(&mut self) -> Option { + let Self { + _postgres: postgres, + } = self; + match postgres { + Some(guard) => Some(guard.backend_pid().await), + None => None, + } + } +} + /// Trait for inferring an object type string from a message type. pub trait ObjectType { fn object_type() -> &'static str; @@ -285,15 +313,20 @@ impl Store { /// Serialize mutations whose invariants span multiple persisted objects. /// /// `SQLite` deployments are single-replica and use only the caller's local - /// mutex. `PostgreSQL` deployments additionally hold a session-level - /// advisory lock so concurrent gateway replicas cannot validate and write - /// the same cross-object invariant independently. + /// locks. `PostgreSQL` deployments additionally hold `locks` as + /// session-level advisory locks, taken in ascending key order on one + /// connection from the dedicated lock pool, so concurrent gateway replicas + /// cannot validate and write the same cross-object invariant + /// independently. Fails with [`PersistenceError::LockTimeout`] when the + /// locks are not acquired by `deadline`. pub async fn acquire_distributed_mutation_guard( &self, + locks: &MutationLockSet, + deadline: tokio::time::Instant, ) -> PersistenceResult { match self { Self::Postgres(store) => Ok(DistributedMutationGuard { - _postgres: Some(store.acquire_cross_object_lock().await?), + _postgres: Some(store.acquire_mutation_locks(locks, deadline).await?), }), Self::Sqlite(_) => Ok(DistributedMutationGuard { _postgres: None }), } @@ -594,6 +627,32 @@ impl Store { store_dispatch_traced!(self.delete_many(object_type, ids)) } + /// Read the authoritative `resource_version` of each listed object without + /// reading or decoding payloads. + /// + /// Ids that do not exist, or that belong to another object type, are absent + /// from the result. Backends issue one statement per + /// `RESOURCE_VERSION_BATCH_SIZE` ids and none for an empty list. When `ids` + /// spans several batches, each batch reads its own snapshot. + #[tracing::instrument( + name = "store", + skip_all, + fields( + otel.name = "store.get_resource_versions", + otel.status_code = tracing::field::Empty, + object_type = %object_type, + object_count = ids.len(), + batch_count = ids.len().div_ceil(RESOURCE_VERSION_BATCH_SIZE), + ) + )] + pub async fn get_resource_versions( + &self, + object_type: &str, + ids: &[String], + ) -> PersistenceResult> { + store_dispatch_traced!(self.get_resource_versions(object_type, ids)) + } + /// Count objects of a given type within a workspace. #[tracing::instrument( name = "store", @@ -663,21 +722,6 @@ impl Store { store_dispatch_traced!(self.list(object_type, workspace, limit, offset)) } - /// List objects by type across all workspaces. - #[tracing::instrument( - name = "store", - skip_all, - fields(otel.name = "store.list_by_type", otel.status_code = tracing::field::Empty, object_type = %object_type) - )] - pub async fn list_by_type( - &self, - object_type: &str, - limit: u32, - offset: u32, - ) -> PersistenceResult> { - store_dispatch_traced!(self.list_by_type(object_type, limit, offset)) - } - /// List workspace objects after a stable cursor, without offset drift. #[tracing::instrument( name = "store", @@ -1007,20 +1051,6 @@ impl Store { .collect() } - /// List and decode protobuf messages across all workspaces, hydrating - /// `resource_version` from the authoritative DB row. - pub async fn list_all_messages( - &self, - limit: u32, - offset: u32, - ) -> PersistenceResult> { - self.list_by_type(T::object_type(), limit, offset) - .await? - .into_iter() - .map(decode_record) - .collect() - } - /// List and decode objects that have a related membership record, with /// pagination. See [`Store::list_with_membership`] for details. pub async fn list_messages_with_membership< @@ -1231,6 +1261,12 @@ pub fn current_time_ms() -> i64 { } fn map_db_error(error: &sqlx::Error) -> PersistenceError { + // 55P03 lock_not_available: only lock-pool sessions set `lock_timeout`. + if let sqlx::Error::Database(db) = error + && db.code().as_deref() == Some("55P03") + { + return PersistenceError::LockTimeout(db.message().to_string()); + } if let sqlx::Error::Database(db) = error && db.is_unique_violation() { @@ -1375,5 +1411,11 @@ pub async fn test_store() -> Store { .expect("in-memory SQLite store should connect") } +#[cfg(test)] +pub mod test_postgres; + +#[cfg(test)] +mod mutation_lock_pg_tests; + #[cfg(test)] mod tests; diff --git a/crates/openshell-server/src/persistence/mutation_lock.rs b/crates/openshell-server/src/persistence/mutation_lock.rs new file mode 100644 index 0000000000..8ba5a0550e --- /dev/null +++ b/crates/openshell-server/src/persistence/mutation_lock.rs @@ -0,0 +1,267 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Keys, modes, and deadlines of the mutation locks that serialize +//! cross-object mutations across gateway replicas. +//! +//! The locks form a hierarchy of intention locks. Each key is held shared +//! (S) or exclusive (X): +//! +//! | Mutation | Keys | +//! |---|---| +//! | global policy and settings, platform-scope profiles | X(global) | +//! | providers and workspace-scoped profiles | S(global) X(workspace) | +//! | one sandbox, admin or supervisor | S(global) S(workspace) X(sandbox) | +//! | lifecycle, driver watch, reconcile (process-local only) | S(global) X(sandbox) | +//! | provisioning-deadline reconcile (process-local only) | S(global) S(workspace) X(sandbox) | +//! +//! Ordering rules, which make the scheme deadlock-free: +//! +//! 1. A per-sandbox lifecycle gate, where a path uses one, comes first. +//! 2. Process-local keys follow in ascending `i64` order. +//! 3. On `PostgreSQL`, the same keys follow as session-level advisory locks in +//! ascending order, all on one lock-pool connection. +//! 4. A task never acquires a mutation guard or a local lifecycle lock while +//! it holds one: no nesting and no upgrade. +//! +//! Within each layer every waiter on a key holds only smaller keys, and the +//! local phase ends before the `PostgreSQL` phase starts, so no wait-for cycle +//! can form. +//! +//! The global key is the legacy cross-object key. Gateways from earlier +//! releases hold it exclusively for every mutation, which conflicts with every +//! scope of this release, so mixed-version fleets stay mutually exclusive +//! during a rolling upgrade. + +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; +use std::time::Duration; + +/// Advisory-lock key of the global mutation lock. +/// +/// Never change this value: gateways from earlier releases hold it +/// exclusively for every cross-object mutation, and a rolling upgrade relies +/// on old and new replicas excluding each other through it. The bytes spell +/// "OPENSHLL" and stay within `PostgreSQL`'s signed 64-bit key space. +pub const GLOBAL_MUTATION_LOCK_KEY: i64 = 0x4f50_454e_5348_4c4c; + +/// Upper bound on acquiring one mutation lock set. +/// +/// The holder only validates and writes, so a wait this long means a stuck +/// replica or an overloaded database; failing beats blocking mutations +/// indefinitely. Keep [`MUTATION_LOCK_TIMEOUT_SETTING`] in sync. +pub const MUTATION_LOCK_TIMEOUT: Duration = Duration::from_secs(10); + +/// [`MUTATION_LOCK_TIMEOUT`] as a `PostgreSQL` `lock_timeout` value. +pub const MUTATION_LOCK_TIMEOUT_SETTING: &str = "10s"; + +/// Size of the dedicated `PostgreSQL` lock pool. +/// +/// Lock connections come from their own pool so that guard holders can never +/// starve the data pool their critical sections need. Each replica opens at +/// most 10 data plus 4 lock connections, so size `max_connections` for +/// rollouts as the high-availability guide describes +/// (`(2 × replicas + surge) × 14`). Each guard holds +/// one lock connection, so a replica sustains about 4 / c guarded operations +/// per second, where c is how long one guard is held. +pub(super) const MUTATION_LOCK_POOL_MAX_CONNECTIONS: u32 = 4; + +/// Domain separator hashed into every derived key. +const KEY_DOMAIN: &[u8] = b"openshell/mutation-lock/v1"; + +/// Advisory-lock key of the one-time time-payload migration +/// (`PostgresStore::migrate_legacy_time_payloads`). Derived keys never use it. +const TIME_PAYLOAD_MIGRATION_LOCK_KEY: i64 = 3052; + +/// Mode in which a mutation lock key is held. `Shared` sorts first, so the +/// maximum of two modes is the stronger one. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum LockMode { + Shared, + Exclusive, +} + +/// A mutation lock key. +#[derive(Clone, Copy, Debug)] +pub enum MutationLockKey<'a> { + /// The fleet-wide key, [`GLOBAL_MUTATION_LOCK_KEY`]. + Global, + /// One workspace, by name. + Workspace(&'a str), + /// One sandbox, by stable id. + Sandbox(&'a str), +} + +impl MutationLockKey<'_> { + /// The `PostgreSQL` advisory-lock key, also used by the process-local lock + /// table. + /// + /// Derived keys are the first 8 bytes, as a big-endian `i64`, of + /// `SHA-256(KEY_DOMAIN || 0 || kind || 0 || value)`. They are computed in + /// Rust so every replica and every `PostgreSQL` version agrees on them. A + /// hash collision only over-serializes. + pub fn advisory_key(self) -> i64 { + match self { + Self::Global => GLOBAL_MUTATION_LOCK_KEY, + Self::Workspace(workspace) => derived_key(b"workspace", workspace), + Self::Sandbox(sandbox_id) => derived_key(b"sandbox", sandbox_id), + } + } +} + +fn derived_key(kind: &[u8], value: &str) -> i64 { + let digest = Sha256::new() + .chain_update(KEY_DOMAIN) + .chain_update([0]) + .chain_update(kind) + .chain_update([0]) + .chain_update(value.as_bytes()) + .finalize(); + let mut prefix = [0_u8; 8]; + prefix.copy_from_slice(&digest[..8]); + avoid_reserved(i64::from_be_bytes(prefix)) +} + +/// Keep derived keys off the global and migration keys. +const fn avoid_reserved(key: i64) -> i64 { + if key == GLOBAL_MUTATION_LOCK_KEY || key == TIME_PAYLOAD_MIGRATION_LOCK_KEY { + key ^ 1 + } else { + key + } +} + +/// The keys one mutation holds, each in its strongest requested mode. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct MutationLockSet { + entries: BTreeMap, +} + +impl MutationLockSet { + /// Process-local lock set of a lifecycle, driver-watch, or reconcile path: + /// S(global) X(sandbox). + pub fn sandbox_lifecycle(sandbox_id: &str) -> Self { + let mut set = Self::default(); + set.insert(MutationLockKey::Global, LockMode::Shared); + set.insert(MutationLockKey::Sandbox(sandbox_id), LockMode::Exclusive); + set + } + + pub fn insert(&mut self, key: MutationLockKey<'_>, mode: LockMode) { + self.insert_raw(key.advisory_key(), mode); + } + + fn insert_raw(&mut self, key: i64, mode: LockMode) { + self.entries + .entry(key) + .and_modify(|held| *held = (*held).max(mode)) + .or_insert(mode); + } + + /// Keys in ascending order, the only acquisition order. + pub fn iter(&self) -> impl Iterator + '_ { + self.entries.iter().map(|(key, mode)| (*key, *mode)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn global_key_is_the_legacy_cross_object_key() { + assert_eq!( + MutationLockKey::Global.advisory_key(), + 0x4f50_454e_5348_4c4c + ); + } + + #[test] + fn timeout_setting_matches_duration() { + assert_eq!( + format!("{}s", MUTATION_LOCK_TIMEOUT.as_secs()), + MUTATION_LOCK_TIMEOUT_SETTING + ); + } + + #[test] + fn lock_set_iterates_in_ascending_key_order() { + let mut set = MutationLockSet::default(); + set.insert_raw(5, LockMode::Exclusive); + set.insert_raw(-3, LockMode::Exclusive); + + assert_eq!( + set.iter().collect::>(), + vec![(-3, LockMode::Exclusive), (5, LockMode::Exclusive)] + ); + } + + #[test] + fn derived_keys_match_golden_values() { + // Computed independently from the documented byte layout. Changing + // any of them breaks mutual exclusion with running replicas. + assert_eq!( + MutationLockKey::Workspace("default").advisory_key(), + 4_171_374_605_116_754_083 + ); + assert_eq!( + MutationLockKey::Workspace("team-a").advisory_key(), + 4_635_337_207_968_654_063 + ); + assert_eq!( + MutationLockKey::Sandbox("00000000-0000-0000-0000-000000000001").advisory_key(), + -542_384_872_970_356_635 + ); + assert_eq!( + MutationLockKey::Sandbox("sb-1").advisory_key(), + -7_385_842_969_463_770_825 + ); + } + + #[test] + fn derived_keys_separate_kinds() { + assert_ne!( + MutationLockKey::Workspace("x").advisory_key(), + MutationLockKey::Sandbox("x").advisory_key() + ); + } + + #[test] + fn reserved_keys_are_remapped() { + assert_ne!( + avoid_reserved(GLOBAL_MUTATION_LOCK_KEY), + GLOBAL_MUTATION_LOCK_KEY + ); + assert_eq!(avoid_reserved(TIME_PAYLOAD_MIGRATION_LOCK_KEY), 3053); + assert_eq!(avoid_reserved(42), 42); + } + + #[test] + fn lock_set_keeps_strongest_mode() { + let mut set = MutationLockSet::default(); + set.insert_raw(5, LockMode::Shared); + set.insert_raw(-3, LockMode::Exclusive); + set.insert_raw(5, LockMode::Exclusive); + set.insert_raw(-3, LockMode::Shared); + + assert_eq!( + set.iter().collect::>(), + vec![(-3, LockMode::Exclusive), (5, LockMode::Exclusive)] + ); + } + + #[test] + fn sandbox_lifecycle_set_is_shared_global_exclusive_sandbox() { + let set = MutationLockSet::sandbox_lifecycle("sb-1"); + + let mut expected = vec![ + (GLOBAL_MUTATION_LOCK_KEY, LockMode::Shared), + ( + MutationLockKey::Sandbox("sb-1").advisory_key(), + LockMode::Exclusive, + ), + ]; + expected.sort_unstable(); + assert_eq!(set.iter().collect::>(), expected); + } +} diff --git a/crates/openshell-server/src/persistence/mutation_lock_pg_tests.rs b/crates/openshell-server/src/persistence/mutation_lock_pg_tests.rs new file mode 100644 index 0000000000..5c52c0046c --- /dev/null +++ b/crates/openshell-server/src/persistence/mutation_lock_pg_tests.rs @@ -0,0 +1,872 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `PostgreSQL` tests of the mutation advisory locks: exclusion between two +//! stores (as between two gateway replicas), exclusion against the legacy +//! global key, cleanup after timed-out and cancelled acquisitions, and the +//! lock-pool bound. +//! +//! Advisory locks are database-wide, not per schema, so every test uses +//! random workspace and sandbox ids, and `mise run test:rust:postgres` runs +//! the tests one at a time. The stores here run no migrations: the schema only +//! scopes their connections. + +use super::mutation_lock::{ + GLOBAL_MUTATION_LOCK_KEY, MUTATION_LOCK_POOL_MAX_CONNECTIONS, MUTATION_LOCK_TIMEOUT, + MutationLockKey, MutationLockSet, +}; +use super::postgres::LOCK_CONNECTION_RELEASE_TIMEOUT; +use super::test_postgres::TestSchema; +use super::{ + DistributedMutationGuard, PersistenceError, PersistenceResult, PostgresStore, Store, + map_db_error, +}; +use crate::compute::MutationScope; +use sqlx::{Connection, PgConnection, PgPool}; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; +use tokio::sync::watch; +use tokio::task::{JoinHandle, JoinSet}; +use tokio::time::Instant; + +/// Deadline of an acquisition that must time out. +const EXPECTED_TIMEOUT: Duration = Duration::from_millis(300); +/// Deadline of an acquisition that must succeed, possibly after a conflicting +/// holder releases. +const PROCEEDS_WITHIN: Duration = Duration::from_secs(5); +/// Deadline of an acquisition that a test interrupts while it waits. +const INTERRUPTED_DEADLINE: Duration = Duration::from_secs(1); +/// How long after an interrupted acquisition's deadline its keys may stay held. +const RELEASED_WITHIN: Duration = Duration::from_secs(1); +/// The client-side backstop fires this long after a lock statement's deadline. +const CLIENT_BACKSTOP_GRACE: Duration = Duration::from_millis(500); +/// How long each holder keeps its locks in the throughput tests. +const HOLD: Duration = Duration::from_millis(20); +const POLL_INTERVAL: Duration = Duration::from_millis(10); + +fn random_id(kind: &str) -> String { + format!("{kind}-{}", uuid::Uuid::new_v4()) +} + +/// Drop client traffic while keeping sockets open, like a failed network path. +/// Client EOF still closes the upstream socket so `PostgreSQL` can release locks. +struct StallingProxy { + url: String, + stalled: watch::Sender, + task: JoinHandle<()>, +} + +impl StallingProxy { + async fn start(database_url: &str) -> Self { + let mut url = url::Url::parse(database_url).unwrap(); + let host = url.host_str().expect("TCP PostgreSQL host").to_owned(); + let port = url.port().unwrap_or(5432); + let upstream_addresses: Vec<_> = tokio::net::lookup_host((host.as_str(), port)) + .await + .unwrap() + .collect(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + url.set_host(Some("127.0.0.1")).unwrap(); + url.set_port(Some(listener.local_addr().unwrap().port())) + .unwrap(); + let (stalled, receiver) = watch::channel(false); + let task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (client, _) = accepted.unwrap(); + openshell_core::net::set_tcp_nodelay_best_effort(&client); + let addresses = upstream_addresses.clone(); + let stalled = receiver.clone(); + connections.spawn(async move { + let upstream = openshell_core::net::connect_tcp_nodelay_best_effort( + &addresses, + ).await?; + let (mut client_read, mut client_write) = client.into_split(); + let (mut upstream_read, mut upstream_write) = upstream.into_split(); + let requests = async { + let mut buffer = [0_u8; 8192]; + loop { + let read = client_read.read(&mut buffer).await?; + if read == 0 { + return Ok::<_, std::io::Error>(()); + } + if !*stalled.borrow() { + upstream_write.write_all(&buffer[..read]).await?; + } + } + }; + tokio::select! { + result = requests => result, + result = tokio::io::copy(&mut upstream_read, &mut client_write) => { + result.map(|_| ()) + } + } + }); + } + _ = connections.join_next(), if !connections.is_empty() => {} + } + } + }); + Self { + url: url.into(), + stalled, + task, + } + } +} + +impl Drop for StallingProxy { + fn drop(&mut self) { + self.task.abort(); + } +} + +/// A random sandbox id whose key sorts after the global key and the +/// workspace key. Keys are taken in ascending order, so an acquisition of its +/// sandbox scope already holds S(global) and S(workspace) while it waits for +/// the sandbox key. +fn sandbox_id_locked_last(workspace: &str) -> String { + let taken_first = + GLOBAL_MUTATION_LOCK_KEY.max(MutationLockKey::Workspace(workspace).advisory_key()); + loop { + let sandbox = random_id("sb"); + if MutationLockKey::Sandbox(&sandbox).advisory_key() > taken_first { + return sandbox; + } + } +} + +/// A disposable schema plus an observer pool for `pg_locks`. +struct LockFixture { + schema: TestSchema, + observer: PgPool, +} + +impl LockFixture { + async fn new() -> Self { + let schema = TestSchema::create("lock").await; + let observer = PgPool::connect(schema.url()) + .await + .expect("connect the pg_locks observer"); + Self { schema, observer } + } + + /// A store with its own data and lock pools, like one gateway replica. + /// + /// Its lock pool starts with one idle, connected session, so the first + /// acquisition spends its deadline on locks rather than on connecting. + /// Warming takes S(global), so create stores before any test holder + /// locks. + async fn store(&self, lock_pool_size: u32) -> Store { + let store = Store::Postgres( + PostgresStore::connect_with_lock_pool_size(self.schema.url(), lock_pool_size) + .await + .expect("connect a lock store"), + ); + drop( + acquire_proceeds( + &store, + MutationScope::sandbox(&random_id("ws"), &random_id("sb")), + "warm the lock pool", + ) + .await, + ); + wait_for_idle_lock_connection(&store).await; + store + } + + /// A plain session, like a gateway from an earlier release or a test + /// holder. + async fn raw_session(&self) -> PgConnection { + PgConnection::connect(self.schema.url()) + .await + .expect("connect a raw session") + } + + /// Granted and waiting holders of one bigint advisory key in this + /// database. + async fn lock_count(&self, key: i64) -> i64 { + let (high, low) = key_halves(key); + sqlx::query_scalar( + "SELECT count(*) FROM pg_locks \ + WHERE locktype = 'advisory' AND objsubid = 1 \ + AND classid = $1::bigint::oid AND objid = $2::bigint::oid \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database())", + ) + .bind(high) + .bind(low) + .fetch_one(&self.observer) + .await + .expect("count advisory locks") + } + + /// Advisory locks held or awaited by one backend. + async fn session_lock_count(&self, pid: i32) -> i64 { + sqlx::query_scalar("SELECT count(*) FROM pg_locks WHERE locktype = 'advisory' AND pid = $1") + .bind(pid) + .fetch_one(&self.observer) + .await + .expect("count a session's advisory locks") + } + + async fn wait_for_lock_count(&self, key: i64, expected: i64, until: Instant, what: &str) { + loop { + let count = self.lock_count(key).await; + if count == expected { + return; + } + assert!( + Instant::now() < until, + "{what}: {count} holders remain, expected {expected}" + ); + tokio::time::sleep(POLL_INTERVAL).await; + } + } + + /// Backend that waits for `key`, once an acquisition blocks on it. + async fn wait_for_waiting_backend(&self, key: i64, until: Instant) -> i32 { + let (high, low) = key_halves(key); + loop { + let waiting: Option = sqlx::query_scalar( + "SELECT pid FROM pg_locks \ + WHERE locktype = 'advisory' AND NOT granted AND objsubid = 1 \ + AND classid = $1::bigint::oid AND objid = $2::bigint::oid \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database())", + ) + .bind(high) + .bind(low) + .fetch_optional(&self.observer) + .await + .expect("find the waiting backend"); + if let Some(pid) = waiting { + return pid; + } + assert!( + Instant::now() < until, + "no backend waited for the sandbox key before the deadline" + ); + tokio::time::sleep(POLL_INTERVAL).await; + } + } + + /// Spawn an acquisition of `Sandbox { workspace, sandbox }` on `store`, + /// where another session holds the sandbox key, and wait until its + /// backend holds S(global) and S(workspace) and waits for the sandbox + /// key. Returns the task and the waiting backend's pid. + async fn spawn_waiting_acquisition( + &self, + store: &Store, + workspace: &str, + sandbox: &str, + deadline: Instant, + ) -> (JoinHandle>, i32) { + let acquisition = { + let store = store.clone(); + let set = MutationScope::sandbox(workspace, sandbox).lock_set(); + tokio::spawn(async move { + store + .acquire_distributed_mutation_guard(&set, deadline) + .await + .map(drop) + }) + }; + let pid = self + .wait_for_waiting_backend(MutationLockKey::Sandbox(sandbox).advisory_key(), deadline) + .await; + assert_eq!( + self.lock_count(MutationLockKey::Workspace(workspace).advisory_key()) + .await, + 1, + "the waiting acquisition holds the workspace key" + ); + assert_eq!( + self.lock_count(GLOBAL_MUTATION_LOCK_KEY).await, + 1, + "the waiting acquisition holds the global key" + ); + (acquisition, pid) + } + + async fn finish(self, stores: Vec) { + for store in stores { + store.close().await; + } + self.observer.close().await; + self.schema.drop_schema().await; + } +} + +/// `pg_locks` shows a bigint advisory key as its high and low 32 bits. +fn key_halves(key: i64) -> (i64, i64) { + let [b0, b1, b2, b3, b4, b5, b6, b7] = key.to_be_bytes(); + ( + i64::from(u32::from_be_bytes([b0, b1, b2, b3])), + i64::from(u32::from_be_bytes([b4, b5, b6, b7])), + ) +} + +async fn acquire( + store: &Store, + scope: MutationScope<'_>, + within: Duration, +) -> PersistenceResult { + store + .acquire_distributed_mutation_guard(&scope.lock_set(), Instant::now() + within) + .await +} + +async fn acquire_proceeds( + store: &Store, + scope: MutationScope<'_>, + what: &str, +) -> DistributedMutationGuard { + acquire(store, scope, PROCEEDS_WITHIN) + .await + .unwrap_or_else(|error| panic!("{what}: expected the locks, got {error:?}")) +} + +/// Wait until `store`'s lock pool has an idle, connected session. A released +/// lock connection returns to the pool from a background task, and an +/// acquisition that starts before then opens a new connection within its own +/// deadline. +async fn wait_for_idle_lock_connection(store: &Store) { + let Store::Postgres(postgres) = store else { + panic!("the lock tests use PostgreSQL stores"); + }; + let until = Instant::now() + PROCEEDS_WITHIN; + while postgres.lock_pool_idle() == 0 { + assert!( + Instant::now() < until, + "no lock connection returned to the pool" + ); + tokio::time::sleep(POLL_INTERVAL).await; + } +} + +/// Acquire `scope` on an idle lock connection with a deadline that +/// `PostgreSQL` must end with its lock timeout (55P03, "canceling statement +/// due to lock timeout"). A client-side timeout fails the test: a connect +/// that misses the deadline never reaches the advisory locks. +async fn acquire_times_out(store: &Store, scope: MutationScope<'_>, what: &str) { + wait_for_idle_lock_connection(store).await; + match acquire(store, scope, EXPECTED_TIMEOUT).await { + Err(PersistenceError::LockTimeout(detail)) if detail.contains("lock timeout") => {} + Err(error) => panic!("{what}: expected PostgreSQL's lock timeout, got {error:?}"), + Ok(_) => panic!("{what}: expected a lock timeout, but the locks were acquired"), + } +} + +async fn raw_lock(session: &mut PgConnection, key: i64) -> sqlx::Result<()> { + sqlx::query("SELECT pg_advisory_lock($1)") + .bind(key) + .execute(session) + .await + .map(drop) +} + +async fn raw_unlock(session: &mut PgConnection, key: i64) { + let released: bool = sqlx::query_scalar("SELECT pg_advisory_unlock($1)") + .bind(key) + .fetch_one(session) + .await + .expect("unlock the raw session's key"); + assert!(released, "the raw session held the key"); +} + +/// Spawn one holder per lock set, spread over `stores`, each keeping its +/// locks for [`HOLD`]; returns how long all of them took. +async fn hold_concurrently(stores: &[Store], sets: Vec) -> Duration { + let started = Instant::now(); + let holders: Vec<_> = sets + .into_iter() + .zip(stores.iter().cycle()) + .map(|(set, store)| { + let store = store.clone(); + tokio::spawn(async move { + let guard = store + .acquire_distributed_mutation_guard( + &set, + Instant::now() + MUTATION_LOCK_TIMEOUT, + ) + .await?; + tokio::time::sleep(HOLD).await; + drop(guard); + Ok::<_, PersistenceError>(()) + }) + }) + .collect(); + for holder in holders { + holder + .await + .expect("holder task") + .expect("holder acquires its locks"); + } + started.elapsed() +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_disjoint_sandbox_scopes_hold_concurrently_across_stores() { + let fixture = LockFixture::new().await; + let store_a = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let store_b = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let workspace = random_id("ws"); + let (sandbox_1, sandbox_2) = (random_id("sb"), random_id("sb")); + + let held_a = acquire_proceeds( + &store_a, + MutationScope::sandbox(&workspace, &sandbox_1), + "store A", + ) + .await; + let held_b = acquire_proceeds( + &store_b, + MutationScope::sandbox(&workspace, &sandbox_2), + "store B, another sandbox in the same workspace", + ) + .await; + assert_eq!( + fixture + .lock_count(MutationLockKey::Workspace(&workspace).advisory_key()) + .await, + 2, + "both sessions hold the workspace key shared at once" + ); + + drop(held_a); + drop(held_b); + fixture.finish(vec![store_a, store_b]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_stalled_release_closes_session_and_recovers_pool() { + let fixture = LockFixture::new().await; + let proxy = StallingProxy::start(fixture.schema.url()).await; + let store = Store::Postgres( + PostgresStore::connect_with_lock_pool_size(&proxy.url, 1) + .await + .unwrap(), + ); + let workspace = random_id("ws"); + let sandbox = random_id("sb"); + let scope = MutationScope::sandbox(&workspace, &sandbox); + let held = acquire_proceeds(&store, scope, "before the network stalls").await; + let key = MutationLockKey::Sandbox(&sandbox).advisory_key(); + assert_eq!(fixture.lock_count(key).await, 1); + + proxy.stalled.send_replace(true); + drop(held); + + // The old pool return waits forever for pg_advisory_unlock_all. Dropping + // a guard must instead bound cleanup and close the stalled connection. + fixture + .wait_for_lock_count( + key, + 0, + Instant::now() + LOCK_CONNECTION_RELEASE_TIMEOUT + Duration::from_secs(2), + "a stalled release must close its lock session", + ) + .await; + let Store::Postgres(postgres) = &store else { + unreachable!(); + }; + assert_eq!(postgres.lock_pool_size(), 0, "the pool permit is recovered"); + + proxy.stalled.send_replace(false); + drop(acquire_proceeds(&store, scope, "after network recovery").await); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_workspace_exclusive_blocks_same_workspace_sandbox_only() { + let fixture = LockFixture::new().await; + let store_a = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let store_b = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let (workspace_1, workspace_2) = (random_id("ws"), random_id("ws")); + let (sandbox_x, sandbox_y) = (random_id("sb"), random_id("sb")); + + let held = acquire_proceeds(&store_a, MutationScope::Workspace(&workspace_1), "store A").await; + acquire_times_out( + &store_b, + MutationScope::sandbox(&workspace_1, &sandbox_x), + "a sandbox in the held workspace", + ) + .await; + drop( + acquire_proceeds( + &store_b, + MutationScope::sandbox(&workspace_2, &sandbox_y), + "a sandbox in another workspace", + ) + .await, + ); + + drop(held); + drop( + acquire_proceeds( + &store_b, + MutationScope::sandbox(&workspace_1, &sandbox_x), + "the sandbox after the workspace holder releases", + ) + .await, + ); + fixture.finish(vec![store_a, store_b]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_global_exclusive_blocks_every_scope() { + let fixture = LockFixture::new().await; + let store_a = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let store_b = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let workspace = random_id("ws"); + let sandbox = random_id("sb"); + let scopes = [ + MutationScope::Global, + MutationScope::Workspace(""), + MutationScope::Workspace(&workspace), + MutationScope::sandbox(&workspace, &sandbox), + ]; + + let held = acquire_proceeds(&store_a, MutationScope::Global, "store A").await; + for scope in scopes { + acquire_times_out(&store_b, scope, &format!("{scope:?} behind the global key")).await; + } + + drop(held); + for scope in scopes { + drop(acquire_proceeds(&store_b, scope, &format!("{scope:?} after release")).await); + } + fixture.finish(vec![store_a, store_b]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_legacy_global_key_holder_excludes_new_scopes() { + let fixture = LockFixture::new().await; + let store = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let workspace = random_id("ws"); + let sandbox = random_id("sb"); + let scope = MutationScope::sandbox(&workspace, &sandbox); + + // A gateway from an earlier release holds the legacy key exclusively. + let mut legacy = fixture.raw_session().await; + raw_lock(&mut legacy, GLOBAL_MUTATION_LOCK_KEY) + .await + .expect("the legacy holder takes the global key"); + acquire_times_out(&store, scope, "a sandbox scope behind a legacy holder").await; + + raw_unlock(&mut legacy, GLOBAL_MUTATION_LOCK_KEY).await; + drop( + acquire_proceeds( + &store, + scope, + "a sandbox scope after the legacy holder releases", + ) + .await, + ); + + legacy.close().await.expect("close the legacy session"); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_new_scope_holder_excludes_legacy_global_key() { + let fixture = LockFixture::new().await; + let store = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let workspace = random_id("ws"); + let sandbox = random_id("sb"); + + let held = acquire_proceeds( + &store, + MutationScope::sandbox(&workspace, &sandbox), + "the new-release holder", + ) + .await; + let mut legacy = fixture.raw_session().await; + sqlx::query("SET lock_timeout = '200ms'") + .execute(&mut legacy) + .await + .expect("bound the legacy lock wait"); + let Err(error) = raw_lock(&mut legacy, GLOBAL_MUTATION_LOCK_KEY).await else { + panic!("the legacy global lock should wait behind a sandbox scope"); + }; + assert_eq!( + error + .as_database_error() + .and_then(sqlx::error::DatabaseError::code) + .as_deref(), + Some("55P03"), + "{error}" + ); + assert!( + matches!(map_db_error(&error), PersistenceError::LockTimeout(_)), + "55P03 maps to a lock timeout" + ); + + drop(held); + sqlx::query("SET lock_timeout = '5s'") + .execute(&mut legacy) + .await + .expect("bound the legacy lock wait"); + raw_lock(&mut legacy, GLOBAL_MUTATION_LOCK_KEY) + .await + .expect("the legacy lock after the new-release holder releases"); + raw_unlock(&mut legacy, GLOBAL_MUTATION_LOCK_KEY).await; + + legacy.close().await.expect("close the legacy session"); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_timeout_releases_held_keys_before_the_holder_does() { + let fixture = LockFixture::new().await; + // One lock connection, so a reused session keeps its backend pid. + let store = fixture.store(1).await; + let workspace = random_id("ws"); + let sandbox = sandbox_id_locked_last(&workspace); + let workspace_key = MutationLockKey::Workspace(&workspace).advisory_key(); + let sandbox_key = MutationLockKey::Sandbox(&sandbox).advisory_key(); + + // The raw session holds only the sandbox key, so the acquisition's own + // backend is the only holder of the global and workspace keys. + let mut holder = fixture.raw_session().await; + raw_lock(&mut holder, sandbox_key) + .await + .expect("the raw session takes the sandbox key"); + + let deadline = Instant::now() + INTERRUPTED_DEADLINE; + let (acquisition, waiting_pid) = fixture + .spawn_waiting_acquisition(&store, &workspace, &sandbox, deadline) + .await; + + let result = tokio::time::timeout_at( + deadline + CLIENT_BACKSTOP_GRACE + RELEASED_WITHIN, + acquisition, + ) + .await + .expect("the acquisition returns by its deadline") + .expect("acquisition task"); + assert!( + matches!(result, Err(PersistenceError::LockTimeout(_))), + "{result:?}" + ); + + // The raw session still holds the sandbox key, yet the global and + // workspace keys the acquisition took are already released. + fixture + .wait_for_lock_count( + workspace_key, + 0, + deadline + RELEASED_WITHIN, + "workspace key after the timeout", + ) + .await; + fixture + .wait_for_lock_count( + GLOBAL_MUTATION_LOCK_KEY, + 0, + deadline + RELEASED_WITHIN, + "global key after the timeout", + ) + .await; + assert_eq!(fixture.lock_count(sandbox_key).await, 1); + + // A server-side timeout returns the healthy session to the pool instead + // of closing it. + let mut reused = acquire_proceeds( + &store, + MutationScope::sandbox(&random_id("ws"), &random_id("sb")), + "a disjoint scope on the same lock connection", + ) + .await; + assert_eq!(reused.postgres_backend_pid().await, Some(waiting_pid)); + drop(reused); + + raw_unlock(&mut holder, sandbox_key).await; + holder.close().await.expect("close the raw session"); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_cancelled_acquisition_releases_held_keys_by_its_deadline() { + let fixture = LockFixture::new().await; + let store = fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await; + let workspace = random_id("ws"); + let sandbox = sandbox_id_locked_last(&workspace); + let workspace_key = MutationLockKey::Workspace(&workspace).advisory_key(); + let sandbox_key = MutationLockKey::Sandbox(&sandbox).advisory_key(); + + let mut holder = fixture.raw_session().await; + raw_lock(&mut holder, sandbox_key) + .await + .expect("the raw session takes the sandbox key"); + + // Cancel mid-wait, while the backend holds the global and workspace keys. + let deadline = Instant::now() + INTERRUPTED_DEADLINE; + let (acquisition, _) = fixture + .spawn_waiting_acquisition(&store, &workspace, &sandbox, deadline) + .await; + acquisition.abort(); + let Err(error) = acquisition.await else { + panic!("the cancelled acquisition should not finish"); + }; + assert!(error.is_cancelled()); + + // The closed session's backend does not notice the disconnect while it + // waits, but the statement's own lock_timeout ends the wait at the + // original deadline. Without it the keys would stay held for the full + // 10 s session backstop. + fixture + .wait_for_lock_count( + workspace_key, + 0, + deadline + RELEASED_WITHIN, + "workspace key after the cancellation", + ) + .await; + fixture + .wait_for_lock_count( + GLOBAL_MUTATION_LOCK_KEY, + 0, + deadline + RELEASED_WITHIN, + "global key after the cancellation", + ) + .await; + + raw_unlock(&mut holder, sandbox_key).await; + fixture + .wait_for_lock_count( + sandbox_key, + 0, + Instant::now() + PROCEEDS_WITHIN, + "sandbox key after the raw session unlocks", + ) + .await; + holder.close().await.expect("close the raw session"); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_released_connection_is_reused_without_residual_locks() { + let fixture = LockFixture::new().await; + let store = fixture.store(1).await; + let workspace = random_id("ws"); + let sandbox = random_id("sb"); + let scope = MutationScope::sandbox(&workspace, &sandbox); + + let mut first = acquire_proceeds(&store, scope, "the first acquisition").await; + let pid = first + .postgres_backend_pid() + .await + .expect("a PostgreSQL guard"); + assert_eq!(fixture.session_lock_count(pid).await, 3); + drop(first); + + // Returning the connection unlocks everything it held. + let until = Instant::now() + PROCEEDS_WITHIN; + while fixture.session_lock_count(pid).await != 0 { + assert!( + Instant::now() < until, + "the released session still holds advisory locks" + ); + tokio::time::sleep(POLL_INTERVAL).await; + } + + let mut second = acquire_proceeds(&store, scope, "the second acquisition").await; + assert_eq!(second.postgres_backend_pid().await, Some(pid)); + drop(second); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_pool_never_exceeds_configured_size() { + let fixture = LockFixture::new().await; + let postgres = PostgresStore::connect(fixture.schema.url()) + .await + .expect("connect a store with the production lock pool"); + let store = Store::Postgres(postgres.clone()); + let workspace = random_id("ws"); + + let holders: Vec<_> = (0..32) + .map(|_| { + let store = store.clone(); + let workspace = workspace.clone(); + tokio::spawn(async move { + let sandbox = random_id("sb"); + let guard = acquire( + &store, + MutationScope::sandbox(&workspace, &sandbox), + PROCEEDS_WITHIN, + ) + .await?; + tokio::time::sleep(HOLD).await; + drop(guard); + Ok::<_, PersistenceError>(()) + }) + }) + .collect(); + let mut largest = 0; + while !holders.iter().all(JoinHandle::is_finished) { + largest = largest.max(postgres.lock_pool_size()); + tokio::time::sleep(Duration::from_millis(1)).await; + } + for holder in holders { + holder + .await + .expect("holder task") + .expect("every holder acquires its locks"); + } + largest = largest.max(postgres.lock_pool_size()); + + assert_eq!( + largest, MUTATION_LOCK_POOL_MAX_CONNECTIONS, + "32 concurrent holders fill the lock pool without exceeding it" + ); + fixture.finish(vec![store]).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_mutation_lock_unrelated_sandboxes_outpace_global_serialization() { + const HOLDERS: usize = 64; + let fixture = LockFixture::new().await; + let stores = vec![ + fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await, + fixture.store(MUTATION_LOCK_POOL_MAX_CONNECTIONS).await, + ]; + let workspace = random_id("ws"); + + let global = hold_concurrently( + &stores, + (0..HOLDERS) + .map(|_| MutationScope::Global.lock_set()) + .collect(), + ) + .await; + let sandboxes = hold_concurrently( + &stores, + (0..HOLDERS) + .map(|_| MutationScope::sandbox(&workspace, &random_id("sb")).lock_set()) + .collect(), + ) + .await; + + // Global holders serialize (about HOLDERS x HOLD); distinct sandboxes + // share the two lock pools. Compare the runs, not absolute times. + assert!( + sandboxes * 3 < global, + "distinct sandboxes took {sandboxes:?}, global serialization took {global:?}" + ); + fixture.finish(stores).await; +} diff --git a/crates/openshell-server/src/persistence/postgres.rs b/crates/openshell-server/src/persistence/postgres.rs index 14bae6dfb1..6724e8769e 100644 --- a/crates/openshell-server/src/persistence/postgres.rs +++ b/crates/openshell-server/src/persistence/postgres.rs @@ -1,6 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use super::mutation_lock::{ + LockMode, MUTATION_LOCK_POOL_MAX_CONNECTIONS, MUTATION_LOCK_TIMEOUT, + MUTATION_LOCK_TIMEOUT_SETTING, MutationLockSet, +}; use super::{ DraftChunkRecord, ObjectCursor, ObjectListQuery, ObjectRecord, PersistenceError, PersistenceResult, PolicyRecord, WriteCondition, WriteResult, current_time_ms, map_db_error, @@ -17,6 +21,8 @@ use prost::Message; use sqlx::pool::PoolConnection; use sqlx::postgres::PgPoolOptions; use sqlx::{Connection, PgPool, Postgres, QueryBuilder, Row}; +use std::collections::HashMap; +use std::time::Duration; static POSTGRES_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations/postgres"); @@ -28,39 +34,160 @@ pub(super) fn embedded_migration_sql(version: i64) -> Option<&'static str> { .map(|migration| migration.sql.as_ref()) } -use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; +use super::{ + DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE, + RESOURCE_VERSION_BATCH_SIZE, +}; #[derive(Debug, Clone)] pub struct PostgresStore { pool: PgPool, + /// Dedicated connections for session-level mutation advisory locks. + lock_pool: PgPool, } -// Stable cluster-wide key for serializing sandbox/provider cross-object -// mutations. The bytes spell "OPENSHLL" and stay within PostgreSQL's signed -// 64-bit advisory-lock key space. -const CROSS_OBJECT_ADVISORY_LOCK_KEY: i64 = 0x4f50_454e_5348_4c4c; +/// How long the client waits past the caller's deadline for a lock statement. +/// +/// Each statement sets `lock_timeout` to the remaining deadline, so +/// `PostgreSQL` ends the wait on time; this timer only catches a stalled +/// connection. +const LOCK_STATEMENT_CLIENT_GRACE: Duration = Duration::from_millis(500); + +/// Bound the complete pool return, including the unlock hook and `SQLx`'s ping. +/// A stalled connection must not retain a lock-pool permit indefinitely. +pub(super) const LOCK_CONNECTION_RELEASE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Holds the mutation advisory locks of one acquisition. +/// +/// Returned to the lock pool on drop; `after_release` runs +/// `pg_advisory_unlock_all()` before reuse. The entire return is bounded by +/// [`LOCK_CONNECTION_RELEASE_TIMEOUT`]. An acquisition that `PostgreSQL` +/// times out (`55P03`) is returned the same way. A cancelled or otherwise failed +/// acquisition closes its session instead (see [`PendingLockConnection`]). +pub(super) struct PostgresAdvisoryLockGuard { + connection: PoolConnection, +} -// Bounds the wait for the cross-object lock. The holder only validates and -// writes, so a wait this long means a stuck replica; failing beats blocking -// every sandbox and provider mutation in the fleet indefinitely. -const CROSS_OBJECT_ADVISORY_LOCK_TIMEOUT: &str = "10s"; +impl Drop for PostgresAdvisoryLockGuard { + fn drop(&mut self) { + return_lock_connection(&mut self.connection); + } +} -pub(super) struct PostgresAdvisoryLockGuard { - // `close_on_drop` is set before this guard is constructed. Closing the - // dedicated session releases the session-level advisory lock even when a - // request is cancelled or returns early. - _connection: PoolConnection, +fn return_lock_connection(connection: &mut PoolConnection) { + // SQLx transfers the connection and pool permit into this owned future + // immediately. Dropping it on timeout closes the socket and releases the + // permit, including when the unlock hook succeeded but the final ping stalls. + // This relies on sqlx-core 0.9's doc-hidden `PoolConnection::return_to_pool`; + // rerun the `postgres_mutation_lock_stalled_release_*` test on any sqlx bump. + let returning = connection.return_to_pool(); + tokio::spawn(async move { + if tokio::time::timeout(LOCK_CONNECTION_RELEASE_TIMEOUT, returning) + .await + .is_err() + { + tracing::warn!( + "timed out returning PostgreSQL mutation lock connection; discarded connection" + ); + } + }); +} + +#[cfg(test)] +impl PostgresAdvisoryLockGuard { + /// Backend process id of the session that holds the locks. + pub(super) async fn backend_pid(&mut self) -> i32 { + let Self { connection } = self; + sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut **connection) + .await + .expect("read the lock session's backend pid") + } +} + +/// A lock-pool connection whose acquisition is still in progress. +/// +/// Dropping it closes the session, which releases every advisory lock the +/// backend holds once its current statement ends. A backend blocked on a lock +/// does not notice the closed socket, so each lock statement carries its own +/// `lock_timeout` to bound that wait by the caller's deadline. +struct PendingLockConnection(Option>); + +impl PendingLockConnection { + fn connection(&mut self) -> &mut PoolConnection { + self.0 + .as_mut() + .expect("pending lock connection is present until disarmed") + } + + /// Keep the session: every lock was acquired. + fn into_connection(mut self) -> PoolConnection { + self.0 + .take() + .expect("pending lock connection is present until disarmed") + } + + /// Return the healthy session to the pool, whose `after_release` unlocks + /// any keys it already holds. + fn release(mut self) { + if let Some(mut connection) = self.0.take() { + return_lock_connection(&mut connection); + } + } +} + +impl Drop for PendingLockConnection { + fn drop(&mut self) { + if let Some(connection) = self.0.as_mut() { + connection.close_on_drop(); + } + } } impl PostgresStore { pub async fn connect(url: &str) -> PersistenceResult { + Self::connect_with_lock_pool_size(url, MUTATION_LOCK_POOL_MAX_CONNECTIONS).await + } + + pub(super) async fn connect_with_lock_pool_size( + url: &str, + lock_pool_size: u32, + ) -> PersistenceResult { let pool = PgPoolOptions::new() .max_connections(10) .connect(url) .await .map_err(|e| map_db_error(&e))?; + let lock_pool = PgPoolOptions::new() + .max_connections(lock_pool_size) + .min_connections(0) + // Backstop only; callers bound the acquire by their own deadline. + .acquire_timeout(MUTATION_LOCK_TIMEOUT) + .after_connect(|connection, _metadata| { + Box::pin(async move { + // Backstop only: every lock statement sets the remaining + // deadline of its own acquisition. + sqlx::query("SELECT set_config('lock_timeout', $1, false)") + .bind(MUTATION_LOCK_TIMEOUT_SETTING) + .execute(&mut *connection) + .await?; + Ok(()) + }) + }) + .after_release(|connection, _metadata| { + Box::pin(async move { + // Scrub every returned lock connection. On error sqlx closes + // it, and the backend exit releases whatever it held. + sqlx::query("SELECT pg_advisory_unlock_all()") + .execute(&mut *connection) + .await?; + Ok(true) + }) + }) + .connect_lazy(url) + .map_err(|e| map_db_error(&e))?; - Ok(Self { pool }) + Ok(Self { pool, lock_pool }) } pub async fn migrate(&self) -> PersistenceResult<()> { @@ -114,32 +241,97 @@ impl PostgresStore { conn.ping().await.map_err(|e| map_db_error(&e)) } - pub(super) async fn acquire_cross_object_lock( + /// Acquire `locks` as session-level advisory locks, in ascending key + /// order, on one lock-pool connection. + /// + /// Fails with [`PersistenceError::LockTimeout`] when the connection or a + /// lock is not available by `deadline`. + pub(super) async fn acquire_mutation_locks( &self, + locks: &MutationLockSet, + deadline: tokio::time::Instant, ) -> PersistenceResult { - let mut connection = self.pool.acquire().await.map_err(|e| map_db_error(&e))?; - connection.close_on_drop(); - sqlx::query("SELECT set_config('lock_timeout', $1, false)") - .bind(CROSS_OBJECT_ADVISORY_LOCK_TIMEOUT) - .execute(&mut *connection) - .await - .map_err(|e| map_db_error(&e))?; - sqlx::query("SELECT pg_advisory_lock($1)") - .bind(CROSS_OBJECT_ADVISORY_LOCK_KEY) - .execute(&mut *connection) - .await - .map_err(|e| map_db_error(&e))?; + let connection = match tokio::time::timeout_at(deadline, self.lock_pool.acquire()).await { + Err(_) | Ok(Err(sqlx::Error::PoolTimedOut)) => { + return Err(PersistenceError::LockTimeout( + "waiting for a mutation lock connection".into(), + )); + } + Ok(Err(error)) => return Err(map_db_error(&error)), + Ok(Ok(connection)) => connection, + }; + let mut pending = PendingLockConnection(Some(connection)); + for (key, mode) in locks.iter() { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining < Duration::from_millis(1) { + // Nothing waits server-side; `after_release` unlocks the keys + // taken so far. + pending.release(); + return Err(PersistenceError::LockTimeout( + "waiting for a PostgreSQL advisory lock".into(), + )); + } + // `set_config` and the lock run in one statement. A CTE that calls + // a volatile function is never inlined, and the outer projection + // needs its row, so `lock_timeout` is set before the lock wait + // starts. PostgreSQL then abandons the wait at the caller's + // deadline even if this future is cancelled and the socket closed. + let sql = match mode { + LockMode::Shared => { + "WITH timeout AS (SELECT set_config('lock_timeout', $1, false)) \ + SELECT pg_advisory_lock_shared($2) FROM timeout" + } + LockMode::Exclusive => { + "WITH timeout AS (SELECT set_config('lock_timeout', $1, false)) \ + SELECT pg_advisory_lock($2) FROM timeout" + } + }; + let statement = sqlx::query(sql) + .bind(format!("{}ms", remaining.as_millis())) + .bind(key) + .execute(&mut **pending.connection()); + match tokio::time::timeout_at(deadline + LOCK_STATEMENT_CLIENT_GRACE, statement).await { + Err(_) => { + return Err(PersistenceError::LockTimeout( + "waiting for a PostgreSQL advisory lock".into(), + )); + } + Ok(Err(error)) => { + let error = map_db_error(&error); + if matches!(error, PersistenceError::LockTimeout(_)) { + // Healthy session: return it; `after_release` unlocks + // the keys already held. + pending.release(); + } + return Err(error); + } + Ok(Ok(_)) => {} + } + } Ok(PostgresAdvisoryLockGuard { - _connection: connection, + connection: pending.into_connection(), }) } - /// Test support only: close the underlying connection pool. + /// Connections the lock pool holds, idle or in use. + #[cfg(test)] + pub(super) fn lock_pool_size(&self) -> u32 { + self.lock_pool.size() + } + + /// Lock-pool connections that are connected and ready for reuse. + #[cfg(test)] + pub(super) fn lock_pool_idle(&self) -> usize { + self.lock_pool.num_idle() + } + + /// Test support only: close the underlying connection pools. /// - /// Do not call from runtime code; this tears down the active pool. + /// Do not call from runtime code; this tears down the active pools. #[cfg(any(test, feature = "test-support"))] pub async fn close(&self) { self.pool.close().await; + self.lock_pool.close().await; } pub async fn put( @@ -587,6 +779,36 @@ WHERE object_type = $1 AND workspace = $2 AND name = $3 Ok(deleted) } + pub async fn get_resource_versions( + &self, + object_type: &str, + ids: &[String], + ) -> PersistenceResult> { + let mut versions = HashMap::with_capacity(ids.len()); + for ids in ids.chunks(RESOURCE_VERSION_BATCH_SIZE) { + // The chunk binds as one TEXT[] parameter, so every chunk length + // shares one prepared statement and uses the primary-key index. + let rows = sqlx::query( + "SELECT id, resource_version FROM objects WHERE object_type = $1 AND id = ANY($2)", + ) + .bind(object_type) + .bind(ids) + .fetch_all(&self.pool) + .await + .map_err(|e| map_db_error(&e))?; + for row in rows { + let id: String = row.try_get("id").map_err(|e| map_db_error(&e))?; + let resource_version: i64 = row + .try_get("resource_version") + .map_err(|e| map_db_error(&e))?; + // Same normalization as row_to_object_record, so the poller + // compares exactly what get_message would have returned. + versions.insert(id, resource_version.max(1).cast_unsigned()); + } + } + Ok(versions) + } + pub async fn count_in_workspace( &self, object_type: &str, @@ -672,30 +894,6 @@ LIMIT $3 OFFSET $4 Ok(rows.into_iter().map(row_to_object_record).collect()) } - pub async fn list_by_type( - &self, - object_type: &str, - limit: u32, - offset: u32, - ) -> PersistenceResult> { - let rows = sqlx::query( - r" -SELECT object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version -FROM objects -WHERE object_type = $1 -ORDER BY created_at_ms ASC, name ASC, workspace ASC, id ASC -LIMIT $2 OFFSET $3 -", - ) - .bind(object_type) - .bind(i64::from(limit)) - .bind(i64::from(offset)) - .fetch_all(&self.pool) - .await - .map_err(|e| map_db_error(&e))?; - - Ok(rows.into_iter().map(row_to_object_record).collect()) - } pub async fn list_after( &self, object_type: &str, @@ -1017,7 +1215,7 @@ LIMIT $3 OFFSET $4 load_error: None, created_at_ms: now_ms, loaded_at_ms: None, - provenance: std::collections::HashMap::default(), + provenance: HashMap::default(), }; let wrapped_payload = policy_payload_from_record(&record)?; diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index a2a5e74f60..26fd077d64 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -19,6 +19,7 @@ use sqlx::sqlite::{ SqliteConnectOptions, SqliteConnection, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous, }; use sqlx::{Connection, QueryBuilder, Row, Sqlite, SqlitePool}; +use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; @@ -36,7 +37,10 @@ pub(super) fn embedded_migration_sql(version: i64) -> Option<&'static str> { } static IN_MEMORY_DB_SEQUENCE: AtomicU64 = AtomicU64::new(0); -use super::{DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE}; +use super::{ + DELETE_MANY_BATCH_SIZE, DRAFT_CHUNK_OBJECT_TYPE, POLICY_OBJECT_TYPE, + RESOURCE_VERSION_BATCH_SIZE, +}; #[derive(Debug, Clone)] pub struct SqliteStore { @@ -757,6 +761,38 @@ WHERE "object_type" = ?1 AND "id" = ?2 Ok(deleted) } + pub async fn get_resource_versions( + &self, + object_type: &str, + ids: &[String], + ) -> PersistenceResult> { + let mut versions = HashMap::with_capacity(ids.len()); + for ids in ids.chunks(RESOURCE_VERSION_BATCH_SIZE) { + let mut query = QueryBuilder::::new( + "SELECT id, resource_version FROM objects WHERE object_type = ", + ); + query.push_bind(object_type).push(" AND id IN ("); + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + let rows = query + .build() + .fetch_all(&self.pool) + .await + .map_err(|e| map_db_error(&e))?; + for row in rows { + let id: String = row.try_get("id").map_err(|e| map_db_error(&e))?; + let resource_version: i64 = row + .try_get("resource_version") + .map_err(|e| map_db_error(&e))?; + versions.insert(id, resource_version.max(1).cast_unsigned()); + } + } + Ok(versions) + } + pub async fn count_in_workspace( &self, object_type: &str, @@ -858,30 +894,6 @@ LIMIT ?3 OFFSET ?4 Ok(rows.into_iter().map(row_to_object_record).collect()) } - pub async fn list_by_type( - &self, - object_type: &str, - limit: u32, - offset: u32, - ) -> PersistenceResult> { - let rows = sqlx::query( - r#" -SELECT "object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version" -FROM "objects" -WHERE "object_type" = ?1 -ORDER BY "created_at_ms" ASC, "name" ASC, "workspace" ASC, "id" ASC -LIMIT ?2 OFFSET ?3 -"#, - ) - .bind(object_type) - .bind(i64::from(limit)) - .bind(i64::from(offset)) - .fetch_all(&self.pool) - .await - .map_err(|e| map_db_error(&e))?; - - Ok(rows.into_iter().map(row_to_object_record).collect()) - } pub async fn list_after( &self, object_type: &str, @@ -1247,7 +1259,7 @@ WHERE o."object_type" = "#, load_error: None, created_at_ms: now_ms, loaded_at_ms: None, - provenance: std::collections::HashMap::default(), + provenance: HashMap::default(), }; let wrapped_payload = policy_payload_from_record(&record)?; diff --git a/crates/openshell-server/src/persistence/test_postgres.rs b/crates/openshell-server/src/persistence/test_postgres.rs new file mode 100644 index 0000000000..8e31d8d80d --- /dev/null +++ b/crates/openshell-server/src/persistence/test_postgres.rs @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Disposable `PostgreSQL` schemas for `#[ignore]`d backend tests. + +use super::Store; + +/// Names the disposable `PostgreSQL` database for `#[ignore]`d tests. +/// `mise run test:rust:postgres` always exports it. +pub const TEST_POSTGRES_URL_ENV: &str = "OPENSHELL_TEST_POSTGRES_URL"; + +/// A uniquely named schema. Stores connected through it see only that schema. +pub struct TestSchema { + admin: sqlx::PgPool, + schema: String, + url: String, +} + +impl TestSchema { + /// Create a schema named `_` in the database named by + /// [`TEST_POSTGRES_URL_ENV`]. + /// + /// Panics when the variable is unset; callers are `#[ignore]`d. + pub async fn create(prefix: &str) -> Self { + let base = std::env::var(TEST_POSTGRES_URL_ENV).unwrap_or_else(|_| { + panic!( + "{TEST_POSTGRES_URL_ENV} must name a disposable PostgreSQL database; \ + run mise run test:rust:postgres" + ) + }); + assert!( + base.starts_with("postgres"), + "{TEST_POSTGRES_URL_ENV} must be a postgres:// URL" + ); + let schema = format!("{prefix}_{}", uuid::Uuid::new_v4().simple()); + let admin = sqlx::PgPool::connect(&base) + .await + .expect("connect to the disposable PostgreSQL database"); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE SCHEMA {schema}"))) + .execute(&admin) + .await + .expect("create the test schema"); + let mut url = url::Url::parse(&base).expect("parse the PostgreSQL URL"); + url.query_pairs_mut() + .append_pair("options", &format!("-csearch_path={schema}")); + Self { + admin, + schema, + url: url.into(), + } + } + + /// Connection URL scoped to this schema. + pub fn url(&self) -> &str { + &self.url + } + + /// A new store with its own pool, as a separate gateway replica would + /// have. Runs migrations. + pub async fn connect_store(&self) -> Store { + Store::connect(&self.url) + .await + .expect("connect a store to the test schema") + } + + /// Drops only this test's schema. + pub async fn drop_schema(self) { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DROP SCHEMA {} CASCADE", + self.schema + ))) + .execute(&self.admin) + .await + .expect("drop the test schema"); + self.admin.close().await; + } +} diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 7ba0bd1245..20f647f1ac 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -804,6 +804,130 @@ async fn delete_many_is_bounded_idempotent_and_type_scoped() { assert!(store.get("provider", "other-type").await.unwrap().is_some()); } +/// Shared by the `SQLite` and `PostgreSQL` tests: spans two lookup batches, +/// mixes in missing, other-type and duplicate ids, and checks that versions +/// follow updates and deletions. +async fn assert_get_resource_versions_contract(store: &Store) { + let batch = super::RESOURCE_VERSION_BATCH_SIZE; + let mut ids = Vec::new(); + for idx in 0..(batch + 12) { + let id = format!("sandbox-{idx}"); + store + .put( + "sandbox", + &id, + &format!("name-{idx}"), + "default", + b"payload", + None, + ) + .await + .unwrap(); + ids.push(id); + } + store + .put( + "provider", + "other-type", + "other-type", + "default", + b"payload", + None, + ) + .await + .unwrap(); + + let mut queried = ids.clone(); + queried.extend([ + "missing".to_string(), + "other-type".to_string(), + "sandbox-0".to_string(), + ]); + let versions = store + .get_resource_versions("sandbox", &queried) + .await + .unwrap(); + assert_eq!(versions.len(), batch + 12); + for id in &ids { + let record = store.get("sandbox", id).await.unwrap().unwrap(); + assert_eq!(versions[id], record.resource_version, "{id}"); + assert_eq!(versions[id], 1, "{id}"); + } + assert!(!versions.contains_key("missing")); + assert!(!versions.contains_key("other-type")); + + // `chunks` follows input order, so these land in the first and second batch. + for idx in [0, batch] { + store + .put_if( + "sandbox", + &format!("sandbox-{idx}"), + &format!("name-{idx}"), + "default", + b"payload-2", + None, + super::WriteCondition::MatchResourceVersion(1), + ) + .await + .unwrap(); + } + let versions = store.get_resource_versions("sandbox", &ids).await.unwrap(); + assert_eq!(versions["sandbox-0"], 2); + assert_eq!(versions[&format!("sandbox-{batch}")], 2); + assert_eq!(versions["sandbox-1"], 1); + + assert!(store.delete("sandbox", "sandbox-1").await.unwrap()); + let versions = store.get_resource_versions("sandbox", &ids).await.unwrap(); + assert!(!versions.contains_key("sandbox-1")); + assert_eq!(versions.len(), batch + 11); + + let versions = store + .get_resource_versions( + "provider", + &["other-type".to_string(), "sandbox-0".to_string()], + ) + .await + .unwrap(); + assert_eq!( + versions, + StdHashMap::from([("other-type".to_string(), 1_u64)]) + ); + + assert!( + store + .get_resource_versions("sandbox", &[]) + .await + .unwrap() + .is_empty() + ); + + // More ids than SQLite accepts as bound variables in one statement (32766), + // so an unchunked lookup fails with "too many SQL variables". + let mut many: Vec = (0..40_000).map(|i| format!("missing-{i}")).collect(); + many.push("sandbox-0".to_string()); + let versions = store.get_resource_versions("sandbox", &many).await.unwrap(); + assert_eq!( + versions, + StdHashMap::from([("sandbox-0".to_string(), 2_u64)]) + ); +} + +#[tokio::test] +async fn get_resource_versions_is_batched_type_scoped_and_omits_missing() { + let store = test_store().await; + assert_get_resource_versions_contract(&store).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_get_resource_versions_binds_text_array_across_batches() { + let schema = super::test_postgres::TestSchema::create("rv").await; + let store = schema.connect_store().await; + assert_get_resource_versions_contract(&store).await; + store.close().await; + schema.drop_schema().await; +} + #[tokio::test] async fn file_backed_sqlite_bulk_delete_allows_concurrent_control_reads() { use std::sync::Arc; diff --git a/crates/openshell-server/src/provider_refresh.rs b/crates/openshell-server/src/provider_refresh.rs index e8b39cbac4..7ee295595d 100644 --- a/crates/openshell-server/src/provider_refresh.rs +++ b/crates/openshell-server/src/provider_refresh.rs @@ -1158,6 +1158,11 @@ async fn apply_minted_credential( credential_key: &str, minted: &MintedCredential, ) -> Result<(), Status> { + // Validate the expiration before staging anything, so this conversion can + // never leave staged credential handles behind. + let credential_expiration_time = + openshell_core::time::optional_timestamp_from_legacy_millis(minted.expires_at_ms) + .map_err(|error| Status::internal(error.to_string()))?; let mut updated = provider.clone(); let staging_id = format!("{}-refresh-{}", provider.object_id(), uuid::Uuid::new_v4()); let staged_handles = if let Some(credentials) = credentials @@ -1205,9 +1210,6 @@ async fn apply_minted_credential( } None }; - let credential_expiration_time = - openshell_core::time::optional_timestamp_from_legacy_millis(minted.expires_at_ms) - .map_err(|error| Status::internal(error.to_string()))?; if let Some(expiration_time) = credential_expiration_time.as_ref() { updated .credential_expiration_times @@ -1223,14 +1225,28 @@ async fn apply_minted_credential( updated.credential_expiration_times.remove(key); } } - // Acquire the shared sandbox mutation boundary only around validation and + // Acquire the workspace mutation key only around validation and // persistence, after any remote minting or credential staging. This // prevents route status from committing against the old provider revision // after the rotation writes, without holding the guard across network I/O. - let _sandbox_sync_guard = if let Some(compute) = compute { - Some(compute.sandbox_sync_guard().await.map_err(|error| { - Status::internal(format!("acquire provider mutation lock: {error}")) - })?) + let _mutation_guard = if let Some(compute) = compute { + match compute + .mutation_guard(crate::compute::MutationScope::Workspace(workspace)) + .await + { + Ok(guard) => Some(guard), + Err(error) => { + if let Some(credentials) = credentials + && let Some(handles) = &staged_handles + { + cleanup_staged_refresh_handles(credentials, provider, handles).await; + } + return Err(crate::grpc::persistence_error_to_status( + error, + "acquire provider mutation lock", + )); + } + } } else { None }; @@ -4103,6 +4119,83 @@ mod tests { assert_eq!(credentials.stored_credential_count(), Some(0)); } + #[tokio::test] + async fn apply_minted_credential_returns_unavailable_when_guard_times_out() { + use super::apply_minted_credential; + + let state = crate::grpc::test_support::test_server_state().await; + state + .compute + .set_mutation_lock_timeout_for_tests(std::time::Duration::from_millis(50)); + let credentials = test_credentials(); + let mut prov = provider("guarded-aws", "aws"); + let original_handles = credentials + .store_provider_credentials( + prov.object_name(), + prov.object_workspace(), + prov.object_id(), + &HashMap::from([( + "AWS_ACCESS_KEY_ID".to_string(), + "old-access-key".to_string(), + )]), + &HashMap::new(), + ) + .await + .unwrap(); + prov.credential_handles.clone_from(&original_handles); + let stored_credential_count = credentials.stored_credential_count(); + state.store.put_message(&prov).await.unwrap(); + let provider_writer = state + .compute + .mutation_guard(crate::compute::MutationScope::Workspace("default")) + .await + .unwrap(); + + let minted = super::MintedCredential { + access_token: "AKIAIOSFODNN7EXAMPLE".to_string(), + expires_at_ms: 4_000_000_000_000, + refresh_token: None, + additional_credentials: HashMap::new(), + }; + let err = apply_minted_credential( + &state.store, + "default", + Some(&credentials), + Some(&state.compute), + &prov, + "AWS_ACCESS_KEY_ID", + &minted, + ) + .await + .unwrap_err(); + assert_eq!(err.code(), tonic::Code::Unavailable); + let details = openshell_core::rpc_error::decode_details(&err).expect("error details"); + assert_eq!( + details.error_info().expect("error info").reason, + "MUTATION_LOCK_TIMEOUT" + ); + let stored = state + .store + .get_message_by_name::("default", "guarded-aws") + .await + .unwrap() + .unwrap(); + assert_eq!(stored.credential_handles, original_handles); + let resolved = credentials + .resolve_provider_handles(&stored, current_time_ms()) + .await + .unwrap(); + assert_eq!( + resolved.values.get("AWS_ACCESS_KEY_ID"), + Some(&"old-access-key".to_string()) + ); + assert_eq!( + credentials.stored_credential_count(), + stored_credential_count + ); + drop(provider_writer); + } + // A wiremock responder that blocks the STS response until the test releases // it, so a delete-refresh can be interleaved deterministically while the // rotation is parked awaiting STS. diff --git a/crates/openshell-server/src/sandbox_watch.rs b/crates/openshell-server/src/sandbox_watch.rs index 679912d352..091382965a 100644 --- a/crates/openshell-server/src/sandbox_watch.rs +++ b/crates/openshell-server/src/sandbox_watch.rs @@ -5,12 +5,14 @@ use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::time::{Duration, Instant}; +use metrics::{counter, gauge, histogram}; use openshell_core::proto::SandboxStreamWarning; use tokio::sync::{broadcast, watch}; -use crate::persistence::Store; +use crate::gateway_metrics; +use crate::persistence::{ObjectType, PersistenceResult, Store}; use openshell_core::proto::Sandbox; /// How often [`spawn_store_poller`] rechecks watched sandboxes for writes made @@ -78,11 +80,148 @@ impl SandboxWatchBus { } } +/// Source of authoritative sandbox resource versions for the cross-replica +/// watch poller. [`Store`] is the production implementation; tests supply +/// fakes that count lookups and inject failures. +trait SandboxVersionSource: Sync { + /// Current `resource_version` of each listed sandbox that exists. Missing + /// sandboxes are absent from the map. + fn sandbox_versions( + &self, + ids: &[String], + ) -> impl Future>> + Send; +} + +impl SandboxVersionSource for Store { + async fn sandbox_versions(&self, ids: &[String]) -> PersistenceResult> { + self.get_resource_versions(Sandbox::object_type(), ids) + .await + } +} + +/// State the poller carries between ticks. +#[derive(Debug, Default)] +struct PollerState { + /// Last observed version per actively watched sandbox. `None` records a + /// sandbox that was missing (deleted) at the last successful lookup. + known_versions: HashMap>, + /// Consecutive ticks whose lookup failed. Only the first failure of a + /// streak is logged as a warning. + consecutive_failures: u32, +} + +/// What one poller tick did. Returned for tests; the loop ignores it. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct PollOutcome { + /// Watched sandboxes included in this tick's lookup. + polled: usize, + /// Sandboxes whose local watchers were notified. + notified: usize, + /// Whether the lookup failed. Known versions were left untouched. + failed: bool, +} + +/// Run one poller tick. +async fn poll_once( + source: &S, + bus: &SandboxWatchBus, + state: &mut PollerState, +) -> PollOutcome { + let active = bus.active_sandbox_ids(); + state + .known_versions + .retain(|sandbox_id, _| active.contains(sandbox_id)); + gauge!(gateway_metrics::SANDBOX_WATCH_POLLED_SANDBOXES) + .set(gateway_metrics::count_as_f64(active.len())); + if active.is_empty() { + return PollOutcome::default(); + } + check_versions(source, bus, state, active.into_iter().collect()).await +} + +#[tracing::instrument( + name = "sandbox_watch", + skip_all, + fields( + otel.name = "sandbox_watch.poll", + otel.status_code = tracing::field::Empty, + watched_count = ids.len(), + notified_count = tracing::field::Empty, + ) +)] +async fn check_versions( + source: &S, + bus: &SandboxWatchBus, + state: &mut PollerState, + ids: Vec, +) -> PollOutcome { + let started = Instant::now(); + let result = source.sandbox_versions(&ids).await; + histogram!(gateway_metrics::SANDBOX_WATCH_POLL_DURATION_SECONDS) + .record(started.elapsed().as_secs_f64()); + + let versions = match result { + Ok(versions) => versions, + Err(err) => { + crate::otel_tracing::mark_error(&tracing::Span::current()); + counter!(gateway_metrics::SANDBOX_WATCH_POLL_ERRORS_TOTAL).increment(1); + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + if state.consecutive_failures == 1 { + tracing::warn!( + watched_count = ids.len(), + error = %err, + "sandbox watch poller: failed to read persisted sandbox versions; retrying every interval" + ); + } + // Keep known versions so the next successful tick still sees + // every change and deletion made during the outage. + return PollOutcome { + polled: ids.len(), + notified: 0, + failed: true, + }; + } + }; + if state.consecutive_failures > 0 { + tracing::info!( + failed_ticks = state.consecutive_failures, + "sandbox watch poller: persisted sandbox version reads recovered" + ); + state.consecutive_failures = 0; + } + + let polled = ids.len(); + let mut notified = 0_usize; + for sandbox_id in ids { + let current = versions.get(&sandbox_id).copied(); + // A first observation always notifies: WatchSandbox subscribes before + // reading its snapshot, and only this catches a remote write between + // those two steps. + let changed = state + .known_versions + .get(&sandbox_id) + .is_none_or(|previous| *previous != current); + if changed { + bus.notify(&sandbox_id); + notified += 1; + } + state.known_versions.insert(sandbox_id, current); + } + tracing::Span::current().record("notified_count", notified); + PollOutcome { + polled, + notified, + failed: false, + } +} + /// Poll persisted sandbox resource versions once per gateway and notify the /// existing in-memory watch bus when another replica changes a record. /// -/// The poller performs at most one lookup per actively watched sandbox per -/// interval, regardless of how many clients are watching that sandbox. +/// Each tick makes one batched lookup of the id and `resource_version` of +/// every actively watched sandbox, regardless of how many clients watch it. +/// The store issues one statement per `RESOURCE_VERSION_BATCH_SIZE` ids and +/// reads no payloads. pub fn spawn_store_poller( store: Arc, bus: SandboxWatchBus, @@ -90,7 +229,7 @@ pub fn spawn_store_poller( mut shutdown_rx: watch::Receiver, ) { tokio::spawn(async move { - let mut known_versions: HashMap> = HashMap::new(); + let mut state = PollerState::default(); let mut timer = tokio::time::interval(interval); timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -102,31 +241,7 @@ pub fn spawn_store_poller( } } _ = timer.tick() => { - let active = bus.active_sandbox_ids(); - known_versions.retain(|sandbox_id, _| active.contains(sandbox_id)); - - for sandbox_id in active { - let current = match store.get_message::(&sandbox_id).await { - Ok(sandbox) => sandbox.map(|sandbox| { - sandbox.metadata.as_ref().map_or(0, |metadata| metadata.resource_version) - }), - Err(err) => { - tracing::warn!( - sandbox_id, - error = %err, - "sandbox watch poller: failed to read persisted sandbox" - ); - continue; - } - }; - - let changed = known_versions - .insert(sandbox_id.clone(), current) - .is_none_or(|previous| previous != current); - if changed { - bus.notify(&sandbox_id); - } - } + poll_once(store.as_ref(), &bus, &mut state).await; } } } @@ -157,7 +272,9 @@ pub fn lag_warning_event(n: u64) -> openshell_core::proto::SandboxStreamEvent { #[cfg(test)] mod tests { use super::*; + use crate::persistence::{PersistenceError, RESOURCE_VERSION_BATCH_SIZE, WriteCondition}; use openshell_core::proto::datamodel::v1::ObjectMeta; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; #[test] fn sandbox_watch_bus_remove_cleans_up() { @@ -281,4 +398,437 @@ mod tests { shutdown_tx.send(true).unwrap(); } + + /// In-memory version source with injectable failures. + #[derive(Default)] + struct FakeSource { + versions: Mutex>, + fail: AtomicBool, + calls: AtomicUsize, + last_ids: Mutex>, + } + + impl FakeSource { + fn with_versions(entries: &[(&str, u64)]) -> Self { + let source = Self::default(); + for (id, version) in entries { + source.set(id, *version); + } + source + } + + fn set(&self, id: &str, version: u64) { + self.versions + .lock() + .unwrap() + .insert(id.to_string(), version); + } + + fn remove(&self, id: &str) { + self.versions.lock().unwrap().remove(id); + } + + fn set_failing(&self, failing: bool) { + self.fail.store(failing, Ordering::Relaxed); + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::Relaxed) + } + + fn last_ids_sorted(&self) -> Vec { + let mut ids = self.last_ids.lock().unwrap().clone(); + ids.sort(); + ids + } + } + + impl SandboxVersionSource for FakeSource { + async fn sandbox_versions( + &self, + ids: &[String], + ) -> PersistenceResult> { + self.calls.fetch_add(1, Ordering::Relaxed); + ids.clone_into(&mut self.last_ids.lock().unwrap()); + if self.fail.load(Ordering::Relaxed) { + return Err(PersistenceError::Database( + "injected lookup failure".to_string(), + )); + } + let versions = self.versions.lock().unwrap(); + Ok(ids + .iter() + .filter_map(|id| versions.get(id).map(|version| (id.clone(), *version))) + .collect()) + } + } + + /// Wraps a real store and counts poller lookups. + struct CountingSource<'a> { + store: &'a Store, + calls: AtomicUsize, + request_sizes: Mutex>, + } + + impl<'a> CountingSource<'a> { + const fn new(store: &'a Store) -> Self { + Self { + store, + calls: AtomicUsize::new(0), + request_sizes: Mutex::new(Vec::new()), + } + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::Relaxed) + } + + fn request_sizes(&self) -> Vec { + self.request_sizes.lock().unwrap().clone() + } + } + + impl SandboxVersionSource for CountingSource<'_> { + async fn sandbox_versions( + &self, + ids: &[String], + ) -> PersistenceResult> { + self.calls.fetch_add(1, Ordering::Relaxed); + self.request_sizes.lock().unwrap().push(ids.len()); + self.store.sandbox_versions(ids).await + } + } + + fn drain(rx: &mut broadcast::Receiver<()>) -> usize { + let mut received = 0; + while rx.try_recv().is_ok() { + received += 1; + } + received + } + + async fn put_sandbox_row(store: &Store, idx: usize) { + store + .put( + "sandbox", + &format!("sb-{idx}"), + &format!("sandbox-{idx}"), + "default", + b"payload", + None, + ) + .await + .unwrap(); + } + + async fn bump_sandbox_row(store: &Store, idx: usize) { + store + .put_if( + "sandbox", + &format!("sb-{idx}"), + &format!("sandbox-{idx}"), + "default", + b"payload-2", + None, + WriteCondition::MatchResourceVersion(1), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn poll_once_uses_one_store_lookup_for_large_watch_set() { + const WATCHED: usize = 5_000; + const { assert!(WATCHED > 4 * RESOURCE_VERSION_BATCH_SIZE) }; + + let store = crate::persistence::test_store().await; + let bus = SandboxWatchBus::new(); + let mut receivers = Vec::with_capacity(WATCHED); + for idx in 0..WATCHED { + put_sandbox_row(&store, idx).await; + receivers.push(bus.subscribe(&format!("sb-{idx}"))); + } + let source = CountingSource::new(&store); + let mut state = PollerState::default(); + + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!( + outcome, + PollOutcome { + polled: WATCHED, + notified: WATCHED, + failed: false, + } + ); + assert_eq!(source.calls(), 1); + assert_eq!(source.request_sizes(), vec![WATCHED]); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 1)); + + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, 0); + assert_eq!(source.calls(), 2); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 0)); + + let changed = [0, 2_500, WATCHED - 1]; + for idx in changed { + bump_sandbox_row(&store, idx).await; + } + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, 3); + assert_eq!(source.calls(), 3); + for (idx, rx) in receivers.iter_mut().enumerate() { + let expected = usize::from(changed.contains(&idx)); + assert_eq!(drain(rx), expected, "sb-{idx}"); + } + + assert!(store.delete("sandbox", "sb-7").await.unwrap()); + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, 1); + assert_eq!(source.calls(), 4); + for (idx, rx) in receivers.iter_mut().enumerate() { + assert_eq!(drain(rx), usize::from(idx == 7), "sb-{idx}"); + } + assert_eq!(state.known_versions["sb-7"], None); + + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, 0); + assert_eq!(source.calls(), 5); + assert_eq!(source.request_sizes(), vec![WATCHED; 5]); + } + + #[tokio::test] + async fn poll_once_notifies_first_observation_change_and_disappearance() { + let source = FakeSource::with_versions(&[("sb-1", 1)]); + let bus = SandboxWatchBus::new(); + let mut present = bus.subscribe("sb-1"); + let mut never = bus.subscribe("sb-never"); + let mut state = PollerState::default(); + + // A first observation notifies whether or not the sandbox exists. + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 2); + assert_eq!((drain(&mut present), drain(&mut never)), (1, 1)); + + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 0); + + source.set("sb-1", 2); + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 1); + assert_eq!((drain(&mut present), drain(&mut never)), (1, 0)); + + source.remove("sb-1"); + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 1); + assert_eq!((drain(&mut present), drain(&mut never)), (1, 0)); + + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 0); + assert_eq!( + state.known_versions, + HashMap::from([("sb-1".to_string(), None), ("sb-never".to_string(), None)]) + ); + } + + #[tokio::test] + async fn poll_once_skips_lookup_when_nothing_is_watched() { + let source = FakeSource::with_versions(&[("sb-1", 1)]); + let bus = SandboxWatchBus::new(); + drop(bus.subscribe("sb-1")); + let mut state = PollerState::default(); + + assert_eq!( + poll_once(&source, &bus, &mut state).await, + PollOutcome::default() + ); + assert_eq!(source.calls(), 0); + assert!(state.known_versions.is_empty()); + } + + #[tokio::test] + async fn poll_once_prunes_unwatched_and_renotifies_rewatched_sandboxes() { + let source = FakeSource::with_versions(&[("sb-1", 1), ("sb-2", 1)]); + let bus = SandboxWatchBus::new(); + let mut first = bus.subscribe("sb-1"); + let second = bus.subscribe("sb-2"); + let mut state = PollerState::default(); + + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 2); + drain(&mut first); + + drop(second); + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 0); + assert_eq!(source.last_ids_sorted(), vec!["sb-1".to_string()]); + assert_eq!( + state.known_versions.keys().collect::>(), + vec!["sb-1"] + ); + + // The version did not change, but a re-watched sandbox is a first + // observation again. + let mut second = bus.subscribe("sb-2"); + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 1); + assert_eq!((drain(&mut first), drain(&mut second)), (0, 1)); + } + + #[tokio::test] + async fn poll_once_keeps_known_versions_when_lookup_fails() { + let source = FakeSource::with_versions(&[("sb-1", 1), ("sb-2", 1)]); + let bus = SandboxWatchBus::new(); + let mut receivers = vec![bus.subscribe("sb-1"), bus.subscribe("sb-2")]; + let mut state = PollerState::default(); + + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 2); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 1)); + + source.set_failing(true); + source.set("sb-1", 2); + source.remove("sb-2"); + receivers.push(bus.subscribe("sb-3")); + source.set("sb-3", 1); + + assert_eq!( + poll_once(&source, &bus, &mut state).await, + PollOutcome { + polled: 3, + notified: 0, + failed: true, + } + ); + let before_outage = + HashMap::from([("sb-1".to_string(), Some(1)), ("sb-2".to_string(), Some(1))]); + assert_eq!(state.known_versions, before_outage); + assert_eq!(state.consecutive_failures, 1); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 0)); + + assert!(poll_once(&source, &bus, &mut state).await.failed); + assert_eq!(state.consecutive_failures, 2); + assert_eq!(state.known_versions, before_outage); + + source.set_failing(false); + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!( + outcome, + PollOutcome { + polled: 3, + notified: 3, + failed: false, + } + ); + assert_eq!(state.consecutive_failures, 0); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 1)); + assert_eq!( + state.known_versions, + HashMap::from([ + ("sb-1".to_string(), Some(2)), + ("sb-2".to_string(), None), + ("sb-3".to_string(), Some(1)), + ]) + ); + } + + #[tokio::test] + async fn poll_once_records_watch_metrics() { + let capture = gateway_metrics::MetricsCapture::install(); + let source = FakeSource::with_versions(&[("sb-1", 1), ("sb-2", 1), ("sb-3", 1)]); + let bus = SandboxWatchBus::new(); + let receivers = vec![ + bus.subscribe("sb-1"), + bus.subscribe("sb-2"), + bus.subscribe("sb-3"), + ]; + let mut state = PollerState::default(); + + assert!(!poll_once(&source, &bus, &mut state).await.failed); + assert_eq!( + capture.value("openshell_server_sandbox_watch_polled_sandboxes"), + Some(3) + ); + source.set_failing(true); + assert!(poll_once(&source, &bus, &mut state).await.failed); + drop(receivers); + assert_eq!( + poll_once(&source, &bus, &mut state).await, + PollOutcome::default() + ); + + assert_eq!( + capture.value("openshell_server_sandbox_watch_polled_sandboxes"), + Some(0) + ); + assert_eq!( + capture.value("openshell_server_sandbox_watch_poll_errors_total"), + Some(1) + ); + // The empty tick issued no lookup, so it records no duration. + assert_eq!( + capture.value("openshell_server_sandbox_watch_poll_duration_seconds_count"), + Some(2) + ); + assert!( + capture + .render() + .contains("openshell_server_sandbox_watch_poll_duration_seconds_bucket{le="), + "the watch poll duration is a bucketed histogram" + ); + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_poller_observes_writes_from_another_replica() { + let schema = crate::persistence::test_postgres::TestSchema::create("watch").await; + // Separate pools on one schema, like two gateway replicas. + let replica_a = schema.connect_store().await; + let replica_b = Arc::new(schema.connect_store().await); + let watched = RESOURCE_VERSION_BATCH_SIZE + 5; + let bus = SandboxWatchBus::new(); + let mut receivers = Vec::with_capacity(watched); + for idx in 0..watched { + put_sandbox_row(&replica_a, idx).await; + receivers.push(bus.subscribe(&format!("sb-{idx}"))); + } + + let source = CountingSource::new(replica_b.as_ref()); + let mut state = PollerState::default(); + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, watched); + assert_eq!(source.calls(), 1); + assert!(receivers.iter_mut().all(|rx| drain(rx) == 1)); + + // HashSet order decides which lookup batch each change lands in; the + // store contract test pins cross-batch reads. + let last = watched - 1; + bump_sandbox_row(&replica_a, 0).await; + bump_sandbox_row(&replica_a, last).await; + assert!(replica_a.delete("sandbox", "sb-1").await.unwrap()); + let outcome = poll_once(&source, &bus, &mut state).await; + assert_eq!(outcome.notified, 3); + for (idx, rx) in receivers.iter_mut().enumerate() { + let expected = usize::from([0, 1, last].contains(&idx)); + assert_eq!(drain(rx), expected, "sb-{idx}"); + } + assert_eq!(state.known_versions["sb-1"], None); + assert_eq!(state.known_versions["sb-0"], Some(2)); + assert_eq!(poll_once(&source, &bus, &mut state).await.notified, 0); + assert_eq!(source.calls(), 3); + + // The real loop delivers a remote write to a local watcher. + let (shutdown_tx, shutdown_rx) = watch::channel(false); + spawn_store_poller( + replica_b.clone(), + bus.clone(), + Duration::from_millis(20), + shutdown_rx, + ); + let mut rx = bus.subscribe("sb-2"); + tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .expect("poller should publish its initial observation") + .unwrap(); + bump_sandbox_row(&replica_a, 2).await; + tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("poller should observe a write made through another store") + .unwrap(); + shutdown_tx.send(true).unwrap(); + + replica_a.close().await; + replica_b.close().await; + schema.drop_schema().await; + } } diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 710ecc1406..86be46fef1 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -27,6 +27,9 @@ use openshell_core::transport_errors::is_expected_transport_close_status; use crate::ServerState; use crate::auth::principal::Principal; +use crate::gateway_metrics::{ + self, GaugeSlot, PeerRequestTimer, PeerRpc, RelayCapacity, RelayRejection, +}; use crate::grpc::provider_readiness::ProviderReadinessEvidence; use crate::persistence::ObjectId; use crate::supervisor_owner::{OWNER_TTL, OwnerError, OwnerGuard, SupervisorOwnerIndex}; @@ -48,6 +51,33 @@ const MAX_PENDING_RELAYS: usize = 256; /// consume the entire global budget. Sits above the SSH-tunnel per-sandbox /// cap (20) so tunnel-specific limits still fire first for that caller. const MAX_PENDING_RELAYS_PER_SANDBOX: usize = 32; +/// The relay caps above, published as capacity gauges when the metrics recorder is installed. +pub(crate) const RELAY_CAPACITY: RelayCapacity = RelayCapacity { + global: MAX_PENDING_RELAYS, + per_sandbox: MAX_PENDING_RELAYS_PER_SANDBOX, +}; +/// Serve normally this long after SIGTERM before closing sessions, so endpoint +/// removal reaches kube-proxy and ingress and reconnects land on other replicas. +pub(crate) const DRAIN_PROPAGATION_DELAY: Duration = Duration::from_secs(3); +/// Upper bound on the paced session-close window. +pub(crate) const DRAIN_CLOSE_WINDOW: Duration = Duration::from_secs(12); +/// Largest gap between two paced closes. Small drains finish in +/// `sessions x 100ms` instead of stretching to the window. +pub(crate) const DRAIN_MAX_CLOSE_INTERVAL: Duration = Duration::from_millis(100); +/// Final wait for supervisor session ownership cleanup during shutdown. +pub(crate) const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(10); +/// The drain and session cleanup budgets fit the chart's default 30s grace +/// with margin. Compute-driver cleanup and the OTLP flush are outside this sum. +const GATEWAY_SHUTDOWN_BUDGET: Duration = Duration::from_secs(25); +const _: () = assert!( + DRAIN_PROPAGATION_DELAY.as_millis() + + DRAIN_CLOSE_WINDOW.as_millis() + + SESSION_CLEANUP_TIMEOUT.as_millis() + <= GATEWAY_SHUTDOWN_BUDGET.as_millis(), + "gateway drain plus cleanup must fit the shutdown budget" +); +/// How long an owner waits for a local session before failing a peer relay. +const PEER_RELAY_SESSION_WAIT: Duration = Duration::from_secs(5); const PEER_TLS_CA_FILE_ENV: &str = "OPENSHELL_PEER_TLS_CA_FILE"; const PEER_TLS_CERT_FILE_ENV: &str = "OPENSHELL_PEER_TLS_CERT_FILE"; const PEER_TLS_KEY_FILE_ENV: &str = "OPENSHELL_PEER_TLS_KEY_FILE"; @@ -274,12 +304,19 @@ struct LiveSession { /// removing a session that has since been superseded by a reconnect. session_id: String, tx: mpsc::Sender, - /// Fires when this session is superseded by a reconnect so the old session - /// task can exit promptly — dropping its own `tx` clone and closing the - /// outbound stream. Without this, a concurrent `open_relay` that grabbed - /// the old session's `tx` just before supersede could still enqueue a - /// `RelayOpen` onto the stale stream and sit until the relay timeout. - shutdown: oneshot::Sender<()>, + /// Fires on supersede, lifecycle disconnect, or a drain slot so the + /// session task can exit promptly — dropping its own `tx` clone and + /// closing the outbound stream. Without this, a concurrent `open_relay` + /// that grabbed the old session's `tx` just before supersede could still + /// enqueue a `RelayOpen` onto the stale stream and sit until the relay + /// timeout. A drain slot takes it and leaves the entry registered, so + /// `None` means the session is closing for a drain. + shutdown: Option>, + /// True while `RelayOpen` may be queued on `tx`: set after `SessionAccepted` + /// is queued (so the supervisor always reads it first) and cleared when a + /// drain closes the session. `has_session` ignores it; compute readiness + /// (`compute::supervisor_session_ready`) must keep seeing the session. + accepts_relays: bool, /// Set after the supervisor confirms that every expected foreground /// attachment has closed and terminal output delivery is complete. terminal_delivery_finalized: bool, @@ -297,6 +334,9 @@ struct LiveSession { provider_readiness: Option, #[allow(dead_code)] connected_at: Instant, + /// This session's share of `openshell_server_supervisor_sessions`, released when the entry + /// leaves the registry by any path (supersede, remove, disconnect, cleanup). + _gauge_slot: GaugeSlot, } /// Idempotency state for tool server endpoint-status reports from one live supervisor. @@ -313,6 +353,28 @@ pub(crate) struct EndpointReportCursor { pub(crate) report_digest: [u8; 32], } +/// Whether this replica can route a relay to its local session for a sandbox. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum LocalSessionRoute { + /// No session for the sandbox on this replica. + Absent, + /// Registered but not accepting relays yet (before `SessionAccepted`), or + /// any more (drain close in progress). + Settling, + /// `RelayOpen` can be queued now. + Ready, +} + +/// Counts from a paced drain, for logging. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DrainSummary { + /// Sessions registered when the drain took its snapshot, after the + /// propagation delay. + pub(crate) planned: usize, + /// Sessions signaled at their slot. The rest had already ended or were replaced. + pub(crate) signaled: usize, +} + /// Holds a oneshot sender that will deliver the upgraded relay stream or a /// target-open failure reported by the supervisor. type RelayStreamSender = oneshot::Sender>; @@ -336,6 +398,8 @@ struct PendingRelay { sandbox_id: String, relay_open: RelayOpen, created_at: Instant, + /// This relay's share of `openshell_server_relay_pending`. + _gauge_slot: GaugeSlot, } #[derive(Debug)] @@ -377,6 +441,16 @@ impl SupervisorSessionRegistry { self.admission_closed.store(true, Ordering::Release); } + /// True once shutdown has closed admission: no new session can register. + pub(crate) fn admission_closed(&self) -> bool { + self.admission_closed.load(Ordering::Acquire) + } + + /// Number of registered sessions, including ones still being accepted. + pub(crate) fn session_count(&self) -> usize { + self.sessions.lock().unwrap().len() + } + /// Close control sessions and wait for tracked ownership cleanup. Call /// after compute shutdown so supervisors can finish normal stop reporting. pub(crate) async fn shutdown(&self, timeout: Duration) -> Result<(), String> { @@ -390,17 +464,118 @@ impl SupervisorSessionRegistry { }) } - /// Register a live supervisor session for the given sandbox. + /// Close every current session on a fixed schedule so its supervisor + /// reconnects to another replica. Each session serves relays and + /// heartbeats until its slot and stays registered until its own task runs + /// the normal cleanup (owner release, disconnect bookkeeping). + pub(crate) async fn drain_sessions( + &self, + window: Duration, + max_interval: Duration, + ) -> DrainSummary { + let targets = self.drain_targets(); + let interval = drain_close_interval(targets.len(), window, max_interval); + let start = tokio::time::Instant::now(); + let mut signaled = 0; + for (slot, (sandbox_id, session_id)) in targets.iter().enumerate() { + let offset = interval.saturating_mul(u32::try_from(slot).unwrap_or(u32::MAX)); + tokio::time::sleep_until(start + offset).await; + if self.close_for_drain(sandbox_id, session_id) { + signaled += 1; + } + } + DrainSummary { + planned: targets.len(), + signaled, + } + } + + /// Snapshot of `(sandbox_id, session_id)` pairs sorted by sandbox id. The + /// order carries no priority; sorting keeps tests deterministic. + fn drain_targets(&self) -> Vec<(String, String)> { + let mut targets: Vec<(String, String)> = self + .sessions + .lock() + .unwrap() + .iter() + .map(|(sandbox_id, session)| (sandbox_id.clone(), session.session_id.clone())) + .collect(); + targets.sort_unstable(); + targets + } + + /// Signal one session to end for a drain without removing it, and stop + /// routing new relays to it. Returns `false` when it already ended, was + /// replaced, or was already signaled. + fn close_for_drain(&self, sandbox_id: &str, session_id: &str) -> bool { + let mut sessions = self.sessions.lock().unwrap(); + let Some(session) = sessions + .get_mut(sandbox_id) + .filter(|session| session.session_id == session_id) + else { + return false; + }; + session.accepts_relays = false; + session + .shutdown + .take() + .is_some_and(|shutdown| shutdown.send(()).is_ok()) + } + + /// Register a live supervisor session for the given sandbox that can + /// receive `RelayOpen` immediately. /// /// If a previous session exists for the same sandbox, its shutdown signal /// is fired so the old session task exits promptly. Returns `true` iff a - /// previous session was superseded. + /// previous session was superseded. The gateway's own `ConnectSupervisor` + /// path uses `register_awaiting_accept` instead. pub fn register( &self, sandbox_id: String, session_id: String, tx: mpsc::Sender, shutdown: oneshot::Sender<()>, + ) -> bool { + self.insert_session(sandbox_id, session_id, tx, shutdown, true) + } + + /// Register a session whose `SessionAccepted` is not queued yet. + /// + /// Relay routing skips it until `mark_accepts_relays`. Supersede handling + /// and the return value match `register`. + pub(crate) fn register_awaiting_accept( + &self, + sandbox_id: String, + session_id: String, + tx: mpsc::Sender, + shutdown: oneshot::Sender<()>, + ) -> bool { + self.insert_session(sandbox_id, session_id, tx, shutdown, false) + } + + /// Open relay routing to this session. Call only after `SessionAccepted` + /// and any replayed relays are queued on its sender. Returns `false` when + /// the session was replaced or removed, or already closed for a drain. + pub(crate) fn mark_accepts_relays(&self, sandbox_id: &str, session_id: &str) -> bool { + let mut sessions = self.sessions.lock().unwrap(); + let Some(session) = sessions + .get_mut(sandbox_id) + .filter(|session| session.session_id == session_id && session.shutdown.is_some()) + else { + return false; + }; + session.accepts_relays = true; + true + } + + /// Insert a session and fire the shutdown signal of the one it replaces. + fn insert_session( + &self, + sandbox_id: String, + session_id: String, + tx: mpsc::Sender, + shutdown: oneshot::Sender<()>, + accepts_relays: bool, ) -> bool { let mut sessions = self.sessions.lock().unwrap(); let previous = sessions.remove(&sandbox_id); @@ -410,18 +585,23 @@ impl SupervisorSessionRegistry { sandbox_id, session_id, tx, - shutdown, + shutdown: Some(shutdown), + accepts_relays, terminal_delivery_finalized: false, endpoint_status_initialized: false, endpoint_report_cursor: None, provider_readiness: None, connected_at: Instant::now(), + _gauge_slot: GaugeSlot::supervisor_session(), }, ); match previous { Some(prev) => { - // Best-effort — the old task may have already exited. - let _ = prev.shutdown.send(()); + // Best-effort — the old task may have already exited, or a + // drain slot already signaled it. + if let Some(shutdown) = prev.shutdown { + let _ = shutdown.send(()); + } true } None => false, @@ -440,7 +620,9 @@ impl SupervisorSessionRegistry { pub fn disconnect(&self, sandbox_id: &str) -> bool { let session = self.sessions.lock().unwrap().remove(sandbox_id); if let Some(session) = session { - let _ = session.shutdown.send(()); + if let Some(shutdown) = session.shutdown { + let _ = shutdown.send(()); + } true } else { false @@ -495,6 +677,7 @@ impl SupervisorSessionRegistry { .lock() .unwrap() .get(sandbox_id) + .filter(|s| s.accepts_relays) .map(|s| s.tx.clone()) } @@ -502,6 +685,23 @@ impl SupervisorSessionRegistry { self.sessions.lock().unwrap().contains_key(sandbox_id) } + /// Classify the local session for relay routing. Unlike `has_session`, + /// this separates sessions that can take `RelayOpen` from those still + /// being accepted. + pub(crate) fn local_session_route(&self, sandbox_id: &str) -> LocalSessionRoute { + self.sessions + .lock() + .unwrap() + .get(sandbox_id) + .map_or(LocalSessionRoute::Absent, |session| { + if session.accepts_relays { + LocalSessionRoute::Ready + } else { + LocalSessionRoute::Settling + } + }) + } + pub fn terminal_delivery_finalized(&self, sandbox_id: &str) -> bool { self.sessions .lock() @@ -805,6 +1005,7 @@ impl SupervisorSessionRegistry { { let mut pending = self.pending_relays.lock().unwrap(); if pending.len() >= MAX_PENDING_RELAYS { + gateway_metrics::record_relay_rejected(RelayRejection::GlobalCapacity); return Err(Status::resource_exhausted(format!( "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)" ))); @@ -814,6 +1015,7 @@ impl SupervisorSessionRegistry { .filter(|p| p.sandbox_id == sandbox_id) .count(); if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX { + gateway_metrics::record_relay_rejected(RelayRejection::SandboxCapacity); return Err(Status::resource_exhausted(format!( "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})" ))); @@ -825,6 +1027,7 @@ impl SupervisorSessionRegistry { sandbox_id: sandbox_id.to_string(), relay_open: relay_open.clone(), created_at: Instant::now(), + _gauge_slot: GaugeSlot::relay_pending(), }, ); } @@ -843,13 +1046,19 @@ impl SupervisorSessionRegistry { } pub fn fail_pending_relay(&self, channel_id: &str, error: String) -> bool { - let pending = self.pending_relays.lock().unwrap().remove(channel_id); - if let Some(pending) = pending { - let _ = pending.sender.send(Err(Status::unavailable(error))); - true - } else { - false - } + // The rest of the entry, including its gauge slot, drops inside this statement while the + // lock is still held, so `relay_pending` never exceeds capacity. + let Some(sender) = self + .pending_relays + .lock() + .unwrap() + .remove(channel_id) + .map(|pending| pending.sender) + else { + return false; + }; + let _ = sender.send(Err(Status::unavailable(error))); + true } /// Claim a pending relay channel. Called by the `/relay/{channel_id}` HTTP handler @@ -863,7 +1072,7 @@ impl SupervisorSessionRegistry { channel_id: &str, principal: Option<&Principal>, ) -> Result { - let pending = { + let (sender, sandbox_id) = { let mut map = self.pending_relays.lock().unwrap(); let pending = map .get(channel_id) @@ -883,13 +1092,22 @@ impl SupervisorSessionRegistry { return Err(status); } - if pending.created_at.elapsed() > RELAY_PENDING_TIMEOUT { + let waited = pending.created_at.elapsed(); + if waited > RELAY_PENDING_TIMEOUT { map.remove(channel_id); + gateway_metrics::record_relay_expired(1); return Err(Status::deadline_exceeded("relay channel timed out")); } - - map.remove(channel_id) - .expect("pending relay existed before removal") + gateway_metrics::record_relay_claimed(waited); + + // The rest of the entry, including its gauge slot, drops at the end of this + // statement while the lock is still held, so `relay_pending` never exceeds capacity. + let PendingRelay { + sender, sandbox_id, .. + } = map + .remove(channel_id) + .expect("pending relay existed before removal"); + (sender, sandbox_id) }; // Create a duplex stream pair: one end for the gateway bridge, one for @@ -897,20 +1115,25 @@ impl SupervisorSessionRegistry { let (gateway_stream, supervisor_stream) = tokio::io::duplex(64 * 1024); // Send the gateway-side stream to the waiter (exec handler or forward handler). - if pending.sender.send(Ok(gateway_stream)).is_err() { + if sender.send(Ok(gateway_stream)).is_err() { return Err(Status::internal("relay requester dropped")); } Ok(ClaimedRelay { stream: supervisor_stream, - sandbox_id: pending.sandbox_id, + sandbox_id, }) } /// Remove all pending relays that have exceeded the timeout. pub fn reap_expired_relays(&self) { - let mut map = self.pending_relays.lock().unwrap(); - map.retain(|_, pending| pending.created_at.elapsed() <= RELAY_PENDING_TIMEOUT); + let reaped = { + let mut map = self.pending_relays.lock().unwrap(); + let before = map.len(); + map.retain(|_, pending| pending.created_at.elapsed() <= RELAY_PENDING_TIMEOUT); + before - map.len() + }; + gateway_metrics::record_relay_expired(reaped); } /// Clean up all state for a sandbox (session + pending relays). @@ -940,6 +1163,38 @@ impl SupervisorSessionRegistry { } } +/// Gap between two paced closes: the window split evenly over the sessions, +/// capped at `max_interval`, so the last close happens before the window ends. +fn drain_close_interval(sessions: usize, window: Duration, max_interval: Duration) -> Duration { + let count = u32::try_from(sessions).unwrap_or(u32::MAX).max(1); + (window / count).min(max_interval) +} + +/// Keep serving for `propagation`, then snapshot and close sessions on a paced +/// schedule. The snapshot is taken after the delay so setups that were in +/// flight at SIGTERM are included. Tests call this with millisecond timings. +pub(crate) async fn drain_with( + registry: &SupervisorSessionRegistry, + propagation: Duration, + window: Duration, + max_interval: Duration, +) -> DrainSummary { + tokio::time::sleep(propagation).await; + registry.drain_sessions(window, max_interval).await +} + +/// Shutdown drain with the production timings. The caller keeps the listener +/// open until this returns. +pub(crate) async fn drain_for_shutdown(registry: &SupervisorSessionRegistry) -> DrainSummary { + drain_with( + registry, + DRAIN_PROPAGATION_DELAY, + DRAIN_CLOSE_WINDOW, + DRAIN_MAX_CLOSE_INTERVAL, + ) + .await +} + /// Spawn a background task that periodically reaps expired pending relay /// entries. /// @@ -1311,15 +1566,16 @@ pub(crate) async fn forward_provider_readiness_to_owner( request: ReportProviderReadinessRequest, ) -> Result { let sandbox_id = request.sandbox_id.clone(); - let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint).await?; - client - .peer_report_provider_readiness(request) + let mut timer = PeerRequestTimer::start(PeerRpc::ReportProviderReadiness); + let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint) .await - .map(Response::into_inner) - .inspect_err(|_| { - state.peer_routes.evict_channel(&owner.owner_peer_endpoint); - state.peer_routes.evict_owner(&sandbox_id); - }) + .inspect_err(|status| timer.client_error(status))?; + let result = client.peer_report_provider_readiness(request).await; + timer.finish(&result); + result.map(Response::into_inner).inspect_err(|_| { + state.peer_routes.evict_channel(&owner.owner_peer_endpoint); + state.peer_routes.evict_owner(&sandbox_id); + }) } pub(crate) async fn forward_endpoint_status_to_owner( @@ -1328,15 +1584,16 @@ pub(crate) async fn forward_endpoint_status_to_owner( request: ReportEndpointStatusRequest, ) -> Result { let sandbox_id = request.sandbox_id.clone(); - let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint).await?; - client - .peer_report_endpoint_status(request) + let mut timer = PeerRequestTimer::start(PeerRpc::ReportEndpointStatus); + let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint) .await - .map(Response::into_inner) - .inspect_err(|_| { - state.peer_routes.evict_channel(&owner.owner_peer_endpoint); - state.peer_routes.evict_owner(&sandbox_id); - }) + .inspect_err(|status| timer.client_error(status))?; + let result = client.peer_report_endpoint_status(request).await; + timer.finish(&result); + result.map(Response::into_inner).inspect_err(|_| { + state.peer_routes.evict_channel(&owner.owner_peer_endpoint); + state.peer_routes.evict_owner(&sandbox_id); + }) } pub(crate) async fn forward_provider_status_query_to_owner( @@ -1345,15 +1602,16 @@ pub(crate) async fn forward_provider_status_query_to_owner( sandbox_id: &str, request: GetSandboxProviderStatusRequest, ) -> Result { - let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint).await?; - client - .peer_get_sandbox_provider_status(request) + let mut timer = PeerRequestTimer::start(PeerRpc::GetSandboxProviderStatus); + let mut client = peer_rpc_client(state, &owner.owner_peer_endpoint) .await - .map(Response::into_inner) - .inspect_err(|_| { - state.peer_routes.evict_channel(&owner.owner_peer_endpoint); - state.peer_routes.evict_owner(sandbox_id); - }) + .inspect_err(|status| timer.client_error(status))?; + let result = client.peer_get_sandbox_provider_status(request).await; + timer.finish(&result); + result.map(Response::into_inner).inspect_err(|_| { + state.peer_routes.evict_channel(&owner.owner_peer_endpoint); + state.peer_routes.evict_owner(sandbox_id); + }) } pub async fn open_routed_relay_with_target( @@ -1394,25 +1652,40 @@ pub async fn open_routed_relay_with_message( let mut backoff = SESSION_WAIT_INITIAL_BACKOFF; let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); loop { - if state.supervisor_sessions.has_session(sandbox_id) { - match state - .supervisor_sessions - .open_relay_with_message(sandbox_id, relay_open.clone(), Duration::ZERO) - .await - { - Ok(relay) => return Ok(relay), - Err(status) if status.code() == tonic::Code::Unavailable => { - // The session can migrate after `has_session` but before - // RelayOpen reaches its sender. Fall through and reread the - // persisted owner instead of surfacing a handoff race. - warn!( - sandbox_id, - error = %status, - "local supervisor relay disappeared during open; resolving owner again" - ); + match state.supervisor_sessions.local_session_route(sandbox_id) { + LocalSessionRoute::Ready => { + match state + .supervisor_sessions + .open_relay_with_message(sandbox_id, relay_open.clone(), Duration::ZERO) + .await + { + Ok(relay) => return Ok(relay), + Err(status) if status.code() == tonic::Code::Unavailable => { + // The session can migrate after the route check but before + // RelayOpen reaches its sender. Fall through and reread the + // persisted owner instead of surfacing a handoff race. + warn!( + sandbox_id, + error = %status, + "local supervisor relay disappeared during open; resolving owner again" + ); + } + Err(status) => return Err(status), + } + } + LocalSessionRoute::Settling => { + // This replica holds the session but it is still being + // accepted, or a drain is closing it. Its owner record names + // this replica until the session task releases it, so wait + // locally instead of logging owner-mismatch retries. + if Instant::now() + backoff > deadline { + return Err(Status::unavailable("supervisor session not connected")); } - Err(status) => return Err(status), + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(SESSION_WAIT_MAX_BACKOFF); + continue; } + LocalSessionRoute::Absent => {} } if let Some(owner) = resolve_owner(state, &owner_index, sandbox_id).await? @@ -1524,15 +1797,30 @@ async fn open_peer_relay( Ok((channel_id, relay_rx)) } +/// Open a `PeerRelay` stream to the owner replica and bridge it to a local duplex stream. +/// +/// The peer request metrics count attempts, so the routed-relay retry loop spikes `unavailable` +/// during rollouts. `ok` means the owner's supervisor claimed the relay (response headers +/// arrived); later bridge failures are not counted. async fn connect_peer_relay( state: &Arc, owner_peer_endpoint: &str, sandbox_id: &str, relay_open: RelayOpen, ) -> Result { - let token = state.peer_routes.peer_token().await?; - let channel = state.peer_routes.channel(owner_peer_endpoint).await?; - let interceptor = PeerAuthInterceptor::new(&token, &state.replica_id)?; + let mut timer = PeerRequestTimer::start(PeerRpc::Relay); + let token = state + .peer_routes + .peer_token() + .await + .inspect_err(|s| timer.client_error(s))?; + let channel = state + .peer_routes + .channel(owner_peer_endpoint) + .await + .inspect_err(|s| timer.client_error(s))?; + let interceptor = PeerAuthInterceptor::new(&token, &state.replica_id) + .inspect_err(|s| timer.client_error(s))?; let mut client = open_shell_client::OpenShellClient::with_interceptor(channel, interceptor); let (out_tx, out_rx) = mpsc::channel::(16); @@ -1545,21 +1833,33 @@ async fn connect_peer_relay( })), }) .await - .map_err(|_| Status::internal("failed to initialize peer relay stream"))?; - - let response = client - .peer_relay(ReceiverStream::new(out_rx)) - .await - .map_err(|err| { - state.peer_routes.evict_channel(owner_peer_endpoint); - Status::unavailable(format!("gateway peer relay RPC failed: {err}")) - })?; + .map_err(|_| Status::internal("failed to initialize peer relay stream")) + .inspect_err(|s| timer.client_error(s))?; + + let result = client.peer_relay(ReceiverStream::new(out_rx)).await; + // Record the owner's code before the remap below hides it as `unavailable`. + timer.finish(&result); + let response = result.map_err(|err| { + state.peer_routes.evict_channel(owner_peer_endpoint); + Status::unavailable(format!("gateway peer relay RPC failed: {err}")) + })?; let inbound = response.into_inner(); let (gateway_stream, bridge_stream) = tokio::io::duplex(64 * 1024); spawn_peer_bridge(bridge_stream, inbound, out_tx, sandbox_id.to_string()); Ok(gateway_stream) } +/// How long the owner waits for its local session before failing a peer +/// relay. A draining replica admits no sessions and closes the ones it has, +/// so waiting only delays the requester's owner re-read. +fn peer_relay_session_wait(admission_closed: bool) -> Duration { + if admission_closed { + Duration::ZERO + } else { + PEER_RELAY_SESSION_WAIT + } +} + pub async fn handle_peer_relay( state: &Arc, request: Request>, @@ -1603,6 +1903,18 @@ pub async fn handle_peer_relay( return Err(Status::invalid_argument("relay channel_id is required")); } + let admission_closed = state.supervisor_sessions.admission_closed(); + if admission_closed && !state.supervisor_sessions.has_session(&init.sandbox_id) { + debug!( + sandbox_id = %init.sandbox_id, + requester = %peer.replica_id, + "gateway peer relay: draining replica holds no session for this sandbox" + ); + return Err(Status::unavailable( + "gateway replica is draining and holds no supervisor session for this sandbox", + )); + } + info!( sandbox_id = %init.sandbox_id, channel_id = %relay_open.channel_id, @@ -1612,7 +1924,11 @@ pub async fn handle_peer_relay( let (channel_id, relay_rx) = state .supervisor_sessions - .open_relay_with_message(&init.sandbox_id, relay_open, Duration::from_secs(5)) + .open_relay_with_message( + &init.sandbox_id, + relay_open, + peer_relay_session_wait(admission_closed), + ) .await?; let supervisor_stream = match tokio::time::timeout(Duration::from_secs(10), relay_rx).await { Ok(Ok(Ok(stream))) => stream, @@ -1853,7 +2169,7 @@ async fn establish_supervisor_session( // Step 2: Create and register the outbound channel. let (tx, rx) = mpsc::channel::(64); let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); - let superseded = state.supervisor_sessions.register( + let superseded = state.supervisor_sessions.register_awaiting_accept( sandbox_id.clone(), session_id.clone(), tx.clone(), @@ -1967,6 +2283,20 @@ async fn establish_supervisor_session( .await; } + // Relays may use this session only now: SessionAccepted and any replayed + // relays are already queued, so the supervisor never reads RelayOpen first, + // and relays opened during setup were not also replayed (no duplicate opens). + if !state + .supervisor_sessions + .mark_accepts_relays(&sandbox_id, &session_id) + { + debug!( + sandbox_id = %sandbox_id, + session_id = %session_id, + "supervisor session: replaced or closed before relay routing opened" + ); + } + // Step 4: Spawn the session loop that reads inbound messages. let state_clone = Arc::clone(&state); let sandbox_id_clone = sandbox_id.clone(); @@ -2107,7 +2437,12 @@ async fn run_session_loop( break; } _ = &mut shutdown_rx => { - info!(sandbox_id = %sandbox_id, session_id = %session_id, "supervisor session: superseded by reconnect, shutting down"); + if state.supervisor_sessions.is_current_session(sandbox_id, session_id) { + // Still registered: a drain slot, not a replacement. + info!(sandbox_id = %sandbox_id, session_id = %session_id, "supervisor session: closing for gateway drain"); + } else { + info!(sandbox_id = %sandbox_id, session_id = %session_id, "supervisor session: superseded by reconnect, shutting down"); + } break; } msg = inbound.message() => { @@ -2269,7 +2604,12 @@ mod tests { use super::*; use crate::auth::identity::{Identity, IdentityProvider}; use crate::auth::principal::{SandboxIdentitySource, SandboxPrincipal, UserPrincipal}; + use crate::gateway_metrics::MetricsCapture; use crate::persistence::Store; + use bytes::Bytes; + use http_body::Frame; + use http_body_util::{BodyExt, Empty, StreamBody}; + use std::convert::Infallible; use tokio::io::{AsyncReadExt, AsyncWriteExt}; async fn test_store() -> Arc { @@ -2496,6 +2836,7 @@ mod tests { service_id: String::new(), }, created_at, + _gauge_slot: GaugeSlot::relay_pending(), } } @@ -2697,6 +3038,41 @@ mod tests { assert_eq!(registry.remove_if_current("sbx", "s1"), Some(true)); } + #[test] + fn session_gauge_tracks_register_supersede_and_removal() { + let metrics = MetricsCapture::install(); + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel(1); + + registry.register( + "sbx-a".to_string(), + "s1".to_string(), + tx.clone(), + make_shutdown(), + ); + assert_eq!(metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), Some(1)); + registry.register( + "sbx-a".to_string(), + "s2".to_string(), + tx.clone(), + make_shutdown(), + ); + assert_eq!( + metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), + Some(1), + "a supersede on the same replica nets zero" + ); + registry.register("sbx-b".to_string(), "s3".to_string(), tx, make_shutdown()); + assert_eq!(metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), Some(2)); + + assert_eq!(registry.remove_if_current("sbx-a", "s1"), None); + assert_eq!(metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), Some(2)); + assert_eq!(registry.remove_if_current("sbx-a", "s2"), Some(false)); + assert_eq!(metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), Some(1)); + assert!(registry.disconnect("sbx-b")); + assert_eq!(metrics.value(gateway_metrics::SUPERVISOR_SESSIONS), Some(0)); + } + // ---- open_relay: happy path and wait semantics ---- #[tokio::test] @@ -2758,6 +3134,7 @@ mod tests { #[tokio::test] async fn open_relay_fails_when_session_receiver_dropped() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (tx, rx) = mpsc::channel::(4); registry.register("sbx".to_string(), "s1".to_string(), tx, make_shutdown()); @@ -2773,10 +3150,13 @@ mod tests { assert_eq!(err.code(), tonic::Code::Unavailable); // The pending-relay entry must have been cleaned up on failure. assert!(registry.pending_relays.lock().unwrap().is_empty()); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), None); } #[tokio::test] async fn open_relay_rejects_when_global_cap_reached() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (tx, _rx) = mpsc::channel::(8); registry.register( @@ -2807,10 +3187,20 @@ mod tests { .expect_err("open_relay should reject once global cap is reached"); assert_eq!(err.code(), tonic::Code::ResourceExhausted); assert!(err.message().contains("gateway relay capacity")); + assert_eq!( + metrics.value("openshell_server_relay_rejected_total{reason=\"global_capacity\"}"), + Some(1) + ); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(256)); + assert_eq!( + metrics.value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}"), + None + ); } #[tokio::test] async fn open_relay_rejects_when_per_sandbox_cap_reached() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (tx, _rx) = mpsc::channel::(8); registry.register("sbx".to_string(), "s".to_string(), tx, make_shutdown()); @@ -2832,6 +3222,11 @@ mod tests { .expect_err("open_relay should reject when per-sandbox cap is reached"); assert_eq!(err.code(), tonic::Code::ResourceExhausted); assert!(err.message().contains("per-sandbox relay limit")); + assert_eq!( + metrics.value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}"), + Some(1) + ); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(32)); // A different sandbox still has headroom. let (tx2, _rx2) = mpsc::channel::(8); @@ -2845,6 +3240,7 @@ mod tests { .open_relay("sbx-other", Duration::from_millis(50)) .await .expect("different sandbox should still accept new relays"); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(33)); } #[tokio::test] @@ -2971,64 +3367,845 @@ mod tests { } } - #[tokio::test] - async fn require_persisted_sandbox_rejects_missing_sandbox() { - let store = test_store().await; - - let err = require_persisted_sandbox(&store, "missing") - .await - .expect_err("missing sandbox should be rejected"); + fn session_accepted_message(session_id: &str) -> GatewayMessage { + GatewayMessage { + payload: Some(gateway_message::Payload::SessionAccepted(SessionAccepted { + session_id: session_id.to_string(), + heartbeat_interval: None, + })), + } + } - assert_eq!(err.code(), tonic::Code::NotFound); + /// Assert that `rx` holds `SessionAccepted` followed by `RelayOpen` for + /// `channel_id`, the order a supervisor requires. + async fn assert_accepted_then_relay_open( + rx: &mut mpsc::Receiver, + channel_id: &str, + ) { + let first = rx.recv().await.expect("SessionAccepted should be queued"); + assert!( + matches!( + first.payload, + Some(gateway_message::Payload::SessionAccepted(_)) + ), + "expected SessionAccepted first, got {:?}", + first.payload + ); + let second = rx.recv().await.expect("RelayOpen should be queued"); + match second.payload { + Some(gateway_message::Payload::RelayOpen(open)) => { + assert_eq!(open.channel_id, channel_id); + } + other => panic!("expected RelayOpen after SessionAccepted, got {other:?}"), + } } #[tokio::test] - async fn require_persisted_sandbox_accepts_existing_sandbox() { - let store = test_store().await; - store - .put_message(&sandbox_record("sbx-1", "sandbox-one")) - .await - .expect("sandbox should persist"); + async fn relay_open_waits_until_session_accepted_is_queued() { + use tokio::sync::mpsc::error::TryRecvError; - require_persisted_sandbox(&store, "sbx-1") - .await - .expect("persisted sandbox should be accepted"); - } + let registry = Arc::new(SupervisorSessionRegistry::new()); + let (tx, mut rx) = mpsc::channel::(4); + assert!(!registry.register_awaiting_accept( + "sbx".to_string(), + "s1".to_string(), + tx.clone(), + make_shutdown(), + )); + // Compute readiness keeps seeing the session before it is accepted. + assert!(registry.has_session("sbx")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + assert_eq!( + registry.local_session_route("missing"), + LocalSessionRoute::Absent + ); - #[test] - fn expected_transport_close_is_nonfatal_only_during_shutdown() { - let status = Status::unknown("h2 protocol error: error reading a body from connection"); + let relay_registry = Arc::clone(®istry); + let relay = tokio::spawn(async move { + relay_registry + .open_relay("sbx", Duration::from_secs(2)) + .await + }); - assert!(expected_transport_close_during_shutdown(&status, true)); - assert!(!expected_transport_close_during_shutdown(&status, false)); - } + tokio::time::sleep(Duration::from_millis(250)).await; + assert!( + matches!(rx.try_recv(), Err(TryRecvError::Empty)), + "RelayOpen must not be queued before SessionAccepted" + ); - #[test] - fn unexpected_transport_error_stays_fatal_during_shutdown() { - let status = Status::internal("policy evaluation failed"); + tx.send(session_accepted_message("s1")).await.unwrap(); + assert!(registry.mark_accepts_relays("sbx", "s1")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Ready + ); - assert!(!expected_transport_close_during_shutdown(&status, true)); + let (channel_id, _relay_rx) = relay + .await + .unwrap() + .expect("relay should open once the session accepts relays"); + assert_accepted_then_relay_open(&mut rx, &channel_id).await; } #[test] - fn gateway_shutdown_makes_session_transport_close_nonfatal() { - let status = - Status::unknown("h2 protocol error: error reading a body from connection: broken pipe"); + fn mark_accepts_relays_rejects_replaced_session() { + let registry = SupervisorSessionRegistry::new(); + let (tx_old, _rx_old) = mpsc::channel::(1); + let (tx_new, _rx_new) = mpsc::channel::(1); - assert!(expected_transport_close_during_session_state( - &status, true, false, false, + registry.register_awaiting_accept( + "sbx".to_string(), + "s-old".to_string(), + tx_old, + make_shutdown(), + ); + assert!(!registry.mark_accepts_relays("sbx", "s-other")); + assert!(!registry.mark_accepts_relays("missing", "s-old")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + + assert!(registry.register_awaiting_accept( + "sbx".to_string(), + "s-new".to_string(), + tx_new, + make_shutdown(), )); + assert!(!registry.mark_accepts_relays("sbx", "s-old")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + assert!(registry.mark_accepts_relays("sbx", "s-new")); } - #[test] - fn sandbox_proto_terminating_detects_deleting_phase() { - let mut sandbox = sandbox_record("sbx-1", "sandbox-one"); - sandbox.set_phase(SandboxPhase::Deleting as i32); - - assert!(sandbox_proto_is_terminating(&sandbox)); - } + #[tokio::test] + async fn routed_relay_does_not_send_relay_open_before_session_accepted() { + use tokio::sync::mpsc::error::TryRecvError; - #[test] + let state = crate::grpc::test_support::test_server_state().await; + // A fresh owner record naming another replica that has no peer + // endpoint, as a stale route would. Consulting it fails the relay at + // once, so the relay only succeeds if it waits for the local session. + SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .publish( + "sbx", + "s0", + "inst", + 1, + "other-replica", + &local_owner_endpoint("other-replica"), + ) + .await + .unwrap(); + let (tx, mut rx) = mpsc::channel::(4); + state.supervisor_sessions.register_awaiting_accept( + "sbx".to_string(), + "s1".to_string(), + tx.clone(), + make_shutdown(), + ); + + let relay_state = Arc::clone(&state); + let relay = tokio::spawn(async move { + open_routed_relay_with_target( + &relay_state, + "sbx", + relay_open::Target::Ssh(SshRelayTarget {}), + String::new(), + Duration::from_secs(2), + ) + .await + }); + + tokio::time::sleep(Duration::from_millis(250)).await; + assert!( + !relay.is_finished(), + "routed relay must wait for the local session instead of using the owner record" + ); + assert!( + matches!(rx.try_recv(), Err(TryRecvError::Empty)), + "routed RelayOpen must not be queued before SessionAccepted" + ); + + tx.send(session_accepted_message("s1")).await.unwrap(); + assert!(state.supervisor_sessions.mark_accepts_relays("sbx", "s1")); + + let (channel_id, _relay_rx) = relay + .await + .unwrap() + .expect("routed relay should open on the local session once accepted"); + assert_accepted_then_relay_open(&mut rx, &channel_id).await; + } + + #[tokio::test] + async fn connect_supervisor_sends_session_accepted_before_relay_open() { + use crate::grpc::OpenShellService; + use openshell_core::proto::open_shell_server::OpenShellServer; + use tokio_stream::wrappers::TcpListenerStream; + + let state = crate::grpc::test_support::test_server_state().await; + state + .store + .put_message(&sandbox_record("sbx", "sandbox-one")) + .await + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let service = OpenShellServer::new(OpenShellService::new(Arc::clone(&state))); + let server = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + let mut client = open_shell_client::OpenShellClient::connect(format!("http://{address}")) + .await + .unwrap(); + + // Session setup registers the session, then parks in the endpoint + // status reset until this sandbox's mutation guard is released. + let sync_guard = state + .compute + .mutation_guard(crate::compute::MutationScope::sandbox("default", "sbx")) + .await + .unwrap(); + let (outbound_tx, outbound_rx) = mpsc::channel::(4); + outbound_tx + .send(SupervisorMessage { + payload: Some(supervisor_message::Payload::Hello(SupervisorHello { + sandbox_id: "sbx".to_string(), + instance_id: "inst".to_string(), + connection_epoch: 1, + ..Default::default() + })), + }) + .await + .unwrap(); + let connect = tokio::spawn(async move { + client + .connect_supervisor(ReceiverStream::new(outbound_rx)) + .await + }); + tokio::time::timeout(Duration::from_secs(5), async { + while !state.supervisor_sessions.has_session("sbx") { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("session setup should register the session"); + assert_eq!( + state.supervisor_sessions.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + + let relay_state = Arc::clone(&state); + let relay = tokio::spawn(async move { + open_routed_relay_with_target( + &relay_state, + "sbx", + relay_open::Target::Ssh(SshRelayTarget {}), + String::new(), + Duration::from_secs(5), + ) + .await + }); + tokio::time::sleep(Duration::from_millis(250)).await; + assert!( + !relay.is_finished(), + "relay must not open while session setup is in progress" + ); + drop(sync_guard); + + let mut inbound = connect + .await + .unwrap() + .expect("ConnectSupervisor should accept the session") + .into_inner(); + let (channel_id, _relay_rx) = relay + .await + .unwrap() + .expect("relay should open once the session is accepted"); + let first = inbound + .message() + .await + .unwrap() + .expect("SessionAccepted should be sent"); + assert!( + matches!( + first.payload, + Some(gateway_message::Payload::SessionAccepted(_)) + ), + "expected SessionAccepted first, got {:?}", + first.payload + ); + let second = inbound + .message() + .await + .unwrap() + .expect("RelayOpen should be sent"); + match second.payload { + Some(gateway_message::Payload::RelayOpen(open)) => { + assert_eq!(open.channel_id, channel_id); + } + other => panic!("expected RelayOpen after SessionAccepted, got {other:?}"), + } + + drop(outbound_tx); + server.abort(); + } + + /// Register `count` Ready sessions `sb-000`, `sb-001`, ... (sorted in slot + /// order) and return their shutdown receivers. + fn register_drain_targets( + registry: &SupervisorSessionRegistry, + count: usize, + ) -> Vec> { + (0..count) + .map(|index| { + let (tx, _rx) = mpsc::channel::(1); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + registry.register( + format!("sb-{index:03}"), + format!("s-{index}"), + tx, + shutdown_tx, + ); + shutdown_rx + }) + .collect() + } + + #[test] + fn drain_close_interval_scales_with_session_count() { + let window = Duration::from_secs(12); + let max_interval = Duration::from_millis(100); + assert_eq!(drain_close_interval(0, window, max_interval), max_interval); + assert_eq!(drain_close_interval(1, window, max_interval), max_interval); + assert_eq!(drain_close_interval(10, window, max_interval), max_interval); + assert_eq!( + drain_close_interval(120, window, max_interval), + max_interval + ); + assert_eq!( + drain_close_interval(1000, window, max_interval), + Duration::from_millis(12) + ); + assert!(drain_close_interval(usize::MAX, window, max_interval) <= max_interval); + + // 10 sessions close within a second, 120 use the whole window, and the + // last close always lands inside the window. + assert_eq!( + drain_close_interval(10, window, max_interval) * 9, + Duration::from_millis(900) + ); + assert_eq!( + drain_close_interval(120, window, max_interval) * 119, + Duration::from_millis(11_900) + ); + for sessions in [1_usize, 10, 120, 121, 1000, 100_000] { + let interval = drain_close_interval(sessions, window, max_interval); + let last_close = interval * u32::try_from(sessions - 1).unwrap(); + assert!( + last_close < window, + "{sessions} sessions: last close at {last_close:?}" + ); + } + } + + #[tokio::test] + async fn drain_with_waits_for_propagation_then_paces_closes() { + let registry = SupervisorSessionRegistry::new(); + let start = tokio::time::Instant::now(); + let observers: Vec<_> = register_drain_targets(®istry, 4) + .into_iter() + .map(|closed| { + tokio::spawn(async move { + closed + .await + .expect("the drain slot should signal the session"); + start.elapsed() + }) + }) + .collect(); + + // Δ = min(20s / 4, 50ms) = 50ms. Without the cap the slots would be 5s + // apart and the timeout below would trip. + let propagation = Duration::from_millis(50); + let summary = tokio::time::timeout( + Duration::from_secs(5), + drain_with( + ®istry, + propagation, + Duration::from_secs(20), + Duration::from_millis(50), + ), + ) + .await + .expect("a four-session drain should finish quickly"); + assert_eq!( + summary, + DrainSummary { + planned: 4, + signaled: 4 + } + ); + for (slot, observer) in (0_u32..).zip(observers) { + let closed_at = observer.await.unwrap(); + let earliest = propagation + Duration::from_millis(50) * slot; + assert!( + closed_at >= earliest, + "slot {slot} closed at {closed_at:?}, before {earliest:?}" + ); + } + // The entry stays for the session task's own cleanup. + assert!(registry.has_session("sb-000")); + assert_eq!( + registry.local_session_route("sb-000"), + LocalSessionRoute::Settling + ); + } + + #[tokio::test] + async fn drain_sessions_signals_every_session_for_large_counts() { + let registry = SupervisorSessionRegistry::new(); + let closed = register_drain_targets(®istry, 200); + + let start = tokio::time::Instant::now(); + // Δ = min(200ms / 200, 100ms) = 1ms. At the 100ms cap the drain would + // take 19.9s and trip the timeout. + let summary = tokio::time::timeout( + Duration::from_secs(5), + registry.drain_sessions(Duration::from_millis(200), Duration::from_millis(100)), + ) + .await + .expect("a 1ms-paced drain of 200 sessions should finish in about 200ms"); + let elapsed = start.elapsed(); + + assert_eq!( + summary, + DrainSummary { + planned: 200, + signaled: 200 + } + ); + assert!( + elapsed >= Duration::from_millis(199), + "the last of 200 slots starts at 199ms, drain took {elapsed:?}" + ); + for mut receiver in closed { + assert!(receiver.try_recv().is_ok()); + } + } + + #[tokio::test] + async fn drain_sessions_keeps_serving_sessions_until_their_slot() { + use tokio::sync::oneshot::error::TryRecvError; + + let registry = Arc::new(SupervisorSessionRegistry::new()); + let (first_tx, _first_rx) = mpsc::channel::(4); + let (first_shutdown, first_closed) = oneshot::channel(); + registry.register("sb-0".into(), "s-0".into(), first_tx, first_shutdown); + let (second_tx, mut second_rx) = mpsc::channel::(4); + let (second_shutdown, mut second_closed) = oneshot::channel(); + registry.register("sb-1".into(), "s-1".into(), second_tx, second_shutdown); + + // Δ = min(20s / 2, 10s) = 10s, so the second slot is far away. + let drain_registry = Arc::clone(®istry); + let drain = tokio::spawn(async move { + drain_registry + .drain_sessions(Duration::from_secs(20), Duration::from_secs(10)) + .await + }); + first_closed + .await + .expect("the first slot should fire at once"); + assert_eq!( + registry.local_session_route("sb-0"), + LocalSessionRoute::Settling + ); + + assert_eq!( + registry.local_session_route("sb-1"), + LocalSessionRoute::Ready + ); + let (channel_id, _relay_rx) = registry + .open_relay("sb-1", Duration::ZERO) + .await + .expect("a session should serve relays until its slot"); + match second_rx.recv().await.and_then(|message| message.payload) { + Some(gateway_message::Payload::RelayOpen(open)) => { + assert_eq!(open.channel_id, channel_id); + } + other => panic!("expected RelayOpen before the slot, got {other:?}"), + } + assert!(matches!(second_closed.try_recv(), Err(TryRecvError::Empty))); + drain.abort(); + } + + #[tokio::test] + async fn drain_sessions_skips_sessions_that_end_after_the_snapshot() { + use tokio::sync::oneshot::error::TryRecvError; + + let registry = SupervisorSessionRegistry::new(); + let _closed = register_drain_targets(®istry, 3); + + let drain = registry.drain_sessions(Duration::from_millis(150), Duration::from_millis(50)); + tokio::pin!(drain); + // The first poll takes the snapshot before the drain first sleeps. + assert!(futures_util::poll!(&mut drain).is_pending()); + assert_eq!(registry.remove_if_current("sb-001", "s-1"), Some(false)); + let (tx, _rx) = mpsc::channel::(1); + let (replacement_shutdown, mut replacement_closed) = oneshot::channel(); + assert!(registry.register("sb-002".into(), "s-2b".into(), tx, replacement_shutdown)); + + assert_eq!( + drain.await, + DrainSummary { + planned: 3, + signaled: 1 + } + ); + assert!(matches!( + replacement_closed.try_recv(), + Err(TryRecvError::Empty) + )); + assert_eq!( + registry.local_session_route("sb-002"), + LocalSessionRoute::Ready + ); + } + + #[test] + fn close_for_drain_skips_sessions_that_ended_or_were_replaced() { + use tokio::sync::oneshot::error::TryRecvError; + + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel::(1); + let (ended_shutdown, _ended_closed) = oneshot::channel(); + registry.register("sb-1".into(), "s-1".into(), tx.clone(), ended_shutdown); + assert_eq!(registry.remove_if_current("sb-1", "s-1"), Some(false)); + assert!(!registry.close_for_drain("sb-1", "s-1")); + + let (old_shutdown, _old_closed) = oneshot::channel(); + registry.register("sb-2".into(), "s-old".into(), tx.clone(), old_shutdown); + let (new_shutdown, mut new_closed) = oneshot::channel(); + assert!(registry.register("sb-2".into(), "s-new".into(), tx, new_shutdown)); + assert!(!registry.close_for_drain("sb-2", "s-old")); + assert!(matches!(new_closed.try_recv(), Err(TryRecvError::Empty))); + assert_eq!( + registry.local_session_route("sb-2"), + LocalSessionRoute::Ready + ); + } + + #[tokio::test] + async fn close_for_drain_is_single_shot_and_stops_relay_routing() { + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel::(4); + let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); + registry.register("sbx".into(), "s1".into(), tx, shutdown_tx); + + assert!(registry.close_for_drain("sbx", "s1")); + assert!(shutdown_rx.try_recv().is_ok()); + assert!( + !registry.close_for_drain("sbx", "s1"), + "a session is signaled once" + ); + + let err = registry + .open_relay("sbx", Duration::ZERO) + .await + .expect_err("a closing session must not take new relays"); + assert_eq!(err.code(), tonic::Code::Unavailable); + // Compute readiness keeps seeing the session until its task cleans up. + assert!(registry.has_session("sbx")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + } + + #[test] + fn register_after_drain_close_does_not_panic() { + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel::(1); + let (old_shutdown, _old_closed) = oneshot::channel(); + registry.register("sbx".into(), "s-old".into(), tx.clone(), old_shutdown); + assert!(registry.close_for_drain("sbx", "s-old")); + assert!(registry.register("sbx".into(), "s-new".into(), tx.clone(), make_shutdown())); + assert!(registry.is_current_session("sbx", "s-new")); + + let (other_shutdown, _other_closed) = oneshot::channel(); + registry.register("sb-other".into(), "s-1".into(), tx, other_shutdown); + assert!(registry.close_for_drain("sb-other", "s-1")); + assert!(registry.disconnect("sb-other")); + assert!(!registry.has_session("sb-other")); + } + + #[test] + fn mark_accepts_relays_rejects_drained_session() { + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel::(1); + let (shutdown_tx, _shutdown_rx) = oneshot::channel(); + registry.register_awaiting_accept("sbx".into(), "s1".into(), tx, shutdown_tx); + + assert!(registry.close_for_drain("sbx", "s1")); + assert!(!registry.mark_accepts_relays("sbx", "s1")); + assert_eq!( + registry.local_session_route("sbx"), + LocalSessionRoute::Settling + ); + } + + #[test] + fn peer_relay_session_wait_is_zero_while_draining() { + assert_eq!(peer_relay_session_wait(true), Duration::ZERO); + assert_eq!(peer_relay_session_wait(false), Duration::from_secs(5)); + } + + fn peer_relay_init(sandbox_id: &str, channel_id: &str) -> PeerRelayFrame { + PeerRelayFrame { + payload: Some(peer_relay_frame::Payload::Init(PeerRelayInit { + sandbox_id: sandbox_id.to_string(), + relay_open: Some(peer_relay_open(channel_id)), + requester_replica_id: "replica-requester".to_string(), + })), + } + } + + #[tokio::test] + async fn draining_owner_fails_peer_relay_at_once_without_a_session() { + use crate::auth::principal::PeerPrincipal; + use crate::grpc::OpenShellService; + use openshell_core::proto::open_shell_server::OpenShellServer; + use tokio_stream::wrappers::TcpListenerStream; + + let state = crate::grpc::test_support::test_server_state().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + // Stands in for peer authentication, which runs in the multiplexer. + let service = OpenShellServer::with_interceptor( + OpenShellService::new(Arc::clone(&state)), + |mut request: Request<()>| { + request + .extensions_mut() + .insert(Principal::Peer(PeerPrincipal { + replica_id: "replica-requester".to_string(), + pod_uid: "pod-uid".to_string(), + })); + Ok(request) + }, + ); + let server = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + let mut client = open_shell_client::OpenShellClient::connect(format!("http://{address}")) + .await + .unwrap(); + + state.supervisor_sessions.close_admission(); + let started = Instant::now(); + let err = client + .peer_relay(tokio_stream::iter([peer_relay_init("sbx", "ch-1")])) + .await + .expect_err("a draining replica without the session must fail the relay"); + assert_eq!(err.code(), tonic::Code::Unavailable); + assert!(err.message().contains("draining"), "{err}"); + assert!( + started.elapsed() < PEER_RELAY_SESSION_WAIT, + "the relay must fail before the normal session wait" + ); + + // A draining owner still serves the sessions it holds. + let (tx, mut rx) = mpsc::channel::(4); + state + .supervisor_sessions + .register("sbx".into(), "s1".into(), tx, make_shutdown()); + let relay = tokio::spawn(async move { + client + .peer_relay(tokio_stream::iter([peer_relay_init("sbx", "ch-2")])) + .await + }); + let message = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("the owner should queue RelayOpen on its session"); + match message.and_then(|message| message.payload) { + Some(gateway_message::Payload::RelayOpen(open)) => { + assert_eq!(open.channel_id, "ch-2"); + } + other => panic!("expected RelayOpen, got {other:?}"), + } + relay.abort(); + server.abort(); + } + + #[tokio::test] + async fn drain_closes_a_connected_supervisor_session_after_its_cleanup() { + use crate::grpc::OpenShellService; + use openshell_core::proto::open_shell_server::OpenShellServer; + use tokio_stream::wrappers::TcpListenerStream; + + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + state + .store + .put_message(&sandbox_record("sbx", "sandbox-one")) + .await + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let service = OpenShellServer::new(OpenShellService::new(Arc::clone(&state))); + let server = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + .unwrap(); + }); + let mut client = open_shell_client::OpenShellClient::connect(format!("http://{address}")) + .await + .unwrap(); + let (outbound_tx, outbound_rx) = mpsc::channel::(4); + outbound_tx + .send(SupervisorMessage { + payload: Some(supervisor_message::Payload::Hello(SupervisorHello { + sandbox_id: "sbx".to_string(), + instance_id: "inst".to_string(), + connection_epoch: 1, + ..Default::default() + })), + }) + .await + .unwrap(); + let mut inbound = client + .connect_supervisor(ReceiverStream::new(outbound_rx)) + .await + .expect("ConnectSupervisor should accept the session") + .into_inner(); + let accepted = inbound.message().await.unwrap(); + assert!(matches!( + accepted.and_then(|message| message.payload), + Some(gateway_message::Payload::SessionAccepted(_)) + )); + assert_eq!( + state.supervisor_sessions.local_session_route("sbx"), + LocalSessionRoute::Ready + ); + assert_eq!( + metrics.value("openshell_server_supervisor_sessions"), + Some(1) + ); + + state.supervisor_sessions.close_admission(); + let summary = drain_with( + &state.supervisor_sessions, + Duration::from_millis(10), + Duration::from_millis(100), + Duration::from_millis(50), + ) + .await; + assert_eq!( + summary, + DrainSummary { + planned: 1, + signaled: 1 + } + ); + + // The session task owns the stream sender, so the supervisor sees the + // end of the stream only after ownership cleanup finished. + let end = tokio::time::timeout(Duration::from_secs(5), inbound.message()) + .await + .expect("the stream should end after the drain slot"); + assert!( + matches!(end, Ok(None)), + "expected a clean end of stream, got {end:?}" + ); + assert!(!state.supervisor_sessions.has_session("sbx")); + assert!( + SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .read("sbx") + .await + .unwrap() + .is_none() + ); + assert_eq!( + metrics.value("openshell_server_supervisor_sessions"), + Some(0) + ); + + drop(outbound_tx); + server.abort(); + } + + #[tokio::test] + async fn require_persisted_sandbox_rejects_missing_sandbox() { + let store = test_store().await; + + let err = require_persisted_sandbox(&store, "missing") + .await + .expect_err("missing sandbox should be rejected"); + + assert_eq!(err.code(), tonic::Code::NotFound); + } + + #[tokio::test] + async fn require_persisted_sandbox_accepts_existing_sandbox() { + let store = test_store().await; + store + .put_message(&sandbox_record("sbx-1", "sandbox-one")) + .await + .expect("sandbox should persist"); + + require_persisted_sandbox(&store, "sbx-1") + .await + .expect("persisted sandbox should be accepted"); + } + + #[test] + fn expected_transport_close_is_nonfatal_only_during_shutdown() { + let status = Status::unknown("h2 protocol error: error reading a body from connection"); + + assert!(expected_transport_close_during_shutdown(&status, true)); + assert!(!expected_transport_close_during_shutdown(&status, false)); + } + + #[test] + fn unexpected_transport_error_stays_fatal_during_shutdown() { + let status = Status::internal("policy evaluation failed"); + + assert!(!expected_transport_close_during_shutdown(&status, true)); + } + + #[test] + fn gateway_shutdown_makes_session_transport_close_nonfatal() { + let status = + Status::unknown("h2 protocol error: error reading a body from connection: broken pipe"); + + assert!(expected_transport_close_during_session_state( + &status, true, false, false, + )); + } + + #[test] + fn sandbox_proto_terminating_detects_deleting_phase() { + let mut sandbox = sandbox_record("sbx-1", "sandbox-one"); + sandbox.set_phase(SandboxPhase::Deleting as i32); + + assert!(sandbox_proto_is_terminating(&sandbox)); + } + + #[test] fn sandbox_proto_terminating_detects_deletion_timestamp() { let mut sandbox = sandbox_record("sbx-1", "sandbox-one"); sandbox.metadata.as_mut().unwrap().deletion_time = @@ -3059,6 +4236,7 @@ mod tests { #[test] fn claim_relay_success() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, _relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3070,10 +4248,86 @@ mod tests { let result = registry.claim_relay("ch-1", Some(&principal)); assert!(result.is_ok()); assert!(!registry.pending_relays.lock().unwrap().contains_key("ch-1")); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); + assert_eq!( + metrics.value("openshell_server_relay_claim_duration_seconds_count"), + Some(1) + ); + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), None); + } + + /// Waker that reads `relay_pending` each time it is woken. `oneshot::Sender::send` wakes a + /// registered receiver synchronously, so the probe sees the gauge exactly as a waiter on + /// another worker thread could at that instant. + struct PendingGaugeProbe { + read: Box Option + Send + Sync>, + seen: Mutex>>, + } + + impl PendingGaugeProbe { + fn register(metrics: &MetricsCapture, rx: &mut oneshot::Receiver) -> Arc { + let probe = Arc::new(Self { + read: metrics.value_reader(gateway_metrics::RELAY_PENDING), + seen: Mutex::new(Vec::new()), + }); + let waker = std::task::Waker::from(Arc::clone(&probe)); + let mut cx = std::task::Context::from_waker(&waker); + assert!(Pin::new(rx).poll(&mut cx).is_pending()); + probe + } + + fn seen(&self) -> Vec> { + self.seen.lock().unwrap().clone() + } + } + + impl std::task::Wake for PendingGaugeProbe { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.seen.lock().unwrap().push((self.read)()); + } + } + + #[test] + fn claim_relay_releases_pending_slot_before_waking_waiter() { + let metrics = MetricsCapture::install(); + let registry = SupervisorSessionRegistry::new(); + let (relay_tx, mut relay_rx) = oneshot::channel(); + registry.pending_relays.lock().unwrap().insert( + "ch-1".to_string(), + pending_relay("sbx-test", relay_tx, Instant::now()), + ); + let probe = PendingGaugeProbe::register(&metrics, &mut relay_rx); + + registry + .claim_relay("ch-1", Some(&sandbox_principal("sbx-test"))) + .expect("claim should succeed"); + // The slot is released under the pending lock, before the waiter is woken, so a + // concurrent open can never push `relay_pending` above capacity. + assert_eq!(probe.seen(), vec![Some(0)]); + } + + #[test] + fn fail_pending_relay_releases_pending_slot_before_waking_waiter() { + let metrics = MetricsCapture::install(); + let registry = SupervisorSessionRegistry::new(); + let (relay_tx, mut relay_rx) = oneshot::channel(); + registry.pending_relays.lock().unwrap().insert( + "ch-fail".to_string(), + pending_relay("sbx-test", relay_tx, Instant::now()), + ); + let probe = PendingGaugeProbe::register(&metrics, &mut relay_rx); + + assert!(registry.fail_pending_relay("ch-fail", "target refused".to_string())); + assert_eq!(probe.seen(), vec![Some(0)]); } #[test] fn claim_relay_rejects_cross_sandbox_principal_without_consuming_channel() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, _relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3094,6 +4348,11 @@ mod tests { .contains_key("ch-cross"), "failed cross-sandbox claim must not consume the channel" ); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(1)); + assert_eq!( + metrics.value("openshell_server_relay_claim_duration_seconds_count"), + None + ); } #[test] @@ -3113,6 +4372,7 @@ mod tests { #[tokio::test] async fn relay_open_failure_completes_pending_waiter() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3133,10 +4393,13 @@ mod tests { let status = result.expect_err("waiter should receive status failure"); assert_eq!(status.code(), tonic::Code::Unavailable); assert_eq!(status.message(), "target refused"); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), None); } #[test] fn claim_relay_expired_returns_deadline_exceeded() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, _relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3162,10 +4425,17 @@ mod tests { .unwrap() .contains_key("ch-old") ); + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), Some(1)); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); + assert_eq!( + metrics.value("openshell_server_relay_claim_duration_seconds_count"), + None + ); } #[test] fn claim_relay_receiver_dropped_returns_internal() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, relay_rx) = oneshot::channel::>(); drop(relay_rx); // Gateway-side waiter has given up already. @@ -3178,6 +4448,11 @@ mod tests { .claim_relay("ch-1", Some(&sandbox_principal("sbx-test"))) .expect_err("should err when receiver is gone"); assert_eq!(err.code(), tonic::Code::Internal); + assert_eq!( + metrics.value("openshell_server_relay_claim_duration_seconds_count"), + Some(1) + ); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); } #[tokio::test] @@ -3215,6 +4490,7 @@ mod tests { #[test] fn reap_expired_relays_removes_old_entries() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, _relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3236,10 +4512,13 @@ mod tests { .unwrap() .contains_key("ch-old") ); + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), Some(1)); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(0)); } #[test] fn reap_expired_relays_keeps_fresh_entries() { + let metrics = MetricsCapture::install(); let registry = SupervisorSessionRegistry::new(); let (relay_tx, _relay_rx) = oneshot::channel(); registry.pending_relays.lock().unwrap().insert( @@ -3255,6 +4534,9 @@ mod tests { .unwrap() .contains_key("ch-fresh") ); + // Reaping nothing records nothing. + assert_eq!(metrics.value(gateway_metrics::RELAY_EXPIRED_TOTAL), None); + assert_eq!(metrics.value(gateway_metrics::RELAY_PENDING), Some(1)); } fn owner_record(replica: &str) -> crate::supervisor_owner::OwnerRecord { @@ -3333,6 +4615,238 @@ mod tests { assert!(OWNER_CACHE_TTL < OWNER_TTL); } + // ---- peer request metrics (requester side) ---- + + #[derive(Clone, Copy)] + enum FakePeerReply { + Status(tonic::Code), + EmptyOk, + } + + /// Minimal h2c server that answers every gRPC call the same way. It stands in for an owner + /// replica without implementing the full `OpenShell` service. + async fn spawn_fake_peer(reply: FakePeerReply) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + tokio::spawn(async move { + let service = hyper::service::service_fn( + move |_req: http::Request| async move { + Ok::<_, Infallible>(fake_peer_response(reply)) + }, + ); + let _ = hyper_util::server::conn::auto::Builder::new( + hyper_util::rt::TokioExecutor::new(), + ) + .serve_connection(hyper_util::rt::TokioIo::new(stream), service) + .await; + }); + } + }); + format!("http://{addr}") + } + + fn fake_peer_response( + reply: FakePeerReply, + ) -> http::Response> { + let builder = http::Response::builder() + .status(200) + .header("content-type", "application/grpc"); + match reply { + // Trailers-only error: tonic returns Err(status) for unary and streaming calls. + FakePeerReply::Status(code) => builder + .header("grpc-status", i32::from(code).to_string()) + .header("grpc-message", "fake peer") + .body(Empty::new().boxed_unsync()) + .unwrap(), + // One empty message (5-byte frame header, zero length), then grpc-status 0. This + // decodes as a default response for any unary RPC, and gives streaming calls an OK + // header. + FakePeerReply::EmptyOk => { + let mut trailers = http::HeaderMap::new(); + trailers.insert("grpc-status", http::HeaderValue::from_static("0")); + let frames = futures::stream::iter([ + Ok::<_, Infallible>(Frame::data(Bytes::from_static(&[0, 0, 0, 0, 0]))), + Ok(Frame::trailers(trailers)), + ]); + builder + .body(StreamBody::new(frames).boxed_unsync()) + .unwrap() + } + } + } + + fn seed_peer_token(state: &ServerState) { + *state.peer_routes.token.lock().unwrap() = Some(CachedPeerToken { + token: "test-peer-token".to_string(), + refresh_at: Instant::now() + Duration::from_mins(5), + }); + } + + fn owner_at(endpoint: &str) -> crate::supervisor_owner::OwnerRecord { + let mut owner = owner_record("replica-owner"); + owner.owner_peer_endpoint = endpoint.to_string(); + owner + } + + fn closed_local_endpoint() -> String { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + format!("http://{addr}") + } + + fn peer_relay_open(channel_id: &str) -> RelayOpen { + RelayOpen { + channel_id: channel_id.to_string(), + target: Some(relay_open::Target::Ssh(SshRelayTarget {})), + service_id: String::new(), + } + } + + #[tokio::test] + async fn peer_relay_metrics_keep_owner_code_before_unavailable_remap() { + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + seed_peer_token(&state); + let endpoint = spawn_fake_peer(FakePeerReply::Status(tonic::Code::ResourceExhausted)).await; + + let err = connect_peer_relay(&state, &endpoint, "sbx-peer", peer_relay_open("ch-peer")) + .await + .expect_err("the owner rejected the relay"); + assert_eq!(err.code(), tonic::Code::Unavailable); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerRelay\",outcome=\"rpc_error\",code=\"resource_exhausted\"}" + ), + Some(1) + ); + assert_eq!( + metrics.value( + "openshell_server_peer_request_duration_seconds_count{rpc=\"PeerRelay\",outcome=\"rpc_error\"}" + ), + Some(1) + ); + let rendered = metrics.render(); + assert!(rendered.contains( + "openshell_server_peer_request_duration_seconds_bucket{rpc=\"PeerRelay\",outcome=\"rpc_error\",le=\"0.001\"}" + )); + assert!( + !state + .peer_routes + .channels + .lock() + .unwrap() + .contains_key(&endpoint), + "a failed peer relay must evict the channel" + ); + let host_port = endpoint.trim_start_matches("http://"); + assert!( + !rendered.contains(host_port), + "metrics must not carry peer endpoints" + ); + } + + #[tokio::test] + async fn peer_relay_metrics_record_ok_when_owner_accepts() { + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + seed_peer_token(&state); + let endpoint = spawn_fake_peer(FakePeerReply::EmptyOk).await; + + connect_peer_relay(&state, &endpoint, "sbx-peer", peer_relay_open("ch-peer")) + .await + .expect("the owner accepted the relay"); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerRelay\",outcome=\"ok\",code=\"ok\"}" + ), + Some(1) + ); + } + + #[tokio::test] + async fn peer_forward_metrics_record_client_error_when_owner_unreachable() { + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + seed_peer_token(&state); + let endpoint = closed_local_endpoint(); + + let err = forward_provider_status_query_to_owner( + &state, + &owner_at(&endpoint), + "sbx-peer", + GetSandboxProviderStatusRequest::default(), + ) + .await + .expect_err("the owner is unreachable"); + assert_eq!(err.code(), tonic::Code::Unavailable); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerGetSandboxProviderStatus\",outcome=\"client_error\",code=\"unavailable\"}" + ), + Some(1) + ); + assert!( + !metrics + .render() + .contains("rpc=\"PeerGetSandboxProviderStatus\",outcome=\"rpc_error\"") + ); + } + + #[tokio::test] + async fn peer_forward_metrics_record_owner_rpc_error_code() { + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + seed_peer_token(&state); + let endpoint = spawn_fake_peer(FakePeerReply::Status(tonic::Code::PermissionDenied)).await; + + let err = forward_endpoint_status_to_owner( + &state, + &owner_at(&endpoint), + ReportEndpointStatusRequest { + sandbox_id: "sbx-peer".into(), + ..Default::default() + }, + ) + .await + .expect_err("the owner rejected the report"); + // Unary forwarders return the owner's status unchanged. + assert_eq!(err.code(), tonic::Code::PermissionDenied); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerReportEndpointStatus\",outcome=\"rpc_error\",code=\"permission_denied\"}" + ), + Some(1) + ); + } + + #[tokio::test] + async fn peer_forward_metrics_record_ok() { + let metrics = MetricsCapture::install(); + let state = crate::grpc::test_support::test_server_state().await; + seed_peer_token(&state); + let endpoint = spawn_fake_peer(FakePeerReply::EmptyOk).await; + + forward_provider_readiness_to_owner( + &state, + &owner_at(&endpoint), + ReportProviderReadinessRequest { + sandbox_id: "sbx-peer".into(), + ..Default::default() + }, + ) + .await + .expect("the owner accepted the report"); + assert_eq!( + metrics.value( + "openshell_server_peer_requests_total{rpc=\"PeerReportProviderReadiness\",outcome=\"ok\",code=\"ok\"}" + ), + Some(1) + ); + } + #[tokio::test] async fn evict_channel_removes_only_the_named_peer() { let cache = PeerRouteCache::default(); diff --git a/crates/openshell-server/tests/supervisor_relay_integration.rs b/crates/openshell-server/tests/supervisor_relay_integration.rs index 5856acbbda..73d9c909f7 100644 --- a/crates/openshell-server/tests/supervisor_relay_integration.rs +++ b/crates/openshell-server/tests/supervisor_relay_integration.rs @@ -22,6 +22,8 @@ use hyper_util::{ rt::{TokioExecutor, TokioIo}, server::conn::auto::Builder, }; +use metrics::LocalRecorderGuard; +use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle, PrometheusRecorder}; use openshell_core::proto::{ GatewayMessage, PeerRelayFrame, RelayFrame, RelayInit, SupervisorMessage, TcpForwardFrame, open_shell_client::OpenShellClient, @@ -658,6 +660,38 @@ fn register_session_with_capacity( rx } +/// Captures metrics recorded on this test's thread. `#[tokio::test]` is current-thread, so the +/// in-process gateway tasks record here too. Do not pass it into async helper fns (it is !Send). +struct MetricsCapture { + handle: PrometheusHandle, + _guard: LocalRecorderGuard<'static>, +} + +impl MetricsCapture { + fn install() -> Self { + // Leaked (test only) so the guard can borrow the recorder for 'static. + let recorder: &'static PrometheusRecorder = + Box::leak(Box::new(PrometheusBuilder::new().build_recorder())); + let handle = recorder.handle(); + let guard = metrics::set_default_local_recorder(recorder); + Self { + handle, + _guard: guard, + } + } + + fn render(&self) -> String { + self.handle.render() + } + + /// Integer value of one exact series, or `None` when it was never emitted. + fn value(&self, series: &str) -> Option { + self.render() + .lines() + .find_map(|line| line.strip_prefix(series)?.strip_prefix(' ')?.parse().ok()) + } +} + /// Mock supervisor that opens a `RelayStream`, sends `Init`, then echoes every /// data frame it receives. Returns when the gateway drops the stream or when /// the supervisor's own outbound channel closes. @@ -859,6 +893,7 @@ async fn concurrent_relays_multiplex_independently() { /// rather than racing the pending map into an inconsistent state. #[tokio::test] async fn open_relay_enforces_per_sandbox_cap_under_concurrent_burst() { + let metrics = MetricsCapture::install(); let registry = Arc::new(SupervisorSessionRegistry::new()); let _channel = spawn_gateway(Arc::clone(®istry)).await; // Oversized mpsc so the session doesn't backpressure the burst — the cap, @@ -894,6 +929,23 @@ async fn open_relay_enforces_per_sandbox_cap_under_concurrent_burst() { } assert_eq!(ok, 32, "exactly per-sandbox cap should succeed"); assert_eq!(exhausted, 32, "overflow should be rejected, not dropped"); + // The successful receivers were dropped, so their entries stay pending + // until they are claimed or reaped. + assert_eq!(metrics.value("openshell_server_relay_pending"), Some(32)); + assert_eq!( + metrics.value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}"), + Some(32) + ); + assert_eq!( + metrics + .value("openshell_server_relay_rejected_total{reason=\"global_capacity\"}") + .unwrap_or(0), + 0 + ); + assert_eq!( + metrics.value("openshell_server_supervisor_sessions"), + Some(1) + ); // A different sandbox still has headroom — the per-sandbox cap doesn't // leak onto unrelated tenants. @@ -902,6 +954,107 @@ async fn open_relay_enforces_per_sandbox_cap_under_concurrent_burst() { .open_relay("sbx-other", Duration::from_secs(1)) .await .expect("other sandbox should not be affected by sbx cap"); + assert_eq!(metrics.value("openshell_server_relay_pending"), Some(33)); + assert_eq!( + metrics.value("openshell_server_supervisor_sessions"), + Some(2) + ); +} + +/// Bursts more `open_relay` calls than the global cap allows, spread so that +/// no sandbox reaches its own cap, and asserts the global ceiling and its +/// metrics. +#[tokio::test] +async fn open_relay_enforces_global_cap_under_concurrent_burst() { + let metrics = MetricsCapture::install(); + let registry = Arc::new(SupervisorSessionRegistry::new()); + let _channel = spawn_gateway(Arc::clone(®istry)).await; + let sandbox_ids: Vec = (0..9).map(|i| format!("sbx-{i}")).collect(); + let _session_rxs: Vec<_> = sandbox_ids + .iter() + .map(|id| register_session_with_capacity(®istry, id, 64)) + .collect(); + + // 9 x 32 = 288 opens against a global cap of 256. No sandbox gets more + // than its cap of 32 attempts and the global check runs first, so exactly + // 32 opens hit the global cap. + let mut handles = Vec::with_capacity(288); + for id in &sandbox_ids { + for _ in 0..32 { + let r = Arc::clone(®istry); + let id = id.clone(); + handles.push(tokio::spawn(async move { + r.open_relay(&id, Duration::from_secs(1)).await + })); + } + } + + let mut ok = 0usize; + let mut exhausted = 0usize; + for h in handles { + match h.await.expect("task joined") { + Ok(_pair) => ok += 1, + Err(status) if status.code() == tonic::Code::ResourceExhausted => { + assert!( + status.message().contains("gateway relay capacity"), + "expected global capacity error message, got: {}", + status.message() + ); + exhausted += 1; + } + Err(other) => panic!("unexpected open_relay error: {other:?}"), + } + } + assert_eq!(ok, 256, "exactly the global cap should succeed"); + assert_eq!(exhausted, 32, "overflow should be rejected, not dropped"); + + assert_eq!(metrics.value("openshell_server_relay_pending"), Some(256)); + assert_eq!( + metrics.value("openshell_server_relay_rejected_total{reason=\"global_capacity\"}"), + Some(32) + ); + assert_eq!( + metrics + .value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}") + .unwrap_or(0), + 0 + ); + assert_eq!( + metrics.value("openshell_server_supervisor_sessions"), + Some(9) + ); + let rendered = metrics.render(); + for forbidden in ["sandbox_id=", "channel_id=", "sbx-0", "endpoint="] { + assert!( + !rendered.contains(forbidden), + "metric labels must not carry identifiers ({forbidden})" + ); + } +} + +#[tokio::test] +async fn relay_claim_releases_pending_slot_and_records_claim_latency() { + let metrics = MetricsCapture::install(); + let registry = Arc::new(SupervisorSessionRegistry::new()); + let channel = spawn_gateway(Arc::clone(®istry)).await; + let _session_rx = register_session(®istry, "sbx"); + + let (channel_id, relay_rx) = registry + .open_relay("sbx", Duration::from_secs(2)) + .await + .expect("open_relay"); + assert_eq!(metrics.value("openshell_server_relay_pending"), Some(1)); + + tokio::spawn(run_echo_supervisor(channel, channel_id)); + let _relay = relay_rx.await.expect("relay result").expect("relay duplex"); + + // The claim releases the pending slot under the pending lock, before it + // wakes the waiter. + assert_eq!(metrics.value("openshell_server_relay_pending"), Some(0)); + assert_eq!( + metrics.value("openshell_server_relay_claim_duration_seconds_count"), + Some(1) + ); } /// Build an in-memory store sufficient for wiring `health_router` in tests diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index 0b713f9921..c2fc179d7f 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -37,6 +37,50 @@ use openshell_core::transport_errors::is_expected_transport_close_status; const INITIAL_BACKOFF: Duration = Duration::from_secs(1); const MAX_BACKOFF: Duration = Duration::from_secs(30); +/// Reconnect delay policy for the gateway control session. +/// +/// The ceiling doubles after every failed attempt and returns to +/// [`INITIAL_BACKOFF`] once the gateway accepted the attempt that just ended, +/// so a long-lived sandbox reconnects within about a second when its gateway +/// replica drains. Delays use equal jitter, `[ceiling / 2, ceiling]`, so +/// sessions closed together do not redial in lockstep. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ReconnectBackoff { + ceiling: Duration, +} + +impl Default for ReconnectBackoff { + fn default() -> Self { + Self { + ceiling: INITIAL_BACKOFF, + } + } +} + +impl ReconnectBackoff { + /// Delay before the next attempt. `accepted` reports whether the attempt + /// that just ended received `SessionAccepted`; `unit_sample` is uniform in + /// `[0, 1)`. + fn next_delay(&mut self, accepted: bool, unit_sample: f64) -> Duration { + if accepted { + self.ceiling = INITIAL_BACKOFF; + } + let delay = equal_jitter(self.ceiling, unit_sample); + self.ceiling = self.ceiling.saturating_mul(2).min(MAX_BACKOFF); + delay + } +} + +fn equal_jitter(ceiling: Duration, unit_sample: f64) -> Duration { + let sample = if unit_sample.is_finite() { + unit_sample.clamp(0.0, 1.0) + } else { + 1.0 + }; + let half = ceiling / 2; + half + half.mul_f64(sample) +} + /// Runtime identity and status channel shared with a supervisor session task. pub struct SessionRuntimeContext { /// Identifies the local supervisor process across gateway reconnects. @@ -354,7 +398,7 @@ struct SessionConfig { } async fn run_session_loop(config: SessionConfig) { - let mut backoff = INITIAL_BACKOFF; + let mut backoff = ReconnectBackoff::default(); let mut attempt: u64 = 0; loop { @@ -364,9 +408,11 @@ async fn run_session_loop(config: SessionConfig) { if let Some(updates) = &config.session_id_updates { updates.send_replace(None); } + // `ready_tx` turns true only after SessionAccepted (run_single_session), + // so its previous value says whether the gateway accepted this attempt. + let accepted = config.ready_tx.send_replace(false); match result { Ok(()) => { - config.ready_tx.send_replace(false); let event = session_closed_event( openshell_ocsf::ctx::ctx(), &config.endpoint, @@ -376,7 +422,6 @@ async fn run_session_loop(config: SessionConfig) { break; } Err(e) => { - config.ready_tx.send_replace(false); let event = session_failed_event( openshell_ocsf::ctx::ctx(), &config.endpoint, @@ -384,8 +429,14 @@ async fn run_session_loop(config: SessionConfig) { &e.to_string(), ); ocsf_emit!(event); - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(MAX_BACKOFF); + let delay = backoff.next_delay(accepted, rand::random::()); + debug!( + attempt, + accepted, + retry_after_ms = delay.as_millis(), + "supervisor session: reconnecting after backoff" + ); + tokio::time::sleep(delay).await; } } } @@ -1201,3 +1252,43 @@ mod ocsf_event_tests { accept_task.await.unwrap(); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reconnect_backoff_doubles_to_max_without_acceptance() { + let mut backoff = ReconnectBackoff::default(); + let delays: Vec = (0..7).map(|_| backoff.next_delay(false, 1.0)).collect(); + assert_eq!(delays, [1, 2, 4, 8, 16, 30, 30].map(Duration::from_secs)); + } + + #[test] + fn reconnect_backoff_resets_after_accepted_session() { + let mut backoff = ReconnectBackoff::default(); + for _ in 0..6 { + backoff.next_delay(false, 1.0); + } + assert_eq!(backoff.next_delay(true, 1.0), Duration::from_secs(1)); + assert_eq!(backoff.next_delay(false, 1.0), Duration::from_secs(2)); + } + + #[test] + fn equal_jitter_stays_between_half_and_full_ceiling() { + let ceiling = Duration::from_secs(2); + assert_eq!(equal_jitter(ceiling, 0.0), Duration::from_secs(1)); + assert_eq!(equal_jitter(ceiling, 0.5), Duration::from_millis(1500)); + assert_eq!(equal_jitter(ceiling, 1.0), Duration::from_secs(2)); + assert_eq!(equal_jitter(ceiling, f64::NAN), Duration::from_secs(2)); + assert_eq!(equal_jitter(ceiling, -3.0), Duration::from_secs(1)); + + for _ in 0..1000 { + let delay = ReconnectBackoff::default().next_delay(true, rand::random()); + assert!( + (Duration::from_millis(500)..=Duration::from_secs(1)).contains(&delay), + "delay {delay:?} outside [500ms, 1s]" + ); + } + } +} diff --git a/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index 4ab3517b27..a4ffd02629 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -81,6 +81,7 @@ See [`values.yaml`](values.yaml) for source defaults. Selected overlays: - [`ci/values-cert-manager.yaml`](ci/values-cert-manager.yaml) - cert-manager integration - [`ci/values-keycloak.yaml`](ci/values-keycloak.yaml) - Keycloak OIDC integration - [`ci/values-high-availability.yaml`](ci/values-high-availability.yaml) - CI overlay for multi-replica external PostgreSQL testing +- [`ci/values-autoscaling.yaml`](ci/values-autoscaling.yaml) - CI overlay for rendering the optional gateway HorizontalPodAutoscaler - [`ci/values-spire.yaml`](ci/values-spire.yaml) - SPIFFE/SPIRE provider token grants - [`ci/values-spire-stack.yaml`](ci/values-spire-stack.yaml) - SPIRE hardened chart values for local development @@ -167,6 +168,35 @@ DNS name while connecting directly to the owning pod. Custom TLS Secrets must include that Service DNS name in the server certificate and provide the CA and client credentials configured by `server.tls`. +Gateway pods get 30 seconds to stop (`podLifecycle.terminationGracePeriodSeconds`). +A gateway that uses an external PostgreSQL database (`server.externalDbSecret`) +uses that time to hand its supervisor sessions to the other replicas before it +exits. Keep the value at 30 or higher, and pass it explicitly when you upgrade +with `--reuse-values`, which keeps the previous release's value. + +Set `autoscaling.enabled=true` to render an `autoscaling/v2` +HorizontalPodAutoscaler for the gateway workload. The Deployment or +StatefulSet then omits `spec.replicas`, and the chart applies its +multi-replica checks to `autoscaling.maxReplicas`. CPU and memory targets +need a matching `resources.requests` entry, or a `resources.limits` entry, +which Kubernetes copies into the request. Add custom metrics from a metrics +adapter with `autoscaling.metrics`. See the +[High Availability guide](https://docs.nvidia.com/openshell/latest/kubernetes/high-availability) +for the metrics to scale and alert on. + +Enabling autoscaling on an existing release removes `spec.replicas` from the +workload in that upgrade. Kubernetes can reset the workload to one replica +until the HPA scales it back to at least `autoscaling.minReplicas`, which +drains supervisor sessions from the other gateway pods. Enable autoscaling +when you install the chart, or during a maintenance window. On an existing +release, upgrade with `--reset-then-reuse-values`. `--reuse-values` keeps the +previous chart version's defaults, which lack the autoscaling values, +including the scale-down `behavior`, when that release predates them. For background, refer to +[Migrating Deployments and StatefulSets to horizontal autoscaling](https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/#migrating-deployments-and-statefulsets-to-horizontal-autoscaling). +Disabling autoscaling renders `spec.replicas` from `replicaCount` again, +which defaults to 1, so set `replicaCount` to the replica count you want in +that upgrade. + ## Secret bootstrap By default, a pre-install/pre-upgrade hook Job runs `openshell-gateway generate-certs` @@ -203,6 +233,13 @@ discovery endpoint or its TLS CA. |-----|------|---------|-------------| | affinity | object | `{}` | Affinity rules for the gateway pod. | | agentSandbox.preflight.enabled | bool | `true` | Check the live cluster for a supported Agent Sandbox API before rendering gateway resources. Disable only for offline rendering and linting. | +| autoscaling.behavior | object | `{"scaleDown":{"policies":[{"periodSeconds":120,"type":"Pods","value":1}],"stabilizationWindowSeconds":300}}` | HPA scaling behavior. Scale-down drains supervisor sessions, so the default removes at most one replica every two minutes after a five-minute stabilization window. Helm merges maps: set autoscaling.behavior.scaleDown to null to drop the default. | +| autoscaling.enabled | bool | `false` | Render a HorizontalPodAutoscaler and stop rendering spec.replicas. | +| autoscaling.maxReplicas | int | `4` | Maximum gateway replicas. Each replica opens its own PostgreSQL connection pools; size the database for this count. | +| autoscaling.metrics | list | `[]` | Additional autoscaling/v2 MetricSpec entries appended verbatim, such as Pods metrics served by prometheus-adapter. | +| autoscaling.minReplicas | int | `2` | Minimum gateway replicas. Use 2 or more to survive a pod failure. | +| autoscaling.targetCPUUtilizationPercentage | int | `80` | Target average CPU utilization, as a percentage of resources.requests.cpu. Set to null to disable. Requires resources.requests.cpu, or resources.limits.cpu, which Kubernetes copies into the request. | +| autoscaling.targetMemoryUtilizationPercentage | int | `nil` | Target average memory utilization, as a percentage of resources.requests.memory. Null disables it. Requires resources.requests.memory, or resources.limits.memory, which Kubernetes copies into the request. | | certManager.caSecretName | string | `"openshell-ca-tls"` | Secret created for the intermediate CA (Certificate with isCA: true). | | certManager.certificateDuration | string | `"8760h"` | Duration for cert-manager-issued certificates. | | certManager.certificateRenewBefore | string | `"720h"` | Renewal window for cert-manager-issued certificates. | @@ -247,7 +284,7 @@ discovery endpoint or its TLS CA. | pkiInitJob.timeoutSeconds | int | `120` | Maximum time in seconds for the certgen hook to poll for cert-manager certificates. When using cert-manager with BackendTLSPolicy, the hook polls for this many seconds waiting for the certificate to be issued, then creates the backend CA ConfigMap. The Job deadline is set to (timeoutSeconds + 30) to allow time for ConfigMap creation and cleanup. Increase this if cert-manager takes longer than 120 seconds to issue certificates. | | podAnnotations | object | `{}` | Extra annotations to add to the gateway pod. | | podLabels | object | `{}` | Extra labels to add to the gateway pod. | -| podLifecycle.terminationGracePeriodSeconds | int | `5` | Grace period, in seconds, before Kubernetes terminates the gateway pod. | +| podLifecycle.terminationGracePeriodSeconds | int | `30` | Grace period, in seconds, before Kubernetes kills the gateway pod after SIGTERM. A gateway that uses an external PostgreSQL database (server.externalDbSecret) moves its supervisor sessions to the other replicas during this window. The drain is sized to finish within 30 seconds; lower values cut it short. Gateways on the default SQLite database exit as soon as shutdown cleanup finishes. | | podSecurityContext.fsGroup | int | `1000` | fsGroup assigned to the gateway pod. | | probes.liveness.failureThreshold | int | `3` | Liveness probe failure threshold before the container is restarted. | | probes.liveness.initialDelaySeconds | int | `2` | Liveness probe initial delay, in seconds. | diff --git a/deploy/helm/openshell/README.md.gotmpl b/deploy/helm/openshell/README.md.gotmpl index 6e49260c20..0965bc1157 100644 --- a/deploy/helm/openshell/README.md.gotmpl +++ b/deploy/helm/openshell/README.md.gotmpl @@ -81,6 +81,7 @@ See [`values.yaml`](values.yaml) for source defaults. Selected overlays: - [`ci/values-cert-manager.yaml`](ci/values-cert-manager.yaml) - cert-manager integration - [`ci/values-keycloak.yaml`](ci/values-keycloak.yaml) - Keycloak OIDC integration - [`ci/values-high-availability.yaml`](ci/values-high-availability.yaml) - CI overlay for multi-replica external PostgreSQL testing +- [`ci/values-autoscaling.yaml`](ci/values-autoscaling.yaml) - CI overlay for rendering the optional gateway HorizontalPodAutoscaler - [`ci/values-spire.yaml`](ci/values-spire.yaml) - SPIFFE/SPIRE provider token grants - [`ci/values-spire-stack.yaml`](ci/values-spire-stack.yaml) - SPIRE hardened chart values for local development @@ -167,6 +168,35 @@ DNS name while connecting directly to the owning pod. Custom TLS Secrets must include that Service DNS name in the server certificate and provide the CA and client credentials configured by `server.tls`. +Gateway pods get 30 seconds to stop (`podLifecycle.terminationGracePeriodSeconds`). +A gateway that uses an external PostgreSQL database (`server.externalDbSecret`) +uses that time to hand its supervisor sessions to the other replicas before it +exits. Keep the value at 30 or higher, and pass it explicitly when you upgrade +with `--reuse-values`, which keeps the previous release's value. + +Set `autoscaling.enabled=true` to render an `autoscaling/v2` +HorizontalPodAutoscaler for the gateway workload. The Deployment or +StatefulSet then omits `spec.replicas`, and the chart applies its +multi-replica checks to `autoscaling.maxReplicas`. CPU and memory targets +need a matching `resources.requests` entry, or a `resources.limits` entry, +which Kubernetes copies into the request. Add custom metrics from a metrics +adapter with `autoscaling.metrics`. See the +[High Availability guide](https://docs.nvidia.com/openshell/latest/kubernetes/high-availability) +for the metrics to scale and alert on. + +Enabling autoscaling on an existing release removes `spec.replicas` from the +workload in that upgrade. Kubernetes can reset the workload to one replica +until the HPA scales it back to at least `autoscaling.minReplicas`, which +drains supervisor sessions from the other gateway pods. Enable autoscaling +when you install the chart, or during a maintenance window. On an existing +release, upgrade with `--reset-then-reuse-values`. `--reuse-values` keeps the +previous chart version's defaults, which lack the autoscaling values, +including the scale-down `behavior`, when that release predates them. For background, refer to +[Migrating Deployments and StatefulSets to horizontal autoscaling](https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/#migrating-deployments-and-statefulsets-to-horizontal-autoscaling). +Disabling autoscaling renders `spec.replicas` from `replicaCount` again, +which defaults to 1, so set `replicaCount` to the replica count you want in +that upgrade. + ## Secret bootstrap By default, a pre-install/pre-upgrade hook Job runs `openshell-gateway generate-certs` diff --git a/deploy/helm/openshell/ci/values-autoscaling.yaml b/deploy/helm/openshell/ci/values-autoscaling.yaml new file mode 100644 index 0000000000..f563095503 --- /dev/null +++ b/deploy/helm/openshell/ci/values-autoscaling.yaml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# CI overlay for rendering the optional gateway HorizontalPodAutoscaler. Like +# values-high-availability.yaml, it expects a PostgreSQL Secret named +# openshell-ha-pg. It is exercised by helm lint and helm-unittest only; kind +# CI has no metrics-server, so HPA runtime behavior is not tested. +workload: + kind: deployment + +server: + externalDbSecret: openshell-ha-pg + +resources: + requests: + cpu: 250m + memory: 256Mi + +autoscaling: + enabled: true + minReplicas: 2 + maxReplicas: 4 + targetCPUUtilizationPercentage: 80 + # Example Pods metric served by a metrics adapter; the chart passes it through + # unchanged. + metrics: + - type: Pods + pods: + metric: + name: openshell_grpc_requests_per_second + target: + type: AverageValue + averageValue: "100" diff --git a/deploy/helm/openshell/ci/values-high-availability.yaml b/deploy/helm/openshell/ci/values-high-availability.yaml index 407326d675..a1d384174b 100644 --- a/deploy/helm/openshell/ci/values-high-availability.yaml +++ b/deploy/helm/openshell/ci/values-high-availability.yaml @@ -11,3 +11,8 @@ workload: server: externalDbSecret: openshell-ha-pg + +# Match the chart default so the e2e drain runs with the production grace +# period, even if the default changes. +podLifecycle: + terminationGracePeriodSeconds: 30 diff --git a/deploy/helm/openshell/templates/_autoscaling.tpl b/deploy/helm/openshell/templates/_autoscaling.tpl new file mode 100644 index 0000000000..b261946a98 --- /dev/null +++ b/deploy/helm/openshell/templates/_autoscaling.tpl @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +{{/* +Whether the chart renders a HorizontalPodAutoscaler for the gateway workload. +A missing autoscaling map (for example `helm upgrade --reuse-values` from a +release that predates these values) means disabled. +*/}} +{{- define "openshell.autoscalingEnabled" -}} +{{- $autoscaling := .Values.autoscaling | default dict -}} +{{- if and (kindIs "map" $autoscaling) (get $autoscaling "enabled") -}}true{{- end -}} +{{- end }} + +{{/* +Largest replica count the chart can run: autoscaling.maxReplicas when the +HPA is enabled, otherwise replicaCount. +*/}} +{{- define "openshell.maxReplicas" -}} +{{- if eq (include "openshell.autoscalingEnabled" .) "true" -}} +{{- int (get .Values.autoscaling "maxReplicas" | default 1) -}} +{{- else -}} +{{- int (default 1 .Values.replicaCount) -}} +{{- end -}} +{{- end }} + +{{/* +Name of the value that sets the largest replica count, for error messages. +*/}} +{{- define "openshell.maxReplicasSource" -}} +{{- ternary "autoscaling.maxReplicas" "replicaCount" (eq (include "openshell.autoscalingEnabled" .) "true") -}} +{{- end }} + +{{/* +Validate autoscaling values. Called from openshell.validateValues. +*/}} +{{- define "openshell.validateAutoscaling" -}} +{{- if eq (include "openshell.autoscalingEnabled" .) "true" -}} +{{- $a := .Values.autoscaling -}} +{{- if not (and (hasKey $a "minReplicas") (hasKey $a "maxReplicas")) -}} +{{- fail "autoscaling.minReplicas and autoscaling.maxReplicas are not set. helm upgrade --reuse-values keeps the values of a release that predates the chart's autoscaling defaults; upgrade with --reset-then-reuse-values, or set the autoscaling values explicitly, including autoscaling.behavior." -}} +{{- end -}} +{{- $min := int (get $a "minReplicas" | default 0) -}} +{{- $max := int (get $a "maxReplicas" | default 0) -}} +{{- if lt $min 1 -}} +{{- fail "autoscaling.minReplicas must be at least 1." -}} +{{- end -}} +{{- if lt $max $min -}} +{{- fail "autoscaling.maxReplicas must be greater than or equal to autoscaling.minReplicas." -}} +{{- end -}} +{{- $cpu := get $a "targetCPUUtilizationPercentage" -}} +{{- $memory := get $a "targetMemoryUtilizationPercentage" -}} +{{- if not (or $cpu $memory (get $a "metrics")) -}} +{{- fail "autoscaling.enabled requires targetCPUUtilizationPercentage, targetMemoryUtilizationPercentage, or autoscaling.metrics." -}} +{{- end -}} +{{- $requests := (.Values.resources | default dict).requests | default dict -}} +{{- $limits := (.Values.resources | default dict).limits | default dict -}} +{{- if and $cpu (not (or (get $requests "cpu") (get $limits "cpu"))) -}} +{{- fail "autoscaling.targetCPUUtilizationPercentage requires resources.requests.cpu (or resources.limits.cpu, which Kubernetes copies into the request); Kubernetes computes utilization against the container request." -}} +{{- end -}} +{{- if and $memory (not (or (get $requests "memory") (get $limits "memory"))) -}} +{{- fail "autoscaling.targetMemoryUtilizationPercentage requires resources.requests.memory (or resources.limits.memory, which Kubernetes copies into the request); Kubernetes computes utilization against the container request." -}} +{{- end -}} +{{- end -}} +{{- end }} diff --git a/deploy/helm/openshell/templates/_helpers.tpl b/deploy/helm/openshell/templates/_helpers.tpl index 32e6b7cb63..c808451f1e 100644 --- a/deploy/helm/openshell/templates/_helpers.tpl +++ b/deploy/helm/openshell/templates/_helpers.tpl @@ -51,6 +51,18 @@ app.kubernetes.io/name: {{ include "openshell.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} {{- end }} +{{/* +Pod labels for the certgen hook Jobs. They keep the release instance label but +do not match openshell.selectorLabels, so gateway selectors (the workload, +Services, HorizontalPodAutoscaler, anti-affinity, and PodDisruptionBudgets) +never select hook pods. +*/}} +{{- define "openshell.certgenPodLabels" -}} +app.kubernetes.io/name: {{ printf "%s-certgen" (include "openshell.name" . | trunc 55 | trimSuffix "-") }} +app.kubernetes.io/instance: {{ .Release.Name }} +app.kubernetes.io/component: certgen +{{- end }} + {{/* Create the name of the service account to use */}} @@ -325,7 +337,8 @@ Validate chart values that Helm would otherwise accept silently. {{- define "openshell.validateValues" -}} {{- $workloadKind := include "openshell.workloadKind" . -}} {{- $workload := .Values.workload | default dict -}} -{{- $replicaCount := int (default 1 .Values.replicaCount) -}} +{{- $maxReplicas := int (include "openshell.maxReplicas" .) -}} +{{- $maxReplicasSource := include "openshell.maxReplicasSource" . -}} {{- if and (hasKey .Values "postgres") (kindIs "map" .Values.postgres) (hasKey .Values.postgres "enabled") -}} {{- fail "postgres.enabled was removed; the OpenShell chart no longer deploys PostgreSQL. Provision PostgreSQL separately and set server.externalDbSecret to a Secret containing a PostgreSQL URI." -}} {{- end -}} @@ -335,11 +348,12 @@ Validate chart values that Helm would otherwise accept silently. {{- if and (eq $workloadKind "deployment") (not .Values.server.externalDbSecret) -}} {{- fail "workload.kind=deployment requires server.externalDbSecret; use workload.kind=statefulset for the default SQLite database." -}} {{- end -}} -{{- if and (gt $replicaCount 1) (not .Values.server.externalDbSecret) -}} -{{- fail "replicaCount > 1 requires server.externalDbSecret; multiple gateway replicas cannot share the default per-pod SQLite database." -}} +{{- include "openshell.validateAutoscaling" . -}} +{{- if and (gt $maxReplicas 1) (not .Values.server.externalDbSecret) -}} +{{- fail (printf "%s > 1 requires server.externalDbSecret; multiple gateway replicas cannot share the default per-pod SQLite database." $maxReplicasSource) -}} {{- end -}} -{{- if and (eq $workloadKind "statefulset") (gt $replicaCount 1) (not (get $workload "allowMultiReplicaStatefulSet" | default false)) -}} -{{- fail "replicaCount > 1 with workload.kind=statefulset requires workload.allowMultiReplicaStatefulSet=true; use workload.kind=deployment for external database-backed multi-replica gateways." -}} +{{- if and (eq $workloadKind "statefulset") (gt $maxReplicas 1) (not (get $workload "allowMultiReplicaStatefulSet" | default false)) -}} +{{- fail (printf "%s > 1 with workload.kind=statefulset requires workload.allowMultiReplicaStatefulSet=true; use workload.kind=deployment for external database-backed multi-replica gateways." $maxReplicasSource) -}} {{- end -}} {{- $workspaceMode := .Values.server.drivers.kubernetes.workspaceMode | default "shared" -}} {{- if not (has $workspaceMode (list "shared" "managed" "operator")) -}} diff --git a/deploy/helm/openshell/templates/certgen.yaml b/deploy/helm/openshell/templates/certgen.yaml index f7c9a751d3..d5610faf2f 100644 --- a/deploy/helm/openshell/templates/certgen.yaml +++ b/deploy/helm/openshell/templates/certgen.yaml @@ -75,7 +75,7 @@ spec: template: metadata: labels: - {{- include "openshell.selectorLabels" . | nindent 8 }} + {{- include "openshell.certgenPodLabels" . | nindent 8 }} spec: restartPolicy: OnFailure serviceAccountName: {{ $hookName }} @@ -146,7 +146,7 @@ spec: template: metadata: labels: - {{- include "openshell.selectorLabels" . | nindent 8 }} + {{- include "openshell.certgenPodLabels" . | nindent 8 }} spec: restartPolicy: OnFailure serviceAccountName: {{ $hookName }} diff --git a/deploy/helm/openshell/templates/deployment.yaml b/deploy/helm/openshell/templates/deployment.yaml index f94900b136..1c9f589bcf 100644 --- a/deploy/helm/openshell/templates/deployment.yaml +++ b/deploy/helm/openshell/templates/deployment.yaml @@ -9,7 +9,9 @@ metadata: labels: {{- include "openshell.labels" . | nindent 4 }} spec: + {{- if ne (include "openshell.autoscalingEnabled" .) "true" }} replicas: {{ .Values.replicaCount }} + {{- end }} selector: matchLabels: {{- include "openshell.selectorLabels" . | nindent 6 }} diff --git a/deploy/helm/openshell/templates/hpa.yaml b/deploy/helm/openshell/templates/hpa.yaml new file mode 100644 index 0000000000..d88c46290b --- /dev/null +++ b/deploy/helm/openshell/templates/hpa.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +{{- include "openshell.validateValues" . }} +{{- if eq (include "openshell.autoscalingEnabled" .) "true" }} +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: {{ include "openshell.fullname" . }} + labels: + {{- include "openshell.labels" . | nindent 4 }} +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: {{ ternary "Deployment" "StatefulSet" (eq (include "openshell.workloadKind" .) "deployment") }} + name: {{ include "openshell.fullname" . }} + minReplicas: {{ .Values.autoscaling.minReplicas }} + maxReplicas: {{ .Values.autoscaling.maxReplicas }} + metrics: + {{- with .Values.autoscaling.targetCPUUtilizationPercentage }} + - type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: {{ . }} + {{- end }} + {{- with .Values.autoscaling.targetMemoryUtilizationPercentage }} + - type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: {{ . }} + {{- end }} + {{- with .Values.autoscaling.metrics }} + {{- toYaml . | nindent 4 }} + {{- end }} + {{- with .Values.autoscaling.behavior }} + behavior: + {{- toYaml . | nindent 4 }} + {{- end }} +{{- end }} diff --git a/deploy/helm/openshell/templates/statefulset.yaml b/deploy/helm/openshell/templates/statefulset.yaml index 10d0839f60..be65064c19 100644 --- a/deploy/helm/openshell/templates/statefulset.yaml +++ b/deploy/helm/openshell/templates/statefulset.yaml @@ -10,7 +10,9 @@ metadata: {{- include "openshell.labels" . | nindent 4 }} spec: serviceName: {{ include "openshell.peerServiceName" . }} + {{- if ne (include "openshell.autoscalingEnabled" .) "true" }} replicas: {{ .Values.replicaCount }} + {{- end }} selector: matchLabels: {{- include "openshell.selectorLabels" . | nindent 6 }} diff --git a/deploy/helm/openshell/tests/autoscaling_test.yaml b/deploy/helm/openshell/tests/autoscaling_test.yaml new file mode 100644 index 0000000000..b2309653c3 --- /dev/null +++ b/deploy/helm/openshell/tests/autoscaling_test.yaml @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +suite: gateway autoscaling +templates: + - templates/hpa.yaml + - templates/deployment.yaml + - templates/statefulset.yaml + - templates/gateway-config.yaml +release: + name: openshell + namespace: my-namespace + +tests: + - it: renders no HorizontalPodAutoscaler by default + template: templates/hpa.yaml + asserts: + - hasDocuments: + count: 0 + + - it: keeps spec.replicas on the default StatefulSet + template: templates/statefulset.yaml + asserts: + - equal: + path: spec.replicas + value: 1 + + - it: renders an autoscaling/v2 HPA targeting the Deployment from the CI overlay + template: templates/hpa.yaml + values: + - ../ci/values-autoscaling.yaml + asserts: + - isKind: + of: HorizontalPodAutoscaler + - isAPIVersion: + of: autoscaling/v2 + - equal: + path: metadata.name + value: openshell + - equal: + path: spec.scaleTargetRef + value: + apiVersion: apps/v1 + kind: Deployment + name: openshell + - equal: + path: spec.minReplicas + value: 2 + - equal: + path: spec.maxReplicas + value: 4 + - contains: + path: spec.metrics + content: + type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 80 + - contains: + path: spec.metrics + content: + type: Pods + pods: + metric: + name: openshell_grpc_requests_per_second + target: + type: AverageValue + averageValue: "100" + - equal: + path: spec.behavior.scaleDown.stabilizationWindowSeconds + value: 300 + - equal: + path: spec.behavior.scaleDown.policies[0] + value: + type: Pods + value: 1 + periodSeconds: 120 + + - it: omits spec.replicas from the Deployment when autoscaling is enabled + template: templates/deployment.yaml + values: + - ../ci/values-autoscaling.yaml + set: + replicaCount: 3 + asserts: + - isKind: + of: Deployment + - notExists: + path: spec.replicas + + - it: renders memory utilization only when configured + template: templates/hpa.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.targetMemoryUtilizationPercentage: 70 + asserts: + - contains: + path: spec.metrics + content: + type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: 70 + + - it: renders only custom metrics when utilization targets are disabled + template: templates/hpa.yaml + set: + workload.kind: deployment + server.externalDbSecret: my-pg-secret + autoscaling.enabled: true + autoscaling.targetCPUUtilizationPercentage: null + autoscaling.metrics: + - type: Pods + pods: + metric: + name: openshell_grpc_requests_per_second + target: + type: AverageValue + averageValue: "50" + asserts: + - lengthEqual: + path: spec.metrics + count: 1 + - equal: + path: spec.metrics[0].type + value: Pods + + - it: merges a scale-up policy with the default scale-down behavior + template: templates/hpa.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.behavior: + scaleUp: + stabilizationWindowSeconds: 0 + asserts: + - equal: + path: spec.behavior.scaleUp.stabilizationWindowSeconds + value: 0 + - equal: + path: spec.behavior.scaleDown.stabilizationWindowSeconds + value: 300 + + - it: omits behavior when set to null + template: templates/hpa.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.behavior: null + asserts: + - notExists: + path: spec.behavior + + - it: targets a multi-replica StatefulSet with the explicit override + template: templates/hpa.yaml + set: + server.externalDbSecret: my-pg-secret + workload.allowMultiReplicaStatefulSet: true + resources.requests.cpu: 100m + autoscaling.enabled: true + asserts: + - equal: + path: spec.scaleTargetRef.kind + value: StatefulSet + + - it: omits spec.replicas from the StatefulSet when autoscaling is enabled + template: templates/statefulset.yaml + set: + server.externalDbSecret: my-pg-secret + workload.allowMultiReplicaStatefulSet: true + resources.requests.cpu: 100m + autoscaling.enabled: true + asserts: + - notExists: + path: spec.replicas + + - it: treats a missing autoscaling map as disabled + template: templates/statefulset.yaml + set: + autoscaling: null + asserts: + - equal: + path: spec.replicas + value: 1 + + - it: fails when autoscaling can exceed one replica on SQLite + template: templates/statefulset.yaml + set: + autoscaling.enabled: true + resources.requests.cpu: 100m + asserts: + - failedTemplate: + errorPattern: "autoscaling.maxReplicas > 1 requires server.externalDbSecret" + + - it: fails when autoscaling a StatefulSet without the multi-replica override + template: templates/statefulset.yaml + set: + autoscaling.enabled: true + server.externalDbSecret: my-pg-secret + resources.requests.cpu: 100m + asserts: + - failedTemplate: + errorPattern: "autoscaling.maxReplicas > 1 with workload.kind=statefulset requires workload.allowMultiReplicaStatefulSet=true" + + - it: fails when maxReplicas is below minReplicas + template: templates/statefulset.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.maxReplicas: 1 + asserts: + - failedTemplate: + errorPattern: "autoscaling.maxReplicas must be greater than or equal to autoscaling.minReplicas" + + - it: fails when minReplicas is below one + template: templates/statefulset.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.minReplicas: 0 + asserts: + - failedTemplate: + errorPattern: "autoscaling.minReplicas must be at least 1" + + - it: fails when no autoscaling metric is configured + template: templates/statefulset.yaml + set: + workload.kind: deployment + server.externalDbSecret: my-pg-secret + autoscaling.enabled: true + autoscaling.targetCPUUtilizationPercentage: null + asserts: + - failedTemplate: + errorPattern: "autoscaling.enabled requires targetCPUUtilizationPercentage, targetMemoryUtilizationPercentage, or autoscaling.metrics" + + - it: fails when a CPU target has no CPU request + template: templates/statefulset.yaml + set: + workload.kind: deployment + server.externalDbSecret: my-pg-secret + autoscaling.enabled: true + asserts: + - failedTemplate: + errorPattern: "autoscaling.targetCPUUtilizationPercentage requires resources.requests.cpu" + + - it: fails when a memory target has no memory request + template: templates/statefulset.yaml + values: + - ../ci/values-autoscaling.yaml + set: + autoscaling.targetMemoryUtilizationPercentage: 70 + resources.requests.memory: null + asserts: + - failedTemplate: + errorPattern: "autoscaling.targetMemoryUtilizationPercentage requires resources.requests.memory" + + - it: accepts CPU and memory limits in place of requests + template: templates/hpa.yaml + set: + workload.kind: deployment + server.externalDbSecret: my-pg-secret + autoscaling.enabled: true + autoscaling.targetMemoryUtilizationPercentage: 70 + resources.limits.cpu: "1" + resources.limits.memory: 1Gi + asserts: + - isKind: + of: HorizontalPodAutoscaler + - contains: + path: spec.metrics + content: + type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 80 + - contains: + path: spec.metrics + content: + type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: 70 + + - it: explains --reuse-values when the autoscaling defaults are missing + template: templates/statefulset.yaml + set: + workload.kind: deployment + server.externalDbSecret: my-pg-secret + resources.requests.cpu: 100m + # --reuse-values from a release without autoscaling values leaves only + # the enabled flag that the upgrade sets. + autoscaling.enabled: true + autoscaling.minReplicas: null + autoscaling.maxReplicas: null + autoscaling.targetCPUUtilizationPercentage: null + autoscaling.metrics: null + autoscaling.behavior: null + asserts: + - failedTemplate: + errorPattern: "autoscaling.minReplicas and autoscaling.maxReplicas are not set.*--reset-then-reuse-values" + + - it: keeps the replicaCount wording when autoscaling is disabled + template: templates/statefulset.yaml + set: + replicaCount: 2 + asserts: + - failedTemplate: + errorPattern: "replicaCount > 1 requires server.externalDbSecret" diff --git a/deploy/helm/openshell/tests/certgen_test.yaml b/deploy/helm/openshell/tests/certgen_test.yaml index cd88b60e97..7a478bbf76 100644 --- a/deploy/helm/openshell/tests/certgen_test.yaml +++ b/deploy/helm/openshell/tests/certgen_test.yaml @@ -114,3 +114,38 @@ tests: path: spec.template.spec.containers[0].args content: "--server-san=192.0.2.10" documentIndex: 3 + + - it: labels hook pods so that gateway selectors do not match them + template: templates/certgen.yaml + asserts: + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: openshell-certgen + app.kubernetes.io/instance: openshell + app.kubernetes.io/component: certgen + documentIndex: 3 + - notEqual: + path: spec.template.metadata.labels["app.kubernetes.io/name"] + value: openshell + documentIndex: 3 + + - it: labels the backend CA hook pod so that gateway selectors do not match it + template: templates/certgen.yaml + set: + certManager.enabled: true + grpcRoute.backendTLSPolicy.enabled: true + asserts: + - hasDocuments: + count: 5 + - equal: + path: metadata.name + value: openshell-certgen-backend-ca + documentIndex: 4 + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: openshell-certgen + app.kubernetes.io/instance: openshell + app.kubernetes.io/component: certgen + documentIndex: 4 diff --git a/deploy/helm/openshell/tests/gateway_config_test.yaml b/deploy/helm/openshell/tests/gateway_config_test.yaml index 83fbb7aa08..3d4d5428da 100644 --- a/deploy/helm/openshell/tests/gateway_config_test.yaml +++ b/deploy/helm/openshell/tests/gateway_config_test.yaml @@ -210,6 +210,9 @@ tests: - equal: path: kind value: StatefulSet + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 30 - it: treats a null workload map as the default StatefulSet template: templates/statefulset.yaml @@ -1063,6 +1066,9 @@ tests: - equal: path: spec.replicas value: 2 + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 30 - contains: path: spec.template.spec.containers[0].env content: diff --git a/deploy/helm/openshell/values.yaml b/deploy/helm/openshell/values.yaml index b48a7c269f..b01312ce86 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -192,8 +192,55 @@ agentSandbox: # Pod restart behavior and health probe tuning. podLifecycle: - # -- Grace period, in seconds, before Kubernetes terminates the gateway pod. - terminationGracePeriodSeconds: 5 + # -- Grace period, in seconds, before Kubernetes kills the gateway pod after + # SIGTERM. A gateway that uses an external PostgreSQL database + # (server.externalDbSecret) moves its supervisor sessions to the other + # replicas during this window. The drain is sized to finish within 30 + # seconds; lower values cut it short. Gateways on the default SQLite database + # exit as soon as shutdown cleanup finishes. + terminationGracePeriodSeconds: 30 + +# Optional HorizontalPodAutoscaler (autoscaling/v2) for the gateway workload. +# When enabled, the Deployment or StatefulSet omits spec.replicas and the HPA +# owns the replica count; replicaCount is ignored. Multi-replica limits +# (server.externalDbSecret, and workload.allowMultiReplicaStatefulSet for a +# StatefulSet) apply to maxReplicas. Scaling out does not move established +# supervisor sessions; read the High Availability guide before choosing metrics. +# The workload can drop to one replica when you enable this on an existing +# release, until the HPA scales it back up; read the chart README before you +# upgrade. +autoscaling: + # -- Render a HorizontalPodAutoscaler and stop rendering spec.replicas. + enabled: false + # -- Minimum gateway replicas. Use 2 or more to survive a pod failure. + minReplicas: 2 + # -- Maximum gateway replicas. Each replica opens its own PostgreSQL + # connection pools; size the database for this count. + maxReplicas: 4 + # -- Target average CPU utilization, as a percentage of + # resources.requests.cpu. Set to null to disable. Requires + # resources.requests.cpu, or resources.limits.cpu, which Kubernetes copies + # into the request. + targetCPUUtilizationPercentage: 80 + # -- (int) Target average memory utilization, as a percentage of + # resources.requests.memory. Null disables it. Requires + # resources.requests.memory, or resources.limits.memory, which Kubernetes + # copies into the request. + targetMemoryUtilizationPercentage: null + # -- Additional autoscaling/v2 MetricSpec entries appended verbatim, such as + # Pods metrics served by prometheus-adapter. + metrics: [] + # -- HPA scaling behavior. Scale-down drains supervisor sessions, so the + # default removes at most one replica every two minutes after a + # five-minute stabilization window. Helm merges maps: set + # autoscaling.behavior.scaleDown to null to drop the default. + behavior: + scaleDown: + stabilizationWindowSeconds: 300 + policies: + - type: Pods + value: 1 + periodSeconds: 120 probes: startup: diff --git a/docs/kubernetes/high-availability.mdx b/docs/kubernetes/high-availability.mdx index f83b181462..c3d8be629f 100644 --- a/docs/kubernetes/high-availability.mdx +++ b/docs/kubernetes/high-availability.mdx @@ -3,8 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 title: "High Availability" sidebar-title: "High Availability" -description: "Run multiple OpenShell gateway replicas on Kubernetes with shared PostgreSQL and authenticated peer routing." -keywords: "Generative AI, Cybersecurity, Kubernetes, High Availability, HA, Gateway, PostgreSQL, Replicas, Failover" +description: "Run multiple OpenShell gateway replicas on Kubernetes with shared PostgreSQL, authenticated peer routing, session draining during rollouts, capacity metrics, and optional autoscaling." +keywords: "Generative AI, Cybersecurity, Kubernetes, High Availability, HA, Gateway, PostgreSQL, Replicas, Failover, Autoscaling, HPA, Metrics, Draining" position: 3 --- @@ -32,6 +32,9 @@ An HA gateway deployment requires: database are intended for a single gateway replica. - An ingress or load balancer that routes clients to the gateway Service. Refer to [Ingress](/kubernetes/ingress) for a Gateway API configuration. +- PostgreSQL connection capacity for every replica, including pods that are + starting or draining during a rollout. Refer to + [Size PostgreSQL Connections](#size-postgresql-connections). The Helm chart rejects `replicaCount` values above `1` unless `server.externalDbSecret` is set. It also rejects a multi-replica StatefulSet @@ -83,6 +86,11 @@ required rule leaves a replica Pending when the cluster does not have enough matching failure domains; use preferred anti-affinity when scheduling the replica is more important than strict separation. +The chart gives each gateway pod 30 seconds to stop +(`podLifecycle.terminationGracePeriodSeconds`) so it can move its supervisor +sessions to the other replicas. Keep that value when you customize pod +settings. + Install or upgrade OpenShell: ```shell @@ -169,24 +177,205 @@ connections also use the chart's TLS materials. You do not need session affinity at the ingress or client-facing Service. -## Failure and Rollout Behavior +## Drain and Rollout Behavior + +A gateway replica that Kubernetes stops on purpose, for a rolling update, a +scale-down, or a node drain, moves its supervisor sessions to the other +replicas before it exits. + +When a gateway pod receives `SIGTERM`, the gateway takes these steps: + +1. It stops accepting new supervisor sessions and reports not ready. `/readyz` + and `/health` return `503` with status `draining`. `/healthz` keeps + returning `200`, so the liveness probe does not restart the pod during the + drain. +2. It sets `openshell_server_draining` to `1` and keeps serving for 3 seconds + while Service endpoints, kube-proxy, and ingress controllers stop routing + to it. Other replicas can still relay requests to the sessions it holds. +3. It closes its supervisor sessions one at a time, at most 100 milliseconds + apart and within 12 seconds. Each supervisor reconnects through the gateway + Service to a ready replica, usually within a second or two. A session keeps + serving requests until its turn. +4. It stops its listener, waits up to 10 seconds for the remaining ownership + records to be released, and exits. + +A gateway that uses an external PostgreSQL database (`server.externalDbSecret`) +runs all four steps, as a Deployment or as a StatefulSet. A gateway on the +default SQLite database skips steps 2 and 3 and exits as soon as step 4 +finishes. A lone replica on PostgreSQL, such as a single-replica StatefulSet, +has no other replica to take its sessions, so the drain adds up to 15 seconds +to its shutdown. A StatefulSet starts the replacement pod only after the old +one exits, so its sandboxes stay unreachable up to 15 seconds longer on every +restart or upgrade. Supervisors retry with growing backoff while no replica is +ready, so some can take up to 30 more seconds, their maximum retry delay, to +reconnect after the new pod is ready. + +### Size the Termination Grace Period + +The drain schedule and supervisor-session cleanup budgets total 25 seconds, +so the chart default of 30 seconds leaves a margin for the built-in +Kubernetes driver. This is not an overall process-exit deadline: compute-driver +cleanup runs separately. If you export gateway traces to an OTLP collector +that can respond slowly, the final trace export after shutdown can take up to +10 more seconds. Allow additional grace for driver cleanup and trace export +when applicable. + +If Kubernetes kills the gateway before the drain finishes, the remaining +supervisors reconnect all at once when the pod exits instead of one at a time, +and their sandboxes can take longer to become reachable again. + + +`helm upgrade --reuse-values` keeps the previous release's grace period, which +was 5 seconds in earlier charts. Upgrade with `--reset-then-reuse-values`, or +set `podLifecycle.terminationGracePeriodSeconds=30` explicitly. + + +### Sessions per Replica + +Plan for no more than about 500 supervisor sessions (connected sandboxes) per +gateway replica. Each reconnect takes one or two short database locks on the +replica that receives it. A replica that drains more sessions than that hands +them over faster than the remaining replicas can register them without +queuing, so some sandboxes take longer to reconnect and requests to them can +fail during the rollout. + +Add replicas as the number of sandboxes grows, before any replica reaches that +count. New replicas receive only new or reconnecting sessions, so adding them +just before a rollout does not lighten the replicas being drained. During +rollouts, watch `openshell_server_mutation_lock_wait_seconds` and +`openshell_server_mutation_lock_timeouts_total` with `scope="sandbox"` on the +receiving replicas. Rising waits or any timeout mean that the replicas receive +reconnects faster than they can absorb them. + +### What Clients See During a Drain + +Most requests continue to work while sessions move. Expect these effects: + +- New exec, SSH, and port-forward requests for a sandbox that is moving + between replicas fail with `FAILED_PRECONDITION` and the message + `sandbox is not ready` while the sandbox reports `Provisioning`, usually for + a second or two. Sandbox service URLs return `412 Precondition Failed` for + the same window. Retry these requests. Requests that were already routed to + the sandbox wait up to 15 seconds for its supervisor to reconnect, then fail + with `UNAVAILABLE`. +- Exec, port-forwarding, and service streams that pass through the stopping + pod end when it exits and must be retried. File sync retries with a new SSH + session instead of resuming the interrupted transfer. +- Each moved sandbox reports `Provisioning` between leaving one replica and + joining another, then returns to `Ready`. Watchers and + `openshell sandbox list` can show that transition. If a supervisor stays + disconnected for more than a few seconds, the sandbox's `Ready` condition can + also show `DependenciesNotReady` until the session returns. The workload + keeps running. + + +Supervisors from earlier OpenShell releases do not reset their reconnect delay +after a successful connection, so after several disconnects they wait up to 30 +seconds before reconnecting. A sandbox keeps the supervisor it was created +with, so requests to older sandboxes can fail during a rollout until they +reconnect. Recreate long-lived sandboxes after upgrading the gateway. + + +### Session Balance After a Rollout + +The drain moves sessions off a stopping pod but does not balance them. +Supervisors reconnect to whichever replicas are ready at that moment: + +- A Deployment rolling update with the default strategy leaves the first + replacement pod with the most sessions. With two replicas it ends up with + about five of every eight sessions on average. +- Deleting or evicting a pod, or rolling a StatefulSet, moves its sessions to + the replicas that stay up. The replacement pod starts nearly empty. +- Scaling out adds replicas that receive only new or reconnecting sessions. + +Client requests keep working because any replica relays to the owner. +OpenShell does not yet rebalance established sessions, and the imbalance fades +as sandboxes are recreated. Watch `openshell_server_supervisor_sessions` per +pod to see the current distribution. + +### Unplanned Failures + +A gateway pod that is killed or loses its node cannot drain. Its supervisors +detect the broken connection, within about 30 seconds when the node +disappears, then reconnect through the gateway Service and replace the stale +ownership records. Other replicas re-read ownership while they wait to relay a +request, so a request succeeds if the supervisor reconnects before the +request's deadline and fails otherwise. Sandboxes owned by other replicas are +unaffected. Requests that pass through a stopped non-owner pod disconnect and +must be retried. + +### Upgrade from an Earlier Release + +While gateways from the earlier release and the new release run together +during the upgrade, the older replicas serialize every mutation across the +fleet. Mutations can be slower and can briefly fail under load until the +rollout finishes. + +## Monitor Capacity + +Each gateway replica exposes Prometheus metrics on port `9090` +(`service.metricsPort`). Scrape every gateway pod, not the Service, because +these signals describe one replica. Refer to +[Gateway Metrics](/observability/gateway-metrics) for the full catalog, scrape +configuration, and access control. + +| Signal | Metrics | Use | +|---|---|---| +| Sessions per replica | `openshell_server_supervisor_sessions` | Placement, skew, and the per-replica session plan. Do not scale on it. | +| Draining | `openshell_server_draining` | Exclude stopping pods from per-replica views. | +| Relay utilization | `openshell_server_relay_pending`, `openshell_server_relay_pending_capacity` | Graph it. It is transient and usually near 0, so alert on rejections instead. | +| Relay rejections | `openshell_server_relay_rejected_total` | Alert on any increase. | +| Peer requests | `openshell_server_peer_requests_total`, `openshell_server_peer_request_duration_seconds` | Cross-replica routing rate, failures, and latency. Spikes of `code="unavailable"` during rollouts are expected. | +| Lock contention | `openshell_server_mutation_lock_wait_seconds`, `openshell_server_mutation_lock_timeouts_total` | Alert on any timeout. | +| Watch cost | `openshell_server_sandbox_watch_polled_sandboxes`, `openshell_server_sandbox_watch_poll_duration_seconds`, `openshell_server_sandbox_watch_poll_errors_total` | Cross-replica `WatchSandbox` load and failed lookups. | + +These queries assume that Prometheus labels each series with `namespace` and +`pod`: + +```promql +# Relay capacity used on each replica. Usually near 0; graph it and alert +# on rejections instead. +max by (namespace, pod) (openshell_server_relay_pending / openshell_server_relay_pending_capacity) + +# Relays rejected for capacity. Alert on any increase. +sum by (namespace, pod, reason) (increase(openshell_server_relay_rejected_total[5m])) + +# Session skew across replicas that are not draining. 1 means balanced. +max(openshell_server_supervisor_sessions and on (namespace, pod) (openshell_server_draining == 0)) + / clamp_min(avg(openshell_server_supervisor_sessions and on (namespace, pod) (openshell_server_draining == 0)), 1) + +# Failed peer requests as a share of attempts, by RPC. +sum by (rpc) (rate(openshell_server_peer_requests_total{outcome!="ok"}[5m])) + / sum by (rpc) (rate(openshell_server_peer_requests_total[5m])) + +# 99th percentile peer request latency across the fleet. +histogram_quantile(0.99, sum by (le, rpc) (rate(openshell_server_peer_request_duration_seconds_bucket{outcome="ok"}[5m]))) + +# 99th percentile mutation lock wait by scope. Watch scope="sandbox" during rollouts. +histogram_quantile(0.99, sum by (le, scope) (rate(openshell_server_mutation_lock_wait_seconds_bucket[5m]))) + +# Mutation lock timeouts by scope. Alert on any increase. +sum by (scope) (increase(openshell_server_mutation_lock_timeouts_total[10m])) + +# 99th percentile watch-poll time per replica. The poll interval is one second. +histogram_quantile(0.99, sum by (le, pod) (rate(openshell_server_sandbox_watch_poll_duration_seconds_bucket[5m]))) +``` -When a non-owner gateway pod stops, Kubernetes removes it from ready Service -endpoints and new requests go to another replica. Existing requests through the -stopped pod disconnect and must be retried. +Relay capacity is used on the replica that owns the sandbox's supervisor +session, including relays that other replicas request through peers, so +adding replicas does not relieve a busy owner. A relay whose client gave up +stays counted for up to about 40 seconds, the 10-second claim timeout plus the +30-second cleanup interval. Peer request counts include retries during +rollouts. -When an owner gateway pod stops, its connected supervisors reconnect through -the gateway Service and publish new ownership. Other replicas re-read ownership -while waiting to relay a request. A request fails if no fresh, reachable owner -appears before its normal deadline. +## Scale the Gateway -Established streams do not move between gateway replicas. An interrupted exec, -forwarding, or service stream must reconnect. File sync retries with a new SSH -session instead of resuming the interrupted byte stream. Rolling updates can -temporarily concentrate supervisor sessions on the replicas that stayed up; -client requests remain routable through peer relay. +Scale the gateway by changing a fixed replica count or by letting a +HorizontalPodAutoscaler choose one. In both cases, keep at least two ready +replicas when availability must survive one gateway pod failure, and size +PostgreSQL and the cluster nodes for the largest replica count. -## Scale the Gateway +### Scale Manually Change `replicaCount` in `values-ha.yaml`, then apply the release again: @@ -199,9 +388,153 @@ helm upgrade openshell \ --wait ``` -Keep at least two ready replicas when availability must survive one gateway pod -failure. Size PostgreSQL connection capacity and the cluster nodes for the -selected replica count. +### Autoscale with a HorizontalPodAutoscaler + +Set `autoscaling.enabled` to render an `autoscaling/v2` +HorizontalPodAutoscaler for the gateway Deployment, or for a StatefulSet with +`workload.allowMultiReplicaStatefulSet`. Add these values to +`values-ha.yaml`: + +```yaml +resources: + requests: + cpu: 500m + memory: 512Mi + +autoscaling: + enabled: true + minReplicas: 2 + maxReplicas: 6 + targetCPUUtilizationPercentage: 70 +``` + +With autoscaling enabled, the chart stops setting `spec.replicas` and ignores +`replicaCount`. Kubernetes computes CPU and memory utilization against the +container request, so those targets need a matching `resources.requests` +entry. A `resources.limits` entry also works, because Kubernetes copies a +limit into an unset request. The chart rejects an `autoscaling.maxReplicas` +above `1` without `server.externalDbSecret`. Each scale-down drains one +replica, so the default behavior removes at most one pod every two minutes +after a five-minute stabilization window. Override it with +`autoscaling.behavior`, or set `autoscaling.behavior.scaleDown` to `null` to +drop the default. New replicas do not take existing sessions. + +Check the autoscaler after you apply the release: + +```shell +kubectl -n openshell get hpa openshell +kubectl -n openshell describe hpa openshell +``` + + +Enabling autoscaling on an existing release removes `spec.replicas` from the +workload in that upgrade. Kubernetes can reset the workload to one replica +until the HPA scales it back to at least `autoscaling.minReplicas`, which +drains supervisor sessions from the other gateway pods. Enable autoscaling +when you install the chart, or during a maintenance window. On an existing +release, upgrade with `--reset-then-reuse-values`. `--reuse-values` keeps the +previous chart version's defaults, which lack the autoscaling values, +including the scale-down `behavior`, when that release predates them. For background, refer to +[Migrating Deployments and StatefulSets to horizontal autoscaling](https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/#migrating-deployments-and-statefulsets-to-horizontal-autoscaling). +Disabling autoscaling renders `spec.replicas` from `replicaCount` again, +which defaults to 1, so set `replicaCount` to the replica count you want in +that upgrade. + + +### Autoscale on Custom Metrics + +The HPA can also scale on a per-pod metric that a metrics adapter publishes. +The chart passes `autoscaling.metrics` to the HPA unchanged and does not +install Prometheus or an adapter. A custom metric needs these pieces: + +- Prometheus scrapes every gateway pod and labels each series with + `namespace` and `pod`. +- A metrics adapter, such as prometheus-adapter, maps those labels to + Kubernetes resources and publishes a per-pod metric through the + `custom.metrics.k8s.io` API. +- `autoscaling.metrics` references that metric as a `Pods` metric with an + `AverageValue` target. + +Scale on signals that fall when you add replicas, such as CPU utilization. Do +not scale on supervisor sessions, relay utilization, or peer request rate. New +replicas do not take existing sessions, relay load stays on the replica that +owns each session, and peer traffic grows with the replica count. Alert on +those signals instead. + +```yaml +autoscaling: + enabled: true + minReplicas: 2 + maxReplicas: 6 + targetCPUUtilizationPercentage: 70 + metrics: + - type: Pods + pods: + metric: + name: + target: + type: AverageValue + averageValue: "100" +``` + +Replace `` with the name your adapter publishes, and set +`averageValue` to the per-pod value at which you want another replica. Check +that the adapter serves the metric for the gateway pods: + +```shell +kubectl get --raw "/apis/custom.metrics.k8s.io/v1beta1/namespaces/openshell/pods/*/" +``` + +### Size PostgreSQL Connections + +Each gateway pod opens up to 14 PostgreSQL connections on demand, 10 for data +access and 4 for mutation locks. The chart uses the Kubernetes default rolling +update strategy, so a rollout adds up to 25 percent of the replica count, +rounded up, as surge pods. Kubernetes does not count terminating pods against +that surge, and each replaced pod can keep its connections throughout its +termination grace period. A rollout can therefore run up to twice the replica +count at once, and an eviction or an autoscaler scale-down during the rollout adds more +draining pods. Set PostgreSQL `max_connections` to at least +`(2 × replicas + surge) × 14`, which leaves one surge of margin for those +pods, plus headroom for your other clients and administration. Use +`autoscaling.maxReplicas` as the replica count when autoscaling is enabled. + +| Replicas | Surge | Minimum `max_connections` | +|---|---|---| +| 2 | 1 | `(4 + 1) × 14 = 70` | +| 3 | 1 | `(6 + 1) × 14 = 98` | +| 4 | 1 | `(8 + 1) × 14 = 126` | +| 6 | 2 | `(12 + 2) × 14 = 196` | + +Three or more replicas exceed what the PostgreSQL default of 100 leaves for +the gateway, because PostgreSQL reserves 3 connections for superusers. Raise +`max_connections` or use a larger managed instance. Let one rollout finish +before you start another, because each overlapping rollout adds its own +draining pods. + +If a connection pooler sits between the gateway and PostgreSQL, use session +pooling. The gateway holds session-level advisory locks, which transaction +pooling breaks. + +Mutations lock only what they change. Sandbox operations lock their own +sandbox, provider and workspace-profile changes lock their workspace, and +gateway-global policy, settings, and platform-profile changes lock the whole +fleet. A lock wait longer than 10 seconds returns `UNAVAILABLE` with the +reason `MUTATION_LOCK_TIMEOUT`, and clients can retry the request. A request +that carried a `request_id` leaves its admission unresolved, like any other +error, so a retry with the same ID returns `REQUEST_OUTCOME_UNCERTAIN`. Observe +resource state and reconcile effects, then send a new request with a new +`request_id`. Refer to +[Durable Request Admission](/sdk/api-errors#durable-request-admission). + +Each replica also polls PostgreSQL once per second for the sandboxes that its +clients follow with `WatchSandbox`, for example a CLI waiting for a sandbox to +become ready. The poll reads only sandbox IDs and versions, one query per +1,000 watched sandboxes, and uses a connection from the data pool. Changes +made through another replica reach those watchers within about one second. +Track the watched set and query latency per replica with +`openshell_server_sandbox_watch_polled_sandboxes` and +`openshell_server_sandbox_watch_poll_duration_seconds`. ## Next Steps @@ -211,3 +544,5 @@ selected replica count. [Managing Certificates](/kubernetes/managing-certificates). - To configure user authentication and authorization, refer to [Access Control](/kubernetes/access-control). +- To scrape and alert on gateway metrics, refer to + [Gateway Metrics](/observability/gateway-metrics). diff --git a/docs/kubernetes/setup.mdx b/docs/kubernetes/setup.mdx index b43ac2b073..67b9ffbb46 100644 --- a/docs/kubernetes/setup.mdx +++ b/docs/kubernetes/setup.mdx @@ -239,6 +239,8 @@ The most commonly changed values are: | `replicaCount` | Number of gateway replicas. Values above `1` require shared PostgreSQL through `server.externalDbSecret`. | | `workload.kind` | Gateway workload controller. Use `statefulset` for SQLite or `deployment` with `server.externalDbSecret`. | | `workload.allowMultiReplicaStatefulSet` | Allow `replicaCount > 1` with `workload.kind=statefulset`. Prefer Deployment for external database-backed multi-replica gateways. | +| `podLifecycle.terminationGracePeriodSeconds` | Seconds Kubernetes waits before killing a stopping gateway pod. Defaults to `30` so gateways that use an external PostgreSQL database can move their supervisor sessions to other replicas. | +| `autoscaling.enabled` / `autoscaling.minReplicas` / `autoscaling.maxReplicas` | Render a HorizontalPodAutoscaler instead of a fixed `replicaCount`. Refer to [High Availability](/kubernetes/high-availability#scale-the-gateway). | | `server.sandboxNamespace` | Namespace where sandbox pods are created. Defaults to the Helm release namespace when left empty. | | `workspaceResources.enabled` | Create namespace-scoped sandbox prerequisites from the gateway chart. Disable when installing the workspace chart separately. | | `server.externalDbSecret` | Secret containing a PostgreSQL connection URI in the `uri` key. Use when the database is managed outside the chart. | @@ -370,6 +372,8 @@ The gateway exposes `/healthz` for process liveness and `/readyz` for dependency - `startupProbe` and `livenessProbe` use `/healthz`. - `readinessProbe` uses `/readyz`, which reflects the latest result of an in-process background database check. +During shutdown, `/readyz` returns `503` with status `draining` while `/healthz` keeps returning `200`. + ## Next Steps - To run multiple gateway replicas, refer to [High Availability](/kubernetes/high-availability). diff --git a/docs/observability/gateway-metrics.mdx b/docs/observability/gateway-metrics.mdx new file mode 100644 index 0000000000..c45e1ec972 --- /dev/null +++ b/docs/observability/gateway-metrics.mdx @@ -0,0 +1,254 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: "Gateway Metrics" +sidebar-title: "Gateway Metrics" +description: "Scrape Prometheus metrics from the OpenShell gateway, including supervisor session, relay capacity, peer routing, mutation lock, and watch signals for multi-replica deployments." +keywords: "Generative AI, Cybersecurity, Observability, Metrics, Prometheus, Gateway, Kubernetes, High Availability" +--- + +The gateway exports Prometheus metrics for API requests, database readiness, +gateway interceptors, and multi-replica capacity. Use them to alert on +saturation, to follow rollouts, and to choose autoscaling signals. + +## Enable the Metrics Endpoint + +The gateway serves `GET /metrics` in the Prometheus text format on its own +listener, separate from the API and health ports. When the listener is +disabled, the gateway records no metrics. Enable it in one of these ways: + +- In `gateway.toml`, set `metrics_bind_address` under `[openshell.gateway]`, + for example `metrics_bind_address = "0.0.0.0:9090"`. Refer to + [Gateway Configuration](/how-it-works/gateways/configuration). +- On the command line, pass `--metrics-port` or set `OPENSHELL_METRICS_PORT`. + The standalone gateway binary defaults to `0`, which disables the endpoint. +- In the Helm chart, set `service.metricsPort`, which defaults to `9090`. The + chart renders the TOML key, a container port named `metrics`, and a + `metrics` port on the gateway Service. Set it to `0` to disable the + endpoint. + +A stopping gateway keeps serving metrics until the process exits, so you can +observe its drain. + + +The metrics endpoint uses plaintext HTTP with no authentication. Anyone who +can reach the port can read the metrics. The Helm chart publishes the port on +the gateway Service, so a `LoadBalancer` or `NodePort` Service exposes it +outside the cluster. Keep the Service internal, and restrict the port to your +monitoring system. + + +The following NetworkPolicy keeps the API and health ports open and accepts +metrics scrapes only from the `monitoring` namespace. Adjust the ports if you +changed `service.port`, `service.healthPort`, or `service.metricsPort`. + +```yaml +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: openshell-gateway-metrics + namespace: openshell +spec: + podSelector: + matchLabels: + app.kubernetes.io/name: openshell + app.kubernetes.io/instance: openshell + policyTypes: + - Ingress + ingress: + - ports: + - port: 8080 + - port: 8081 + - from: + - namespaceSelector: + matchLabels: + kubernetes.io/metadata.name: monitoring + ports: + - port: 9090 +``` + +## Scrape Each Replica + +Capacity metrics describe one gateway replica, so scrape every gateway pod +rather than the Service, and keep the `namespace` and `pod` labels on each +series. With the Prometheus Operator, a PodMonitor adds those labels +automatically: + +```yaml +apiVersion: monitoring.coreos.com/v1 +kind: PodMonitor +metadata: + name: openshell-gateway + namespace: openshell +spec: + selector: + matchLabels: + app.kubernetes.io/name: openshell + app.kubernetes.io/instance: openshell + podMetricsEndpoints: + - port: metrics + path: /metrics + interval: 30s +``` + +With a plain Prometheus configuration, use pod discovery and copy the labels +yourself: + +```yaml +scrape_configs: + - job_name: openshell-gateway + kubernetes_sd_configs: + - role: pod + namespaces: + names: [openshell] + relabel_configs: + - source_labels: [__meta_kubernetes_pod_label_app_kubernetes_io_instance] + regex: openshell + action: keep + - source_labels: [__meta_kubernetes_pod_container_port_name] + regex: metrics + action: keep + - source_labels: [__meta_kubernetes_namespace] + target_label: namespace + - source_labels: [__meta_kubernetes_pod_name] + target_label: pod +``` + +To check one pod without Prometheus, read its metrics through the Kubernetes +API server proxy: + +```shell +kubectl get --raw "/api/v1/namespaces/openshell/pods/:9090/proxy/metrics" \ + | grep '^openshell_server_supervisor_sessions' +``` + +The API server connects to the pod from the control plane. A NetworkPolicy +like the one above blocks that connection unless it also allows the control +plane on the metrics port. `kubectl port-forward` reaches the pod through the +kubelet instead, and NetworkPolicy does not block it. Start the port-forward: + +```shell +kubectl -n openshell port-forward pod/ 9090:9090 +``` + +Then read the metrics from another terminal: + +```shell +curl -s http://localhost:9090/metrics | grep '^openshell_server_supervisor_sessions' +``` + +## Histograms and Summaries + +The multi-replica latency metrics are Prometheus histograms with `_bucket`, +`_sum`, and `_count` series and buckets from 1 millisecond to 15 seconds, so +`histogram_quantile` can aggregate them across replicas. The older +`*_duration_seconds` metrics and the interceptor latency metric are summaries +with `quantile` labels that each replica computes on its own. Do not average +summary quantiles across pods. + +## Multi-Replica Capacity Metrics + +These metrics describe how much work one replica holds and how the replicas +interact. No metric carries a sandbox, channel, endpoint, or replica +identifier, because the scrape target already identifies the replica. + +Supervisor sessions and draining: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_supervisor_sessions` | Gauge | None | Supervisor control sessions held by this replica. After a rollout settles, the sum across pods equals the number of connected sandboxes. | +| `openshell_server_draining` | Gauge | None | `1` from the start of this replica's drain until the process exits, otherwise `0`. Gateways that use an external PostgreSQL database (`server.externalDbSecret`) drain. Gateways on the default SQLite database skip the drain and always report `0`. | + +Relays: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_relay_pending` | Gauge | None | Relay channels on this replica waiting for the supervisor to connect back, including channels opened for peer replicas. A channel whose caller gave up stays counted until it is claimed or cleaned up, up to about 40 seconds. | +| `openshell_server_relay_pending_capacity` | Gauge | None | Pending relay limit per replica, `256`. | +| `openshell_server_relay_pending_per_sandbox_capacity` | Gauge | None | Pending relay limit per sandbox on one replica, `32`. | +| `openshell_server_relay_rejected_total` | Counter | `reason` | Relay opens rejected at a limit. Rejected relays fail the client request, so use this counter, not client error codes, to detect saturation. | +| `openshell_server_relay_expired_total` | Counter | None | Pending relays dropped because the supervisor did not connect back within 10 seconds. | +| `openshell_server_relay_claim_duration_seconds` | Histogram | None | Time from opening a relay to the supervisor claiming it. | + +Peer requests, which this replica sends to the replica that owns a sandbox's +supervisor session: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_peer_requests_total` | Counter | `rpc`, `outcome`, `code` | Requests sent to the owning replica. Each retry counts separately. | +| `openshell_server_peer_request_duration_seconds` | Histogram | `rpc`, `outcome` | Latency of those requests. For `PeerRelay`, until the owner's supervisor claimed the relay. | + +Mutation locks: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_mutation_lock_wait_seconds` | Histogram | `scope` | Time to acquire a mutation lock, recorded for successful acquisitions. | +| `openshell_server_mutation_lock_timeouts_total` | Counter | `scope` | Acquisitions that waited longer than 10 seconds. The request failed with `UNAVAILABLE` and reason `MUTATION_LOCK_TIMEOUT`. [API Errors](/sdk/api-errors) describes how to retry it. | + +Watch poller: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_sandbox_watch_polled_sandboxes` | Gauge | None | Sandboxes with a local `WatchSandbox` follower checked on the last poll. Always `0` on gateways that use SQLite, which do not run the poller. | +| `openshell_server_sandbox_watch_poll_duration_seconds` | Histogram | None | Duration of the batched version lookup for one poll, including failed lookups. Polls with no watched sandboxes record nothing. | +| `openshell_server_sandbox_watch_poll_errors_total` | Counter | None | Polls whose version lookup failed. The poller keeps its last known versions and retries on the next poll. | + +The labels take these values: + +- `reason` is `global_capacity` or `sandbox_capacity`. +- `rpc` is `PeerRelay`, `PeerReportProviderReadiness`, + `PeerReportEndpointStatus`, or `PeerGetSandboxProviderStatus`. +- `outcome` is `ok`, `client_error`, or `rpc_error`. `client_error` means the + request failed before it reached the owner (token, connection, or stream + setup), or the caller cancelled it at any point, which records + `code="cancelled"` even when the owner was already handling the request. + `rpc_error` means the call to the owner failed, either with the owner's gRPC + status or with a transport error on an already open peer connection. +- `code` is the lowercase gRPC status name, such as `ok`, `unavailable`, or + `resource_exhausted`, as the owner returned it and before the gateway + reports peer failures to clients as `UNAVAILABLE`. It differs from the + numeric `code` label of `openshell_server_grpc_requests_total`. +- `scope` is `global`, `workspace`, or `sandbox`. + +Gauges and counters exist from startup. They start at `0`, except +`openshell_server_relay_pending_capacity` and +`openshell_server_relay_pending_per_sandbox_capacity`, which start at their +limits. `openshell_server_peer_requests_total` starts with only its +`outcome="ok"` series. Other outcomes and codes, and every histogram, appear +after their first sample, so write alert expressions that tolerate absent +series. + +## Request and Readiness Metrics + +The gateway also records every API request and its background database check: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_server_grpc_requests_total` | Counter | `method`, `code` | gRPC requests by method name and numeric status code, including internal peer requests. Streaming RPCs record the status known when the response headers are sent. | +| `openshell_server_grpc_request_duration_seconds` | Summary | `method`, `code` | Time until the response headers are sent. For streaming RPCs this covers setup only. | +| `openshell_server_http_requests_total` | Counter | `path`, `status` | HTTP requests on the API port. `path` is `/_ws_tunnel`, `/auth`, or `unknown`. | +| `openshell_server_http_request_duration_seconds` | Summary | `path`, `status` | Duration of those HTTP requests. | +| `openshell_server_readiness_database_healthy` | Gauge | None | `1` when the last background database check succeeded, otherwise `0`. | +| `openshell_server_readiness_database_probe_duration_seconds` | Summary | `outcome` | Duration of the background database check. `outcome` is `success`, `db_error`, or `timeout`. | + +## Gateway Interceptor Metrics + +Gateways with [gateway interceptors](/extensibility/gateway-interceptors) +record their evaluations: + +| Metric | Type | Labels | Description | +|---|---|---|---| +| `openshell_gateway_interceptor_latency_seconds` | Summary | None | Duration of interceptor calls. | +| `openshell_gateway_interceptor_evaluations_total` | Counter | `decision`, `interceptor`, `binding_id` | Interceptor evaluations by decision. | +| `openshell_gateway_interceptor_patches_total` | Counter | `interceptor`, `binding_id` | Patches that interceptors applied. | +| `openshell_gateway_interceptor_fail_open_total` | Counter | None | Interceptor failures that let the request continue. | +| `openshell_gateway_interceptor_fail_closed_total` | Counter | None | Interceptor failures that rejected the request. | +| `openshell_gateway_interceptor_post_commit_observation_failures_total` | Counter | `stage` | Failed post-commit observations. `stage` is `response_body` or `evaluation`. | + +## Next Steps + +- To monitor capacity and choose autoscaling signals for multiple gateway + replicas, refer to + [High Availability](/kubernetes/high-availability#monitor-capacity). +- To configure the metrics listener in `gateway.toml`, refer to + [Gateway Configuration](/how-it-works/gateways/configuration). diff --git a/docs/sdk/api-errors.mdx b/docs/sdk/api-errors.mdx index 03cd9cf206..23a34c3b96 100644 --- a/docs/sdk/api-errors.mdx +++ b/docs/sdk/api-errors.mdx @@ -31,6 +31,7 @@ Recognized gateway reasons include the following. | `INVALID_ARGUMENT` | `INVALID_ARGUMENT` | Correct the fields listed in `BadRequest`. | | `RESOURCE_VERSION_CONFLICT` | `ABORTED` | Read the resource again and construct a new conditional write. `metadata.recovery` is `REFRESH_STATE`; `current_resource_version` is included when known. | | `PROFILE_SOURCE_UNAVAILABLE` | `UNAVAILABLE` | Retry a profile snapshot read after at least the supplied delay. | +| `MUTATION_LOCK_TIMEOUT` | `UNAVAILABLE` | The request waited more than 10 seconds for a concurrent mutation to release a sandbox, workspace, or gateway-wide lock. Retry after at least the supplied delay. A request that carried a `request_id` leaves its admission unresolved, like any other error, so a retry with that ID returns `REQUEST_OUTCOME_UNCERTAIN`. Observe resource state and reconcile effects before you start a new request with a new ID. | | `REQUEST_ID_PAYLOAD_MISMATCH` | `FAILED_PRECONDITION` | Keep the original payload for that request ID. Inspect the original operation before submitting a different one. | | `REQUEST_OUTCOME_UNCERTAIN` | `FAILED_PRECONDITION` | An attempt is admitted but has no confirmed replayable success. Observe resource state and reconcile effects. Do not switch to a new ID to bypass the admission. | | `REQUEST_REPLAY_UNAVAILABLE` | `FAILED_PRECONDITION` | The original scope, resource, interceptor transformation, or fingerprint key is no longer replayable. Reconcile effects; the gateway does not execute the request again. Missing private-key material can also reject admission before work starts. | diff --git a/e2e/rust/tests/kubernetes_ha_rebalancing.rs b/e2e/rust/tests/kubernetes_ha_rebalancing.rs index ca65f1e3eb..983050b4cc 100644 --- a/e2e/rust/tests/kubernetes_ha_rebalancing.rs +++ b/e2e/rust/tests/kubernetes_ha_rebalancing.rs @@ -3,12 +3,15 @@ #![cfg(feature = "e2e-kubernetes-ha")] +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt::Write as _; use std::fs; use std::io::Write; use std::path::Path; use std::process::Stdio; use std::time::{Duration, Instant}; +use futures_util::future::try_join_all; use openshell_e2e::harness::binary::openshell_cmd; use openshell_e2e::harness::output::strip_ansi; use openshell_e2e::harness::port::{find_free_port, wait_for_port}; @@ -22,6 +25,34 @@ static KUBE_HA_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new const HA_SYNC_PAYLOAD_BYTES: usize = 32 * 1024 * 1024; const HA_SYNC_TIMEOUT: Duration = Duration::from_secs(600); +const HA_REDISTRIBUTION_SANDBOXES: usize = 4; +/// Overall budget for the rollout test. It stays below the 600 s nextest +/// terminate-after limit so a slow step fails with its own diagnostics +/// instead of a bare timeout. +const HA_REDISTRIBUTION_TEST_BUDGET: Duration = Duration::from_secs(540); +// Per-step caps; each step also stops at the overall deadline. +const HA_REDISTRIBUTION_CREATE_TIMEOUT: Duration = Duration::from_secs(180); +const HA_SESSION_ACCOUNTING_TIMEOUT: Duration = Duration::from_secs(90); +const HA_ROLLOUT_TIMEOUT: Duration = Duration::from_secs(180); +/// Old pods can still be draining after `rollout status` returns; the drain +/// plus cleanup takes at most 25 s. +const HA_DRAIN_EXIT_TIMEOUT: Duration = Duration::from_secs(30); +/// A drained gateway pod leaves the pod list within that 25 s bound plus pod +/// teardown and one poll. A pod still listed after this long overran the +/// bound or was killed at the end of its 30 s termination grace period. +const HA_DRAIN_EXIT_BOUND: Duration = Duration::from_secs(28); +const HA_READY_PODS_TIMEOUT: Duration = Duration::from_secs(120); +const HA_EXEC_TIMEOUT: Duration = Duration::from_secs(180); +/// Bounds the wait for a scaled-down pod's sandboxes to report Ready again. +const HA_SETTLE_TIMEOUT: Duration = Duration::from_secs(120); +/// Bounds one metrics scrape or sandbox listing, so a hung call cannot run +/// past a step deadline. +const HA_QUERY_TIMEOUT: Duration = Duration::from_secs(10); +const SUPERVISOR_SESSIONS_METRIC: &str = "openshell_server_supervisor_sessions"; +const DRAINING_METRIC: &str = "openshell_server_draining"; +const RELAY_PENDING_CAPACITY_METRIC: &str = "openshell_server_relay_pending_capacity"; +const RELAY_PENDING_CAPACITY: u64 = 256; + #[derive(Clone)] struct KubeTarget { context: String, @@ -140,7 +171,16 @@ impl KubeTarget { } async fn wait_for_gateway_pods(&self, expected: usize) -> Result, String> { - let deadline = Instant::now() + Duration::from_secs(240); + self.wait_for_gateway_pods_until(expected, Instant::now() + Duration::from_secs(240)) + .await + } + + async fn wait_for_gateway_pods_until( + &self, + expected: usize, + deadline: Instant, + ) -> Result, String> { + let budget = deadline.saturating_duration_since(Instant::now()); let mut last = String::new(); while Instant::now() < deadline { @@ -162,14 +202,28 @@ impl KubeTarget { } Err(format!( - "gateway pods did not reach expected ready count {expected} within 240s; last={last}" + "gateway pods did not reach expected ready count {expected} within {budget:?}; last={last}" )) } async fn gateway_pods(&self) -> Result, String> { + Ok(self + .gateway_pod_states() + .await? + .into_iter() + .filter(|pod| !pod.terminating) + .collect()) + } + + /// Every gateway pod, including pods that are terminating. + async fn gateway_pod_states(&self) -> Result, String> { let selector = format!("app.kubernetes.io/instance={}", self.release); + // Bound each poll so a stalled API server cannot push a step past its + // deadline. + let request_timeout = format!("--request-timeout={}s", HA_QUERY_TIMEOUT.as_secs()); let json = self .kubectl(&[ + &request_timeout, "-n", &self.namespace, "get", @@ -188,9 +242,6 @@ impl KubeTarget { let mut pods = Vec::new(); for item in items { - if !item["metadata"]["deletionTimestamp"].is_null() { - continue; - } let Some(name) = item["metadata"]["name"].as_str() else { continue; }; @@ -202,20 +253,351 @@ impl KubeTarget { && condition["status"].as_str() == Some("True") }) }); + let metrics_port = item["spec"]["containers"] + .as_array() + .and_then(|containers| { + containers + .iter() + .filter_map(|container| container["ports"].as_array()) + .flatten() + .find(|port| port["name"].as_str() == Some("metrics")) + }) + .and_then(|port| port["containerPort"].as_u64()) + .and_then(|port| u16::try_from(port).ok()); pods.push(GatewayPod { name: name.to_string(), ready, + terminating: !item["metadata"]["deletionTimestamp"].is_null(), + metrics_port, }); } pods.sort_by(|a, b| a.name.cmp(&b.name)); Ok(pods) } + + /// Run `kubectl get --raw` and return stdout only, so stderr warnings + /// cannot corrupt the response body. + async fn kubectl_raw(&self, path: &str) -> Result { + let request_timeout = format!("--request-timeout={}s", HA_QUERY_TIMEOUT.as_secs()); + let output = Command::new("kubectl") + .arg("--context") + .arg(&self.context) + .args([request_timeout.as_str(), "get", "--raw", path]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .output() + .await + .map_err(|err| format!("failed to spawn kubectl get --raw {path}: {err}"))?; + + if !output.status.success() { + return Err(format!( + "kubectl get --raw {path} failed with exit {:?}: {}", + output.status.code(), + String::from_utf8_lossy(&output.stderr).trim() + )); + } + + String::from_utf8(output.stdout) + .map_err(|err| format!("kubectl get --raw {path} returned non-UTF-8 output: {err}")) + } + + /// Fetch `/metrics` through the API server's pod proxy, which reaches the + /// plain-HTTP metrics listener without a port-forward. + async fn scrape_gateway_metrics(&self, pod: &GatewayPod) -> Result { + let port = pod.metrics_port.ok_or_else(|| { + format!( + "gateway pod {} has no container port named metrics", + pod.name + ) + })?; + self.kubectl_raw(&format!( + "/api/v1/namespaces/{}/pods/{}:{port}/proxy/metrics", + self.namespace, pod.name + )) + .await + } + + /// No terminating gateway pod, exactly `expected_pods` ready pods, and + /// every `required` sandbox in phase Ready. Returns the pods and the full + /// phase map. + async fn gateway_settled_once( + &self, + expected_pods: usize, + required: &[String], + ) -> Result<(Vec, BTreeMap), String> { + let pods = self.gateway_pod_states().await?; + if pods.len() != expected_pods || pods.iter().any(|pod| pod.terminating || !pod.ready) { + return Err(format!( + "expected {expected_pods} ready gateway pods and none terminating; pods={:?}", + pods.iter() + .map(|pod| format!( + "{} ready={} terminating={}", + pod.name, pod.ready, pod.terminating + )) + .collect::>() + )); + } + + let phases = sandbox_phases().await?; + let not_ready: Vec = required + .iter() + .filter(|name| phases.get(*name).map(String::as_str) != Some("Ready")) + .map(|name| { + format!( + "{name}={}", + phases.get(name).map_or("missing", String::as_str) + ) + }) + .collect(); + if !not_ready.is_empty() { + return Err(format!( + "sandboxes not Ready: {not_ready:?}; phases={phases:?}" + )); + } + Ok((pods, phases)) + } + + async fn wait_for_gateway_settled( + &self, + expected_pods: usize, + required: &[String], + deadline: Instant, + ) -> Result<(), String> { + poll_until( + deadline, + "gateway pods and sandboxes did not settle", + move || self.gateway_settled_once(expected_pods, required), + ) + .await + .map(|_| ()) + } + + /// Settled gateway, and on every pod: the relay capacity gauge, draining + /// at 0, and supervisor sessions that add up to the Ready sandbox count. + async fn session_accounting_once( + &self, + expected_pods: usize, + required: &[String], + ) -> Result { + let (pods, phases) = self.gateway_settled_once(expected_pods, required).await?; + let mut sessions_by_pod = BTreeMap::new(); + for pod in &pods { + let metrics = self.scrape_gateway_metrics(pod).await?; + let capacity = prometheus_count(&metrics, RELAY_PENDING_CAPACITY_METRIC)?; + if capacity != RELAY_PENDING_CAPACITY { + return Err(format!( + "{RELAY_PENDING_CAPACITY_METRIC} on {} is {capacity}, expected {RELAY_PENDING_CAPACITY}", + pod.name + )); + } + let draining = prometheus_count(&metrics, DRAINING_METRIC)?; + if draining != 0 { + return Err(format!( + "{DRAINING_METRIC} on ready pod {} is {draining}", + pod.name + )); + } + sessions_by_pod.insert( + pod.name.clone(), + prometheus_count(&metrics, SUPERVISOR_SESSIONS_METRIC)?, + ); + } + + let ready = phases.values().filter(|phase| *phase == "Ready").count(); + let ready_sandboxes = u64::try_from(ready).unwrap_or(u64::MAX); + let sessions: u64 = sessions_by_pod.values().sum(); + if sessions != ready_sandboxes { + return Err(format!( + "supervisor sessions {sessions_by_pod:?} add up to {sessions}, but {ready_sandboxes} sandboxes are Ready; phases={phases:?}" + )); + } + Ok(SessionAccounting { + sessions_by_pod, + ready_sandboxes, + }) + } + + async fn wait_for_session_accounting( + &self, + expected_pods: usize, + required: &[String], + deadline: Instant, + ) -> Result { + poll_until( + deadline.min(Instant::now() + HA_SESSION_ACCOUNTING_TIMEOUT), + "supervisor session gauges did not account for every Ready sandbox", + move || self.session_accounting_once(expected_pods, required), + ) + .await + } + + /// Run `kubectl rollout restart` on the gateway and wait until the + /// rollout finished and no old pod is left, recording which terminating + /// pods reported `openshell_server_draining` 1 and how long each stayed + /// terminating. `rollout status` returns while old pods may still drain, + /// because terminating pods do not count toward Deployment status. + async fn restart_gateway_and_observe_drain( + &self, + deadline: Instant, + ) -> Result { + let resource = self.gateway_workload_resource().await?; + self.kubectl(&["-n", &self.namespace, "rollout", "restart", &resource]) + .await?; + + let status_timeout = format!( + "--timeout={}s", + step_budget(deadline, HA_ROLLOUT_TIMEOUT).as_secs().max(1) + ); + let status_args = [ + "-n", + &self.namespace, + "rollout", + "status", + &resource, + &status_timeout, + ]; + let status = self.kubectl(&status_args); + tokio::pin!(status); + + let until = deadline.min(Instant::now() + HA_ROLLOUT_TIMEOUT + HA_DRAIN_EXIT_TIMEOUT); + let mut observation = RolloutObservation::default(); + let mut status_done = false; + // Pods still listed as terminating, each with the end of the first + // listing that showed it terminating. + let mut first_seen: BTreeMap = BTreeMap::new(); + loop { + if Instant::now() >= until { + return Err(format!( + "rollout status finished={status_done}; pods still terminating={:?}; {observation:?}", + first_seen.keys().collect::>() + )); + } + if status_done { + tokio::time::sleep(Duration::from_millis(500)).await; + } else { + tokio::select! { + result = &mut status => { + result.map_err(|err| format!("{err}\n{observation:?}"))?; + status_done = true; + } + () = tokio::time::sleep(Duration::from_millis(500)) => {} + } + } + + let listing_started = Instant::now(); + // Pods churn during a rollout; retry on the next tick. + let Ok(pods) = self.gateway_pod_states().await else { + continue; + }; + let listed_at = Instant::now(); + let terminating: Vec<&GatewayPod> = pods.iter().filter(|pod| pod.terminating).collect(); + let names: Vec<&str> = terminating.iter().map(|pod| pod.name.as_str()).collect(); + record_terminating_pods( + &mut first_seen, + &mut observation.exit_times, + &names, + listing_started, + listed_at, + ); + for pod in terminating { + observation.terminating.insert(pod.name.clone()); + // The process may already have exited; only a positive + // reading counts. Stop scraping a pod after one, which keeps + // each poll short and its exit time accurate. + if !observation.draining.contains(&pod.name) + && self + .scrape_gateway_metrics(pod) + .await + .is_ok_and(|metrics| prometheus_count(&metrics, DRAINING_METRIC) == Ok(1)) + { + observation.draining.insert(pod.name.clone()); + } + } + if status_done && first_seen.is_empty() { + return Ok(observation); + } + } + } + + /// Unwrap a step result, or panic with its error followed by a snapshot of + /// the gateway pods and sandboxes. + async fn check(&self, what: &str, result: Result) -> T { + match result { + Ok(value) => value, + Err(err) => panic!("{what}: {err}\n{}", self.diagnostics().await), + } + } + + /// Best-effort snapshot for failure messages: each gateway pod with its + /// state and session and draining gauges, then every sandbox phase. + async fn diagnostics(&self) -> String { + let mut out = String::from("cluster snapshot:"); + match self.gateway_pod_states().await { + Ok(pods) => { + for pod in &pods { + let gauges = match self.scrape_gateway_metrics(pod).await { + Ok(metrics) => { + let gauge = |metric| { + prometheus_count(&metrics, metric) + .map_or_else(|err| err, |value| value.to_string()) + }; + format!( + "sessions={} draining={}", + gauge(SUPERVISOR_SESSIONS_METRIC), + gauge(DRAINING_METRIC) + ) + } + Err(err) => format!("metrics unavailable: {err}"), + }; + let _ = write!( + out, + "\n gateway pod {} ready={} terminating={} {gauges}", + pod.name, pod.ready, pod.terminating + ); + } + } + Err(err) => { + let _ = write!(out, "\n gateway pods unavailable: {err}"); + } + } + match sandbox_phases().await { + Ok(phases) => { + let _ = write!(out, "\n sandbox phases: {phases:?}"); + } + Err(err) => { + let _ = write!(out, "\n sandbox phases unavailable: {err}"); + } + } + out + } } -#[derive(Debug)] +#[derive(Debug, Clone)] struct GatewayPod { name: String, ready: bool, + /// `metadata.deletionTimestamp` is set. + terminating: bool, + /// `containerPort` of the container port named `metrics`. + metrics_port: Option, +} + +#[derive(Debug)] +struct SessionAccounting { + sessions_by_pod: BTreeMap, + ready_sandboxes: u64, +} + +#[derive(Debug, Default)] +struct RolloutObservation { + terminating: BTreeSet, + draining: BTreeSet, + /// For each pod seen terminating: from the end of the first listing that + /// showed it terminating to the start of the first listing without it. A + /// late first sighting only shortens this; the last poll adds at most + /// one poll interval. + exit_times: BTreeMap, } struct PortForward { @@ -371,6 +753,146 @@ async fn assert_exec_through_all_pods( Ok(()) } +/// `name -> phase` for every sandbox in the CLI's workspace. +async fn sandbox_phases() -> Result, String> { + let mut cmd = openshell_cmd(); + // The session gauge counts every workspace, so list every workspace too. + cmd.args(["sandbox", "list", "--all-workspaces", "--output", "json"]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let output = tokio::time::timeout(HA_QUERY_TIMEOUT, cmd.output()) + .await + .map_err(|_| format!("sandbox list did not finish within {HA_QUERY_TIMEOUT:?}"))? + .map_err(|err| format!("failed to spawn openshell sandbox list: {err}"))?; + + let stdout = String::from_utf8_lossy(&output.stdout); + if !output.status.success() { + return Err(format!( + "sandbox list failed with exit {:?}:\n{stdout}{}", + output.status.code(), + String::from_utf8_lossy(&output.stderr) + )); + } + let json = stdout.find('{').map_or("", |start| &stdout[start..]); + let value = serde_json::from_str::(json) + .map_err(|err| format!("failed to parse sandbox list JSON: {err}\n{stdout}"))?; + if value["next_page_token"] + .as_str() + .is_some_and(|token| !token.is_empty()) + { + return Err(format!( + "sandbox list returned more than one page; the check needs every sandbox:\n{stdout}" + )); + } + let sandboxes = value["sandboxes"] + .as_array() + .ok_or_else(|| format!("sandbox list JSON missing sandboxes array: {value}"))?; + + Ok(sandboxes + .iter() + .filter_map(|sandbox| { + Some(( + sandbox["name"].as_str()?.to_string(), + sandbox["phase"].as_str()?.to_string(), + )) + }) + .collect()) +} + +/// Value of an unlabelled sample in Prometheus text exposition output. +fn prometheus_sample(text: &str, metric: &str) -> Option { + text.lines() + .filter(|line| !line.starts_with('#')) + .find_map(|line| { + let mut fields = line.split_whitespace(); + if fields.next()? != metric { + return None; + } + fields.next()?.parse::().ok() + }) +} + +fn prometheus_count(text: &str, metric: &str) -> Result { + let value = prometheus_sample(text, metric) + .ok_or_else(|| format!("metric {metric} missing from /metrics"))?; + if value < 0.0 || value.fract() != 0.0 { + return Err(format!("metric {metric} is not a count: {value}")); + } + format!("{value:.0}") + .parse::() + .map_err(|err| format!("metric {metric} value {value}: {err}")) +} + +/// Record one successful pod listing. A pod newly listed as terminating is +/// timed from `listed_at`, the end of the listing. A timed pod missing from +/// the listing left after the previous one, so its time up to +/// `listing_started` moves to `exit_times`. +fn record_terminating_pods( + first_seen: &mut BTreeMap, + exit_times: &mut BTreeMap, + terminating: &[&str], + listing_started: Instant, + listed_at: Instant, +) { + first_seen.retain(|name, seen| { + let listed = terminating.contains(&name.as_str()); + if !listed { + exit_times.insert( + name.clone(), + listing_started.saturating_duration_since(*seen), + ); + } + listed + }); + for name in terminating { + if !exit_times.contains_key(*name) { + first_seen.entry((*name).to_string()).or_insert(listed_at); + } + } +} + +/// Time left before `deadline`, capped at `step`. +fn step_budget(deadline: Instant, step: Duration) -> Duration { + deadline.saturating_duration_since(Instant::now()).min(step) +} + +/// Run one step under `min(step, time left)`, naming the step on timeout. +async fn within( + deadline: Instant, + step: Duration, + what: &str, + fut: impl Future>, +) -> Result { + let budget = step_budget(deadline, step); + tokio::time::timeout(budget, fut) + .await + .map_err(|_| format!("{what} did not finish within {budget:?}"))? +} + +/// Run `attempt` at least once, then every 2 s until it succeeds or `until` +/// passes. The error names `what` and carries the last failure. +async fn poll_until( + until: Instant, + what: &str, + mut attempt: impl FnMut() -> Fut, +) -> Result +where + Fut: Future>, +{ + let budget = until.saturating_duration_since(Instant::now()); + let interval = Duration::from_secs(2); + loop { + let last = match attempt().await { + Ok(value) => return Ok(value), + Err(err) => err, + }; + if Instant::now() + interval >= until { + return Err(format!("{what} within {budget:?}; last: {last}")); + } + tokio::time::sleep(interval).await; + } +} + fn write_deterministic_payload(path: &Path, size: usize) { let mut file = fs::File::create(path).expect("create HA sync payload"); let mut offset = 0usize; @@ -514,6 +1036,15 @@ async fn sandbox_exec_rebalances_across_gateway_scale_and_rollout() { .wait_for_gateway_pods(2) .await .expect("gateway should scale back to two ready replicas"); + // The removed pod drains after `rollout status` returns, and new execs + // fail while the moved sandbox briefly reports Provisioning. + kube.wait_for_gateway_settled( + 2, + &[sandbox.name.clone()], + Instant::now() + HA_SETTLE_TIMEOUT, + ) + .await + .expect("the removed gateway pod should finish draining and the sandbox should be Ready again"); assert_exec_through_all_pods(&kube, &pods, &sandbox.name, "scale-down") .await .expect("exec should work through every gateway pod after scale-down"); @@ -625,3 +1156,191 @@ async fn sandbox_file_sync_survives_gateway_pod_rolls() { sandbox.cleanup().await; } + +#[tokio::test] +#[allow(clippy::too_many_lines)] // one rollout scenario, each step with its own deadline +async fn supervisor_sessions_redistribute_across_gateway_pod_rolls() { + let _test_lock = KUBE_HA_TEST_LOCK.lock().await; + let kube = KubeTarget::from_env(); + let deadline = Instant::now() + HA_REDISTRIBUTION_TEST_BUDGET; + + let scaled = within( + deadline, + HA_ROLLOUT_TIMEOUT, + "scale gateway to two replicas", + kube.scale_gateway(2), + ) + .await; + kube.check( + "gateway should run two HA replicas before the rollout", + scaled, + ) + .await; + let old_pods = kube + .wait_for_gateway_pods_until(2, deadline.min(Instant::now() + HA_READY_PODS_TIMEOUT)) + .await; + let old_pods = kube + .check("two ready gateway replicas before the rollout", old_pods) + .await; + + let created = within( + deadline, + HA_REDISTRIBUTION_CREATE_TIMEOUT, + "create sandboxes through the configured gateway", + try_join_all((0..HA_REDISTRIBUTION_SANDBOXES).map(|idx| async move { + let phase = format!("redistribute-{idx}"); + create_sandbox_through_configured_gateway(&phase).await + })), + ) + .await; + let mut sandboxes = kube + .check("create sandboxes through the configured gateway", created) + .await; + let names: Vec = sandboxes + .iter() + .map(|sandbox| sandbox.name.clone()) + .collect(); + + let before = kube.wait_for_session_accounting(2, &names, deadline).await; + let before = kube + .check( + "session gauges should account for every Ready sandbox before the rollout", + before, + ) + .await; + eprintln!( + "supervisor sessions before rollout: {:?} for {} Ready sandboxes", + before.sessions_by_pod, before.ready_sandboxes + ); + + let observation = kube.restart_gateway_and_observe_drain(deadline).await; + let observation = kube.check("gateway rollout restart", observation).await; + eprintln!( + "gateway pods terminating during rollout: {:?}; draining: {:?}; left after: {:?}", + observation.terminating, observation.draining, observation.exit_times + ); + assert!( + !observation.draining.is_empty(), + "expected {DRAINING_METRIC}=1 on a terminating gateway pod; terminating pods seen: {:?}", + observation.terminating + ); + assert!( + observation + .draining + .iter() + .all(|pod| old_pods.contains(pod)), + "only replaced pods should drain; old={old_pods:?} draining={:?}", + observation.draining + ); + // The gauge flips before any session closes, so only the exit time tells + // a finished drain from a kill at the end of the grace period. + assert!( + observation + .exit_times + .values() + .all(|elapsed| *elapsed <= HA_DRAIN_EXIT_BOUND), + "every terminating gateway pod should leave within {HA_DRAIN_EXIT_BOUND:?}; a longer stay means the drain overran its 25 s bound or the 30 s grace period cut it off; left after: {:?}", + observation.exit_times + ); + + let new_pods = kube + .wait_for_gateway_pods_until(2, deadline.min(Instant::now() + HA_READY_PODS_TIMEOUT)) + .await; + let new_pods = kube + .check("two ready gateway replicas after the rollout", new_pods) + .await; + assert!( + new_pods.iter().all(|pod| !old_pods.contains(pod)), + "rollout restart should replace every gateway pod; old={old_pods:?} new={new_pods:?}" + ); + + // Exec runs only after this settles: new execs fail with + // FAILED_PRECONDITION while a moved sandbox reports Provisioning. + let after = kube.wait_for_session_accounting(2, &names, deadline).await; + let after = kube + .check( + "every sandbox should be Ready with exactly one supervisor session on the new pods", + after, + ) + .await; + eprintln!( + "supervisor sessions after rollout: {:?} for {} Ready sandboxes", + after.sessions_by_pod, after.ready_sandboxes + ); + + let exec = within( + deadline, + HA_EXEC_TIMEOUT, + "exec through every new gateway pod", + async { + for (idx, sandbox) in sandboxes.iter().enumerate() { + let phase = format!("redistribute-{idx}"); + assert_exec_through_all_pods(&kube, &new_pods, &sandbox.name, &phase) + .await + .map_err(|err| { + format!( + "exec through every new gateway pod for {}: {err}", + sandbox.name + ) + })?; + let marker = format!("ha-redistribute-client-{idx}"); + exec_through_configured_gateway(&sandbox.name, &marker) + .await + .map_err(|err| { + format!( + "exec through the configured gateway for {}: {err}", + sandbox.name + ) + })?; + } + Ok(()) + }, + ) + .await; + kube.check("exec after the rollout", exec).await; + + for sandbox in &mut sandboxes { + sandbox.cleanup().await; + } +} + +#[test] +fn prometheus_count_reads_unlabelled_samples_only() { + let text = "# TYPE openshell_server_supervisor_sessions gauge\n\ + openshell_server_supervisor_sessions 3\n\ + openshell_server_supervisor_sessions_total 9\n\ + openshell_server_relay_rejected_total{reason=\"global_capacity\"} 1\n"; + assert_eq!( + prometheus_count(text, "openshell_server_supervisor_sessions"), + Ok(3) + ); + assert!(prometheus_count(text, "openshell_server_relay_rejected_total").is_err()); + assert!(prometheus_count(text, "openshell_server_missing").is_err()); +} + +#[test] +fn terminating_pod_exit_time_runs_from_first_sighting_to_first_listing_without_it() { + let start = Instant::now(); + let at = |secs| start + Duration::from_secs(secs); + let mut first_seen = BTreeMap::new(); + let mut exit_times = BTreeMap::new(); + + record_terminating_pods(&mut first_seen, &mut exit_times, &["old-a"], at(0), at(1)); + record_terminating_pods( + &mut first_seen, + &mut exit_times, + &["old-a", "old-b"], + at(2), + at(3), + ); + record_terminating_pods(&mut first_seen, &mut exit_times, &["old-b"], at(10), at(11)); + assert_eq!( + exit_times, + BTreeMap::from([("old-a".to_string(), Duration::from_secs(9))]) + ); + assert_eq!(first_seen, BTreeMap::from([("old-b".to_string(), at(3))])); + + record_terminating_pods(&mut first_seen, &mut exit_times, &[], at(40), at(41)); + assert_eq!(exit_times.get("old-b"), Some(&Duration::from_secs(37))); + assert!(first_seen.is_empty()); +} diff --git a/skills/debug-openshell-cluster/SKILL.md b/skills/debug-openshell-cluster/SKILL.md index 8e865f13ec..48b34a5400 100644 --- a/skills/debug-openshell-cluster/SKILL.md +++ b/skills/debug-openshell-cluster/SKILL.md @@ -418,18 +418,86 @@ kubectl -n get deployment,service,pod -l app.kubernetes.io/name= logs deployment/ --tail=200 ``` -Multi-replica gateways serialize cross-object sandbox and provider mutations -with a PostgreSQL advisory lock. If those RPCs stall while ordinary reads and -health checks remain responsive, inspect long-running database sessions and -advisory-lock waiters. Do not print the database URI or Secret contents into -logs: +Gateways that use an external PostgreSQL database guard cross-object +mutations with PostgreSQL advisory locks at three levels: one global key, one +key per workspace, and one key per sandbox, each taken shared or exclusive. +Sandbox operations on different sandboxes do not wait on each other's locks; +provider and workspace-profile changes block sandbox operations in their +workspace; gateway-global policy and settings changes block everything. Each +replica takes its advisory locks on a dedicated pool of 4 PostgreSQL +connections, one per guard, so at most 4 guarded operations per replica hold +or wait for PostgreSQL locks at once. A guard that waits on a PostgreSQL lock +keeps its connection while it waits, so contention on one key can fill the +pool and make unrelated guarded operations on that replica queue for a +connection, within the same 10-second limit. + +A lock wait longer than 10 seconds returns `UNAVAILABLE` with reason +`MUTATION_LOCK_TIMEOUT` and the message "... timed out waiting for a +concurrent mutation; retry the request", logs +`mutation lock acquisition timed out` with `scope`, `waited_ms`, and `detail` +fields, and increments `openshell_server_mutation_lock_timeouts_total`. The +`detail` field says where the wait stopped: + +- `waiting for a local mutation lock`: another operation on the same replica + holds a conflicting key. +- `waiting for a mutation lock connection`: all 4 lock-pool connections of that + replica were in use. `pg_locks` shows no row for the timed-out operation. + Look for the pod's granted or waiting advisory-lock rows by `client_addr`, + and find the holder they wait on. +- `waiting for a PostgreSQL advisory lock`, or PostgreSQL's own + `canceling statement due to lock timeout`: a conflicting key is held on + another connection, usually by another replica or by an older gateway during + an upgrade. `pg_locks` shows the holder and the waiters. + +A gateway log line `timed out returning PostgreSQL mutation lock connection; +discarded connection` means returning a lock session stalled for 5 seconds and +the gateway dropped it; that slot stays busy for up to those 5 seconds. The +warning does not identify the PostgreSQL backend. Healthy guards also hold +idle advisory-lock sessions while validation and writes use separate data +connections. Neither an idle duration nor a matching `client_addr` proves +that a holder is orphaned, even when it blocks a timed-out request. + +Before using `SELECT pg_terminate_backend()`, conclusively map that +backend to its owning gateway process and confirm that process has stopped +or can no longer write. If the owner is still running, stop it first and +verify its exit; draining or failing readiness alone is insufficient. Killing +its lock session while it can still write removes exclusion from an active +mutation. Recheck the backend PID and `backend_start` before terminating the +confirmed orphan: pod IPs and PIDs can be reused, and a database proxy can +hide several gateways behind one `client_addr`. If ownership cannot be +established, investigate connectivity instead of choosing a backend by IP or +idle state. Setting PostgreSQL +`tcp_keepalives_idle`, `tcp_keepalives_interval`, and `tcp_keepalives_count` +(for example 60, 10, and 6) bounds how long such sessions survive. + +If mutations stall or time out while reads and health checks work, inspect +advisory-lock holders and waiters. Do not print the database URI or Secret +contents into logs: ```sql -SELECT pid, granted, waitstart -FROM pg_locks -WHERE locktype = 'advisory'; +SELECT l.pid, l.mode, l.granted, l.waitstart, l.classid, l.objid, + a.client_addr, a.backend_start, a.state, + now() - a.state_change AS in_state_for +FROM pg_locks l +JOIN pg_stat_activity a USING (pid) +WHERE l.locktype = 'advisory' +ORDER BY l.granted, l.waitstart; ``` +Gateways hold these locks at session level outside any transaction, so a +holder (`granted = t`) usually shows `state = idle`. It still holds the lock, +and `in_state_for` approximates how long. `client_addr` is the holding +gateway pod's IP, or the pooler's IP when a connection pooler sits in between. + +The global key appears as `classid = 1330660686` and `objid = 1397247052` +(key `0x4F50454E53484C4C`). Older gateways during a rolling upgrade take that +key exclusively for every guarded mutation, so mutations can queue behind +them until the rollout finishes. + +If a client following a sandbox through one replica misses changes made +through another replica, check that replica's logs for `sandbox watch poller` +warnings and its `openshell_server_sandbox_watch_poll_errors_total` counter. + For multi-replica gateway installs, supervisor and client session traffic may be served by a non-owner gateway replica and relayed to the current supervisor owner over the internal `PeerRelay` RPC. Check the headless peer Service, @@ -466,6 +534,60 @@ name and load the chart CA plus client identity from the `peer-client-tls` volume exists, those files are readable, and the server certificate includes the name in `OPENSHELL_PEER_TLS_SERVER_NAME`. +Gateway pods that use an external PostgreSQL database budget 25 seconds for +supervisor drain and cleanup: a 3-second propagation delay, closes spread +over at most 12 seconds, then up to 10 seconds of session cleanup. This does +not bound other compute-driver cleanup or final trace export; allow for those +when sizing the termination grace period. Gateways on the default SQLite +database skip the drain. A lone replica on PostgreSQL, such as a single-replica StatefulSet, +still drains even though no other replica takes its sessions, so it takes up to +15 seconds longer to stop. A StatefulSet starts the replacement pod only after +the old one exits, so its sandboxes stay unreachable up to 15 seconds longer +on every restart, and supervisors in reconnect backoff can take up to 30 more +seconds, their maximum retry delay, after the new pod is ready. A drain that +ran to completion logs +"Draining supervisor sessions before stopping the gateway listener" and then +"Supervisor session drain finished; stopping gateway listener". If the second +line is missing, the termination grace period cut the drain short. Check the +grace period, the drain, and per-pod capacity: + +```bash +kubectl -n openshell get "${GATEWAY_DEPLOYMENT}" \ + -o jsonpath='{.spec.template.spec.terminationGracePeriodSeconds}{"\n"}' +kubectl -n openshell logs -c openshell-gateway --tail=200 \ + | grep -E 'Draining supervisor sessions|Supervisor session drain finished|mutation lock acquisition timed out' +for pod in $(kubectl -n openshell get pod \ + -l app.kubernetes.io/name=openshell,app.kubernetes.io/instance=openshell \ + -o jsonpath='{range .items[?(@.spec.containers[0].name=="openshell-gateway")]}{.metadata.name}{" "}{end}'); do + echo "${pod}" + kubectl get --raw "/api/v1/namespaces/openshell/pods/${pod}:9090/proxy/metrics" \ + | grep -E '^openshell_server_(supervisor_sessions|draining|relay_pending|relay_rejected_total|mutation_lock_timeouts_total)' +done +kubectl -n openshell get hpa +kubectl -n openshell describe hpa openshell +``` + +The JSONPath filter keeps only pods whose first container is the gateway +(`openshell-gateway`). It skips certificate hook Job pods, which older charts +labeled like gateway pods. The metrics port is `service.metricsPort` (default +`9090`). Read a stopping pod's logs while it is terminating, because a deleted +pod's logs are gone. + +The API server proxy connects to each pod from the control plane. A +NetworkPolicy that accepts the metrics port only from a monitoring namespace +blocks it unless the policy also allows the control plane. In that case, +read one pod at a time through `kubectl port-forward`, which reaches the pod +through the kubelet and is not blocked by NetworkPolicy: + +```bash +kubectl -n openshell port-forward pod/ 9090:9090 >/dev/null & +pf_pid=$! +sleep 2 +curl -s http://localhost:9090/metrics \ + | grep -E '^openshell_server_(supervisor_sessions|draining|relay_pending|relay_rejected_total|mutation_lock_timeouts_total)' +kill "${pf_pid}" +``` + Check required Helm deployment secrets: ```bash @@ -909,6 +1031,15 @@ credential failures. | Kubernetes gateway pod pending | PVC unbound, taint, selector, or insufficient resources | `kubectl -n openshell describe pod ` | | Kubernetes sandbox pod stuck pending, workspace PVC unbound | Cluster has no default `StorageClass` and OpenShell does not set `storageClassName` on the workspace PVC (clusters with a default `StorageClass` bind fine without it) | `kubectl -n openshell describe pvc`; set `server.workspaceStorageClass` (gateway config `workspace_storage_class`) to a valid `StorageClass` | | Kubernetes gateway pod crash loops | Missing secret, bad DB URL, bad TLS config | `kubectl -n openshell logs deployment/openshell -c openshell-gateway` or `kubectl -n openshell logs statefulset/openshell -c openshell-gateway` | +| Sandboxes switch to `Provisioning` for a second or two during a gateway rollout, and new exec, SSH, or forward requests to them fail with `FAILED_PRECONDITION` `sandbox is not ready` (service URLs return `412`) | Expected: the stopping replica moved their supervisor sessions, and requests during the move fail the readiness check. A reconnect that takes more than a few seconds can also show the `Ready` condition as `DependenciesNotReady` while the workload keeps running | Retry. The sandboxes should return to `Ready`; if not, check Service endpoints, gateway logs, and `openshell_server_supervisor_sessions` per pod | +| Exec or forwarding to older sandboxes fails for up to 30 seconds during rollouts | Supervisors from earlier releases wait longer between reconnect attempts | Recreate those sandboxes to pick up the current supervisor | +| Gateway pod exits with code 137 during a rollout, or its log has "Draining supervisor sessions" without "Supervisor session drain finished" | `podLifecycle.terminationGracePeriodSeconds` is shorter than the drain (often 5 after `helm upgrade --reuse-values`) | Workload `terminationGracePeriodSeconds`; set it to 30 | +| A single-replica gateway on PostgreSQL takes up to 15 seconds longer to stop, and a single-replica StatefulSet leaves sandboxes unreachable that much longer on every restart, with some supervisors reconnecting up to 30 seconds after the new pod is ready | Expected: every gateway on an external PostgreSQL database drains, even when no other replica can take its sessions, a StatefulSet starts the replacement pod only after the old one exits, and supervisors back off (up to 30 seconds between attempts) while no replica is ready | None; the drain adds at most 15 seconds and the backoff at most 30 seconds after the new pod is ready. Run more than one replica to keep sandboxes reachable during restarts | +| Many sandboxes take longer than 15 seconds to reconnect during a rollout | A draining replica held more sessions than the others can absorb without lock queuing (about 500) | `openshell_server_supervisor_sessions` per pod before the rollout, `openshell_server_mutation_lock_wait_seconds`; add replicas | +| Mutating RPCs return `UNAVAILABLE` with "timed out waiting for a concurrent mutation" | Advisory-lock contention, a full 4-connection lock pool on one replica, a slow PostgreSQL, or older gateways still running during an upgrade | `openshell_server_mutation_lock_timeouts_total`, the `detail` field of `mutation lock acquisition timed out` in gateway logs, the `pg_locks` query in Step 6 | +| `helm upgrade` fails with an `autoscaling.*` message | HPA values invalid: missing `resources.requests` (or `resources.limits`), `maxReplicas` above 1 without `server.externalDbSecret` (or on a StatefulSet without `workload.allowMultiReplicaStatefulSet`), no metric target, or min/max out of order. "`minReplicas` and `maxReplicas` are not set" means `--reuse-values` kept a release without the chart's autoscaling defaults | Fix the values named in the error; upgrade with `--reset-then-reuse-values` instead of `--reuse-values` | +| HPA shows `` targets | No metrics-server for CPU/memory, or the metrics adapter does not serve the custom metric | `kubectl -n openshell describe hpa openshell`, `kubectl get --raw /apis/custom.metrics.k8s.io/v1beta1` | +| One replica holds most sessions after a rollout | Expected: sessions stay where they reconnected | `openshell_server_supervisor_sessions` per pod; it fades as sandboxes are recreated | | OpenShift gateway pod fails to start with an SCC/`runAsUser` error (e.g. `unable to validate against any security context constraint`) | Chart's default `podSecurityContext`/`securityContext` hardcodes `runAsUser`/`fsGroup`, which the restricted-v2 SCC rejects; it must instead inject the namespace-assigned UID/GID range | `oc -n openshell describe pod `; deploy with `podSecurityContext: null` and clear `securityContext.runAsUser` (see `deploy/helm/openshell/ci/values-openshift-scc.yaml`) | | OpenShift sandbox pod fails to start (`unable to validate against any security context constraint`) | The `openshell-sandbox` service account lacks the privileged SCC it needs | `oc adm policy add-scc-to-user privileged -z openshell-sandbox -n openshell`; remove with `remove-scc-from-user` when done | | OpenShift self-hosted Vault/OpenBao credential store pod never schedules (waits time out with `no matching resources found`) | The store's Helm chart pins `runAsUser`/`fsGroup`/seccomp, which restricted-v2 rejects, so the StatefulSet controller never creates the pod | Deploy the store's chart in its OpenShift mode (`--set global.openshift=true` for the OpenBao/Vault chart) so the namespace SCC assigns a compliant security context — no manual SCC grant needed | diff --git a/tasks/scripts/run-postgres-tests.sh b/tasks/scripts/run-postgres-tests.sh new file mode 100755 index 0000000000..cf27525464 --- /dev/null +++ b/tasks/scripts/run-postgres-tests.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run the PostgreSQL-backed openshell-server tests: #[ignore] tests whose names +# start with `postgres_`. Point OPENSHELL_TEST_POSTGRES_URL at a disposable +# database, or leave it unset to start a throwaway PostgreSQL container with +# the local container engine (Docker or Podman). Never point it at a database +# that a running gateway uses: the tests take fleet-wide advisory locks. +# +# Extra arguments are passed to cargo nextest. Use test-name filters to run a +# subset, for example `postgres_get_resource_versions`; a second -E filterset +# would widen the selection instead of narrowing it. + +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${ROOT}" + +# Same pinned image as the Kubernetes e2e fixture (e2e/kubernetes/postgres-fixture.yaml). +POSTGRES_IMAGE="${OPENSHELL_TEST_POSTGRES_IMAGE:-mirror.gcr.io/library/postgres:17.10-alpine3.23@sha256:979c4379dd698aba0b890599a6104e082035f98ef31d9b9291ec22f2b13059ca}" +CONTAINER_NAME="" + +cleanup() { + if [ -n "${CONTAINER_NAME}" ]; then + ce rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true + fi +} +trap cleanup EXIT + +start_postgres() { + local password port ready=0 + # shellcheck source=tasks/scripts/container-engine.sh + source "${ROOT}/tasks/scripts/container-engine.sh" + + password="$(od -An -N16 -tx1 /dev/urandom | tr -d ' \n')" + CONTAINER_NAME="openshell-test-postgres-$$" + echo "Starting disposable PostgreSQL (${POSTGRES_IMAGE})..." + # No --rm: the EXIT trap removes the container, and keeping it until then + # preserves its logs when PostgreSQL fails to start. + ce run -d --name "${CONTAINER_NAME}" \ + -e POSTGRES_USER=openshell \ + -e POSTGRES_PASSWORD="${password}" \ + -e POSTGRES_DB=openshell \ + -p 127.0.0.1::5432 \ + "${POSTGRES_IMAGE}" >/dev/null + + # The image's init phase runs a socket-only server, so a TCP probe succeeds + # only once the final server accepts connections. + for _ in $(seq 1 60); do + if ce exec "${CONTAINER_NAME}" pg_isready -h 127.0.0.1 -U openshell -d openshell >/dev/null 2>&1; then + ready=1 + break + fi + if [ "$(ce inspect -f '{{.State.Running}}' "${CONTAINER_NAME}" 2>/dev/null)" != "true" ]; then + break + fi + sleep 1 + done + if [ "${ready}" != "1" ]; then + echo "ERROR: PostgreSQL did not become ready within 60s or its container exited" >&2 + ce logs "${CONTAINER_NAME}" >&2 || true + exit 1 + fi + + port="$(ce port "${CONTAINER_NAME}" 5432/tcp | head -n1 | awk -F: '{print $NF}')" + echo "PostgreSQL is ready on 127.0.0.1:${port} (container ${CONTAINER_NAME})" + export OPENSHELL_TEST_POSTGRES_URL="postgres://openshell:${password}@127.0.0.1:${port}/openshell" +} + +if [ -z "${OPENSHELL_TEST_POSTGRES_URL:-}" ]; then + start_postgres +fi + +# The mutation-replay test predates OPENSHELL_TEST_POSTGRES_URL. +export OPENSHELL_REPLAY_TEST_DATABASE_URL="${OPENSHELL_REPLAY_TEST_DATABASE_URL:-${OPENSHELL_TEST_POSTGRES_URL}}" +export OPENSHELL_TELEMETRY_ENABLED=false + +# Advisory locks are database-wide, so run the tests one at a time. +cargo nextest run -p openshell-server --features test-support \ + --run-ignored only --test-threads 1 \ + -E 'test(/(^|::)postgres_/)' "$@" diff --git a/tasks/test.toml b/tasks/test.toml index 39dd070204..75ea00b026 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -104,6 +104,11 @@ run = [ run_windows = "powershell -NoProfile -ExecutionPolicy Bypass -File tasks/scripts/windows-msvc.ps1 test-precommit native" hide = true +["test:rust:postgres"] +description = "Run PostgreSQL-backed openshell-server tests against OPENSHELL_TEST_POSTGRES_URL or a disposable Docker/Podman PostgreSQL container" +run = "tasks/scripts/run-postgres-tests.sh" +run_windows = "echo Skipping test:rust:postgres: PostgreSQL integration tests need a Linux or macOS container engine." + ["test:python"] description = "Run Python tests" depends = ["python:proto"]