diff --git a/.agents/skills/helm-dev-environment/SKILL.md b/.agents/skills/helm-dev-environment/SKILL.md index 0b5a069eeb..df927931c6 100644 --- a/.agents/skills/helm-dev-environment/SKILL.md +++ b/.agents/skills/helm-dev-environment/SKILL.md @@ -68,8 +68,9 @@ gatewayConfig: supervisor_topology: gateway ``` -The gateway chart creates a single-replica shared supervisor Deployment, -Service, NetworkPolicy, and an empty `-supervisor-token` Secret in +The gateway chart creates a shared supervisor Deployment (`supervisor.replicas`, +default one), Service, NetworkPolicy, and an empty +`-supervisor-token` Secret in the release namespace, regardless of shared/managed/operator workspace mode. The gateway mints a supervisor token from its sandbox-JWT signing key, writes it into that Secret (the gateway Role may `get`/`patch` only that Secret), and @@ -96,6 +97,88 @@ OPENSHELL_E2E_KUBE_IMAGE_PULL_SECRET=e2e-regcred \ mise run e2e:kubernetes:shared-supervisor ``` +### Shared supervisor high availability + +The shared supervisor is always a pool of `supervisor.replicas` members +(default 1). With the gateway topology the chart renders the gateway's +`[openshell.gateway.supervisor_pool]` tuning table (only +`recovery_timeout_secs` from `supervisor.recoveryTimeoutSeconds`, default 60; +the driver's `shared_supervisor` table stays unchanged), a +RollingUpdate Deployment (`maxSurge: 1`, `maxUnavailable: 0`) with preferred +hostname anti-affinity and a 60-second grace period, the pool identity env vars +(`OPENSHELL_SUPERVISOR_POOL_MEMBER_ID`, `_MEMBER_UID`, `_ADVERTISE_IP` from the +Pod name, UID, and IP), and a `maxUnavailable: 1` PodDisruptionBudget. +`ci/values-supervisor-ha.yaml` selects two replicas and a 60-second recovery +timeout. The topology needs gateway TLS: the e2e harness keeps it on with +`ci/values-tls-port-forward.yaml` and `OPENSHELL_E2E_KUBE_TLS=1`, which reaches +the gateway over mTLS through its port-forward. `mise run +e2e:kubernetes:shared-supervisor-ha` sets both, creates a task-scoped k3d +cluster, kills the supervisor Pod that owns a sandbox, and checks that the +sandbox keeps its Pod, generation, main process, and in-memory state. + +The rendered gateway configuration and supervisor Pod template do not +depend on `supervisor.replicas`; `deploy/helm/test-supervisor-pool-scaling.sh` +(part of `mise run helm:test`) renders a pool at 1, 2, and 3 replicas and +fails if anything but the replica count differs. `mise run +e2e:kubernetes:shared-supervisor-scaling` installs a one-member pool +(`ci/values-supervisor-ha-scaling.yaml`), then the test runs `helm upgrade +--reuse-values --set supervisor.replicas=2` and back to 1 against the chart +the harness installed (`OPENSHELL_E2E_KUBE_CHART`), probing exec and +policy-mediated networking throughout. It checks that scaling up moves +nothing, that the member removed by scaling down hands its sandboxes to the +other (`cause=leaving`), that no workload, gateway Pod, or remaining member +restarts, and that a new sandbox runs after each step. It upgrades its +release, so it runs on a release of its own; the two pool tests also share a +nextest test group that never runs them at once. Against an existing cluster +whose kubeconfig the shell environment may override, set +`OPENSHELL_E2E_KUBECONFIG` (passed as `--kubeconfig` to kubectl and helm). + +On k3d, deploy the pool with the Skaffold profile, which layers the overlay and +turns gateway TLS back on over `ci/values-skaffold.yaml`: + +```bash +cd deploy/helm/openshell +KUBECONFIG=../../../kubeconfig skaffold run -p shared-supervisor-high-availability +cd ../../.. +``` + +The plaintext registration from `helm:skaffold:run` does not apply. Keep +`mise run helm:k3s:forward` running and set up the client certificate as in +"With mTLS enabled" under +[Connecting through the forwarding task](#connecting-through-the-forwarding-task). +Both members share the single k3d node, so exercise Pod failures: + +```bash +export KUBECONFIG=kubeconfig OPENSHELL_GATEWAY_ENDPOINT=https://localhost:8090 +openshell sandbox create --detach --name ha-probe -- sh -c 'echo $$ > /sandbox/main.pid; exec sleep infinity' +openshell sandbox exec --name ha-probe --no-tty -- cat /sandbox/main.pid +openshell sandbox get ha-probe # note the sandbox ID +# Supervisor logs name the instance by its route, not its sandbox ID. +kubectl get sandboxes.agents.x-k8s.io -A -l openshell.ai/sandbox-name=ha-probe \ + -o jsonpath='{.items[0].metadata.annotations.openshell\.ai/sandbox-runtime-supervisor-session}' +# The member whose latest admission, removal, or release line names the route owns it. +kubectl -n openshell logs -l app.kubernetes.io/name=openshell-supervisor --prefix --tail=-1 \ + | grep -E '(admitted assigned|removing withdrawn|relinquishing a) sandbox instance.* $' +# Force-delete the owner, the closest local stand-in for a crash. +kubectl -n openshell delete pod --grace-period=0 --force +kubectl -n openshell logs -l app.kubernetes.io/name=openshell,app.kubernetes.io/instance=openshell \ + -c openshell-gateway --prefix --tail=200 | grep 'taking over a shared sandbox instance' +kubectl -n openshell logs -l openshell.ai/sandbox-id=,openshell.ai/boundary-role=workload \ + --tail=200 | grep -E 'sandbox-supervisor-|Supervisor takeover confirmed' +openshell sandbox exec --name ha-probe --no-tty -- cat /sandbox/main.pid # same PID +``` + +A plain `kubectl delete pod` exercises the graceful leave instead: the member +announces it is leaving and the gateway hands its instances over at once. A +member that stops without announcing it, as on node loss, is declared dead +only after the 15-second membership TTL. +`kubectl -n openshell rollout restart deployment/openshell-supervisor` replaces +every member one at a time. The PID must survive each case, while an open +`sandbox connect` or `exec` stream breaks and must reconnect. A +`sandbox-supervisor-recovery-expired` finding in the workload log means no +replacement confirmed within the recovery timeout, so the runtime terminated +the workload. + ### 2. Deploy OpenShell **Iterative dev** (rebuilds on file changes, recommended during active development): @@ -498,6 +581,8 @@ 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-supervisor.yaml`, `ci/values-supervisor-ha.yaml`, `ci/values-supervisor-ha-scaling.yaml` | Shared supervisor pool overlays: one member, two members (Skaffold profile `shared-supervisor-high-availability`), or one member to scale to two and back | +| `deploy/helm/openshell/ci/values-tls-port-forward.yaml` | Keeps gateway TLS on over the Skaffold values for e2e topologies that need it (`OPENSHELL_E2E_KUBE_TLS=1`) | | `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 | diff --git a/.config/nextest.toml b/.config/nextest.toml index 5da0fe19b0..1f3f5d0364 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -26,11 +26,26 @@ slow-timeout = { period = "60s", terminate-after = 5 } # names because binary() fails in workspaces without the HA test binary. [test-groups] kubernetes-ha = { max-threads = 1 } +kubernetes-shared-supervisor-pool = { max-threads = 1 } [[profile.e2e-kubernetes.overrides]] filter = "test(/gateway_(scale_and_rollout|pod_rolls)$/)" test-group = "kubernetes-ha" +# Both change the shared supervisor pool of the release they run against: one +# kills a member, the other scales the pool. Never run them at once. Test +# names, not binary(), for the reason above. +[[profile.e2e-kubernetes.overrides]] +filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running)$/)" +test-group = "kubernetes-shared-supervisor-pool" + +# Scaling the pool runs two bounded Helm upgrades (up to 6 minutes each), waits +# for each member to join or leave, and creates sandboxes until one lands on +# the new member, so it gets 20 minutes instead of five. +[[profile.e2e-kubernetes.overrides]] +filter = "test(/^scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running$/)" +slow-timeout = { period = "60s", terminate-after = 20 } + # Relative to the profile store dir (`e2e/rust/target/nextest/e2e-kubernetes/`). [profile.e2e-kubernetes.junit] path = "../../../../../results/e2e-kubernetes.xml" diff --git a/Cargo.lock b/Cargo.lock index f36af0eb0f..87685d16f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4888,6 +4888,7 @@ dependencies = [ "bytes", "clap", "hex", + "hyper-util", "ipnet", "landlock", "libc", @@ -4916,6 +4917,7 @@ dependencies = [ "tokio-rustls", "tokio-stream", "tonic", + "tower", "tracing", "tracing-subscriber", "uuid", diff --git a/TESTING.md b/TESTING.md index e6754df317..c3fbaace03 100644 --- a/TESTING.md +++ b/TESTING.md @@ -578,6 +578,9 @@ Available task variants: | `e2e:kubernetes:workspace-operator` | Operator workspace mode (pre-provisioned namespaces) | | `e2e:kubernetes:v1alpha1` | Agent Sandbox v1alpha1 compatibility | | `e2e:kubernetes:external-driver` | External Kubernetes driver sidecar | +| `e2e:kubernetes:shared-supervisor` | One-member shared supervisor pool for the gateway's sandboxes (gateway TLS through the port-forward) | +| `e2e:kubernetes:shared-supervisor-ha` | Two-replica shared supervisor pool: kills the member that owns a sandbox and checks the sandbox is adopted, not restarted | +| `e2e:kubernetes:shared-supervisor-scaling` | One-member shared supervisor pool: scales it to two members and back with `helm upgrade`, checking sandboxes stay running and the removed member hands its sandboxes off | Kubernetes e2e environment variables: @@ -586,6 +589,8 @@ Kubernetes e2e environment variables: | `OPENSHELL_E2E_KUBE_CONTEXT` | kubectl context for an existing cluster (skips k3d creation) | | `OPENSHELL_E2E_KUBE_TEST` | Scope to a single test (e.g. `smoke`) | | `OPENSHELL_E2E_KUBE_EXTRA_VALUES` | Colon-separated additional Helm values files | +| `OPENSHELL_E2E_KUBE_CHART`, `OPENSHELL_E2E_KUBECONFIG` | Set by the harness for tests that upgrade the release, such as the supervisor pool scaling test: the chart it installed, and the kubeconfig passed as `--kubeconfig` to kubectl and helm. Set them yourself when running such a test directly against an existing release | +| `OPENSHELL_E2E_KUBE_TLS` | `1` reaches the gateway over mTLS through the port-forward; pair it with an overlay that keeps gateway TLS on, such as `ci/values-tls-port-forward.yaml` | | `OPENSHELL_E2E_KUBERNETES_FEATURES` | Cargo feature flags (default: `e2e,e2e-host-gateway,e2e-kubernetes`) | | `IMAGE_TAG` | Gateway/supervisor image tag (default: `latest` for existing clusters) | | `OPENSHELL_REGISTRY` | Image registry prefix (default: `ghcr.io/nvidia/openshell`) | diff --git a/crates/openshell-core/src/jwt.rs b/crates/openshell-core/src/jwt.rs index 7b1bf57ffd..f79490f832 100644 --- a/crates/openshell-core/src/jwt.rs +++ b/crates/openshell-core/src/jwt.rs @@ -50,6 +50,17 @@ mod session { /// JOSE `typ` of the supervisor token: the credential of the gateway's /// shared supervisor process for the assignment stream, and nothing else. pub const SUPERVISOR_JWT_TYPE: &str = "openshell-supervisor+jwt"; + /// JOSE `typ` of a supervisor takeover grant. + /// + /// The grant is the gateway's authorization for one supervisor process to + /// replace the owner of one runtime generation at a strictly higher + /// credential epoch. Only the sandbox runtime verifies it, and only while + /// staging a replacement attach. + pub const SUPERVISOR_TAKEOVER_JWT_TYPE: &str = "openshell-supervisor-takeover+jwt"; + /// Longest lifetime of a takeover grant. Replay safety comes from the + /// strictly increasing epoch it names; the bound only limits how long an + /// unused grant stays usable. + pub const MAX_SUPERVISOR_TAKEOVER_TTL: Duration = Duration::from_mins(10); pub const SANDBOX_SESSION_AUDIENCE: &str = "openshell-sandbox"; pub const DEFAULT_SESSION_TOKEN_TTL: Duration = Duration::from_hours(1); pub const MIN_SESSION_TOKEN_TTL: Duration = Duration::from_mins(1); @@ -194,6 +205,54 @@ mod session { component: SessionComponent, } + /// Claims of a supervisor takeover grant. The subject and audience are the + /// sandbox's, like its Sandbox Protocol session token, but the JOSE `typ` + /// keeps the two from ever being accepted for each other. + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] + #[serde(deny_unknown_fields)] + struct SupervisorTakeoverClaims { + iss: String, + sub: String, + aud: String, + iat: i64, + exp: i64, + jti: String, + sandbox_id: SandboxId, + runtime_generation: SandboxGenerationId, + auth_epoch: CredentialEpoch, + supervisor_instance: String, + owner: String, + component: SessionComponent, + } + + /// The replacement of the supervisor that owns one runtime generation. + /// + /// The gateway grants it only after durably advancing the sandbox's + /// credential epoch to `auth_epoch`; ordinary credential refresh keeps the + /// epoch and never produces a grant. + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct SupervisorTakeover { + pub sandbox_id: SandboxId, + pub runtime_generation: SandboxGenerationId, + /// Ownership term the grant starts; it must exceed the term the + /// runtime currently accepts. + pub auth_epoch: CredentialEpoch, + /// Supervisor process identity the runtime pins once the replacement + /// is confirmed (a canonical UUID chosen by the gateway). + pub supervisor_instance: String, + /// Supervisor pool member that owns the new term, for audit. + pub owner: String, + } + + /// A takeover grant whose signature, issuer, audience, and lifetime were + /// verified. + #[derive(Clone, Debug, PartialEq, Eq)] + pub struct VerifiedSupervisorTakeover { + pub takeover: SupervisorTakeover, + pub token_id: Uuid, + pub expires_at: i64, + } + /// Durable identity shared by every short-lived token for one sandbox runtime. #[derive(Clone, Debug, PartialEq, Eq)] pub struct SandboxRuntimeIdentity { @@ -721,6 +780,49 @@ mod session { }) } + /// Mint a takeover grant for `takeover`, valid for `ttl` (at most + /// [`MAX_SUPERVISOR_TAKEOVER_TTL`]). The caller must already have + /// advanced the sandbox's persisted credential epoch to + /// `takeover.auth_epoch`. + pub fn mint_supervisor_takeover( + &self, + takeover: &SupervisorTakeover, + ttl: Duration, + ) -> Result { + if ttl.is_zero() || ttl > MAX_SUPERVISOR_TAKEOVER_TTL { + return Err(SessionJwtError::InvalidLifetime); + } + validate_takeover_identity(&takeover.supervisor_instance, &takeover.owner)?; + let issued_at = self.clock.now_unix_seconds(); + let expires_at = + issued_at.saturating_add(i64::try_from(ttl.as_secs()).unwrap_or(i64::MAX)); + let token_id = Uuid::new_v4(); + let claims = SupervisorTakeoverClaims { + iss: self.issuer.clone(), + sub: format!("{SANDBOX_SUBJECT_PREFIX}{}", takeover.sandbox_id), + aud: SANDBOX_SESSION_AUDIENCE.to_string(), + iat: issued_at, + exp: expires_at, + jti: token_id.to_string(), + sandbox_id: takeover.sandbox_id.clone(), + runtime_generation: takeover.runtime_generation.clone(), + auth_epoch: takeover.auth_epoch, + supervisor_instance: takeover.supervisor_instance.clone(), + owner: takeover.owner.clone(), + component: SessionComponent::OpenShellSupervisor, + }; + let mut header = Header::new(Algorithm::EdDSA); + header.kid = Some(self.key_id.clone()); + header.typ = Some(SUPERVISOR_TAKEOVER_JWT_TYPE.to_string()); + let token = encode(&header, &claims, &self.encoding_key) + .map_err(|_| SessionJwtError::SigningFailed)?; + Ok(MintedSessionToken { + token: SecretJwt::parse(token)?, + expires_at, + token_id, + }) + } + fn mint( &self, profile: SessionTokenProfile, @@ -820,6 +922,57 @@ mod session { self.validate_claims(claims) } + /// Verify a supervisor takeover grant. Only the Sandbox profile, which + /// the runtime uses, verifies grants; the gateway profile refuses them. + pub fn verify_supervisor_takeover( + &self, + token: &str, + ) -> Result { + if self.profile != SessionTokenProfile::Sandbox { + return Err(SessionJwtError::WrongTokenType); + } + let claims = decode_signed::( + &self.keys, + token, + SUPERVISOR_TAKEOVER_JWT_TYPE, + &["iss", "aud", "iat", "exp", "sub", "jti"], + )?; + if claims.iss != self.issuer { + return Err(SessionJwtError::WrongIssuer); + } + if claims.aud != SANDBOX_SESSION_AUDIENCE { + return Err(SessionJwtError::WrongAudience); + } + if claims.sub != format!("{SANDBOX_SUBJECT_PREFIX}{}", claims.sandbox_id) { + return Err(SessionJwtError::SubjectMismatch); + } + SandboxGenerationId::parse(claims.runtime_generation.to_string()) + .map_err(|_| SessionJwtError::InvalidRuntimeIdentity)?; + validate_takeover_identity(&claims.supervisor_instance, &claims.owner)?; + let token_id = Uuid::parse_str(&claims.jti).map_err(|_| SessionJwtError::InvalidJti)?; + // A grant must expire, within its own bound. + let max_lifetime = + i64::try_from(MAX_SUPERVISOR_TAKEOVER_TTL.as_secs()).unwrap_or(i64::MAX); + if claims.exp == 0 + || claims.exp <= claims.iat + || claims.exp.saturating_sub(claims.iat) > max_lifetime + { + return Err(SessionJwtError::InvalidLifetime); + } + validate_lifetime(claims.iat, claims.exp, self.clock.now_unix_seconds())?; + Ok(VerifiedSupervisorTakeover { + takeover: SupervisorTakeover { + sandbox_id: claims.sandbox_id, + runtime_generation: claims.runtime_generation, + auth_epoch: claims.auth_epoch, + supervisor_instance: claims.supervisor_instance, + owner: claims.owner, + }, + token_id, + expires_at: claims.exp, + }) + } + fn validate_claims( &self, claims: SessionClaims, @@ -1076,6 +1229,22 @@ mod session { Ok(gateway_id) } + /// A takeover names a non-blank supervisor instance and owner without + /// whitespace; the runtime compares the instance with its attach request. + fn validate_takeover_identity( + supervisor_instance: &str, + owner: &str, + ) -> Result<(), SessionJwtError> { + let valid = |value: &str| { + !value.is_empty() && value.len() <= 253 && !value.chars().any(char::is_whitespace) + }; + if valid(supervisor_instance) && valid(owner) { + Ok(()) + } else { + Err(SessionJwtError::InvalidRuntimeIdentity) + } + } + fn validate_key_id(key_id: String) -> Result { if key_id.is_empty() || key_id.trim() != key_id || key_id.chars().any(char::is_whitespace) { return Err(SessionJwtError::InvalidKeyId); @@ -1461,6 +1630,115 @@ mod tests { bundle.validate().expect("non-expiring auth bundle"); } + fn takeover(epoch: u64) -> SupervisorTakeover { + SupervisorTakeover { + sandbox_id: SandboxId::parse("sandbox-a").expect("sandbox ID"), + runtime_generation: SandboxGenerationId::parse("generation-1") + .expect("runtime generation"), + auth_epoch: CredentialEpoch::new(epoch).expect("auth epoch"), + supervisor_instance: "6a8f1f4e-8c0e-4d55-9d8c-1d6e2f3a4b5c".to_string(), + owner: "supervisor-1".to_string(), + } + } + + #[test] + fn takeover_grants_are_their_own_profile_verified_only_by_the_runtime() { + let (issuer, gateway, sandbox, identity) = fixture(); + let expected = takeover(2); + let minted = issuer + .mint_supervisor_takeover(&expected, std::time::Duration::from_mins(5)) + .expect("mint takeover grant"); + let verified = sandbox + .verify_supervisor_takeover(minted.token.expose_secret()) + .expect("runtime verifies the grant"); + assert_eq!(verified.takeover, expected); + assert_eq!(verified.token_id, minted.token_id); + assert_eq!(verified.expires_at, 1_900_000_000 + 300); + + // A grant is never a session credential, and no other credential is + // a grant: neither session profile nor the shared supervisor token. + assert_eq!( + sandbox.verify(minted.token.expose_secret()), + Err(SessionJwtError::WrongTokenType) + ); + assert_eq!( + gateway.verify_supervisor_takeover(minted.token.expose_secret()), + Err(SessionJwtError::WrongTokenType) + ); + let pair = issuer.mint_pair(&identity).expect("token pair"); + for token in [&pair.sandbox.token, &pair.gateway.token] { + assert_eq!( + sandbox.verify_supervisor_takeover(token.expose_secret()), + Err(SessionJwtError::WrongTokenType) + ); + } + let host = issuer + .mint_supervisor_token(std::time::Duration::from_mins(15)) + .expect("supervisor token"); + assert_eq!( + sandbox.verify_supervisor_takeover(host.token.expose_secret()), + Err(SessionJwtError::WrongTokenType) + ); + } + + #[test] + fn takeover_grants_are_bounded_named_and_expire() { + let (issuer, _gateway, _sandbox, _identity) = fixture(); + for ttl in [ + std::time::Duration::ZERO, + MAX_SUPERVISOR_TAKEOVER_TTL + std::time::Duration::from_secs(1), + ] { + assert_eq!( + issuer.mint_supervisor_takeover(&takeover(2), ttl).err(), + Some(SessionJwtError::InvalidLifetime) + ); + } + for (instance, owner) in [("", "member"), ("instance", ""), ("in stance", "m")] { + let mut unnamed = takeover(2); + unnamed.supervisor_instance = instance.to_string(); + unnamed.owner = owner.to_string(); + assert_eq!( + issuer + .mint_supervisor_takeover(&unnamed, std::time::Duration::from_mins(5)) + .err(), + Some(SessionJwtError::InvalidRuntimeIdentity) + ); + } + + let key = openshell_crypto::pki::generate_keypair_for(&PKCS_ED25519).unwrap(); + let issuer = SessionJwtIssuer::from_ed25519_pem( + key.serialize_pem().unwrap().as_bytes(), + "current", + "test", + None, + Arc::new(FixedClock(1_900_000_000)), + ) + .unwrap(); + let verifier_at = |now| { + SessionJwtVerifier::new( + "test", + SessionTokenProfile::Sandbox, + [SessionVerificationKey { + key_id: "current".to_string(), + public_key_pem: key.public_key_pem().into_bytes(), + }], + Arc::new(FixedClock(now)), + ) + .unwrap() + }; + // A non-expiring session issuer still mints expiring grants. + let minted = issuer + .mint_supervisor_takeover(&takeover(3), std::time::Duration::from_mins(1)) + .unwrap(); + verifier_at(1_900_000_030) + .verify_supervisor_takeover(minted.token.expose_secret()) + .expect("grant within its lifetime"); + assert_eq!( + verifier_at(1_900_000_200).verify_supervisor_takeover(minted.token.expose_secret()), + Err(SessionJwtError::Expired) + ); + } + #[test] fn token_debug_is_redacted() { let (issuer, _gateway, _sandbox, identity) = fixture(); diff --git a/crates/openshell-sandbox-backend/README.md b/crates/openshell-sandbox-backend/README.md index 9045a8edf8..d4aa405b94 100644 --- a/crates/openshell-sandbox-backend/README.md +++ b/crates/openshell-sandbox-backend/README.md @@ -40,6 +40,31 @@ supervisor nor runs workload code. Policy discovery is followed by the existing attach, mediation setup, confirm, and explicit start sequence. The existing connection registry continues to govern supervisor replacement and recovery. +## Supervisor takeover + +`sandbox_auth::SandboxConnectionRegistry` pins the first supervisor process +that attaches. Protocol 1.1 adds `openshell.sandbox.supervisor-takeover`: a +supervisor holding a `BoundaryTakeover` (a gateway-signed grant for a later +credential epoch naming the sandbox, generation, and its process identity) +attaches with the grant, and the runtime stages it and promotes it at +confirmation. A grant is never sent to a runtime that does not advertise the +capability; that supervisor fails closed instead. The registry tracks the +accepted epoch, the pinned process, and the highest epoch any granted attach +reached, so a lower grant cannot promote after a higher one disconnected. +Ordinary credential refresh never changes the epoch. The replacement adopts +the main and exec processes from the attach snapshot instead of starting the +agent, and the snapshot reports the policy the runtime bound at launch. Once +confirmed, it forgets the grant and re-attaches without it to refresh the +snapshot, so reconnects of the same process never depend on the grant's +lifetime. + +Every attach also offers the runtime a recovery deadline +(`supervisor_recovery_timeout_secs`, bounded by +`SUPERVISOR_RECOVERY_TIMEOUT_RANGE`): how long the workload stays frozen after +that supervisor's connection is lost. `OpenShellRuntimeBackend` sends the +deadline it was built with on its first attach, a takeover attach, and every +re-attach; without one the runtime uses `DEFAULT_SUPERVISOR_RECOVERY_TIMEOUT`. + Sandbox and generation identity come from the verified bearer, rather than duplicate fields on `OpenBoundaryRequest`. Attach supplies policy; confirmation reports resource claims from backend-owned provisioned state, which the @@ -141,6 +166,25 @@ afterwards. The launch material and status of one hosted instance `openshell-supervisor-control`, which depends on this crate for the runtime descriptor and route validation. +In a supervisor pool, one member hosts each route, but a runtime reaches an +arbitrary member. A routed runtime therefore sends preface version `2` on every +dial, which the supervisor answers before any TLS with a status byte, a +length byte, and a runtime listener address as ASCII `ip:port` (IPv6 +bracketed): `1` accepts, naming the member itself (empty before it has a +route directory), and the handshake proceeds as for version 1; `2` redirects +to the owner; `3` reports the route unknown and carries no address. A member +redirects routes it does not host according to the gateway's route directory +(`RuntimeRouter::set_directory`), never to its own address, and the directory +also wins over a local registration it no longer agrees with, which is being +released. Replies are unauthenticated hints like the handle: the member must +still complete the pinned TLS handshake, and every request still carries the +bearer token. A lone shared supervisor accepts without naming itself, which +tells the runtime there is no pool to probe. Version 1 prefaces, from runtimes +that predate version 2, receive no reply. A supervisor that predates version 2 +closes a version 2 preface without replying (usually as a reset, since it +reads only the header); the runtime then dials it again with version 1, so +upgrades may roll the runtime and the supervisor in either order. + Runtime-initiated transport supports TCP, protected Unix sockets, and host vsock listeners. These share bounded TLS handshakes, the authenticated spare queue, and the same generation-pinned runtime identity check. The legacy diff --git a/crates/openshell-sandbox-backend/src/boundary_protocol.rs b/crates/openshell-sandbox-backend/src/boundary_protocol.rs index 8402a5e7a5..cc3eb6d69b 100644 --- a/crates/openshell-sandbox-backend/src/boundary_protocol.rs +++ b/crates/openshell-sandbox-backend/src/boundary_protocol.rs @@ -165,8 +165,10 @@ impl NativeLinuxSandboxAuditEvidence { /// /// The supervisor generates this value in memory and presents it on every /// attach, including transport reconnects. The sandbox pins the first value it -/// accepts for its process lifetime, so a replacement supervisor cannot reuse -/// launch credentials to take over an existing runtime generation. +/// accepts, so a replacement supervisor cannot reuse launch credentials to +/// take over an existing runtime generation. Only a gateway-signed takeover +/// grant for a higher credential epoch, naming the replacement's identity, +/// moves the pin (see [`crate::sandbox_auth`]). #[derive(Clone, Copy, PartialEq, Eq, Hash)] pub struct SupervisorInstanceId(uuid::Uuid); @@ -579,6 +581,25 @@ pub struct BoundaryConfig { pub child_env: std::collections::HashMap, } +/// Default time a runtime keeps a workload frozen, after losing its +/// supervisor, for an authenticated supervisor to recover it. +pub const DEFAULT_SUPERVISOR_RECOVERY_TIMEOUT: Duration = Duration::from_secs(30); +/// Accepted bounds, in seconds, of a recovery deadline a supervisor offers. +/// +/// A supervisor offers it at attach (see [`Request::Attach`]). Recovery is +/// always bounded: a workload never runs unsupervised, and a frozen workload +/// is terminated once its deadline passes. +pub const SUPERVISOR_RECOVERY_TIMEOUT_RANGE: std::ops::RangeInclusive = 10..=300; + +/// The recovery deadline of an offer of `seconds`, if +/// [`SUPERVISOR_RECOVERY_TIMEOUT_RANGE`] allows it. +#[must_use] +pub fn supervisor_recovery_timeout(seconds: u32) -> Option { + SUPERVISOR_RECOVERY_TIMEOUT_RANGE + .contains(&seconds) + .then(|| Duration::from_secs(u64::from(seconds))) +} + impl fmt::Debug for BoundaryConfig { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -778,6 +799,16 @@ pub enum Request { Attach { supervisor_instance_id: SupervisorInstanceId, policy: Box, + /// Gateway-signed takeover grant, present only when this process + /// replaces the generation's current owner. + #[serde(default, skip_serializing_if = "Option::is_none")] + supervisor_takeover: Option, + /// Seconds the runtime should keep the workload frozen after losing + /// this connection's supervision, offered by the supervisor and + /// bounded by [`SUPERVISOR_RECOVERY_TIMEOUT_RANGE`]. Absent keeps + /// [`DEFAULT_SUPERVISOR_RECOVERY_TIMEOUT`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + supervisor_recovery_timeout_secs: Option, }, Confirm, StartAgent { @@ -859,9 +890,16 @@ impl fmt::Debug for Request { Self::Attach { supervisor_instance_id: _, policy: _, + supervisor_takeover, + supervisor_recovery_timeout_secs, } => formatter .debug_struct("Attach") .field("policy", &"") + .field("takeover", &supervisor_takeover.is_some()) + .field( + "supervisor_recovery_timeout_secs", + supervisor_recovery_timeout_secs, + ) .finish(), Self::Confirm => formatter.write_str("Confirm"), Self::StartAgent { @@ -1015,6 +1053,10 @@ pub struct MediationTimingWire { pub struct SessionSnapshotWire { pub generation: String, pub processes: Vec, + /// Launch controls bound by the first attach, if any. A replacement + /// supervisor reattaches with exactly these. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub bound_policy: Option>, } /// Stable generation-scoped process state available to a replacement supervisor. diff --git a/crates/openshell-sandbox-backend/src/isolation_protocol.rs b/crates/openshell-sandbox-backend/src/isolation_protocol.rs index 9a740e4191..d1ac0aca3b 100644 --- a/crates/openshell-sandbox-backend/src/isolation_protocol.rs +++ b/crates/openshell-sandbox-backend/src/isolation_protocol.rs @@ -11,25 +11,44 @@ use crate::boundary_protocol::{BoundaryConfig, SandboxRuntimeDescriptor}; use crate::proto::{ExchangeFrame, exchange_frame::Payload}; pub const SANDBOX_PROTOCOL_CAPABILITY: &str = "openshell.sandbox.protocol"; -const PROTOCOL_VERSION: ProtocolVersion = ProtocolVersion { major: 1, minor: 0 }; +/// Optional: the runtime accepts a gateway-granted supervisor takeover. +/// +/// The grant travels on Attach (`supervisor_takeover`). A supervisor holding +/// one must not send it, or resume a workload it does not own, unless the +/// runtime advertises this; an older runtime keeps rejecting replacement +/// supervisors. +pub const SUPERVISOR_TAKEOVER_CAPABILITY: &str = "openshell.sandbox.supervisor-takeover"; +const PROTOCOL_VERSION: ProtocolVersion = ProtocolVersion { major: 1, minor: 1 }; -/// The built-in runtime supports fresh launches; suspend/resume is optional. +/// The built-in runtime supports fresh launches and supervisor takeover; +/// suspend/resume is optional. #[must_use] pub fn metadata(implementation: &str) -> PeerMetadata { - let capabilities = vec![SANDBOX_PROTOCOL_CAPABILITY.to_string()]; + let required = vec![SANDBOX_PROTOCOL_CAPABILITY.to_string()]; PeerMetadata { protocol_version: Some(PROTOCOL_VERSION), implementation_name: implementation.to_string(), implementation_version: openshell_core::VERSION.to_string(), - supported_capabilities: capabilities.clone(), - required_capabilities: capabilities, + supported_capabilities: vec![ + SANDBOX_PROTOCOL_CAPABILITY.to_string(), + SUPERVISOR_TAKEOVER_CAPABILITY.to_string(), + ], + required_capabilities: required, } } /// Use the shared extension rules: equal majors and both peers' requirements. pub fn negotiate(local: &PeerMetadata, remote: Option) -> Result<(), tonic::Status> { + negotiate_capabilities(local, remote).map(drop) +} + +/// Negotiate and return the remote peer's normalized metadata, so optional +/// capabilities it advertises can be checked before they are used. +pub fn negotiate_capabilities( + local: &PeerMetadata, + remote: Option, +) -> Result { negotiate_metadata("isolation", "sandbox-protocol", local, remote) - .map(drop) .map_err(|error| tonic::Status::failed_precondition(error.to_string())) } diff --git a/crates/openshell-sandbox-backend/src/lib.rs b/crates/openshell-sandbox-backend/src/lib.rs index 6e1bc80872..6d0af037d9 100644 --- a/crates/openshell-sandbox-backend/src/lib.rs +++ b/crates/openshell-sandbox-backend/src/lib.rs @@ -24,9 +24,9 @@ pub mod stream_adapter; pub use openshell_core::proto::sandbox::backend::v1 as backend_proto; pub use runtime::{ - BoundBoundary, BoundaryExec, BoundaryLoopbackConnector, BoundaryProcess, BoundaryTerminal, - ConfirmedBoundary, ExecSession, NetworkMediationSource, OpenShellRuntimeBackend, - ProcessAttachment, ReadyBoundary, RunningBoundary, + BoundBoundary, BoundaryExec, BoundaryLoopbackConnector, BoundaryProcess, BoundaryTakeover, + BoundaryTerminal, ConfirmedBoundary, ExecSession, NetworkMediationSource, + OpenShellRuntimeBackend, ProcessAttachment, ReadyBoundary, RunningBoundary, }; pub use runtime::{BoundaryTransport, BoundaryTransportConnector, validate_runtime_coordinates}; diff --git a/crates/openshell-sandbox-backend/src/protobuf.rs b/crates/openshell-sandbox-backend/src/protobuf.rs index 21c52788f1..c057679894 100644 --- a/crates/openshell-sandbox-backend/src/protobuf.rs +++ b/crates/openshell-sandbox-backend/src/protobuf.rs @@ -134,6 +134,14 @@ fn duration(micros: u64) -> Result { nanos: number((micros % 1_000_000) * 1000)?, }) } +/// A whole number of seconds that fits `u32`, such as an offered recovery +/// deadline. +fn whole_seconds(v: prost_types::Duration) -> Result { + if v.nanos != 0 { + return Err(invalid("duration is not a whole number of seconds")); + } + u32::try_from(v.seconds).map_err(|_| invalid("duration outside the accepted range")) +} fn micros(v: Option) -> Result { let v = v.unwrap_or_default(); if !(0..=315_576_000_000).contains(&v.seconds) @@ -385,11 +393,20 @@ impl WireMessage for w::SessionSnapshotWire { }) }) .collect::>()?, + bound_policy: self + .bound_policy + .as_ref() + .map(|policy| policy.to_proto()) + .transpose()?, }) } fn from_proto(v: Self::Proto) -> Result { Ok(Self { generation: v.generation, + bound_policy: v + .bound_policy + .map(|policy| w::SandboxPolicyWire::from_proto(policy).map(Box::new)) + .transpose()?, processes: v .processes .into_iter() @@ -441,9 +458,18 @@ impl WireMessage for w::RequestEnvelope { W::Attach { supervisor_instance_id, policy, + supervisor_takeover, + supervisor_recovery_timeout_secs, } => P::Attach(p::Attach { supervisor_instance_id: supervisor_instance_id.to_string(), policy: Some(policy.to_proto()?), + supervisor_takeover: supervisor_takeover.clone(), + supervisor_recovery_timeout: supervisor_recovery_timeout_secs.map(|seconds| { + prost_types::Duration { + seconds: i64::from(seconds), + nanos: 0, + } + }), }), W::Confirm => P::Confirm(p::Confirm {}), W::StartAgent { @@ -525,6 +551,11 @@ impl WireMessage for w::RequestEnvelope { P::Attach(v) => W::Attach { supervisor_instance_id: parse(&v.supervisor_instance_id)?, policy: Box::new(w::SandboxPolicyWire::from_proto(required(v.policy)?)?), + supervisor_takeover: v.supervisor_takeover, + supervisor_recovery_timeout_secs: v + .supervisor_recovery_timeout + .map(whole_seconds) + .transpose()?, }, P::Confirm(_) => W::Confirm, P::StartAgent(v) => W::StartAgent { @@ -794,6 +825,8 @@ mod tests { request: Some(p::request_envelope::Request::Attach(p::Attach { supervisor_instance_id: uuid::Uuid::new_v4().to_string(), policy: None, + supervisor_takeover: None, + supervisor_recovery_timeout: None, })), }; assert!(w::RequestEnvelope::from_proto(attach).is_err()); @@ -807,6 +840,53 @@ mod tests { assert!(w::RequestEnvelope::from_proto(signal).is_err()); } + #[test] + fn attach_carries_an_offered_recovery_deadline_only_when_present() { + for offered in [None, Some(60)] { + let request = w::RequestEnvelope::new(w::Request::Attach { + supervisor_instance_id: w::SupervisorInstanceId::new(), + policy: Box::new(w::SandboxPolicyWire::from( + openshell_core::policy::SandboxPolicy { + version: 1, + filesystem: openshell_core::policy::FilesystemPolicy::default(), + network: openshell_core::policy::NetworkPolicy::default(), + landlock: openshell_core::policy::LandlockPolicy::default(), + process: openshell_core::policy::ProcessPolicy::default(), + }, + )), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: offered, + }) + .unwrap(); + let encoded = request.to_proto().unwrap(); + let Some(p::request_envelope::Request::Attach(attach)) = &encoded.request else { + unreachable!() + }; + assert_eq!( + attach.supervisor_recovery_timeout, + offered.map(|seconds| prost_types::Duration { + seconds: i64::from(seconds), + nanos: 0, + }) + ); + let mut fractional = encoded.clone(); + let decoded = w::RequestEnvelope::from_proto(encoded).unwrap(); + assert!(matches!( + decoded.request, + w::Request::Attach { supervisor_recovery_timeout_secs, .. } + if supervisor_recovery_timeout_secs == offered + )); + // An offer that is not whole seconds is not a valid frame. + if let Some(p::request_envelope::Request::Attach(attach)) = &mut fractional.request { + attach.supervisor_recovery_timeout = Some(prost_types::Duration { + seconds: 60, + nanos: 1, + }); + } + assert!(w::RequestEnvelope::from_proto(fractional).is_err()); + } + } + #[test] fn exec_requires_the_original_valid_deadline() { let request = w::RequestEnvelope::new(w::Request::Exec { diff --git a/crates/openshell-sandbox-backend/src/runtime.rs b/crates/openshell-sandbox-backend/src/runtime.rs index 8e54aa9832..53cf10e703 100644 --- a/crates/openshell-sandbox-backend/src/runtime.rs +++ b/crates/openshell-sandbox-backend/src/runtime.rs @@ -64,6 +64,9 @@ const _: () = assert!( MAX_NETWORK_STREAMS < crate::boundary_protocol::BOUNDARY_MAX_CONCURRENT_STREAMS as usize ); const NETWORK_STREAM_WAIT_WARNING: Duration = Duration::from_secs(5); +/// Bound on waiting for a replaced owner's process attachment to be released +/// after a takeover retires its connection. +const ADOPTED_ATTACHMENT_RELEASE_TIMEOUT: Duration = Duration::from_secs(10); fn begin_recovery_window( deadline: &mut Option, @@ -193,6 +196,29 @@ impl BoundaryTransport { } } +/// Authority for this supervisor process to replace the owner of a running +/// generation: the identity the gateway granted it and the signed grant. +/// +/// The runtime stages the replacement at attach and promotes it at confirm; +/// a runtime that does not advertise +/// [`crate::isolation_protocol::SUPERVISOR_TAKEOVER_CAPABILITY`] is refused +/// before the grant is sent. +#[derive(Clone)] +pub struct BoundaryTakeover { + pub supervisor_instance_id: crate::boundary_protocol::SupervisorInstanceId, + pub grant: openshell_core::jwt::SecretJwt, +} + +impl std::fmt::Debug for BoundaryTakeover { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BoundaryTakeover") + .field("supervisor_instance_id", &self.supervisor_instance_id) + .field("grant", &self.grant) + .finish() + } +} + /// Host-side `OpenShell` Sandbox Protocol implementation registered with the supervisor. #[derive(Debug)] pub struct OpenShellRuntimeBackend { @@ -201,6 +227,8 @@ pub struct OpenShellRuntimeBackend { provider_credentials: openshell_core::provider_credentials::ProviderCredentialState, sandbox_bearer: openshell_core::jwt::SessionBearerTokenSlot, delegated: Option, + takeover: Option, + recovery_timeout_secs: Option, } #[derive(Clone)] @@ -242,6 +270,8 @@ impl OpenShellRuntimeBackend { provider_credentials, sandbox_bearer, delegated: None, + takeover: None, + recovery_timeout_secs: None, } } @@ -286,6 +316,25 @@ impl OpenShellRuntimeBackend { self } + /// Attach as the gateway-granted replacement of the generation's current + /// owner. The runtime's process snapshot then decides whether the agent + /// is adopted or, for a runtime that never started it, launched. + #[must_use] + pub fn with_takeover(mut self, takeover: BoundaryTakeover) -> Self { + self.takeover = Some(takeover); + self + } + + /// Offer the runtime a recovery deadline on every attach: how long it + /// keeps the workload frozen after losing this supervisor's connection + /// (see [`crate::boundary_protocol::SUPERVISOR_RECOVERY_TIMEOUT_RANGE`]). + /// `None` leaves the runtime's default. + #[must_use] + pub fn with_recovery_timeout_secs(mut self, seconds: Option) -> Self { + self.recovery_timeout_secs = seconds; + self + } + pub fn backend_name(&self) -> &str { self.delegated .as_ref() @@ -327,11 +376,22 @@ impl OpenShellRuntimeBackend { descriptor.payload.clone(), ); client.transport = self.transport.clone(); + if let Some(takeover) = &self.takeover { + client.supervisor_instance_id = takeover.supervisor_instance_id; + client.requires_takeover = true; + } + client.recovery_timeout_secs = self.recovery_timeout_secs; let client = Arc::new(client); + let policy = Box::new(SandboxPolicyWire::from(sandbox.policy.clone())); let response = client .call_idempotent(Request::Attach { supervisor_instance_id: client.supervisor_instance_id, - policy: Box::new(SandboxPolicyWire::from(sandbox.policy.clone())), + policy: policy.clone(), + supervisor_takeover: self + .takeover + .as_ref() + .map(|takeover| takeover.grant.expose_secret().to_string()), + supervisor_recovery_timeout_secs: client.recovery_timeout_secs, }) .await?; let Response::Attached { snapshot } = response else { @@ -342,6 +402,12 @@ impl OpenShellRuntimeBackend { "sandbox session snapshot generation does not match runtime descriptor".to_string(), )); } + let takeover = self.takeover.is_some().then(|| RecoveredSession { + // Later reconnects reattach without the grant, with the launch + // controls the runtime already enforces. + reattach_policy: snapshot.bound_policy.clone().unwrap_or(policy), + processes: RecoveredProcesses::from_snapshot(&snapshot), + }); Ok(Box::new(BoundBoundary { client: client.clone(), agent: sandbox.agent, @@ -355,10 +421,45 @@ impl OpenShellRuntimeBackend { session_id, resource_claims, provisioning_evidence, + takeover, })) } } +/// Processes a replacement supervisor found running in the runtime's +/// snapshot. They belong to the generation, not to the supervisor that +/// started them, so the replacement adopts them instead of launching again. +#[derive(Clone, Debug, Default)] +struct RecoveredProcesses { + /// The canonical process, when the previous owner started it. + main: Option, + /// Exec processes that have not exited. + execs: Vec, +} + +impl RecoveredProcesses { + fn from_snapshot(snapshot: &crate::boundary_protocol::SessionSnapshotWire) -> Self { + use crate::boundary_protocol::ProcessKindWire; + let mut recovered = Self::default(); + for process in &snapshot.processes { + match process.kind { + ProcessKindWire::Main => recovered.main = Some(process.process_id.clone()), + ProcessKindWire::Exec if process.status.is_none() => { + recovered.execs.push(process.process_id.clone()); + } + ProcessKindWire::Exec => {} + } + } + recovered + } +} + +/// What a takeover attach learned from the runtime. +struct RecoveredSession { + reattach_policy: Box, + processes: RecoveredProcesses, +} + fn validate_runtime_descriptor( runtime_descriptor: &SandboxRuntimeDescriptor, sandbox: &SandboxContext, @@ -465,13 +566,32 @@ pub struct BoundBoundary { session_id: openshell_core::SandboxSessionId, resource_claims: std::collections::BTreeMap, provisioning_evidence: openshell_isolation_interface::contract::ProvisioningEvidence, + takeover: Option, } impl BoundBoundary { + /// The runtime's network opens. A source taken before [`Self::confirm`] + /// waits for confirmation before accepting any, so a replacement can + /// load its enforcement before the frozen workload resumes. pub fn network_mediation_source(&self) -> Arc { self.mediation.clone() } + /// Whether this boundary was attached as a gateway-granted replacement. + #[must_use] + pub fn is_takeover(&self) -> bool { + self.takeover.is_some() + } + + /// Whether the runtime already runs the canonical process, which a + /// replacement adopts instead of starting. + #[must_use] + pub fn adopts_running_agent(&self) -> bool { + self.takeover + .as_ref() + .is_some_and(|takeover| takeover.processes.main.is_some()) + } + pub fn host_gateway_ip(&self) -> Option { self.host_gateway_ip } @@ -481,6 +601,16 @@ impl BoundBoundary { let Response::Confirmed { confirmation } = response else { return Err(unexpected_response("confirmed", &response)); }; + // A confirmed takeover pins this process at the granted epoch. From + // now on a reconnect reattaches without the grant, so its expiry + // cannot break transport recovery while confirmation is validated. + let reattach_policy = self + .takeover + .as_ref() + .map(|takeover| takeover.reattach_policy.clone()); + if let Some(policy) = &reattach_policy { + self.client.forget_takeover_grant(policy.clone())?; + } if confirmation.generation != self.generation || confirmation.session_id != self.session_id || confirmation.resource_claims != self.resource_claims @@ -496,6 +626,13 @@ impl BoundBoundary { crate::audit::validate_native_linux_audit(&confirmation.backend_audit)?; } let client = self.client.clone(); + // Adopt from a snapshot taken after the promotion: the replaced owner + // may have started processes after the takeover attach, and it can + // start none once this process is promoted. + let recovered = match reattach_policy { + Some(policy) => Some(client.refresh_recovered_processes(policy).await?), + None => None, + }; let confirmed = ConfirmedBoundary::try_new( Box::new(ReadyBoundary { client: self.client, @@ -503,6 +640,7 @@ impl BoundBoundary { sandbox_id: self.sandbox_id, ca_file_paths: self.ca_file_paths, provider_credentials: self.provider_credentials, + recovered, }), *confirmation, &self.identity, @@ -518,10 +656,48 @@ pub struct ReadyBoundary { sandbox_id: String, ca_file_paths: Arc>>, provider_credentials: openshell_core::provider_credentials::ProviderCredentialState, + /// Processes found running by a takeover attach. + recovered: Option, } impl ReadyBoundary { + /// Start the canonical process, or, for a replacement supervisor whose + /// runtime already runs it, adopt it and the running exec processes + /// without launching anything. pub async fn start_agent(self: Box) -> Result, BackendError> { + if let Some(RecoveredProcesses { + main: Some(process_id), + execs, + }) = &self.recovered + { + let adopted = |process_id: &String| { + Arc::new(BoundaryProcess { + exec: true, + adopted: true, + client: self.client.clone(), + process_id: process_id.clone(), + }) + }; + return Ok(Box::new(RunningBoundary { + process: Arc::new(BoundaryProcess { + exec: false, + adopted: true, + client: self.client.clone(), + process_id: process_id.clone(), + }), + adopted_execs: execs.iter().map(adopted).collect(), + exec: Arc::new(BoundaryExec { + client: self.client.clone(), + provider_credentials: self.provider_credentials, + // Publication restarts from the runtime's installed + // generation, which the first reconciliation learns. + publication_generation: tokio::sync::Mutex::new(0), + }), + loopback_connector: Arc::new(BoundaryLoopbackConnector { + client: self.client, + }), + })); + } let ca_paths = self .ca_file_paths .lock() @@ -571,11 +747,13 @@ impl ReadyBoundary { }; let process = Arc::new(BoundaryProcess { exec: false, + adopted: false, client: self.client.clone(), process_id, }); Ok(Box::new(RunningBoundary { process, + adopted_execs: Vec::new(), exec: Arc::new(BoundaryExec { client: self.client.clone(), provider_credentials: self.provider_credentials, @@ -590,6 +768,8 @@ impl ReadyBoundary { pub struct RunningBoundary { process: Arc, + /// Exec processes a replacement supervisor adopted from the runtime. + adopted_execs: Vec>, exec: Arc, loopback_connector: Arc, } @@ -599,6 +779,20 @@ impl RunningBoundary { self.process.clone() } + /// Whether the canonical process was adopted from a previous owner + /// rather than started by this supervisor. + #[must_use] + pub fn agent_adopted(&self) -> bool { + self.process.adopted + } + + /// Running exec processes adopted from a previous owner. Their original + /// clients are gone; the replacement observes their exit. + #[must_use] + pub fn adopted_exec_processes(&self) -> &[Arc] { + &self.adopted_execs + } + pub fn exec(&self) -> Arc { self.exec.clone() } @@ -619,13 +813,37 @@ impl RunningBoundary { pub struct BoundaryProcess { exec: bool, + /// Started by a previous owner of the generation. + adopted: bool, client: Arc, process_id: String, } impl BoundaryProcess { + /// Generation-scoped process ID. + #[must_use] + pub fn process_id(&self) -> &str { + &self.process_id + } + pub async fn attach(&self) -> Result { - open_process_attachment(self.client.clone(), self.process_id.clone()).await + if !self.adopted { + return open_process_attachment(self.client.clone(), self.process_id.clone()).await; + } + // The replaced owner's attachment ends with its retired connection; + // wait a bounded time for the runtime to release it. + let deadline = tokio::time::Instant::now() + ADOPTED_ATTACHMENT_RELEASE_TIMEOUT; + loop { + match open_process_attachment(self.client.clone(), self.process_id.clone()).await { + Err(BackendError::Denied(message)) + if message.contains("already has a control attachment") + && tokio::time::Instant::now() < deadline => + { + tokio::time::sleep(Duration::from_millis(100)).await; + } + result => return result, + } + } } pub async fn wait(&self) -> Result { @@ -903,6 +1121,7 @@ async fn open_exec_session( let process: Arc = Arc::new(BoundaryProcess { exec: true, + adopted: false, client: client.clone(), process_id: process_id.clone(), }); @@ -959,7 +1178,18 @@ pub struct NetworkMediationSource { } impl NetworkMediationSource { + /// Accepts reach the runtime only from a confirmed connection. + async fn wait_confirmed(&self) { + let _ = self + .client + .confirmed + .subscribe() + .wait_for(|confirmed| *confirmed) + .await; + } + pub async fn accept_tcp(&self) -> Result { + self.wait_confirmed().await; // An idle accept has no deadline. Recover transport loss before exposing // a source error, which permanently stops the proxy. The boundary drops // the interrupted pending open; a retry never replays its decision. @@ -1000,6 +1230,7 @@ impl NetworkMediationSource { } pub async fn accept_dns(&self) -> Result { + self.wait_confirmed().await; loop { let session = self.client.mediation_session().await?; match session.accept_dns().await { @@ -1197,6 +1428,13 @@ struct BoundaryClient { credential_monitor_started: AtomicBool, network_streams: Arc, network_streams_exhausted: AtomicBool, + /// This client attaches with a takeover grant; the runtime must + /// advertise takeover support before any request is sent. + requires_takeover: bool, + /// Recovery deadline offered on every attach. + recovery_timeout_secs: Option, + /// True once a connection of this client was confirmed. + confirmed: tokio::sync::watch::Sender, } #[derive(Clone)] @@ -1390,9 +1628,54 @@ impl BoundaryClient { credential_monitor_started: AtomicBool::new(false), network_streams: Arc::new(tokio::sync::Semaphore::new(MAX_NETWORK_STREAMS)), network_streams_exhausted: AtomicBool::new(false), + requires_takeover: false, + recovery_timeout_secs: None, + confirmed: tokio::sync::watch::Sender::new(false), } } + /// Reattach on the confirmed connection as the promoted owner and read the + /// processes to adopt from that snapshot. + async fn refresh_recovered_processes( + &self, + policy: Box, + ) -> Result { + let response = self + .call_idempotent(Request::Attach { + supervisor_instance_id: self.supervisor_instance_id, + policy, + supervisor_takeover: None, + supervisor_recovery_timeout_secs: self.recovery_timeout_secs, + }) + .await?; + let Response::Attached { snapshot } = response else { + return Err(unexpected_response("attached", &response)); + }; + if snapshot.generation != self.runtime_descriptor.generation { + return Err(BackendError::Confirm( + "sandbox session snapshot generation does not match runtime descriptor".to_string(), + )); + } + Ok(RecoveredProcesses::from_snapshot(&snapshot)) + } + + /// Replace the remembered takeover attach with an ordinary attach of the + /// now-pinned process, so transport recovery never depends on the + /// grant's lifetime. + fn forget_takeover_grant(&self, policy: Box) -> Result<(), BackendError> { + let envelope = Self::prepare_request(Request::Attach { + supervisor_instance_id: self.supervisor_instance_id, + policy, + supervisor_takeover: None, + supervisor_recovery_timeout_secs: self.recovery_timeout_secs, + })?; + *self + .attach_request + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(envelope); + Ok(()) + } + async fn acquire_network_stream(&self) -> tokio::sync::OwnedSemaphorePermit { let acquire = self.network_streams.clone().acquire_owned(); tokio::pin!(acquire); @@ -1442,6 +1725,7 @@ impl BoundaryClient { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(envelope.clone()); + self.confirmed.send_replace(true); } return Ok(response); } @@ -2091,8 +2375,19 @@ impl BoundaryClient { .await .map_err(control_rpc_error)? .into_inner(); - crate::isolation_protocol::negotiate(&supervisor, response.backend) - .map_err(control_rpc_error)?; + let backend = + crate::isolation_protocol::negotiate_capabilities(&supervisor, response.backend) + .map_err(control_rpc_error)?; + if self.requires_takeover + && !backend.supported_capabilities.iter().any(|capability| { + capability == crate::isolation_protocol::SUPERVISOR_TAKEOVER_CAPABILITY + }) + { + return Err(BackendError::Unsupported( + "the sandbox runtime does not support supervisor takeover; refusing to resume it" + .to_string(), + )); + } let mut request = tonic::Request::new(OpenBoundaryRequest { session_id: self.runtime_descriptor.session_id.to_string(), backend_name: self.backend_name.clone(), @@ -2742,6 +3037,14 @@ mod tests { mediation_ready: bool, provider_environment_generation: u64, confirmation: BoundaryConfirmation, + /// Processes reported by successive attaches; empty once drained. + attach_snapshots: Arc< + std::sync::Mutex< + std::collections::VecDeque>, + >, + >, + /// Every request kind received, in order. + received: Arc>>, } type TestGrpcStream = @@ -2768,6 +3071,8 @@ mod tests { let mediation_ready = self.mediation_ready; let provider_environment_generation = self.provider_environment_generation; let confirmation = self.confirmation.clone(); + let attach_snapshots = self.attach_snapshots.clone(); + let received = self.received.clone(); let (outbound, outbound_rx) = tokio::sync::mpsc::channel(1); tokio::spawn(async move { let mut frame = Vec::new(); @@ -2801,6 +3106,25 @@ mod tests { return; } }; + received.lock().unwrap().push(match &envelope.request { + Request::Attach { + supervisor_takeover: Some(_), + supervisor_recovery_timeout_secs: Some(_), + .. + } => "takeover-attach offering recovery", + Request::Attach { + supervisor_takeover: Some(_), + .. + } => "takeover-attach", + Request::Attach { + supervisor_recovery_timeout_secs: Some(_), + .. + } => "attach offering recovery", + Request::Attach { .. } => "attach", + Request::Confirm => "confirm", + Request::StartAgent { .. } => "start-agent", + _ => "other", + }); match encode_frame(&ResponseEnvelope { request_id: envelope.request_id, response: match envelope.request { @@ -2808,12 +3132,20 @@ mod tests { yaml: None, invalid: false, }, - Request::Attach { .. } => Response::Attached { - snapshot: crate::boundary_protocol::SessionSnapshotWire { - generation: "test-generation".to_string(), - processes: Vec::new(), - }, - }, + Request::Attach { .. } => { + let processes = attach_snapshots + .lock() + .unwrap() + .pop_front() + .unwrap_or_default(); + Response::Attached { + snapshot: crate::boundary_protocol::SessionSnapshotWire { + generation: "test-generation".to_string(), + processes, + bound_policy: None, + }, + } + } Request::Confirm => Response::Confirmed { confirmation: Box::new(confirmation), }, @@ -2944,6 +3276,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); @@ -3018,6 +3352,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; tokio::spawn(async move { tonic::transport::Server::builder() @@ -3038,6 +3374,8 @@ mod tests { let attach = Request::Attach { supervisor_instance_id: client.supervisor_instance_id, policy: Box::new(SandboxPolicyWire::from(sandbox().policy)), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, }; assert!(matches!( client.call_idempotent(attach).await.unwrap(), @@ -3129,6 +3467,8 @@ mod tests { mediation_ready: true, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; tokio::spawn(async move { tonic::transport::Server::builder() @@ -3149,6 +3489,8 @@ mod tests { let attach = Request::Attach { supervisor_instance_id: client.supervisor_instance_id, policy: Box::new(SandboxPolicyWire::from(sandbox().policy)), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, }; assert!(matches!( client.call_idempotent(attach).await.unwrap(), @@ -3184,6 +3526,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); @@ -3383,6 +3727,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 0, confirmation, + attach_snapshots: Arc::default(), + received: Arc::default(), }; tonic::transport::Server::builder() .add_service(IsolationBackendServer::new(TestIsolationBackend::default())) @@ -3511,6 +3857,7 @@ mod tests { session_id: test_session_id(), resource_claims: std::collections::BTreeMap::new(), provisioning_evidence, + takeover: None, }; assert!(matches!( @@ -3901,6 +4248,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; let server = tokio::spawn(async move { loop { @@ -3974,6 +4323,8 @@ mod tests { mediation_ready: false, provider_environment_generation: 50, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }; let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); @@ -4005,6 +4356,7 @@ mod tests { sandbox_id: context.sandbox_id, ca_file_paths: Arc::new(std::sync::Mutex::new(None)), provider_credentials: credentials.clone(), + recovered: None, }); let running = ready.start_agent().await.unwrap(); assert_eq!(requests.load(Ordering::Acquire), 1); @@ -4027,6 +4379,316 @@ mod tests { let _ = server.await; } + #[tokio::test] + async fn adopted_boundary_never_restarts_the_agent_and_catches_up_publication() { + let certificate = test_certificate(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let service = TestGrpcBoundary { + wait_for_half_close: false, + expected_token: "a".repeat(32), + requests: requests.clone(), + mediation_failures: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + mediation_ready: false, + provider_environment_generation: 50, + confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), + }; + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let stream = tokio_rustls::TlsAcceptor::from(certificate.server_config) + .accept(stream) + .await + .unwrap(); + tonic::transport::Server::builder() + .add_service(IsolationBackendServer::new(TestIsolationBackend::default())) + .add_service(SandboxProtocolServer::new(service)) + .serve_with_incoming(tokio_stream::iter([Ok::<_, std::io::Error>(TestTlsIo( + Box::new(stream), + ))])) + .await + .unwrap(); + }); + let snapshot = crate::boundary_protocol::SessionSnapshotWire { + generation: "test-generation".to_string(), + bound_policy: None, + processes: vec![ + crate::boundary_protocol::ProcessSnapshotWire { + process_id: "test-generation:main:0".to_string(), + kind: crate::boundary_protocol::ProcessKindWire::Main, + terminal: false, + status: None, + retained_output: crate::boundary_protocol::OutputWindowWire { + first_sequence: 0, + next_sequence: 9, + truncated: false, + }, + }, + crate::boundary_protocol::ProcessSnapshotWire { + process_id: "test-generation:exec:3".to_string(), + kind: crate::boundary_protocol::ProcessKindWire::Exec, + terminal: false, + status: None, + retained_output: crate::boundary_protocol::OutputWindowWire { + first_sequence: 0, + next_sequence: 0, + truncated: false, + }, + }, + crate::boundary_protocol::ProcessSnapshotWire { + process_id: "test-generation:exec:1".to_string(), + kind: crate::boundary_protocol::ProcessKindWire::Exec, + terminal: false, + status: Some(ExitStatusWire::Exited(0)), + retained_output: crate::boundary_protocol::OutputWindowWire { + first_sequence: 0, + next_sequence: 4, + truncated: false, + }, + }, + ], + }; + let context = sandbox(); + let ready = Box::new(ReadyBoundary { + client: Arc::new(BoundaryClient::new( + tls_runtime_descriptor(address, certificate.client_tls), + test_bearer(&"a".repeat(32)), + )), + agent: context.agent, + sandbox_id: context.sandbox_id, + ca_file_paths: Arc::new(std::sync::Mutex::new(None)), + provider_credentials: + openshell_core::provider_credentials::ProviderCredentialState::from_child_env_snapshot( + 6, + HashMap::new(), + ), + recovered: Some(RecoveredProcesses::from_snapshot(&snapshot)), + }); + let running = ready.start_agent().await.unwrap(); + // The running canonical process is adopted: nothing is started. + assert_eq!(requests.load(Ordering::Acquire), 0); + assert!(running.agent_adopted()); + assert_eq!(running.agent().process_id(), "test-generation:main:0"); + // Only the exec process that has not exited is adopted. + assert_eq!( + running + .adopted_exec_processes() + .iter() + .map(|process| process.process_id()) + .collect::>(), + ["test-generation:exec:3"] + ); + // Publication resumes above the runtime's installed generation. + running + .exec() + .synchronize_provider_environment() + .await + .unwrap(); + assert_eq!(requests.load(Ordering::Acquire), 2); + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn takeover_confirm_adopts_from_a_snapshot_taken_after_promotion() { + let certificate = test_certificate(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let process = |process_id: &str, kind| crate::boundary_protocol::ProcessSnapshotWire { + process_id: process_id.to_string(), + kind, + terminal: false, + status: None, + retained_output: crate::boundary_protocol::OutputWindowWire { + first_sequence: 0, + next_sequence: 0, + truncated: false, + }, + }; + // The staged snapshot predates an exec the previous owner started + // before the promotion; only the refreshed snapshot shows it. + let attach_snapshots = Arc::new(std::sync::Mutex::new(std::collections::VecDeque::from([ + vec![process( + "test-generation:main:0", + crate::boundary_protocol::ProcessKindWire::Main, + )], + vec![ + process( + "test-generation:main:0", + crate::boundary_protocol::ProcessKindWire::Main, + ), + process( + "test-generation:exec:4", + crate::boundary_protocol::ProcessKindWire::Exec, + ), + ], + ]))); + let received = Arc::new(std::sync::Mutex::new(Vec::new())); + let service = TestGrpcBoundary { + wait_for_half_close: false, + expected_token: "a".repeat(32), + requests: Arc::default(), + mediation_failures: Arc::default(), + mediation_ready: false, + provider_environment_generation: 0, + confirmation: test_confirmation(), + attach_snapshots, + received: received.clone(), + }; + let server_config = certificate.server_config.clone(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let stream = tokio_rustls::TlsAcceptor::from(server_config) + .accept(stream) + .await + .unwrap(); + tonic::transport::Server::builder() + .add_service(IsolationBackendServer::new(TestIsolationBackend::default())) + .add_service(SandboxProtocolServer::new(service)) + .serve_with_incoming(tokio_stream::iter([Ok::<_, std::io::Error>(TestTlsIo( + Box::new(stream), + ))])) + .await + .unwrap(); + }); + let descriptor = tls_runtime_descriptor(address, certificate.client_tls); + let instance = crate::boundary_protocol::SupervisorInstanceId::new(); + let backend = OpenShellRuntimeBackend::new( + Arc::new(std::sync::Mutex::new(None)), + openshell_core::provider_credentials::ProviderCredentialState::from_child_env_snapshot( + 0, + HashMap::new(), + ), + test_bearer(&"a".repeat(32)), + ) + .with_takeover(BoundaryTakeover { + supervisor_instance_id: instance, + grant: SecretJwt::parse("test-grant").unwrap(), + }) + .with_recovery_timeout_secs(Some(60)); + let bound = backend + .attach(descriptor.backend_descriptor().unwrap(), sandbox()) + .await + .unwrap(); + assert!(bound.is_takeover()); + assert!(bound.adopts_running_agent()); + let running = bound + .confirm() + .await + .unwrap() + .into_boundary() + .start_agent() + .await + .unwrap(); + assert!(running.agent_adopted()); + assert_eq!(running.agent().process_id(), "test-generation:main:0"); + assert_eq!( + running + .adopted_exec_processes() + .iter() + .map(|process| process.process_id()) + .collect::>(), + ["test-generation:exec:4"] + ); + // The grant goes out once; the refresh reattaches as the pinned + // owner, and the agent is never started again. Every attach offers + // the recovery deadline. + assert_eq!( + *received.lock().unwrap(), + [ + "takeover-attach offering recovery", + "confirm", + "attach offering recovery" + ] + ); + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn takeover_grants_reach_only_runtimes_that_support_takeover() { + async fn open( + metadata: openshell_core::proto::extension::v1::PeerMetadata, + ) -> (Result<(), BackendError>, usize, BoundaryClient) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let opens = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let backend = TestIsolationBackend { + metadata: Some(metadata), + opens: opens.clone(), + ..TestIsolationBackend::default() + }; + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + tonic::transport::Server::builder() + .add_service(IsolationBackendServer::new(backend)) + .serve_with_incoming(tokio_stream::iter([Ok::<_, std::io::Error>(stream)])) + .await + .unwrap(); + }); + let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) + .unwrap() + .connect() + .await + .unwrap(); + let mut client = BoundaryClient::new( + tls_runtime_descriptor(address, test_certificate().client_tls), + test_bearer(&"a".repeat(32)), + ); + client.requires_takeover = true; + let credential = BoundaryCredential::capture(&client.sandbox_bearer).unwrap(); + let opened = client.open_boundary(channel, &credential).await; + server.abort(); + (opened, opens.load(Ordering::Acquire), client) + } + + // A runtime from before takeover support never sees the boundary + // opened, let alone the grant. + let mut legacy = crate::isolation_protocol::metadata("legacy/runtime"); + legacy.supported_capabilities.retain(|capability| { + capability != crate::isolation_protocol::SUPERVISOR_TAKEOVER_CAPABILITY + }); + let (opened, opens, _) = open(legacy).await; + assert!( + matches!(opened, Err(BackendError::Unsupported(_))), + "{opened:?}" + ); + assert_eq!(opens, 0); + + let (opened, opens, mut client) = + open(crate::isolation_protocol::metadata("openshell/sandbox")).await; + opened.expect("a takeover-capable runtime opens"); + assert_eq!(opens, 1); + client.recovery_timeout_secs = Some(60); + + // Once confirmed, reconnects reattach as the pinned process without + // the grant, so its expiry cannot break transport recovery. + let policy = Box::new(SandboxPolicyWire::from(sandbox().policy)); + *client.attach_request.lock().unwrap() = Some( + BoundaryClient::prepare_request(Request::Attach { + supervisor_instance_id: client.supervisor_instance_id, + policy: policy.clone(), + supervisor_takeover: Some("grant".to_string()), + supervisor_recovery_timeout_secs: Some(60), + }) + .unwrap(), + ); + client.forget_takeover_grant(policy).unwrap(); + let remembered = client.attach_request.lock().unwrap().clone().unwrap(); + // The pinned process keeps offering its recovery deadline. + assert!(matches!( + remembered.request, + Request::Attach { + supervisor_takeover: None, + supervisor_instance_id, + supervisor_recovery_timeout_secs: Some(60), + .. + } if supervisor_instance_id == client.supervisor_instance_id + )); + } + #[tokio::test] async fn provider_environment_acknowledges_the_exact_local_snapshot_and_checked_generation() { let certificate = test_certificate(); diff --git a/crates/openshell-sandbox-backend/src/runtime/test_support.rs b/crates/openshell-sandbox-backend/src/runtime/test_support.rs index 6e94701548..430d8b79b6 100644 --- a/crates/openshell-sandbox-backend/src/runtime/test_support.rs +++ b/crates/openshell-sandbox-backend/src/runtime/test_support.rs @@ -728,6 +728,9 @@ fn client(peer: Peer) -> Arc { task, })); client.boundary_handle.set("test".into()).unwrap(); + // Fixtures stand for an attached, confirmed boundary, whose mediation + // sources accept at once. + client.confirmed.send_replace(true); Arc::new(client) } pub fn process(process: Arc) -> Arc { @@ -738,6 +741,7 @@ pub fn process(process: Arc) -> Arc }), process_id: "main".into(), exec: false, + adopted: false, }) } pub fn terminal(terminal: Arc) -> Arc { diff --git a/crates/openshell-sandbox-backend/src/runtime/tests/credential_renewal.rs b/crates/openshell-sandbox-backend/src/runtime/tests/credential_renewal.rs index 0c5320a46b..a8f1137be7 100644 --- a/crates/openshell-sandbox-backend/src/runtime/tests/credential_renewal.rs +++ b/crates/openshell-sandbox-backend/src/runtime/tests/credential_renewal.rs @@ -69,6 +69,8 @@ async fn renewing_boundary_client( mediation_ready: false, provider_environment_generation: 0, confirmation: test_confirmation(), + attach_snapshots: Arc::default(), + received: Arc::default(), }, expected_token: Arc::new(std::sync::RwLock::new("a".repeat(32))), failures: Arc::new(std::sync::atomic::AtomicUsize::new(0)), @@ -108,6 +110,8 @@ async fn renewing_boundary_client( .call_idempotent(Request::Attach { supervisor_instance_id: client.supervisor_instance_id, policy: Box::new(SandboxPolicyWire::from(sandbox().policy)), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, }) .await .unwrap(); diff --git a/crates/openshell-sandbox-backend/src/runtime/tests/delegated.rs b/crates/openshell-sandbox-backend/src/runtime/tests/delegated.rs index 86af9ab75e..6275eff8bd 100644 --- a/crates/openshell-sandbox-backend/src/runtime/tests/delegated.rs +++ b/crates/openshell-sandbox-backend/src/runtime/tests/delegated.rs @@ -34,6 +34,8 @@ async fn external_backend( mediation_ready: false, provider_environment_generation: 0, confirmation, + attach_snapshots: Arc::default(), + received: Arc::default(), }; tonic::transport::Server::builder() .add_service(IsolationBackendServer::new(TestIsolationBackend { diff --git a/crates/openshell-sandbox-backend/src/runtime/tests/network_recovery.rs b/crates/openshell-sandbox-backend/src/runtime/tests/network_recovery.rs index 0e550b3032..2a76a307fb 100644 --- a/crates/openshell-sandbox-backend/src/runtime/tests/network_recovery.rs +++ b/crates/openshell-sandbox-backend/src/runtime/tests/network_recovery.rs @@ -135,6 +135,7 @@ impl NetworkPeer { snapshot: SessionSnapshotWire { generation: "test-generation".into(), processes: Vec::new(), + bound_policy: None, }, } } @@ -345,6 +346,8 @@ impl Fixture { .call_idempotent(Request::Attach { supervisor_instance_id: client.supervisor_instance_id, policy: Box::new(SandboxPolicyWire::from(sandbox().policy)), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, }) .await .unwrap(); diff --git a/crates/openshell-sandbox-backend/src/runtime_routed.rs b/crates/openshell-sandbox-backend/src/runtime_routed.rs index e2da5a6742..8f238ebfd2 100644 --- a/crates/openshell-sandbox-backend/src/runtime_routed.rs +++ b/crates/openshell-sandbox-backend/src/runtime_routed.rs @@ -15,6 +15,16 @@ //! replaced handles are closed without a handshake, and authenticated //! connections only fill the target instance's bounded spare queue, so no //! peer can displace a live channel. +//! +//! In a supervisor pool one member hosts each route, but a runtime reaches +//! an arbitrary member. A version 2 preface is answered before any TLS: the +//! member that hosts the route accepts, naming itself, and the TLS handshake +//! follows, while any other member names the owner from its gateway-provided +//! route directory, or reports the route unknown, and closes. The directory +//! wins over a local registration it no longer agrees with, which is being +//! released. Replies are hints like the handle itself: the owner must still +//! complete the pinned TLS handshake, and every request still carries the +//! bearer token. use std::collections::HashMap; use std::net::SocketAddr; @@ -23,18 +33,27 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use openshell_isolation_interface::contract::{BackendError, BoundaryDuplexStream}; -use tokio::io::{AsyncRead, AsyncReadExt as _}; +use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use crate::boundary_protocol::SandboxTlsClientConfig; /// First bytes of a routing preface. pub const ROUTE_PREFACE_MAGIC: [u8; 4] = *b"OSRP"; -/// Routing preface layout version. +/// Routing preface layout version that receives no reply. pub const ROUTE_PREFACE_VERSION: u8 = 1; +/// Routing preface layout version that the supervisor answers with a +/// [`RouteReply`] before the TLS handshake. +pub const ROUTE_PREFACE_VERSION_V2: u8 = 2; /// Longest accepted route handle, in bytes. pub const MAX_ROUTE_HANDLE_LEN: usize = 64; const MIN_ROUTE_HANDLE_LEN: usize = 16; +/// Status bytes of a version 2 reply. An accept and a redirect are followed +/// by a length byte and a runtime listener address as ASCII `ip:port`: the +/// accepting member's own (empty while it knows no directory) or the owner's. +const ROUTE_REPLY_ACCEPT: u8 = 0x01; +const ROUTE_REPLY_REDIRECT: u8 = 0x02; +const ROUTE_REPLY_UNKNOWN: u8 = 0x03; /// Unauthenticated sockets the listener handshakes at once, across sandbox /// instances. @@ -72,26 +91,38 @@ pub fn new_route_handle() -> String { /// Encode the routing preface a runtime sends before its TLS handshake. pub fn encode_route_preface(handle: &str) -> Result, BackendError> { + encode_versioned_preface(ROUTE_PREFACE_VERSION, handle) +} + +/// Encode a version 2 routing preface, which the supervisor answers before +/// the TLS handshake (see [`negotiate_route`]). +pub fn encode_route_preface_v2(handle: &str) -> Result, BackendError> { + encode_versioned_preface(ROUTE_PREFACE_VERSION_V2, handle) +} + +fn encode_versioned_preface(version: u8, handle: &str) -> Result, BackendError> { validate_route_handle(handle)?; let length = u8::try_from(handle.len()) .map_err(|_| BackendError::Descriptor("runtime route handle is too long".to_string()))?; let mut preface = Vec::with_capacity(ROUTE_PREFACE_MAGIC.len() + 2 + handle.len()); preface.extend_from_slice(&ROUTE_PREFACE_MAGIC); - preface.push(ROUTE_PREFACE_VERSION); + preface.push(version); preface.push(length); preface.extend_from_slice(handle.as_bytes()); Ok(preface) } -/// Read and validate one routing preface. Reads only the declared bytes. -async fn read_route_preface(stream: &mut S) -> std::io::Result { +/// Read and validate one routing preface, returning its version and route +/// handle. Reads only the declared bytes. +async fn read_route_preface(stream: &mut S) -> std::io::Result<(u8, String)> { let invalid = |message: &str| std::io::Error::new(std::io::ErrorKind::InvalidData, message); let mut header = [0_u8; 6]; stream.read_exact(&mut header).await?; if header[..4] != ROUTE_PREFACE_MAGIC { return Err(invalid("missing runtime routing preface")); } - if header[4] != ROUTE_PREFACE_VERSION { + let version = header[4]; + if version != ROUTE_PREFACE_VERSION && version != ROUTE_PREFACE_VERSION_V2 { return Err(invalid("unsupported runtime routing preface version")); } let length = usize::from(header[5]); @@ -102,7 +133,90 @@ async fn read_route_preface(stream: &mut S) -> std::io::Re stream.read_exact(&mut handle).await?; let handle = String::from_utf8(handle).map_err(|_| invalid("invalid runtime route handle"))?; validate_route_handle(&handle).map_err(|error| invalid(&error.to_string()))?; - Ok(handle) + Ok((version, handle)) +} + +/// A supervisor's answer to a version 2 routing preface. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RouteReply { + /// This supervisor hosts the route; the TLS handshake follows on the same + /// connection. Names the member's runtime listener as its route directory + /// does, when it has one; unauthenticated, like everything before TLS. + Accept(Option), + /// Another pool member hosts the route, at this runtime listener address. + /// Unauthenticated, like the handle that selected it. + Redirect(SocketAddr), + /// No pool member is known to host the route. + Unknown, +} + +/// Name `handle` in a version 2 routing preface and read the reply. +/// +/// Only the reply's declared bytes are read, so after +/// [`RouteReply::Accept`] the stream is positioned at the TLS handshake. The +/// caller bounds the exchange in time. +pub async fn negotiate_route(stream: &mut S, handle: &str) -> std::io::Result +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let invalid = |message: &str| std::io::Error::new(std::io::ErrorKind::InvalidData, message); + let preface = encode_route_preface_v2(handle).map_err(|error| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, error.to_string()) + })?; + stream.write_all(&preface).await?; + stream.flush().await?; + match stream.read_u8().await? { + ROUTE_REPLY_ACCEPT => { + let member = read_reply_address(stream).await?; + Ok(RouteReply::Accept(member)) + } + ROUTE_REPLY_UNKNOWN => Ok(RouteReply::Unknown), + ROUTE_REPLY_REDIRECT => read_reply_address(stream) + .await? + .map(RouteReply::Redirect) + .ok_or_else(|| invalid("invalid runtime route redirect")), + _ => Err(invalid("unknown runtime route reply")), + } +} + +/// Read a reply's length-prefixed listener address; empty reads as `None`. +async fn read_reply_address( + stream: &mut S, +) -> std::io::Result> { + let mut address = vec![0_u8; usize::from(stream.read_u8().await?)]; + if address.is_empty() { + return Ok(None); + } + stream.read_exact(&mut address).await?; + std::str::from_utf8(&address) + .ok() + .and_then(|address| address.parse().ok()) + .map(Some) + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "invalid runtime route reply address", + ) + }) +} + +/// Encode a version 2 reply carrying a listener address. +fn address_reply(status: u8, address: Option) -> Result, String> { + let address = address + .map(|address| address.to_string()) + .unwrap_or_default(); + let length = u8::try_from(address.len()) + .map_err(|_| "runtime listener address is too long".to_string())?; + let mut reply = Vec::with_capacity(2 + address.len()); + reply.push(status); + reply.push(length); + reply.extend_from_slice(address.as_bytes()); + Ok(reply) +} + +/// Encode a version 2 redirect to `owner`. +fn redirect_reply(owner: SocketAddr) -> Result, String> { + address_reply(ROUTE_REPLY_REDIRECT, Some(owner)) } struct Route { @@ -117,6 +231,51 @@ struct Route { struct Routes { next_id: u64, by_handle: HashMap, + /// This listener's address, as the route directory names it. + own_address: Option, + /// Runtime listener of the pool member that hosts each route. + owners: HashMap, +} + +/// Where a routing preface's handle leads. +enum Destination { + /// A sandbox instance registered on this listener, which the directory + /// names as `own`. + Local { + tls: SandboxTlsClientConfig, + spares: mpsc::Sender, + pending: Arc, + own: Option, + }, + /// Another pool member's runtime listener. + Owner(SocketAddr), + Unknown, +} + +impl Routes { + /// The directory never sends a runtime back to this listener. For a + /// runtime that follows redirects, the directory also wins over a local + /// registration naming another owner: the gateway moved the route, and + /// the registration is being released. A runtime that cannot follow one + /// reaches the registration, as before pools existed. + fn destination(&self, handle: &str, follows_redirects: bool) -> Destination { + let owner = self + .owners + .get(handle) + .copied() + .filter(|owner| Some(*owner) != self.own_address); + match (self.by_handle.get(handle), owner) { + (Some(_), Some(owner)) if follows_redirects => Destination::Owner(owner), + (Some(route), _) => Destination::Local { + tls: route.tls.clone(), + spares: route.spares.clone(), + pending: route.pending.clone(), + own: self.own_address, + }, + (None, Some(owner)) => Destination::Owner(owner), + (None, None) => Destination::Unknown, + } + } } /// Shared listener that routes runtime connections to registered sandbox @@ -202,6 +361,18 @@ impl RuntimeRouter { connected: AtomicBool::new(false), })) } + + /// Replace the route directory of a supervisor pool: the runtime listener + /// of the member that hosts each route. A runtime whose version 2 preface + /// names a route not registered here is redirected to that member. `own` + /// is this listener's address as the directory names it; an entry naming + /// it reports the route unknown instead, so a runtime is never redirected + /// back here. + pub fn set_directory(&self, own: Option, owners: HashMap) { + let mut routes = lock(&self.routes); + routes.own_address = own; + routes.owners = owners; + } } impl Drop for RuntimeRouter { @@ -237,9 +408,9 @@ impl RouteRegistration { &self.handle } - /// Wait for one authenticated connection for `descriptor`'s sandbox - /// instance. - pub(crate) async fn accept( + /// Wait for one authenticated connection for this sandbox instance, + /// whose generation-pinned identity `tls` must name. + pub async fn accept( &self, tls: &SandboxTlsClientConfig, ) -> Result { @@ -340,26 +511,46 @@ async fn route_connection( _listener_permit: OwnedSemaphorePermit, ) -> Result<(), String> { let deadline = tokio::time::Instant::now() + HANDSHAKE_TIMEOUT; - let handle = tokio::time::timeout_at(deadline, read_route_preface(&mut stream)) + let (version, handle) = tokio::time::timeout_at(deadline, read_route_preface(&mut stream)) .await .map_err(|_| "routing preface timed out".to_string())? .map_err(|error| format!("routing preface: {error}"))?; - // The handle only selects which pinned identity to expect. - let (tls, spares, route_pending) = { - let routes = lock(routes); - let route = routes - .by_handle - .get(&handle) - .ok_or_else(|| "unknown or retired runtime route handle".to_string())?; - ( - route.tls.clone(), - route.spares.clone(), - route.pending.clone(), - ) + // Only a version 2 runtime reads a reply. + let negotiated = version == ROUTE_PREFACE_VERSION_V2; + // The handle only selects which pinned identity to expect, or which pool + // member to name. + let destination = lock(routes).destination(&handle, negotiated); + let (tls, spares, route_pending, own) = match destination { + Destination::Local { + tls, + spares, + pending, + own, + } => (tls, spares, pending, own), + Destination::Owner(owner) if negotiated => { + send_reply(&mut stream, deadline, &redirect_reply(owner)?).await?; + tracing::debug!(%owner, "redirected routed runtime connection to the route's owner"); + return Ok(()); + } + Destination::Owner(_) | Destination::Unknown => { + if negotiated { + send_reply(&mut stream, deadline, &[ROUTE_REPLY_UNKNOWN]).await?; + } + return Err("unknown or retired runtime route handle".to_string()); + } }; let Ok(_route_permit) = route_pending.try_acquire_owned() else { return Err("runtime route handshake limit reached".to_string()); }; + // Accept only once the TLS handshake can proceed. + if negotiated { + send_reply( + &mut stream, + deadline, + &address_reply(ROUTE_REPLY_ACCEPT, own)?, + ) + .await?; + } let stream: BoundaryDuplexStream = Box::new(stream); let authenticated = tokio::time::timeout_at( deadline, @@ -375,11 +566,22 @@ async fn route_connection( .map_err(|_| "runtime route has no free spare slot".to_string()) } +/// Write one version 2 reply within the handshake deadline. +async fn send_reply( + stream: &mut tokio::net::TcpStream, + deadline: tokio::time::Instant, + reply: &[u8], +) -> Result<(), String> { + tokio::time::timeout_at(deadline, stream.write_all(reply)) + .await + .map_err(|_| "routing reply timed out".to_string())? + .map_err(|error| format!("routing reply: {error}")) +} + #[cfg(test)] mod tests { use super::*; use crate::boundary_protocol::{SandboxTlsMaterial, generate_sandbox_tls_material}; - use tokio::io::AsyncWriteExt as _; /// A runtime: dial the router, name `route`, and serve `identity`'s TLS. async fn runtime_dial( @@ -391,6 +593,20 @@ mod tests { stream .write_all(&encode_route_preface(route).unwrap()) .await?; + serve_tls(stream, identity).await + } + + /// A redirect-aware runtime: negotiate `route` with the router. + async fn negotiate(router: SocketAddr, route: &str) -> (tokio::net::TcpStream, RouteReply) { + let mut stream = tokio::net::TcpStream::connect(router).await.unwrap(); + let reply = negotiate_route(&mut stream, route).await.unwrap(); + (stream, reply) + } + + async fn serve_tls( + stream: tokio::net::TcpStream, + identity: &SandboxTlsMaterial, + ) -> std::io::Result> { let mut config = rustls::ServerConfig::builder_with_provider( openshell_crypto::tls::configuration_provider(), ) @@ -475,23 +691,178 @@ mod tests { } #[tokio::test] - async fn preface_round_trips_and_rejects_malformed_input() { - let handle = new_route_handle(); - let encoded = encode_route_preface(&handle).unwrap(); + async fn negotiated_runtimes_are_accepted_redirected_or_refused_by_route_directory() { + let router = RuntimeRouter::bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let address = router.local_addr(); + let identity = + generate_sandbox_tls_material(openshell_core::SandboxSessionId::new()).unwrap(); + let tls = client_tls(&identity); + let (hosted, elsewhere, looped) = + (new_route_handle(), new_route_handle(), new_route_handle()); + let registration = router.register(&hosted, &tls).unwrap(); + let owner: SocketAddr = "10.0.0.2:7443".parse().unwrap(); + router.set_directory( + Some(address), + HashMap::from([ + (hosted.clone(), owner), + (elsewhere.clone(), owner), + (looped.clone(), address), + ]), + ); + + // A registration the directory moved elsewhere is being released: + // a runtime that follows redirects goes to the new owner, and one + // that cannot still reaches the registration. assert_eq!( - read_route_preface(&mut encoded.as_slice()).await.unwrap(), - handle + negotiate(address, &hosted).await.1, + RouteReply::Redirect(owner) + ); + let (runtime, accepted) = tokio::join!( + runtime_dial(address, &hosted, &identity), + registration.accept(&tls) ); - let mut wrong_version = encoded.clone(); - wrong_version[4] = 2; + assert!(runtime.is_ok() && accepted.is_ok()); + + // A route registered here that the directory agrees with is + // accepted, naming this member, then authenticated and delivered as + // for version 1. + router.set_directory( + Some(address), + HashMap::from([ + (hosted.clone(), address), + (elsewhere.clone(), owner), + (looped.clone(), address), + ]), + ); + let (stream, reply) = negotiate(address, &hosted).await; + assert_eq!(reply, RouteReply::Accept(Some(address))); + let (runtime, accepted) = + tokio::join!(serve_tls(stream, &identity), registration.accept(&tls)); + assert!(runtime.is_ok() && accepted.is_ok()); + + // A route hosted elsewhere is redirected to its owner. A directory + // entry naming this listener, or no entry at all, is unknown. + assert_eq!( + negotiate(address, &elsewhere).await.1, + RouteReply::Redirect(owner) + ); + assert_eq!(negotiate(address, &looped).await.1, RouteReply::Unknown); + assert_eq!( + negotiate(address, &new_route_handle()).await.1, + RouteReply::Unknown + ); + + // A version 1 runtime cannot follow a redirect and gets no reply. + let mut stream = tokio::net::TcpStream::connect(address).await.unwrap(); + stream + .write_all(&encode_route_preface(&elsewhere).unwrap()) + .await + .unwrap(); + let mut reply = Vec::new(); + stream.read_to_end(&mut reply).await.unwrap(); + assert!(reply.is_empty()); + + // A replacement directory takes effect for the next connection. + let moved: SocketAddr = "[fd00::3]:7443".parse().unwrap(); + router.set_directory(Some(address), HashMap::from([(looped.clone(), moved)])); + assert_eq!( + negotiate(address, &looped).await.1, + RouteReply::Redirect(moved) + ); + assert_eq!(negotiate(address, &elsewhere).await.1, RouteReply::Unknown); + assert!(!delivered(®istration, &tls).await); + } + + #[tokio::test] + async fn preface_round_trips_and_rejects_malformed_input() { + let handle = new_route_handle(); + for (encoded, version) in [ + ( + encode_route_preface(&handle).unwrap(), + ROUTE_PREFACE_VERSION, + ), + ( + encode_route_preface_v2(&handle).unwrap(), + ROUTE_PREFACE_VERSION_V2, + ), + ] { + assert_eq!( + read_route_preface(&mut encoded.as_slice()).await.unwrap(), + (version, handle.clone()) + ); + } + let mut wrong_version = encode_route_preface(&handle).unwrap(); + wrong_version[4] = 3; for malformed in [ b"TLS\x16\x03\x01".to_vec(), wrong_version, [b"OSRP\x01\x41".as_slice(), &[b'a'; 65]].concat(), - [b"OSRP\x01\x10".as_slice(), b"UPPERCASE-HANDLE"].concat(), + [b"OSRP\x02\x10".as_slice(), b"UPPERCASE-HANDLE"].concat(), + [b"OSRP\x02\x20".as_slice(), b"truncated-handle"].concat(), ] { assert!(read_route_preface(&mut malformed.as_slice()).await.is_err()); } assert!(encode_route_preface("short").is_err()); + assert!(encode_route_preface_v2("short").is_err()); + } + + /// Negotiate against a supervisor that answers `reply`, returning the + /// outcome and the bytes left unread after it. + async fn negotiate_scripted(reply: &[u8]) -> (std::io::Result, Vec) { + let handle = new_route_handle(); + let (mut runtime, mut supervisor) = tokio::io::duplex(1024); + supervisor.write_all(reply).await.unwrap(); + supervisor.shutdown().await.unwrap(); + let outcome = negotiate_route(&mut runtime, &handle).await; + let mut preface = vec![0_u8; encode_route_preface_v2(&handle).unwrap().len()]; + supervisor.read_exact(&mut preface).await.unwrap(); + assert_eq!(preface, encode_route_preface_v2(&handle).unwrap()); + let mut rest = Vec::new(); + runtime.read_to_end(&mut rest).await.unwrap(); + (outcome, rest) + } + + #[tokio::test] + async fn negotiation_reads_only_the_declared_reply_and_rejects_malformed_replies() { + // The TLS handshake starts right after an accept. + let (outcome, rest) = negotiate_scripted(b"\x01\x00\x16\x03\x01").await; + assert_eq!(outcome.unwrap(), RouteReply::Accept(None)); + assert_eq!(rest, b"\x16\x03\x01"); + let (outcome, rest) = negotiate_scripted(b"\x01\x0d10.0.0.4:7443\x16").await; + assert_eq!( + outcome.unwrap(), + RouteReply::Accept(Some("10.0.0.4:7443".parse().unwrap())) + ); + assert_eq!(rest, b"\x16"); + let (outcome, rest) = negotiate_scripted(b"\x02\x0d10.0.0.2:7443trailing").await; + assert_eq!( + outcome.unwrap(), + RouteReply::Redirect("10.0.0.2:7443".parse().unwrap()) + ); + assert_eq!(rest, b"trailing"); + assert_eq!( + negotiate_scripted(&redirect_reply("[fd00::2]:7443".parse().unwrap()).unwrap()) + .await + .0 + .unwrap(), + RouteReply::Redirect("[fd00::2]:7443".parse().unwrap()) + ); + assert_eq!( + negotiate_scripted(b"\x03").await.0.unwrap(), + RouteReply::Unknown + ); + for malformed in [ + b"".as_slice(), + b"\x01", + b"\x01\x05local", + b"\x04", + b"\x02", + b"\x02\x00", + b"\x02\x0810.0.0.2", + b"\x02\x0d10.0.0.2", + b"\x02\x0enot-an-address", + ] { + assert!(negotiate_scripted(malformed).await.0.is_err()); + } } } diff --git a/crates/openshell-sandbox-backend/src/sandbox_auth.rs b/crates/openshell-sandbox-backend/src/sandbox_auth.rs index 8f9f3ce960..df7a7e7cc9 100644 --- a/crates/openshell-sandbox-backend/src/sandbox_auth.rs +++ b/crates/openshell-sandbox-backend/src/sandbox_auth.rs @@ -2,16 +2,31 @@ // SPDX-License-Identifier: Apache-2.0 //! Transport-neutral authentication for `OpenShell` Sandbox Protocol connections. +//! +//! A runtime generation is owned by one supervisor process at a time. The +//! owner's term is the credential epoch the runtime currently accepts, and +//! the owner's process identity is pinned at its first attach. A different +//! process can take over only with a gateway-signed takeover grant for a +//! strictly higher epoch, naming that process; the replacement is staged +//! while the current owner stays active and is promoted only when the +//! replacement confirms, which advances the accepted epoch, moves the pin, +//! and retires every connection of the previous term. Credentials of lower +//! epochs are rejected from then on. Ordinary credential renewal keeps the +//! epoch and cannot move ownership. +use std::str::FromStr as _; use std::sync::Mutex; use openshell_core::jwt::{ AuthenticatedSandboxSession, CredentialEpoch, SandboxId, SessionJwtError, SessionJwtVerifier, + VerifiedSupervisorTakeover, }; use openshell_core::sandbox_generation::SandboxGenerationId; use tonic::metadata::MetadataMap; use uuid::Uuid; +use crate::boundary_protocol::SupervisorInstanceId; + /// Server-local identity assigned after a byte stream completes TLS. /// /// It cannot be supplied by a compute driver or workload. Protocol handlers use @@ -56,7 +71,9 @@ pub struct SandboxProtocolAuthenticator { verifier: SessionJwtVerifier, expected_sandbox_id: SandboxId, expected_runtime_generation: SandboxGenerationId, - expected_auth_epoch: CredentialEpoch, + /// Epoch the runtime was launched with. No credential below it is ever + /// valid; [`SandboxConnectionRegistry`] holds the owner term in force. + launch_auth_epoch: CredentialEpoch, } impl std::fmt::Debug for SandboxProtocolAuthenticator { @@ -69,7 +86,7 @@ impl std::fmt::Debug for SandboxProtocolAuthenticator { "expected_runtime_generation", &self.expected_runtime_generation, ) - .field("expected_auth_epoch", &self.expected_auth_epoch) + .field("launch_auth_epoch", &self.launch_auth_epoch) .finish() } } @@ -80,16 +97,35 @@ impl SandboxProtocolAuthenticator { verifier: SessionJwtVerifier, expected_sandbox_id: SandboxId, expected_runtime_generation: SandboxGenerationId, - expected_auth_epoch: CredentialEpoch, + launch_auth_epoch: CredentialEpoch, ) -> Self { Self { verifier, expected_sandbox_id, expected_runtime_generation, - expected_auth_epoch, + launch_auth_epoch, + } + } + + /// Verify a supervisor takeover grant for this runtime generation. The + /// registry decides whether its epoch may replace the current owner. + pub fn verify_takeover( + &self, + grant: &str, + ) -> Result { + let verified = self.verifier.verify_supervisor_takeover(grant)?; + if verified.takeover.sandbox_id != self.expected_sandbox_id { + return Err(SandboxAuthError::WrongSandbox); + } + if verified.takeover.runtime_generation != self.expected_runtime_generation { + return Err(SandboxAuthError::WrongRuntimeGeneration); } + Ok(verified) } + /// Verify the bearer for this sandbox and generation. Its epoch is only + /// checked against the launch epoch here; the connection registry admits + /// the current owner term, stages takeovers, and rejects older terms. pub fn authenticate( &self, connection_id: SandboxConnectionId, @@ -113,7 +149,7 @@ impl SandboxProtocolAuthenticator { if session.runtime_generation != self.expected_runtime_generation { return Err(SandboxAuthError::WrongRuntimeGeneration); } - if session.auth_epoch != self.expected_auth_epoch { + if session.auth_epoch < self.launch_auth_epoch { return Err(SandboxAuthError::StaleCredentialEpoch); } Ok(SandboxProtocolPrincipal { @@ -129,121 +165,205 @@ struct ActiveConnection { epoch: CredentialEpoch, } +/// A connection attached but not yet confirmed. +#[derive(Debug, Clone, PartialEq, Eq)] +struct PendingConnection { + id: SandboxConnectionId, + epoch: CredentialEpoch, + /// Present when the connection replaces the owner under a takeover grant. + takeover: Option, +} + +/// A verified takeover staged on a pending connection. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StagedTakeover { + /// Process identity pinned when the takeover is confirmed. + pub supervisor_instance_id: SupervisorInstanceId, + /// Owner term accepted once the takeover is confirmed. + pub auth_epoch: CredentialEpoch, + /// Pool member that owns the new term, for audit. + pub owner: String, + pub grant_id: Uuid, +} + +/// What a confirmation changed. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct Promotion { + /// The previously active connection, which may now be closed without + /// creating an unsupervised interval. + pub replaced: Option, + /// The takeover this confirmation completed. Every connection of the + /// previous owner term is then stale and must be closed. + pub takeover: Option, +} + #[derive(Debug)] struct ConnectionState { active: Option, - pending: Option, + pending: Option, highest_epoch: Option, - supervisor_instance_id: Option, + /// The owner term in force: credentials below it are stale, and + /// credentials above it act only to stage a granted takeover. + accepted_epoch: CredentialEpoch, + supervisor_instance_id: Option, terminal: bool, } -/// Enforces the one-active-connection and monotonically increasing epoch rules. +/// Enforces the one-active-connection and monotonically increasing epoch +/// rules, and the owner term a takeover must advance. #[derive(Debug)] pub struct SandboxConnectionRegistry { state: Mutex, } impl SandboxConnectionRegistry { + /// Track connections of a runtime launched at `launch_auth_epoch`. #[must_use] pub fn new( _session_id: openshell_core::SandboxSessionId, _session_rotation: openshell_core::jwt::SessionRotation, + launch_auth_epoch: CredentialEpoch, ) -> Self { Self { state: Mutex::new(ConnectionState { active: None, pending: None, highest_epoch: None, + accepted_epoch: launch_auth_epoch, supervisor_instance_id: None, terminal: false, }), } } + fn lock(&self) -> std::sync::MutexGuard<'_, ConnectionState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// The owner term in force. + #[must_use] + pub fn accepted_epoch(&self) -> CredentialEpoch { + self.lock().accepted_epoch + } + + /// The supervisor process allowed to attach without a takeover grant. + #[must_use] + pub fn supervisor_instance(&self) -> Option { + self.lock().supervisor_instance_id + } + + /// Reject credentials of an owner term that has been replaced. Higher + /// epochs pass so a granted replacement can open its boundary and stage + /// its attach; nothing else accepts them until it is confirmed. + pub fn admit(&self, principal: &SandboxProtocolPrincipal) -> Result<(), SandboxAuthError> { + if principal.session.auth_epoch < self.lock().accepted_epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } + Ok(()) + } + + /// Check, without changing anything, whether [`Self::attach`] would + /// accept this attach. Callers that act on the boundary before + /// committing the attach run this first, under the same serialization + /// as the commit, so a rejected candidate never touches the boundary. + pub fn check_attach( + &self, + principal: &SandboxProtocolPrincipal, + supervisor_instance_id: SupervisorInstanceId, + takeover: Option<&VerifiedSupervisorTakeover>, + ) -> Result<(), SandboxAuthError> { + plan_attach(&self.lock(), principal, supervisor_instance_id, takeover).map(drop) + } + /// Stage a fully authenticated connection for confirmation. The returned /// ID identifies an older unconfirmed candidate that may be closed. The /// current active connection remains authoritative until [`Self::confirm`] /// promotes this candidate. + /// + /// Without `takeover`, only the pinned supervisor process (or the first + /// one, which is then pinned) may attach, at the accepted epoch. With a + /// verified grant, a process named by the grant may stage a replacement + /// at the grant's epoch when it exceeds the accepted epoch and is not + /// below any epoch already staged; replaying the attach of an already + /// confirmed takeover is an ordinary attach. pub fn attach( &self, principal: &SandboxProtocolPrincipal, - supervisor_instance_id: crate::boundary_protocol::SupervisorInstanceId, + supervisor_instance_id: SupervisorInstanceId, + takeover: Option<&VerifiedSupervisorTakeover>, ) -> Result, SandboxAuthError> { + let mut state = self.lock(); + let AttachPlan::Stage { takeover: staged } = + plan_attach(&state, principal, supervisor_instance_id, takeover)? + else { + return Ok(None); + }; let epoch = principal.session.auth_epoch; - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.terminal { - return Err(SandboxAuthError::TerminalSession); - } - match state.supervisor_instance_id { - Some(expected) if expected != supervisor_instance_id => { - return Err(SandboxAuthError::WrongSupervisorInstance); - } - Some(_) => {} - None => state.supervisor_instance_id = Some(supervisor_instance_id), - } - if let Some(active) = state.active { - if active.id == principal.connection_id && active.epoch == epoch { - return Ok(None); - } - if epoch == active.epoch { - return Err(SandboxAuthError::ConnectionStillActive); - } - if epoch < active.epoch { - return Err(SandboxAuthError::StaleCredentialEpoch); - } - } - if state.highest_epoch.is_some_and(|highest| epoch < highest) { - return Err(SandboxAuthError::StaleCredentialEpoch); - } - if let Some(pending) = state.pending { - if pending.id == principal.connection_id && pending.epoch == epoch { - return Ok(None); - } - if epoch < pending.epoch { - return Err(SandboxAuthError::StaleCredentialEpoch); - } + let replaced = state + .pending + .take() + .map(|pending| pending.id) + .filter(|id| *id != principal.connection_id); + state.highest_epoch = Some( + state + .highest_epoch + .map_or(epoch, |highest| highest.max(epoch)), + ); + if staged.is_none() { + // The first process to attach is pinned; a replacement then + // needs a grant. + state + .supervisor_instance_id + .get_or_insert(supervisor_instance_id); } - let replaced = state.pending.map(|pending| pending.id); - state.highest_epoch = Some(epoch); - state.pending = Some(ActiveConnection { + state.pending = Some(PendingConnection { id: principal.connection_id, epoch, + takeover: staged, }); Ok(replaced) } - /// Promote an attached, confirmed candidate to the active connection. The - /// returned ID is the previously active connection, which may now be - /// closed without creating an unsupervised interval. + /// Promote an attached, confirmed candidate to the active connection. + /// Confirming a staged takeover advances the accepted epoch and pins the + /// replacement process before the previous owner's connection is + /// reported for closing. pub fn confirm( &self, principal: &SandboxProtocolPrincipal, - ) -> Result, SandboxAuthError> { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + ) -> Result { + let epoch = principal.session.auth_epoch; + let mut state = self.lock(); if state.terminal { return Err(SandboxAuthError::TerminalSession); } if state .active - .is_some_and(|active| active.id == principal.connection_id) + .is_some_and(|active| active.id == principal.connection_id && active.epoch == epoch) { - return Ok(None); + return Ok(Promotion::default()); } let pending = state .pending - .filter(|pending| pending.id == principal.connection_id) + .clone() + .filter(|pending| pending.id == principal.connection_id && pending.epoch == epoch) .ok_or(SandboxAuthError::ConnectionNotAttached)?; + if let Some(takeover) = &pending.takeover { + state.accepted_epoch = takeover.auth_epoch; + state.supervisor_instance_id = Some(takeover.supervisor_instance_id); + } let replaced = state.active.map(|active| active.id); - state.active = Some(pending); + state.active = Some(ActiveConnection { + id: pending.id, + epoch: pending.epoch, + }); state.pending = None; - Ok(replaced) + Ok(Promotion { + replaced, + takeover: pending.takeover, + }) } /// Confirm may run on either the active connection (an idempotent replay) @@ -252,19 +372,17 @@ impl SandboxConnectionRegistry { &self, principal: &SandboxProtocolPrincipal, ) -> Result<(), SandboxAuthError> { - let state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let epoch = principal.session.auth_epoch; + let state = self.lock(); if state.terminal { return Err(SandboxAuthError::TerminalSession); } if state .active - .is_some_and(|active| active.id == principal.connection_id) - || state - .pending - .is_some_and(|pending| pending.id == principal.connection_id) + .is_some_and(|active| active.id == principal.connection_id && active.epoch == epoch) + || state.pending.as_ref().is_some_and(|pending| { + pending.id == principal.connection_id && pending.epoch == epoch + }) { Ok(()) } else { @@ -276,16 +394,17 @@ impl SandboxConnectionRegistry { &self, principal: &SandboxProtocolPrincipal, ) -> Result<(), SandboxAuthError> { - let state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let epoch = principal.session.auth_epoch; + let state = self.lock(); if state.terminal { return Err(SandboxAuthError::TerminalSession); } + if epoch != state.accepted_epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } if state .active - .is_none_or(|active| active.id != principal.connection_id) + .is_none_or(|active| active.id != principal.connection_id || active.epoch != epoch) { return Err(SandboxAuthError::ConnectionNotAttached); } @@ -295,10 +414,7 @@ impl SandboxConnectionRegistry { /// Remove a physical connection. Returns `true` only when it was the /// confirmed active connection and recovery must begin. pub fn disconnect(&self, connection_id: SandboxConnectionId) -> bool { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut state = self.lock(); let was_active = state .active .is_some_and(|active| active.id == connection_id); @@ -307,6 +423,7 @@ impl SandboxConnectionRegistry { } if state .pending + .as_ref() .is_some_and(|pending| pending.id == connection_id) { state.pending = None; @@ -315,16 +432,110 @@ impl SandboxConnectionRegistry { } pub fn mark_terminal(&self) { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut state = self.lock(); state.terminal = true; state.active = None; state.pending = None; } } +/// What an accepted attach does to the registry. +enum AttachPlan { + /// The connection is already attached at this epoch. + Unchanged, + /// Stage the connection, as a takeover when one is given. + Stage { takeover: Option }, +} + +/// Decide an attach against the current registry state without changing it. +fn plan_attach( + state: &ConnectionState, + principal: &SandboxProtocolPrincipal, + supervisor_instance_id: SupervisorInstanceId, + takeover: Option<&VerifiedSupervisorTakeover>, +) -> Result { + let epoch = principal.session.auth_epoch; + let connection_id = principal.connection_id; + if state.terminal { + return Err(SandboxAuthError::TerminalSession); + } + if epoch < state.accepted_epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } + let staged = if let Some(grant) = takeover { + let staged = staged_takeover(grant, epoch, supervisor_instance_id)?; + let promoted = epoch == state.accepted_epoch + && state.supervisor_instance_id == Some(supervisor_instance_id); + if promoted { + // A replayed attach of the confirmed takeover is an ordinary + // attach by the owner it pinned. + None + } else if epoch <= state.accepted_epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } else { + Some(staged) + } + } else { + if epoch > state.accepted_epoch { + return Err(SandboxAuthError::TakeoverRequired); + } + if state + .supervisor_instance_id + .is_some_and(|expected| expected != supervisor_instance_id) + { + return Err(SandboxAuthError::WrongSupervisorInstance); + } + None + }; + if staged.is_none() + && let Some(active) = state.active + { + if active.id == connection_id && active.epoch == epoch { + return Ok(AttachPlan::Unchanged); + } + if epoch == active.epoch { + return Err(SandboxAuthError::ConnectionStillActive); + } + if epoch < active.epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } + } + // A term below the highest one ever staged can no longer win, even after + // that candidate disconnected: the gateway granted the higher term + // later and has already fenced the lower one. + if state.highest_epoch.is_some_and(|highest| epoch < highest) { + return Err(SandboxAuthError::StaleCredentialEpoch); + } + if let Some(pending) = &state.pending { + if pending.id == connection_id && pending.epoch == epoch && pending.takeover == staged { + return Ok(AttachPlan::Unchanged); + } + if epoch < pending.epoch { + return Err(SandboxAuthError::StaleCredentialEpoch); + } + } + Ok(AttachPlan::Stage { takeover: staged }) +} + +/// Bind a verified grant to the attach that presents it. +fn staged_takeover( + grant: &VerifiedSupervisorTakeover, + epoch: CredentialEpoch, + supervisor_instance_id: SupervisorInstanceId, +) -> Result { + let granted = SupervisorInstanceId::from_str(&grant.takeover.supervisor_instance) + .map_err(|_| SandboxAuthError::TakeoverMismatch)?; + if grant.takeover.auth_epoch != epoch || granted != supervisor_instance_id { + return Err(SandboxAuthError::TakeoverMismatch); + } + Ok(StagedTakeover { + supervisor_instance_id, + auth_epoch: epoch, + owner: grant.takeover.owner.clone(), + grant_id: grant.token_id, + }) +} + #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] pub enum SandboxAuthError { #[error("authorization metadata is missing")] @@ -350,6 +561,12 @@ pub enum SandboxAuthError { ConnectionNotAttached, #[error("sandbox session is terminal")] TerminalSession, + /// The bearer is ahead of the owner term in force, which only a + /// takeover grant may advance. + #[error("Sandbox Protocol credential epoch requires a supervisor takeover grant")] + TakeoverRequired, + #[error("supervisor takeover grant does not match this attach")] + TakeoverMismatch, } #[cfg(test)] @@ -358,7 +575,7 @@ mod tests { use openshell_core::jwt::{ DEFAULT_SESSION_TOKEN_TTL, JwtClock, SandboxRuntimeIdentity, SessionJwtIssuer, - SessionTokenProfile, SessionVerificationKey, + SessionTokenProfile, SessionVerificationKey, SupervisorTakeover, }; use openshell_core::sandbox_generation::SandboxGenerationId; use rcgen::PKCS_ED25519; @@ -374,59 +591,102 @@ mod tests { } } - fn fixture( - epoch: u64, - ) -> ( - SandboxProtocolAuthenticator, - openshell_core::jwt::MintedSessionToken, - ) { - let key = openshell_crypto::pki::generate_keypair_for(&PKCS_ED25519).expect("generate key"); - let public_key_pem = key.public_key_pem().into_bytes(); - let clock: Arc = Arc::new(FixedClock); - let sandbox_id = SandboxId::parse("sandbox-a").expect("sandbox ID"); - let issuer = SessionJwtIssuer::from_ed25519_pem( - key.serialize_pem().unwrap().as_bytes(), - "current", - "test", - Some(DEFAULT_SESSION_TOKEN_TTL), - clock.clone(), - ) - .expect("issuer"); - let verifier = SessionJwtVerifier::new( - "test", - SessionTokenProfile::Sandbox, - [SessionVerificationKey { - key_id: "current".to_string(), - public_key_pem, - }], - clock, - ) - .expect("verifier"); - let token = issuer - .mint_pair(&SandboxRuntimeIdentity { - sandbox_id: sandbox_id.clone(), - runtime_generation: SandboxGenerationId::parse("generation-1") - .expect("runtime generation"), - auth_epoch: CredentialEpoch::new(epoch).expect("epoch"), - }) - .expect("token pair") - .sandbox; - ( - SandboxProtocolAuthenticator::new( - verifier, - sandbox_id, - SandboxGenerationId::parse("generation-1").expect("runtime generation"), - CredentialEpoch::new(epoch).expect("epoch"), - ), - token, - ) + /// One gateway issuer and the runtime that trusts it, launched at epoch 1. + struct Gateway { + issuer: SessionJwtIssuer, + authenticator: SandboxProtocolAuthenticator, } - fn registry_for(_principal: &SandboxProtocolPrincipal) -> SandboxConnectionRegistry { - SandboxConnectionRegistry::new( - openshell_core::SandboxSessionId::new(), - openshell_core::jwt::SessionRotation::new(1).expect("rotation"), - ) + impl Gateway { + fn new() -> Self { + Self::launched_at(1) + } + + fn launched_at(launch_epoch: u64) -> Self { + let key = + openshell_crypto::pki::generate_keypair_for(&PKCS_ED25519).expect("generate key"); + let public_key_pem = key.public_key_pem().into_bytes(); + let clock: Arc = Arc::new(FixedClock); + let issuer = SessionJwtIssuer::from_ed25519_pem( + key.serialize_pem().unwrap().as_bytes(), + "current", + "test", + Some(DEFAULT_SESSION_TOKEN_TTL), + clock.clone(), + ) + .expect("issuer"); + let verifier = SessionJwtVerifier::new( + "test", + SessionTokenProfile::Sandbox, + [SessionVerificationKey { + key_id: "current".to_string(), + public_key_pem, + }], + clock, + ) + .expect("verifier"); + Self { + issuer, + authenticator: SandboxProtocolAuthenticator::new( + verifier, + SandboxId::parse("sandbox-a").expect("sandbox ID"), + SandboxGenerationId::parse("generation-1").expect("runtime generation"), + CredentialEpoch::new(launch_epoch).expect("epoch"), + ), + } + } + + fn registry(&self) -> SandboxConnectionRegistry { + SandboxConnectionRegistry::new( + openshell_core::SandboxSessionId::new(), + openshell_core::jwt::SessionRotation::new(1).expect("rotation"), + self.authenticator.launch_auth_epoch, + ) + } + + /// A Sandbox Protocol bearer for `epoch`. + fn token(&self, epoch: u64) -> String { + self.issuer + .mint_pair(&SandboxRuntimeIdentity { + sandbox_id: SandboxId::parse("sandbox-a").expect("sandbox ID"), + runtime_generation: SandboxGenerationId::parse("generation-1") + .expect("runtime generation"), + auth_epoch: CredentialEpoch::new(epoch).expect("epoch"), + }) + .expect("token pair") + .sandbox + .token + .expose_secret() + .to_string() + } + + /// Authenticate `token` on a new physical connection. + fn connect(&self, token: &str) -> SandboxProtocolPrincipal { + self.authenticator + .authenticate(SandboxConnectionId::new(), &metadata(token)) + .expect("principal") + } + + /// A verified grant for `instance` to own `epoch`. + fn grant(&self, epoch: u64, instance: SupervisorInstanceId) -> VerifiedSupervisorTakeover { + let minted = self + .issuer + .mint_supervisor_takeover( + &SupervisorTakeover { + sandbox_id: SandboxId::parse("sandbox-a").expect("sandbox ID"), + runtime_generation: SandboxGenerationId::parse("generation-1") + .expect("runtime generation"), + auth_epoch: CredentialEpoch::new(epoch).expect("epoch"), + supervisor_instance: instance.to_string(), + owner: "supervisor-b".to_string(), + }, + std::time::Duration::from_mins(5), + ) + .expect("mint grant"); + self.authenticator + .verify_takeover(minted.token.expose_secret()) + .expect("verify grant") + } } fn metadata(token: &str) -> MetadataMap { @@ -438,57 +698,79 @@ mod tests { metadata } + fn epoch(value: u64) -> CredentialEpoch { + CredentialEpoch::new(value).expect("epoch") + } + #[test] fn bearer_identity_must_match_the_provisioned_runtime() { - let (mut authenticator, token) = fixture(1); - let metadata = metadata(token.token.expose_secret()); - authenticator.expected_sandbox_id = SandboxId::parse("sandbox-b").unwrap(); + let mut gateway = Gateway::new(); + let metadata = metadata(&gateway.token(1)); + gateway.authenticator.expected_sandbox_id = SandboxId::parse("sandbox-b").unwrap(); assert_eq!( - authenticator.authenticate(SandboxConnectionId::new(), &metadata), + gateway + .authenticator + .authenticate(SandboxConnectionId::new(), &metadata), Err(SandboxAuthError::WrongSandbox) ); - authenticator.expected_sandbox_id = SandboxId::parse("sandbox-a").unwrap(); - authenticator.expected_runtime_generation = + gateway.authenticator.expected_sandbox_id = SandboxId::parse("sandbox-a").unwrap(); + gateway.authenticator.expected_runtime_generation = SandboxGenerationId::parse("generation-2").unwrap(); assert_eq!( - authenticator.authenticate(SandboxConnectionId::new(), &metadata), + gateway + .authenticator + .authenticate(SandboxConnectionId::new(), &metadata), Err(SandboxAuthError::WrongRuntimeGeneration) ); } + #[test] + fn credentials_below_the_launch_epoch_are_never_valid() { + let gateway = Gateway::launched_at(3); + assert_eq!( + gateway + .authenticator + .authenticate(SandboxConnectionId::new(), &metadata(&gateway.token(2))), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + } + #[test] fn bearer_metadata_must_occur_exactly_once() { - let (authenticator, token) = fixture(1); + let gateway = Gateway::new(); + let token = gateway.token(1); assert_eq!( - authenticator.authenticate(SandboxConnectionId::new(), &MetadataMap::new()), + gateway + .authenticator + .authenticate(SandboxConnectionId::new(), &MetadataMap::new()), Err(SandboxAuthError::MissingBearer) ); - let mut duplicate = metadata(token.token.expose_secret()); + let mut duplicate = metadata(&token); duplicate.append( "authorization", - format!("Bearer {}", token.token.expose_secret()) - .parse() - .expect("metadata value"), + format!("Bearer {token}").parse().expect("metadata value"), ); assert_eq!( - authenticator.authenticate(SandboxConnectionId::new(), &duplicate), + gateway + .authenticator + .authenticate(SandboxConnectionId::new(), &duplicate), Err(SandboxAuthError::DuplicateBearer) ); } #[test] fn bearer_scheme_is_case_insensitive() { - let (authenticator, token) = fixture(1); + let gateway = Gateway::new(); + let token = gateway.token(1); for scheme in ["bearer", "BEARER"] { let mut metadata = MetadataMap::new(); metadata.insert( "authorization", - format!("{scheme} {}", token.token.expose_secret()) - .parse() - .expect("metadata value"), + format!("{scheme} {token}").parse().expect("metadata value"), ); assert!( - authenticator + gateway + .authenticator .authenticate(SandboxConnectionId::new(), &metadata) .is_ok(), "{scheme}" @@ -498,111 +780,405 @@ mod tests { #[test] fn reconnect_requires_disconnect_and_terminal_is_final() { - let (first_authenticator, first_token) = fixture(1); - let first_id = SandboxConnectionId::new(); - let first = first_authenticator - .authenticate(first_id, &metadata(first_token.token.expose_secret())) - .expect("first principal"); - let registry = registry_for(&first); - let instance = crate::boundary_protocol::SupervisorInstanceId::new(); - assert_eq!(registry.attach(&first, instance), Ok(None)); - assert_eq!(registry.attach(&first, instance), Ok(None)); - assert_eq!(registry.confirm(&first), Ok(None)); - assert_eq!(registry.confirm(&first), Ok(None)); - - let second_id = SandboxConnectionId::new(); - let second = first_authenticator - .authenticate(second_id, &metadata(first_token.token.expose_secret())) - .expect("second principal"); - assert_eq!( - registry.attach(&second, instance), + let gateway = Gateway::new(); + let token = gateway.token(1); + let first = gateway.connect(&token); + let registry = gateway.registry(); + let instance = SupervisorInstanceId::new(); + assert_eq!(registry.attach(&first, instance, None), Ok(None)); + assert_eq!(registry.attach(&first, instance, None), Ok(None)); + assert_eq!(registry.confirm(&first), Ok(Promotion::default())); + assert_eq!(registry.confirm(&first), Ok(Promotion::default())); + + let second = gateway.connect(&token); + assert_eq!( + registry.attach(&second, instance, None), Err(SandboxAuthError::ConnectionStillActive) ); - assert!(registry.disconnect(first_id)); - assert_eq!(registry.attach(&second, instance), Ok(None)); - assert_eq!(registry.confirm(&second), Ok(None)); - assert!(registry.disconnect(second_id)); - assert_eq!(registry.attach(&first, instance), Ok(None)); - assert_eq!(registry.confirm(&first), Ok(None)); + assert!(registry.disconnect(first.connection_id())); + assert_eq!(registry.attach(&second, instance, None), Ok(None)); + assert_eq!(registry.confirm(&second), Ok(Promotion::default())); + assert!(registry.disconnect(second.connection_id())); + assert_eq!(registry.attach(&first, instance, None), Ok(None)); + assert_eq!(registry.confirm(&first), Ok(Promotion::default())); registry.mark_terminal(); assert_eq!( registry.require_active(&first), Err(SandboxAuthError::TerminalSession) ); assert_eq!( - registry.attach(&first, instance), + registry.attach(&first, instance, None), Err(SandboxAuthError::TerminalSession) ); } #[test] fn replacement_supervisor_cannot_resume_existing_runtime_generation() { - let (authenticator, token) = fixture(1); - let first_id = SandboxConnectionId::new(); - let first = authenticator - .authenticate(first_id, &metadata(token.token.expose_secret())) - .expect("first principal"); - let registry = registry_for(&first); - let first_instance = crate::boundary_protocol::SupervisorInstanceId::new(); + let gateway = Gateway::new(); + let token = gateway.token(1); + let first = gateway.connect(&token); + let registry = gateway.registry(); + let first_instance = SupervisorInstanceId::new(); registry - .attach(&first, first_instance) + .attach(&first, first_instance, None) .expect("attach first supervisor"); registry.confirm(&first).expect("confirm first supervisor"); - assert!(registry.disconnect(first_id)); + assert!(registry.disconnect(first.connection_id())); - let replacement_id = SandboxConnectionId::new(); - let replacement = authenticator - .authenticate(replacement_id, &metadata(token.token.expose_secret())) - .expect("replacement principal"); + // Launch credentials, a renewed bearer, or a bearer from a later + // epoch never move ownership without a grant. + let replacement = gateway.connect(&token); assert_eq!( - registry.attach( - &replacement, - crate::boundary_protocol::SupervisorInstanceId::new(), - ), + registry.attach(&replacement, SupervisorInstanceId::new(), None), + Err(SandboxAuthError::WrongSupervisorInstance) + ); + let renewed = gateway.connect(&gateway.token(1)); + assert_eq!( + registry.attach(&renewed, SupervisorInstanceId::new(), None), Err(SandboxAuthError::WrongSupervisorInstance) ); - assert_eq!(registry.attach(&replacement, first_instance), Ok(None)); + let ahead = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&ahead, SupervisorInstanceId::new(), None), + Err(SandboxAuthError::TakeoverRequired) + ); + assert_eq!( + registry.attach(&ahead, first_instance, None), + Err(SandboxAuthError::TakeoverRequired) + ); + assert_eq!( + registry.attach(&replacement, first_instance, None), + Ok(None) + ); + assert_eq!(registry.accepted_epoch(), epoch(1)); } #[test] fn replacement_does_not_displace_active_connection_before_confirm() { - let (first_authenticator, first_token) = fixture(1); - let first_id = SandboxConnectionId::new(); - let first = first_authenticator - .authenticate(first_id, &metadata(first_token.token.expose_secret())) - .expect("first principal"); - let (replacement_authenticator, replacement_token) = fixture(2); - let replacement_id = SandboxConnectionId::new(); - let replacement = replacement_authenticator - .authenticate( - replacement_id, - &metadata(replacement_token.token.expose_secret()), - ) - .expect("replacement principal"); - let registry = registry_for(&first); - let instance = crate::boundary_protocol::SupervisorInstanceId::new(); - registry.attach(&first, instance).expect("attach first"); + let gateway = Gateway::new(); + let first = gateway.connect(&gateway.token(1)); + let registry = gateway.registry(); + let owner = SupervisorInstanceId::new(); + registry.attach(&first, owner, None).expect("attach first"); registry.confirm(&first).expect("confirm first"); - assert_eq!(registry.attach(&replacement, instance), Ok(None)); + let candidate = SupervisorInstanceId::new(); + let replacement = gateway.connect(&gateway.token(2)); + let grant = gateway.grant(2, candidate); + assert_eq!( + registry.attach(&replacement, candidate, Some(&grant)), + Ok(None) + ); registry .require_active(&first) .expect("first remains active"); assert_eq!( registry.require_active(&replacement), - Err(SandboxAuthError::ConnectionNotAttached) + Err(SandboxAuthError::StaleCredentialEpoch) ); registry .require_attached(&replacement) .expect("replacement may confirm"); - assert_eq!(registry.confirm(&replacement), Ok(Some(first_id))); + let promotion = registry.confirm(&replacement).expect("confirm"); + assert_eq!(promotion.replaced, Some(first.connection_id())); + assert_eq!( + promotion.takeover.map(|takeover| takeover.auth_epoch), + Some(epoch(2)) + ); assert_eq!( registry.require_active(&first), - Err(SandboxAuthError::ConnectionNotAttached) + Err(SandboxAuthError::StaleCredentialEpoch) ); registry .require_active(&replacement) .expect("replacement became active"); } + + #[test] + fn confirmed_takeover_fences_every_older_credential_and_process() { + let gateway = Gateway::new(); + let registry = gateway.registry(); + let old_owner = SupervisorInstanceId::new(); + let old = gateway.connect(&gateway.token(1)); + registry.attach(&old, old_owner, None).expect("attach old"); + registry.confirm(&old).expect("confirm old"); + // The old owner's process dies: the runtime freezes. + assert!(registry.disconnect(old.connection_id())); + + let candidate = SupervisorInstanceId::new(); + let replacement = gateway.connect(&gateway.token(2)); + let grant = gateway.grant(2, candidate); + assert_eq!( + registry.attach(&replacement, candidate, Some(&grant)), + Ok(None) + ); + // A staged takeover already outranks the old term: the old owner can + // no longer reattach while the replacement prepares. + let returning = gateway.connect(&gateway.token(1)); + assert_eq!( + registry.attach(&returning, old_owner, None), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + let promotion = registry.confirm(&replacement).expect("confirm takeover"); + assert_eq!(promotion.replaced, None); + assert_eq!(registry.accepted_epoch(), epoch(2)); + assert_eq!(registry.supervisor_instance(), Some(candidate)); + + // Every credential of the replaced term is now stale, including a + // renewal minted for it and an attempt to reuse the old process pin. + for token in [gateway.token(1), gateway.token(1)] { + let stale = gateway.connect(&token); + assert_eq!( + registry.admit(&stale), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + assert_eq!( + registry.attach(&stale, old_owner, None), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + assert_eq!( + registry.attach(&stale, candidate, None), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + } + // The same grant cannot be replayed by another process, and a grant + // for the term in force cannot move it. + let thief = SupervisorInstanceId::new(); + let replayed = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&replayed, thief, Some(&grant)), + Err(SandboxAuthError::TakeoverMismatch) + ); + assert_eq!( + registry.attach(&replayed, thief, Some(&gateway.grant(2, thief))), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + // The new owner's own transport recovery attaches without the grant, + // and replaying its takeover attach is harmless. + assert!(registry.disconnect(replacement.connection_id())); + let reconnect = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&reconnect, candidate, Some(&grant)), + Ok(None) + ); + let plain = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&plain, candidate, None), + Ok(Some(reconnect.connection_id())) + ); + assert_eq!(registry.confirm(&plain), Ok(Promotion::default())); + registry + .require_active(&plain) + .expect("same owner recovered"); + } + + #[test] + fn concurrent_candidates_resolve_to_the_highest_granted_term() { + let gateway = Gateway::new(); + let registry = gateway.registry(); + let owner = SupervisorInstanceId::new(); + let launch = gateway.connect(&gateway.token(1)); + registry.attach(&launch, owner, None).expect("attach"); + registry.confirm(&launch).expect("confirm"); + + let (second, third) = (SupervisorInstanceId::new(), SupervisorInstanceId::new()); + let candidate_two = gateway.connect(&gateway.token(2)); + let candidate_three = gateway.connect(&gateway.token(3)); + assert_eq!( + registry.attach(&candidate_two, second, Some(&gateway.grant(2, second))), + Ok(None) + ); + // A higher term replaces the lower staged candidate, which is closed. + assert_eq!( + registry.attach(&candidate_three, third, Some(&gateway.grant(3, third))), + Ok(Some(candidate_two.connection_id())) + ); + assert_eq!( + registry.confirm(&candidate_two), + Err(SandboxAuthError::ConnectionNotAttached) + ); + // The lower candidate cannot stage again behind the higher one. + let retry = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&retry, second, Some(&gateway.grant(2, second))), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + let promotion = registry + .confirm(&candidate_three) + .expect("highest term wins"); + assert_eq!(promotion.replaced, Some(launch.connection_id())); + assert_eq!(registry.accepted_epoch(), epoch(3)); + // Grants with a mismatched epoch or instance never stage. + let fourth = SupervisorInstanceId::new(); + let ahead = gateway.connect(&gateway.token(4)); + assert_eq!( + registry.attach(&ahead, fourth, Some(&gateway.grant(5, fourth))), + Err(SandboxAuthError::TakeoverMismatch) + ); + assert_eq!( + registry.attach( + &ahead, + fourth, + Some(&gateway.grant(4, SupervisorInstanceId::new())) + ), + Err(SandboxAuthError::TakeoverMismatch) + ); + } + + #[test] + fn a_dropped_higher_candidate_still_outranks_lower_grants() { + // The gateway granted term 2 to a slow candidate, then term 3 to + // another that staged and lost its transport before confirming. The + // slow candidate's still-valid term-2 grant must not win. + let gateway = Gateway::new(); + let registry = gateway.registry(); + let owner = SupervisorInstanceId::new(); + let launch = gateway.connect(&gateway.token(1)); + registry.attach(&launch, owner, None).expect("attach"); + registry.confirm(&launch).expect("confirm"); + + let (slow, third) = (SupervisorInstanceId::new(), SupervisorInstanceId::new()); + let stage3 = gateway.connect(&gateway.token(3)); + let grant3 = gateway.grant(3, third); + assert_eq!(registry.attach(&stage3, third, Some(&grant3)), Ok(None)); + assert!(!registry.disconnect(stage3.connection_id())); + + let attempt2 = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.check_attach(&attempt2, slow, Some(&gateway.grant(2, slow))), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + assert_eq!( + registry.attach(&attempt2, slow, Some(&gateway.grant(2, slow))), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + assert_eq!(registry.accepted_epoch(), epoch(1)); + + // The term-3 candidate may still stage again and win. + let retry3 = gateway.connect(&gateway.token(3)); + assert_eq!(registry.attach(&retry3, third, Some(&grant3)), Ok(None)); + let promotion = registry.confirm(&retry3).expect("term 3 confirms"); + assert_eq!(promotion.replaced, Some(launch.connection_id())); + assert_eq!(registry.accepted_epoch(), epoch(3)); + assert_eq!(registry.supervisor_instance(), Some(third)); + } + + #[test] + fn attach_checks_match_commits_without_changing_the_registry() { + let gateway = Gateway::new(); + let registry = gateway.registry(); + let owner = SupervisorInstanceId::new(); + let first = gateway.connect(&gateway.token(1)); + // A check never pins or stages: a different first process still wins. + registry + .check_attach(&first, SupervisorInstanceId::new(), None) + .expect("check"); + assert_eq!(registry.supervisor_instance(), None); + registry.attach(&first, owner, None).expect("attach"); + assert_eq!(registry.supervisor_instance(), Some(owner)); + let candidate = SupervisorInstanceId::new(); + let ahead = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.check_attach(&ahead, candidate, None), + Err(SandboxAuthError::TakeoverRequired) + ); + let grant = gateway.grant(2, candidate); + registry + .check_attach(&ahead, candidate, Some(&grant)) + .expect("granted candidate passes the check"); + assert_eq!( + registry.confirm(&ahead), + Err(SandboxAuthError::ConnectionNotAttached), + "a check alone stages nothing" + ); + assert_eq!( + registry.check_attach(&ahead, SupervisorInstanceId::new(), Some(&grant)), + Err(SandboxAuthError::TakeoverMismatch) + ); + } + + #[test] + fn takeover_grants_bind_this_runtime_and_terminal_sessions_reject_them() { + let gateway = Gateway::new(); + let other = Gateway::new(); + let candidate = SupervisorInstanceId::new(); + // A grant signed by another gateway's key is not verifiable here. + let minted = other + .issuer + .mint_supervisor_takeover( + &SupervisorTakeover { + sandbox_id: SandboxId::parse("sandbox-a").unwrap(), + runtime_generation: SandboxGenerationId::parse("generation-1").unwrap(), + auth_epoch: epoch(2), + supervisor_instance: candidate.to_string(), + owner: "supervisor-b".to_string(), + }, + std::time::Duration::from_mins(5), + ) + .unwrap(); + assert!(matches!( + gateway + .authenticator + .verify_takeover(minted.token.expose_secret()), + Err(SandboxAuthError::Jwt(_)) + )); + // A grant for another sandbox or generation is refused. + for (sandbox, generation, expected) in [ + ("sandbox-b", "generation-1", SandboxAuthError::WrongSandbox), + ( + "sandbox-a", + "generation-2", + SandboxAuthError::WrongRuntimeGeneration, + ), + ] { + let minted = gateway + .issuer + .mint_supervisor_takeover( + &SupervisorTakeover { + sandbox_id: SandboxId::parse(sandbox).unwrap(), + runtime_generation: SandboxGenerationId::parse(generation).unwrap(), + auth_epoch: epoch(2), + supervisor_instance: candidate.to_string(), + owner: "supervisor-b".to_string(), + }, + std::time::Duration::from_mins(5), + ) + .unwrap(); + assert_eq!( + gateway + .authenticator + .verify_takeover(minted.token.expose_secret()), + Err(expected) + ); + } + let registry = gateway.registry(); + registry.mark_terminal(); + let replacement = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&replacement, candidate, Some(&gateway.grant(2, candidate))), + Err(SandboxAuthError::TerminalSession) + ); + } + + #[test] + fn a_fresh_runtime_accepts_a_granted_first_owner() { + // The launch owner died before attaching; its replacement starts the + // runtime under a grant, which pins the granted process. + let gateway = Gateway::new(); + let registry = gateway.registry(); + let candidate = SupervisorInstanceId::new(); + let replacement = gateway.connect(&gateway.token(2)); + assert_eq!( + registry.attach(&replacement, candidate, Some(&gateway.grant(2, candidate))), + Ok(None) + ); + registry.confirm(&replacement).expect("confirm"); + assert_eq!(registry.supervisor_instance(), Some(candidate)); + let launch_owner = gateway.connect(&gateway.token(1)); + assert_eq!( + registry.attach(&launch_owner, SupervisorInstanceId::new(), None), + Err(SandboxAuthError::StaleCredentialEpoch) + ); + } } diff --git a/crates/openshell-sandbox/Cargo.toml b/crates/openshell-sandbox/Cargo.toml index 78414b267a..88f3380580 100644 --- a/crates/openshell-sandbox/Cargo.toml +++ b/crates/openshell-sandbox/Cargo.toml @@ -80,8 +80,11 @@ socket2 = { workspace = true, features = ["all"] } tempfile = "3" [dev-dependencies] +hyper-util = { workspace = true } +openshell-ocsf = { path = "../openshell-ocsf", features = ["test-support"] } rcgen = { workspace = true } tempfile = "3" +tower = { workspace = true } [lints] workspace = true diff --git a/crates/openshell-sandbox/README.md b/crates/openshell-sandbox/README.md index 1cde277dd8..9ebd3a1d44 100644 --- a/crates/openshell-sandbox/README.md +++ b/crates/openshell-sandbox/README.md @@ -26,3 +26,56 @@ and dials a replacement when the supervisor starts HTTP/2 on it or the spare closes. Live connections keep the same disconnect, freeze, and recovery-deadline handling as listener mode, and the workload still cannot connect to the control port. + +A runtime of a supervisor shared by several sandboxes (one with a route handle) +names its route in a version 2 preface on every dial and follows the +supervisors' redirects, on the same port and never on loopback, at most two +hops. A supervisor that closes the negotiation without replying, as one from +before version 2 does, is dialed again with a version 1 preface and treated +as a lone supervisor; this changes only routing, since TLS pinning, request +authentication, and takeover grants are unaffected. It keeps a spare at the member whose connection it last confirmed, its +authenticated owner; an unauthenticated accept never makes a member the +owner. Once a supervisor names itself a pool member or redirects it, the +runtime also dials the configured addresses past every known member every 15 +seconds, up to three times, so a member that hosts the route without a spare +gets one: a previous owner that stays reachable after the gateway granted the +route elsewhere, holding an unused spare, cannot keep the new owner from the +runtime. A lone shared supervisor never names itself, so it holds the one +spare and is never probed. The runtime keeps at most three unused spares, one +per member, and never closes an established connection to refresh routing. A +member's address grants nothing: a later term at a fenced owner's address +authenticates with its own grant. A retired connection that does not shut down +gracefully within two seconds, such as an unused spare, has its transport +dropped. + +## Supervisor takeover + +The boundary advertises `openshell.sandbox.supervisor-takeover`. A replacement +supervisor process may attach to a running generation only with a +gateway-signed takeover grant for a later credential epoch, its bearer of that +epoch, and the grant's supervisor process identity. The attach is staged and +promoted at confirmation, which advances the accepted epoch, pins the new +process, and closes every connection of the previous owner; credentials of +earlier epochs are refused from then on. Refused attempts emit an OCSF +configuration-state event and a `sandbox-supervisor-takeover-denied` finding. + +Attach and confirmation are serialized, and the registry checks an attach +before it binds anything. Owner operations (exec, signals, termination, +provider updates, network decisions, DNS answers) hold an ownership guard from +authorization through their effect, and a promotion takes it exclusively, so +none of the replaced owner's operations takes effect afterwards. The guard is +asynchronous: tasks on the runtime's workers wait for it without blocking +them. A confirmation replay on the connection in force promotes nothing and +skips the guard. Promotion and the supervisory state change together, so a +replacement lost mid-promotion leaves the workload frozen under a fresh +recovery deadline. Each supervisor offers a recovery deadline when it attaches +(10 to 300 seconds; the runtime refuses any other before binding anything), and +the deadline of the connection in force when supervision is lost bounds the +freeze: 30 seconds without an offer. A supervisor pool offers 60 seconds unless +the gateway configures another. + +The runtime keeps the trust-on-first-use executable pins of the generation +(`binary_pins`): the first digest it offers for each path, from its own +measurements. A contradicting identity is offered as a resolution failure, +which every supervisor denies, so a replacement with an empty cache cannot pin +a binary the workload replaced. diff --git a/crates/openshell-sandbox/src/binary_pins.rs b/crates/openshell-sandbox/src/binary_pins.rs new file mode 100644 index 0000000000..32a6c94868 --- /dev/null +++ b/crates/openshell-sandbox/src/binary_pins.rs @@ -0,0 +1,191 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Trust-on-first-use executable pins kept by the sandbox runtime itself. +//! +//! A supervisor pins the digest of every executable in an identity chain the +//! first time a mediated TCP open offers it, before it evaluates policy, and +//! denies a later open whose executable at that path has another digest. Those +//! pins live in the supervisor process, so a replacement supervisor that takes +//! over a running generation would start without them and pin a binary the +//! workload replaced in the meantime. +//! +//! The runtime therefore records the first digest it ever offered for each +//! path, from its own measurements, for the lifetime of the generation, and +//! refuses to offer an identity that contradicts one: such an open reaches +//! every supervisor, old or new, as an identity failure, which it denies. The +//! pins are never relayed by a supervisor. The bound matches the supervisor's +//! identity cache, so a full journal refuses new paths as a full cache does. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Mutex; + +use openshell_isolation_interface::contract::{BinaryIdentity, ResolveError, Sha256Digest}; + +/// Most pinned paths, matching the supervisor's binary identity cache. +pub const MAX_BINARY_IDENTITY_PINS: usize = 4096; + +/// First digest offered for each executable path in this generation. +#[derive(Debug, Default)] +pub struct BinaryIdentityPins { + pins: Mutex>, +} + +impl BinaryIdentityPins { + /// Check an identity about to be offered for a TCP open against the + /// pins, recording the paths it introduces. + /// + /// An identity a supervisor could not pin (an unresolved one, a relative + /// path, a missing digest, or conflicting evidence) passes unchanged: the + /// supervisor rejects it itself. An executable whose digest differs from + /// its pin, or a new path once the journal is full, turns the identity + /// into a resolution failure. + pub fn admit( + &self, + identity: Result, + ) -> Result { + let Ok(identity) = identity else { + return identity; + }; + let mut offered = HashMap::new(); + for executable in std::iter::once(&identity.executable).chain(&identity.ancestors) { + let Some(digest) = executable.digest else { + return Ok(identity); + }; + if !executable.path.is_absolute() { + return Ok(identity); + } + if offered + .insert(executable.path.clone(), digest) + .is_some_and(|existing| existing != digest) + { + return Ok(identity); + } + } + let mut pins = self + .pins + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for (path, digest) in &offered { + if pins.get(path).is_some_and(|pinned| pinned != digest) { + return Err(ResolveError::Failed(format!( + "binary integrity violation: {} executable changed", + path.display() + ))); + } + } + let new_paths = offered + .keys() + .filter(|path| !pins.contains_key(*path)) + .count(); + if new_paths > MAX_BINARY_IDENTITY_PINS.saturating_sub(pins.len()) { + return Err(ResolveError::Failed(format!( + "binary identity pins exhausted (maximum {MAX_BINARY_IDENTITY_PINS} pinned paths)" + ))); + } + for (path, digest) in offered { + pins.entry(path).or_insert(digest); + } + Ok(identity) + } + + /// Number of pinned paths. + #[must_use] + pub fn len(&self) -> usize { + self.pins + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() + } + + /// Whether nothing is pinned yet. + #[must_use] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use openshell_isolation_interface::contract::ExecutableIdentity; + + use super::*; + + fn executable(path: &str, digest: char) -> ExecutableIdentity { + ExecutableIdentity { + path: PathBuf::from(path), + digest: Some(digest.to_string().repeat(64).parse().expect("digest")), + } + } + + /// A resolved identity, as offered for a TCP open. + fn identity(leaf: ExecutableIdentity, ancestors: Vec) -> BinaryIdentity { + BinaryIdentity { + executable: leaf, + ancestors, + cmdline_paths: Vec::new(), + } + } + + #[test] + fn first_offered_digests_are_pinned_and_contradictions_refused() { + let pins = BinaryIdentityPins::default(); + let curl = || executable("/usr/bin/curl", 'a'); + let shell = || executable("/bin/sh", 'b'); + assert!(pins.admit(Ok(identity(curl(), vec![shell()]))).is_ok()); + assert_eq!(pins.len(), 2); + assert!(pins.admit(Ok(identity(curl(), vec![]))).is_ok()); + // A replaced leaf or ancestor is refused, and pins never move. + for replaced in [ + identity(executable("/usr/bin/curl", 'c'), vec![shell()]), + identity(curl(), vec![executable("/bin/sh", 'd')]), + ] { + assert!(matches!( + pins.admit(Ok(replaced)), + Err(ResolveError::Failed(_)) + )); + } + assert!(pins.admit(Ok(identity(curl(), vec![shell()]))).is_ok()); + assert_eq!(pins.len(), 2); + } + + #[test] + fn identities_a_supervisor_rejects_itself_pass_unpinned() { + let pins = BinaryIdentityPins::default(); + assert!(matches!( + pins.admit(Err(ResolveError::NotFound)), + Err(ResolveError::NotFound) + )); + let missing_digest = ExecutableIdentity { + path: PathBuf::from("/usr/bin/curl"), + digest: None, + }; + let conflicting = vec![executable("/bin/sh", 'b'), executable("/bin/sh", 'c')]; + for unpinnable in [ + identity(missing_digest, vec![]), + identity(executable("relative/tool", 'a'), vec![]), + identity(executable("/usr/bin/env", 'a'), conflicting), + ] { + assert!(pins.admit(Ok(unpinnable)).is_ok()); + } + assert!(pins.is_empty()); + } + + #[test] + fn a_full_journal_refuses_new_paths_but_keeps_known_ones() { + let pins = BinaryIdentityPins::default(); + for index in 0..MAX_BINARY_IDENTITY_PINS { + let tool = identity(executable(&format!("/bin/tool-{index}"), 'a'), vec![]); + assert!(pins.admit(Ok(tool)).is_ok()); + } + assert!(matches!( + pins.admit(Ok(identity(executable("/bin/one-more", 'a'), vec![]))), + Err(ResolveError::Failed(_)) + )); + assert!( + pins.admit(Ok(identity(executable("/bin/tool-7", 'a'), vec![]))) + .is_ok() + ); + } +} diff --git a/crates/openshell-sandbox/src/boundary_server.rs b/crates/openshell-sandbox/src/boundary_server.rs index 56a974b9a1..e5efa25ceb 100644 --- a/crates/openshell-sandbox/src/boundary_server.rs +++ b/crates/openshell-sandbox/src/boundary_server.rs @@ -60,7 +60,7 @@ mod linux { sandbox_protocol_server::{SandboxProtocol, SandboxProtocolServer}, }; use openshell_sandbox_backend::sandbox_auth::{ - SandboxAuthError, SandboxConnectionId, SandboxConnectionRegistry, + Promotion, SandboxAuthError, SandboxConnectionId, SandboxConnectionRegistry, SandboxProtocolAuthenticator, SandboxProtocolPrincipal, }; use openshell_sandbox_backend::{ @@ -86,7 +86,6 @@ mod linux { const CONTROL_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5); const CONTROL_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10); const MEDIATION_REPLACEMENT_TIMEOUT: Duration = Duration::from_secs(20); - const AUTHENTICATED_RECONNECT_TIMEOUT: Duration = Duration::from_secs(30); const ENFORCEMENT_LOSS_TERMINATION_GRACE: Duration = Duration::from_secs(openshell_core::config::DEFAULT_STOP_TIMEOUT_SECS as u64); const FORCE_KILL_REAP_TIMEOUT: Duration = Duration::from_secs(2); @@ -95,12 +94,32 @@ mod linux { const FILESYSTEM_FLUSH_TIMEOUT: Duration = Duration::from_secs(3); const MAX_PENDING_HANDSHAKES: usize = 32; const MAX_CONTROL_CONNECTIONS: usize = 128; - /// Bounds one runtime-initiated TCP connect attempt. + /// Bounds one runtime-initiated dial, including finding the supervisor + /// pool member that hosts the route. const RUNTIME_DIAL_TIMEOUT: Duration = Duration::from_secs(5); + /// Bounds connecting to one supervisor address and negotiating the route + /// there, so an unreachable cached owner leaves time for the configured + /// addresses. + const RUNTIME_ROUTE_HOP_TIMEOUT: Duration = Duration::from_secs(2); + /// Redirects followed from each supervisor address a dial starts at. + const MAX_RUNTIME_ROUTE_REDIRECTS: usize = 2; + /// Spares a routed runtime keeps unused at once: one at its owner, and in + /// a supervisor pool one at each member that discovery finds hosting its + /// route. A lone shared supervisor is one member, so it holds one. + const POOL_MAX_UNCLAIMED_SPARES: usize = 3; + /// Interval of route discovery through the configured supervisor + /// addresses, which finds a new owner while the previous one stays + /// reachable. Only runs once a supervisor showed it is a pool member. + const POOL_ROUTE_DISCOVERY_INTERVAL: Duration = Duration::from_secs(15); + /// Dials in one discovery round. + const POOL_ROUTE_DISCOVERY_PROBES: usize = 3; const RUNTIME_DIAL_BACKOFF_MIN: Duration = Duration::from_millis(100); const RUNTIME_DIAL_BACKOFF_MAX: Duration = Duration::from_secs(2); /// A spare that dies this quickly without being claimed counts as a failure. const RUNTIME_SPARE_MIN_LIFETIME: Duration = Duration::from_secs(1); + /// How long a retired connection may take to shut down gracefully before + /// its transport is dropped. + const RETIRED_CONNECTION_GRACE: Duration = Duration::from_secs(2); /// Consecutive dial failures are expected while the supervisor starts; /// report a persistent failure at most this often. const RUNTIME_DIAL_WARNING_INTERVAL: Duration = Duration::from_secs(30); @@ -679,13 +698,33 @@ mod linux { /// connection before retiring the old one, never wait on this loop. TLS /// failures back off like dial failures. Established spares are served /// like accepted connections: same identity, slot limit, and disconnect, - /// freeze, and termination handling. + /// freeze, and termination handling. A runtime of a supervisor shared by + /// several sandboxes keeps its spares as [`run_routed_runtime_control`] + /// describes. async fn run_runtime_initiated_control( config: BoundaryListenerConfig, server_config: Arc, runtime: Arc, active_connections: Arc, ) { + if let BoundaryListenerConfig::RuntimeInitiatedTlsTcp { + addresses, + route: Some(route), + .. + } = &config + { + return run_routed_runtime_control( + SupervisorRoute { + addresses: addresses.clone(), + route: route.clone(), + }, + server_config, + runtime, + active_connections, + POOL_ROUTE_DISCOVERY_INTERVAL, + ) + .await; + } let mut backoff = RUNTIME_DIAL_BACKOFF_MIN; let mut warn_at: Option = None; loop { @@ -717,6 +756,7 @@ mod linux { runtime.clone(), active_connections.clone(), Some(claimed), + None, ); runtime.process_runtime.spawn(async move { if let Err(error) = served.await { @@ -727,15 +767,7 @@ mod linux { claim.await.is_ok() || established_at.elapsed() >= RUNTIME_SPARE_MIN_LIFETIME } Err(error) => { - // Failures are expected while the supervisor starts; report - // a persistent one at most once per interval. - let now = tokio::time::Instant::now(); - if now >= *warn_at.get_or_insert(now + RUNTIME_DIAL_WARNING_INTERVAL) { - warn_at = Some(now + RUNTIME_DIAL_WARNING_INTERVAL); - tracing::warn!(%error, "Runtime-initiated control cannot reach the supervisor; retrying"); - } else { - tracing::debug!(%error, "Runtime-initiated control dial failed; retrying"); - } + report_runtime_dial_failure(&mut warn_at, &error); false } }; @@ -748,31 +780,28 @@ mod linux { } } - /// Dial only driver-provided supervisor addresses from the bootstrap. - /// - /// A shared supervisor first receives this generation's route handle so - /// it can expect the right pinned identity; this runtime then serves TLS - /// exactly as it does for a dedicated supervisor. + /// Failures are expected while the supervisor starts; report a persistent + /// one at most once per interval. + fn report_runtime_dial_failure( + warn_at: &mut Option, + error: &dyn std::fmt::Display, + ) { + let now = tokio::time::Instant::now(); + if now >= *warn_at.get_or_insert(now + RUNTIME_DIAL_WARNING_INTERVAL) { + *warn_at = Some(now + RUNTIME_DIAL_WARNING_INTERVAL); + tracing::warn!(%error, "Runtime-initiated control cannot reach the supervisor; retrying"); + } else { + tracing::debug!(%error, "Runtime-initiated control dial failed; retrying"); + } + } + + /// Dial only driver-provided addresses of a dedicated supervisor from the + /// bootstrap. A shared supervisor is reached through + /// [`run_routed_runtime_control`] instead. async fn dial_supervisor(config: &BoundaryListenerConfig) -> io::Result { match config { - BoundaryListenerConfig::RuntimeInitiatedTlsTcp { - addresses, route, .. - } => { - let mut stream = - openshell_core::net::connect_tcp_nodelay_best_effort(addresses).await?; - let keepalive = socket2::TcpKeepalive::new() - .with_time(Duration::from_secs(10)) - .with_interval(Duration::from_secs(5)) - .with_retries(3); - let _ = socket2::SockRef::from(&stream).set_tcp_keepalive(&keepalive); - if let Some(route) = route { - let preface = - openshell_sandbox_backend::runtime_routed::encode_route_preface(route) - .map_err(|error| { - io::Error::new(io::ErrorKind::InvalidInput, error.to_string()) - })?; - stream.write_all(&preface).await?; - } + BoundaryListenerConfig::RuntimeInitiatedTlsTcp { addresses, .. } => { + let stream = connect_supervisor(addresses).await?; Ok(PlainControlStream::Tcp(stream.into_std()?)) } BoundaryListenerConfig::RuntimeInitiatedUnix { socket_path, .. } => { @@ -806,6 +835,434 @@ mod linux { } } + /// Connect to the first reachable supervisor address. Keepalive detects a + /// supervisor that vanishes while a spare sits idle. + async fn connect_supervisor( + addresses: &[std::net::SocketAddr], + ) -> io::Result { + let stream = openshell_core::net::connect_tcp_nodelay_best_effort(addresses).await?; + let keepalive = socket2::TcpKeepalive::new() + .with_time(Duration::from_secs(10)) + .with_interval(Duration::from_secs(5)) + .with_retries(3); + let _ = socket2::SockRef::from(&stream).set_tcp_keepalive(&keepalive); + Ok(stream) + } + + /// Where a runtime of a shared supervisor dials: the configured supervisor + /// addresses (the shared supervisor's Service) and its route handle. + struct SupervisorRoute { + addresses: Vec, + route: String, + } + + /// An authenticated connection a supervisor has not used yet. + struct RouteSpare { + member: std::net::SocketAddr, + claimed: tokio::sync::oneshot::Receiver<()>, + established_at: tokio::time::Instant, + } + + /// Unclaimed spares of a routed runtime, at most one per supervisor pool + /// member; a lone shared supervisor is one member. A member's address + /// grants nothing: a member whose earlier term was fenced may hold a + /// later one, which only its signed grant proves, so every member may be + /// offered a spare and stale credentials are refused on it. + #[derive(Default)] + struct RouteSpares { + unclaimed: Vec, + /// The member whose connection the runtime confirmed last. + owner: Option, + /// A supervisor named itself as a pool member or redirected this + /// runtime: its route can move to another member, so discovery runs. + pooled: bool, + } + + impl RouteSpares { + fn holds(&self, member: std::net::SocketAddr) -> bool { + self.unclaimed.iter().any(|spare| spare.member == member) + } + + /// Whether `member` may be given another spare. + fn admits(&self, member: std::net::SocketAddr) -> bool { + self.unclaimed.len() < POOL_MAX_UNCLAIMED_SPARES && !self.holds(member) + } + + /// Whether the confirmed owner, or before one is confirmed any + /// member, lacks a spare. + fn wants_spare(&self) -> bool { + self.unclaimed.len() < POOL_MAX_UNCLAIMED_SPARES + && self + .owner + .map_or(self.unclaimed.is_empty(), |owner| !self.holds(owner)) + } + + /// The next spare its member used (`true`) or that closed unused. + async fn next_settled(&mut self) -> (RouteSpare, bool) { + if self.unclaimed.is_empty() { + return std::future::pending().await; + } + let (index, claimed) = std::future::poll_fn(|context| { + for (index, spare) in self.unclaimed.iter_mut().enumerate() { + if let Poll::Ready(result) = Future::poll(Pin::new(&mut spare.claimed), context) + { + return Poll::Ready((index, result.is_ok())); + } + } + Poll::Pending + }) + .await; + (self.unclaimed.swap_remove(index), claimed) + } + } + + /// Keep a runtime of a shared supervisor reachable by whichever + /// supervisor owns its route. + /// + /// Every dial names the route in a version 2 preface. A spare waits at + /// the member whose connection the runtime confirmed last, its + /// authenticated owner, or wherever the configured addresses lead before + /// one is confirmed; an unauthenticated accept never makes a member the + /// owner. In a supervisor pool a previous owner can stay reachable, and + /// keep its spare unused, after the gateway granted the route to another + /// member, so once a supervisor names itself a pool member or redirects + /// this runtime, discovery rounds also dial the configured addresses past + /// every known member: a member found hosting the route without a spare + /// gets one, on which it can authenticate its takeover. A lone shared + /// supervisor is never probed. Spares are bounded per member and in + /// total, and established connections are never closed to refresh + /// routing. + async fn run_routed_runtime_control( + pool: SupervisorRoute, + server_config: Arc, + runtime: Arc, + active_connections: Arc, + discovery_interval: Duration, + ) { + let mut spares = RouteSpares::default(); + let mut backoff = RUNTIME_DIAL_BACKOFF_MIN; + let mut retry_at = tokio::time::Instant::now(); + let mut warn_at: Option = None; + let mut discovery = tokio::time::interval_at( + tokio::time::Instant::now() + discovery_interval, + discovery_interval, + ); + discovery.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + spares.owner = runtime.confirmed_route_owner(); + if spares.wants_spare() && tokio::time::Instant::now() >= retry_at { + let dialed = tokio::time::timeout( + RUNTIME_DIAL_TIMEOUT, + dial_route_owner(&pool.addresses, &pool.route, spares.owner), + ) + .await + .unwrap_or_else(|_| { + Err(io::Error::new( + io::ErrorKind::TimedOut, + "supervisor dial timed out", + )) + }); + let started = match dialed { + Ok(dialed) => { + spares.pooled |= dialed.shows_pool(); + if spares.admits(dialed.member) { + start_route_spare(dialed, &server_config, &runtime, &active_connections) + .await + .map(Some) + } else { + // The addresses led to a member that already + // holds a spare; discovery tries again. + Ok(None) + } + } + Err(error) => Err(error), + }; + match started { + Ok(Some(spare)) => { + warn_at = None; + spares.unclaimed.push(spare); + } + Ok(None) => retry_at = tokio::time::Instant::now() + discovery_interval, + Err(error) => { + report_runtime_dial_failure(&mut warn_at, &error); + retry_at = tokio::time::Instant::now() + backoff; + backoff = (backoff * 2).min(RUNTIME_DIAL_BACKOFF_MAX); + } + } + } + let waiting = spares.wants_spare() && retry_at > tokio::time::Instant::now(); + tokio::select! { + (spare, claimed) = spares.next_settled() => { + if claimed { + // Replace a used spare at once. + backoff = RUNTIME_DIAL_BACKOFF_MIN; + retry_at = tokio::time::Instant::now(); + } else if spare.established_at.elapsed() >= RUNTIME_SPARE_MIN_LIFETIME { + backoff = RUNTIME_DIAL_BACKOFF_MIN; + } else { + retry_at = tokio::time::Instant::now() + backoff; + backoff = (backoff * 2).min(RUNTIME_DIAL_BACKOFF_MAX); + } + } + _ = discovery.tick(), if spares.pooled => { + discover_route_owner(&pool, &mut spares, &server_config, &runtime, &active_connections) + .await; + } + () = tokio::time::sleep_until(retry_at), if waiting => {} + } + } + } + + /// One discovery round: dial the configured addresses, past any known + /// member, so a member that hosts the route without a spare gets one. A + /// redirect to a member that already holds a spare means the pool's + /// directory agrees with this runtime and ends the round. + async fn discover_route_owner( + pool: &SupervisorRoute, + spares: &mut RouteSpares, + server_config: &Arc, + runtime: &Arc, + active_connections: &Arc, + ) { + for _ in 0..POOL_ROUTE_DISCOVERY_PROBES { + if spares.unclaimed.len() >= POOL_MAX_UNCLAIMED_SPARES { + return; + } + let Ok(Ok(dialed)) = tokio::time::timeout( + RUNTIME_DIAL_TIMEOUT, + dial_route_owner(&pool.addresses, &pool.route, None), + ) + .await + else { + return; + }; + spares.pooled |= dialed.shows_pool(); + if spares.admits(dialed.member) { + tracing::info!( + member = %dialed.member, + "Supervisor pool member hosts this runtime's route without a spare; opening one" + ); + match start_route_spare(dialed, server_config, runtime, active_connections).await { + Ok(spare) => spares.unclaimed.push(spare), + Err(error) => { + tracing::debug!(%error, "Could not open a spare to a supervisor pool member"); + } + } + return; + } + if dialed.redirected { + return; + } + } + } + + /// Complete the runtime's TLS handshake with a reached supervisor and + /// serve the connection until the supervisor uses it or it closes. + async fn start_route_spare( + dialed: DialedMember, + server_config: &Arc, + runtime: &Arc, + active_connections: &Arc, + ) -> io::Result { + let stream = ControlStream::PendingTls { + stream: PlainControlStream::Tcp(dialed.stream.into_std()?), + server_config: server_config.clone(), + } + .establish_async(&runtime.process_runtime) + .await?; + let (claimed_tx, claimed) = tokio::sync::oneshot::channel(); + let served = serve_authenticated_connection( + stream, + runtime.clone(), + active_connections.clone(), + Some(claimed_tx), + Some(dialed.member), + ); + runtime.process_runtime.spawn(async move { + if let Err(error) = served.await { + tracing::debug!(%error, "Boundary control connection ended"); + } + }); + Ok(RouteSpare { + member: dialed.member, + claimed, + established_at: tokio::time::Instant::now(), + }) + } + + /// A supervisor that accepted this runtime's route. + struct DialedMember { + /// The member's runtime listener as it names itself, or the address + /// dialed when it names none. + member: std::net::SocketAddr, + stream: tokio::net::TcpStream, + /// The supervisor named itself, as only a pool member does. + named: bool, + /// A member's route directory led here. + redirected: bool, + } + + impl DialedMember { + /// Whether reaching this supervisor showed it belongs to a pool. + fn shows_pool(&self) -> bool { + self.named || self.redirected + } + } + + /// Find a supervisor pool member that hosts `route`: `preferred` first, + /// then each configured address, following at most + /// [`MAX_RUNTIME_ROUTE_REDIRECTS`] redirects from each. + /// + /// Replies are unauthenticated hints, so they only lead where a + /// configured address could (see [`route_redirect_allowed`]), and an + /// accept names a member without making it the owner. The member still + /// has to complete the generation-pinned TLS handshake, and every request + /// still carries the instance's bearer token. + async fn dial_route_owner( + addresses: &[std::net::SocketAddr], + route: &str, + preferred: Option, + ) -> io::Result { + use openshell_sandbox_backend::runtime_routed::RouteReply; + + let candidates = preferred.into_iter().chain( + addresses + .iter() + .copied() + .filter(|address| Some(*address) != preferred), + ); + let mut failure = io::Error::new( + io::ErrorKind::NotFound, + "no supervisor address is configured", + ); + for candidate in candidates { + let mut target = candidate; + let mut redirects = 0; + failure = loop { + let (stream, reply) = match negotiate_route_at(target, route).await { + Ok(negotiated) => negotiated, + Err(error) => break error, + }; + match reply { + RouteReply::Accept(named) => { + let member = named + .filter(|named| route_redirect_allowed(*named, addresses)) + .unwrap_or(target); + return Ok(DialedMember { + member, + stream, + named: named.is_some(), + redirected: redirects > 0, + }); + } + RouteReply::Redirect(_) if redirects == MAX_RUNTIME_ROUTE_REDIRECTS => { + break io::Error::new( + io::ErrorKind::InvalidData, + "supervisor route redirects exceeded the hop limit", + ); + } + RouteReply::Redirect(owner) if !route_redirect_allowed(owner, addresses) => { + break io::Error::new( + io::ErrorKind::PermissionDenied, + format!("refused supervisor route redirect to {owner}"), + ); + } + RouteReply::Redirect(owner) => { + tracing::debug!(%target, %owner, "Following supervisor route redirect"); + redirects += 1; + target = owner; + } + RouteReply::Unknown => { + break io::Error::new( + io::ErrorKind::NotFound, + format!("supervisor {target} does not know this runtime's route"), + ); + } + } + }; + tracing::debug!( + error = %failure, + %candidate, + "Supervisor address did not accept this runtime's route" + ); + } + Err(failure) + } + + /// Connect to one supervisor address and negotiate `route` there. + /// + /// A shared supervisor from before route negotiation closes a version 2 + /// preface without replying; it reads only the preface header, so the + /// close usually arrives as a reset. That address is then reached again with the + /// version 1 preface, which such a supervisor routes to its local + /// registration of the route, and the dial reads as an accept that names + /// no member: no pool evidence, so no discovery. The fallback changes only + /// routing; the supervisor still has to complete the generation-pinned + /// TLS handshake and authenticate every request, and a takeover still + /// needs its signed grant. + async fn negotiate_route_at( + address: std::net::SocketAddr, + route: &str, + ) -> io::Result<( + tokio::net::TcpStream, + openshell_sandbox_backend::runtime_routed::RouteReply, + )> { + use openshell_sandbox_backend::runtime_routed::{ + RouteReply, encode_route_preface, negotiate_route, + }; + + tokio::time::timeout(RUNTIME_ROUTE_HOP_TIMEOUT, async { + let mut stream = connect_supervisor(&[address]).await?; + match negotiate_route(&mut stream, route).await { + Err(error) + if matches!( + error.kind(), + io::ErrorKind::UnexpectedEof + | io::ErrorKind::ConnectionReset + | io::ErrorKind::ConnectionAborted + | io::ErrorKind::BrokenPipe + ) => + { + tracing::debug!( + %address, + "Supervisor closed the route negotiation without replying; using the version 1 preface" + ); + let mut stream = connect_supervisor(&[address]).await?; + let preface = encode_route_preface(route).map_err(|error| { + io::Error::new(io::ErrorKind::InvalidInput, error.to_string()) + })?; + stream.write_all(&preface).await?; + Ok((stream, RouteReply::Accept(None))) + } + negotiated => negotiated.map(|reply| (stream, reply)), + } + }) + .await + .unwrap_or_else(|_| { + Err(io::Error::new( + io::ErrorKind::TimedOut, + format!("supervisor {address} did not answer the route negotiation"), + )) + }) + } + + /// A redirect may only lead where a configured address could: to the + /// configured supervisor port, and to a loopback or unspecified address + /// only when a configured address is one. Validation refuses those for + /// configured addresses because workloads share the loopback interface. + fn route_redirect_allowed( + target: std::net::SocketAddr, + addresses: &[std::net::SocketAddr], + ) -> bool { + let is_local = |address: std::net::SocketAddr| { + let ip = address.ip().to_canonical(); + ip.is_loopback() || ip.is_unspecified() + }; + addresses.iter().any(|address| { + address.port() == target.port() && (!is_local(target) || is_local(*address)) + }) + } + /// Signals once the supervisor sends its first application bytes. struct ClaimOnFirstRead { inner: openshell_isolation_interface::contract::BoundaryDuplexStream, @@ -870,17 +1327,19 @@ mod linux { .await .map_err(|error| format!("authenticate boundary transport: {error}"))?; drop(pending); - serve_authenticated_connection(stream, runtime, active_connections, None).await + serve_authenticated_connection(stream, runtime, active_connections, None, None).await } /// Serve one TLS-established connection, accepted or runtime-initiated. /// `claimed`, used by runtime-initiated spares, fires when the supervisor - /// first sends data. + /// first sends data. `origin` names the supervisor pool member that + /// accepted a pool runtime's connection. async fn serve_authenticated_connection( stream: ControlStream, runtime: Arc, active_connections: Arc, claimed: Option>, + origin: Option, ) -> Result<(), String> { let Some(_slot) = acquire_control_connection_slot(&active_connections) else { return Err("authenticated control connection limit reached".to_string()); @@ -893,7 +1352,39 @@ mod linux { }), None => stream, }; - serve_grpc(stream, runtime, SandboxConnectionId::new()).await + let connection_id = SandboxConnectionId::new(); + let _origin = + origin.map(|origin| ConnectionOrigin::record(&runtime, connection_id, origin)); + serve_grpc(stream, runtime, connection_id).await + } + + /// The pool member a served connection came from, for as long as it is + /// served. + struct ConnectionOrigin { + runtime: std::sync::Weak, + connection_id: SandboxConnectionId, + } + + impl ConnectionOrigin { + fn record( + runtime: &Arc, + connection_id: SandboxConnectionId, + origin: std::net::SocketAddr, + ) -> Self { + lock(&runtime.connection_origins).insert(connection_id, origin); + Self { + runtime: Arc::downgrade(runtime), + connection_id, + } + } + } + + impl Drop for ConnectionOrigin { + fn drop(&mut self) { + if let Some(runtime) = self.runtime.upgrade() { + lock(&runtime.connection_origins).remove(&self.connection_id); + } + } } async fn serve_grpc( @@ -912,6 +1403,14 @@ mod linux { runtime: Arc::downgrade(&runtime), connection_id, }, + retired: Box::pin({ + let mut retired = connection_closed.clone(); + async move { + let _ = retired.changed().await; + tokio::time::sleep(RETIRED_CONNECTION_GRACE).await; + } + }), + ended: false, })]), tokio_stream::pending(), ); @@ -958,6 +1457,12 @@ mod linux { // bridges, including on keepalive failure or task cancellation. _connection_alive: tokio::sync::watch::Sender<()>, _disconnect: TransportDisconnectGuard, + /// Resolves a short grace after the connection is retired. A retired + /// connection shuts down gracefully, but one whose peer never started + /// HTTP/2, such as an unused spare, would wait for it forever; reads + /// then end so the transport is dropped. + retired: Pin + Send>>, + ended: bool, } struct TransportDisconnectGuard { @@ -979,6 +1484,10 @@ mod linux { context: &mut Context<'_>, buffer: &mut tokio::io::ReadBuf<'_>, ) -> Poll> { + if self.ended || self.retired.as_mut().poll(context).is_ready() { + self.ended = true; + return Poll::Ready(Ok(())); + } Pin::new(&mut self.stream).poll_read(context, buffer) } } @@ -1337,11 +1846,14 @@ mod linux { .map_err(|_| "mediation attach timed out".to_string())? .map_err(|error| format!("read mediation attach: {error}"))?; let request_id = request.request_id.clone(); - runtime.authorize_request(&principal, &request.request)?; + let owner = runtime + .authorize_request_async(&principal, &request.request) + .await?; if !matches!(request.request, Request::OpenMediation) { return Err("persistent mediation stream omitted OpenMediation".to_string()); } let mut response = runtime.dispatch(request); + drop(owner); let lease = if matches!(response, Response::MediationReady) { tokio::time::timeout( MEDIATION_REPLACEMENT_TIMEOUT, @@ -1378,13 +1890,14 @@ mod linux { return Ok(()); }; let broker = runtime.network_accept_context()?; - run_boundary_mediation(stream, runtime.clone(), broker).await + run_boundary_mediation(stream, runtime.clone(), broker, principal).await } async fn run_boundary_mediation( stream: tokio::io::DuplexStream, runtime: Arc, broker: NetworkBroker, + principal: SandboxProtocolPrincipal, ) -> Result<(), String> { let (mut reader, mut writer) = tokio::io::split(stream); let (outbound_tx, mut outbound_rx) = @@ -1412,8 +1925,13 @@ mod linux { } Ok::<(), String>(()) }; - let accept_task = - run_boundary_accepts(runtime, broker, outbound_tx.clone(), routes.clone()); + let accept_task = run_boundary_accepts( + runtime, + broker, + outbound_tx.clone(), + routes.clone(), + principal, + ); tokio::pin!(writer_task); tokio::pin!(reader_task); tokio::pin!(accept_task); @@ -1431,6 +1949,7 @@ mod linux { broker: NetworkBroker, outbound: tokio::sync::mpsc::Sender, routes: BoundaryMediationRoutes, + principal: SandboxProtocolPrincipal, ) -> Result<(), String> { loop { let pending = broker @@ -1448,6 +1967,7 @@ mod linux { route_rx, outbound.clone(), routes.clone(), + (runtime.clone(), principal.clone()), )); } } @@ -1458,6 +1978,7 @@ mod linux { mut inbound: tokio::sync::mpsc::Receiver, outbound: tokio::sync::mpsc::Sender, routes: BoundaryMediationRoutes, + (runtime, principal): (Arc, SandboxProtocolPrincipal), ) { let query = DnsQueryWire { request: pending.request.clone(), @@ -1500,7 +2021,19 @@ mod linux { "mediation session ended before DNS response", )), }; - let _ = pending.complete(result); + // Answer only for the owner in force; a replaced owner's answer fails. + { + let current = runtime.owner_operation_async(&principal).await; + let result = if current.is_ok() { + result + } else { + Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "the supervisor that answered is no longer the owner", + )) + }; + let _ = pending.complete(result); + } routes.lock().await.remove(&stream_id); } @@ -1514,7 +2047,8 @@ mod linux { .map_err(|error| format!("set control timeout: {error}"))?; let request: RequestEnvelope = read_frame(&mut stream).map_err(|error| format!("read control frame: {error}"))?; - runtime.authorize_request(principal, &request.request)?; + // Held until the operation takes effect; see `BoundaryRuntime::ownership`. + let owner = runtime.authorize_request(principal, &request.request)?; if request.validate_payload_digest().is_err() { let response = ResponseEnvelope { request_id: request.request_id, @@ -1531,6 +2065,7 @@ mod linux { let response = runtime .process_runtime .block_on(runtime.terminate_boundary()); + drop(owner); return write_frame( &mut stream, &ResponseEnvelope { @@ -1541,12 +2076,16 @@ mod linux { .map_err(|error| format!("write boundary termination response: {error}")); } Request::Exec { spec } => { - let started = match runtime.start_exec( + let started = runtime.start_exec( &request.request_id, &request.payload_digest, request.exec_admission_deadline_unix_ms, spec, - ) { + ); + // The exec has started or failed; its output streams outside + // the ownership guard. + drop(owner); + let started = match started { Ok(started) => started, Err(response) => { return write_frame( @@ -1574,7 +2113,9 @@ mod linux { return runtime.stream_process(stream, started.attachment); } Request::AttachProcess { process_id } => { - let (attachment, terminal) = match runtime.attach_process(&process_id) { + let attachment = runtime.attach_process(&process_id); + drop(owner); + let (attachment, terminal) = match attachment { Ok(attachment) => attachment, Err(response) => { return write_frame( @@ -1598,13 +2139,15 @@ mod linux { return runtime.stream_process(stream, attachment); } Request::LoopbackConnect { host, port } => { - let target = match LoopbackTarget::new(host, port) + let target = LoopbackTarget::new(host, port) .map_err(|error| format!("validate port-forward target: {error}")) .and_then(|target| { runtime .connect_port(target) .map_err(|error| format!("connect boundary loopback port: {error}")) - }) { + }); + drop(owner); + let target = match target { Ok(target) => target, Err(error) => { write_frame( @@ -1636,6 +2179,9 @@ mod linux { return Ok(()); } Request::AcceptNetwork => { + // Waiting for an open does not hold the ownership guard; the + // decision is admitted again when it is applied. + drop(owner); let broker = runtime.network_accept_context()?; let request_id = request.request_id; runtime.process_runtime.block_on(async move { @@ -1653,10 +2199,13 @@ mod linux { pending = broker.accept() => pending .map_err(|error| format!("accept sandbox network open: {error}"))?, }; + // Executables keep their first offered digest across + // supervisor owners; a contradiction is never offered. + let identity = runtime.binary_pins.admit(pending.identity.clone()); let response = encode_frame(&ResponseEnvelope { request_id, response: Response::NetworkConnected { - identity: BinaryIdentityWire::from(pending.identity.clone()), + identity: BinaryIdentityWire::from(identity), destination: pending.destination, socket: Some(pending.socket), policy_generation: 0, @@ -1686,8 +2235,22 @@ mod linux { } let decision = openshell_sandbox_backend::protobuf::decode(&payload) .map_err(|error| format!("decode network-open decision: {error}"))?; - let Some(target) = pending - .complete(decision) + // Apply the decision only while its supervisor is the + // owner in force; a replaced owner's decision is denied. + let decided = { + let current = runtime.owner_operation_async(principal).await; + let decision = if current.is_ok() { + decision + } else { + openshell_isolation_interface::contract::TcpOpenDecision::Denied( + openshell_isolation_interface::contract::TcpOpenDenial::MediationUnavailable, + ) + }; + pending.decide(decision) + } + .map_err(|error| format!("complete sandbox network open: {error}"))?; + let Some(target) = decided + .established() .await .map_err(|error| format!("complete sandbox network open: {error}"))? else { @@ -1708,25 +2271,119 @@ mod linux { } _ => {} } - let supervisor_instance_id = match &request.request { + let (supervisor_instance_id, supervisor_takeover, recovery_offer) = match &request.request { Request::Attach { supervisor_instance_id, + supervisor_takeover, + supervisor_recovery_timeout_secs, .. - } => Some(*supervisor_instance_id), - _ => None, + } => ( + Some(*supervisor_instance_id), + supervisor_takeover.clone(), + *supervisor_recovery_timeout_secs, + ), + _ => (None, None, None), }; let is_attach = supervisor_instance_id.is_some(); let is_confirm = matches!(&request.request, Request::Confirm); + // A recovery deadline outside the bounds is refused before the attach + // can touch the boundary: recovery is always bounded. + let recovery_offer = match recovery_offer { + None => None, + Some(seconds) => { + let Some(offer) = + openshell_sandbox_backend::boundary_protocol::supervisor_recovery_timeout( + seconds, + ) + else { + let range = + openshell_sandbox_backend::boundary_protocol::SUPERVISOR_RECOVERY_TIMEOUT_RANGE; + return write_frame( + &mut stream, + &ResponseEnvelope { + request_id: request.request_id, + response: guest_error( + BoundaryErrorKind::Invalid, + format!( + "supervisor recovery timeout must be between {} and {} seconds", + range.start(), + range.end() + ), + ), + }, + ) + .map_err(|error| format!("write control frame: {error}")); + }; + Some(offer) + } + }; + // Attach and confirm are serialized, so what an attach's precheck + // admits still holds when it commits. + let _admission = (is_attach || is_confirm).then(|| lock(&runtime.admission)); + let denied = |stream: &mut ControlStream, request_id: String, message: String| { + write_frame( + stream, + &ResponseEnvelope { + request_id, + response: guest_error(BoundaryErrorKind::Denied, message), + }, + ) + .map_err(|error| format!("write control frame: {error}")) + }; + // A takeover grant is verified before the attach can touch the + // boundary; an invalid grant is refused like any other credential. + let takeover = match supervisor_takeover + .as_deref() + .map(|grant| runtime.authenticator.verify_takeover(grant)) + .transpose() + { + Ok(takeover) => takeover, + Err(error) => { + emit_takeover_denial(&error.to_string(), None); + return denied(&mut stream, request.request_id, error.to_string()); + } + }; + // Refuse an attach the registry would refuse before the boundary + // binds anything for it. + if let Some(supervisor_instance_id) = supervisor_instance_id + && let Err(error) = runtime.connections.check_attach( + principal, + supervisor_instance_id, + takeover.as_ref(), + ) + { + if takeover.is_some() || attach_denial_is_takeover(&error) { + emit_takeover_denial(&error.to_string(), takeover.as_ref()); + } + let (kind, message) = connection_auth_error(error); + return write_frame( + &mut stream, + &ResponseEnvelope { + request_id: request.request_id, + response: guest_error(kind, message), + }, + ) + .map_err(|error| format!("write control frame: {error}")); + } let mut response = ResponseEnvelope { request_id: request.request_id.clone(), response: runtime.dispatch(request), }; + drop(owner); // Report commit failures to the supervisor: a silently closed stream // is indistinguishable from transport loss. if is_attach && matches!(&response.response, Response::Attached { .. }) { let supervisor_instance_id = supervisor_instance_id .ok_or_else(|| "attach request lost supervisor instance identity".to_string())?; - if let Err((kind, message)) = runtime.commit_attach(principal, supervisor_instance_id) { + if let Err((kind, message)) = runtime.commit_attach( + principal, + supervisor_instance_id, + takeover.as_ref(), + recovery_offer, + ) { + if takeover.is_some() { + emit_takeover_denial(&message, takeover.as_ref()); + } response.response = guest_error(kind, message); } } @@ -1781,6 +2438,9 @@ mod linux { nix::unistd::syncfs(directory.as_raw_fd()).map_err(io::Error::from) } + /// Admission of one owner operation; see `BoundaryRuntime::ownership`. + type OwnershipGuard<'a> = tokio::sync::RwLockReadGuard<'a, ()>; + struct BoundaryRuntime { config: BoundaryConfig, routing_handle: String, @@ -1793,6 +2453,42 @@ mod linux { supervisor_connection: Mutex, /// VM-only termination flush. filesystem_flush: Option, + /// How long a workload stays frozen without supervision before it is + /// terminated, unless the connection in force negotiated a deadline. + default_recovery_timeout: Duration, + /// Recovery deadlines supervisors offered when their connections + /// attached; the deadline of the connection in force governs the + /// freeze that follows its loss. Entries live as long as their + /// connections. + recovery_offers: Mutex>, + /// Recovery deadlines of the freezes this runtime started. + #[cfg(test)] + scheduled_recoveries: Mutex>, + /// Serializes attach and confirm, so a candidate's registry checks + /// hold from its precheck, before it touches the boundary, through + /// its commit. + admission: Mutex<()>, + /// Owner operations hold this for reading from their authorization + /// through their effect; promoting a connection takes it for writing. + /// A promoted replacement therefore waits for the replaced owner's + /// authorized operations, and none of them takes effect after the + /// promotion. Asynchronous tasks wait for it without blocking a + /// runtime worker: an operation may hold it across `block_on`, which + /// needs the workers to make progress. + ownership: tokio::sync::RwLock<()>, + /// Runs between a promotion and the supervisory state recording it. + #[cfg(test)] + promotion_hook: Mutex>>, + /// Executable pins of this generation, which outlive every supervisor + /// owner. + binary_pins: crate::binary_pins::BinaryIdentityPins, + /// Pool runtimes only: the supervisor pool member each served + /// connection came from. + connection_origins: + Mutex>, + /// Pool runtimes only: the member whose connection was confirmed + /// last, where the next spare goes first. + route_owner: Mutex>, next_recovery_id: AtomicU64, /// The wire policy bound at first attach, so an idempotent attach retry /// carrying a different policy is denied instead of silently keeping @@ -2089,7 +2785,20 @@ mod linux { connections: SandboxConnectionRegistry::new( config.session_id, config.session_rotation, + config.auth_epoch, ), + default_recovery_timeout: + openshell_sandbox_backend::boundary_protocol::DEFAULT_SUPERVISOR_RECOVERY_TIMEOUT, + recovery_offers: Mutex::new(std::collections::HashMap::new()), + #[cfg(test)] + scheduled_recoveries: Mutex::new(Vec::new()), + admission: Mutex::new(()), + ownership: tokio::sync::RwLock::new(()), + #[cfg(test)] + promotion_hook: Mutex::new(None), + binary_pins: crate::binary_pins::BinaryIdentityPins::default(), + connection_origins: Mutex::new(std::collections::HashMap::new()), + route_owner: Mutex::new(None), connection_shutdowns: Mutex::new(std::collections::HashMap::new()), routing_handle: uuid::Uuid::new_v4().to_string(), config, @@ -2118,14 +2827,20 @@ mod linux { self } + /// Authenticate a bearer and reject credentials of a replaced owner term. fn authenticate_request( &self, connection_id: SandboxConnectionId, metadata: &tonic::metadata::MetadataMap, ) -> Result { - self.authenticator + let principal = self + .authenticator .authenticate(connection_id, metadata) - .map_err(|error| tonic::Status::unauthenticated(error.to_string())) + .map_err(|error| tonic::Status::unauthenticated(error.to_string()))?; + self.connections + .admit(&principal) + .map_err(|error| tonic::Status::unauthenticated(error.to_string()))?; + Ok(principal) } fn discover_image_policy() -> Response { @@ -2169,45 +2884,177 @@ mod linux { } } - fn authorize_request( + /// Check a request that runs without the ownership guard; `false` + /// when it is an owner operation, which must take the guard. + fn admits_unguarded( &self, principal: &SandboxProtocolPrincipal, request: &Request, - ) -> Result<(), String> { - if matches!(request, Request::Attach { .. } | Request::DiscoverPolicy) { - return Ok(()); - } - if matches!(request, Request::Confirm) { - self.connections + ) -> Result { + match request { + Request::Attach { .. } | Request::DiscoverPolicy => Ok(true), + Request::Confirm => self + .connections .require_attached(principal) - .map_err(|error| error.to_string()) - } else { - self.connections + .map(|()| true) + .map_err(|error| error.to_string()), + // Waiting reads status and may block for the process lifetime. + Request::Wait { .. } => self + .connections .require_active(principal) - .map_err(|error| error.to_string()) + .map(|()| true) + .map_err(|error| error.to_string()), + _ => Ok(false), + } + } + + /// Authorize a request on a blocking thread. Owner operations return + /// the ownership guard, which the caller holds until the operation + /// has taken effect. + fn authorize_request( + &self, + principal: &SandboxProtocolPrincipal, + request: &Request, + ) -> Result>, String> { + if self.admits_unguarded(principal, request)? { + return Ok(None); + } + self.owner_operation(principal) + .map(Some) + .map_err(|error| error.to_string()) + } + + /// [`Self::authorize_request`] for an asynchronous task. + async fn authorize_request_async( + &self, + principal: &SandboxProtocolPrincipal, + request: &Request, + ) -> Result>, String> { + if self.admits_unguarded(principal, request)? { + return Ok(None); } + self.owner_operation_async(principal) + .await + .map(Some) + .map_err(|error| error.to_string()) + } + + /// Admit one operation of the owner in force, on a blocking thread. + /// While the returned guard lives, no confirmation can promote + /// another connection, so the operation cannot take effect for an + /// owner that was replaced. + fn owner_operation( + &self, + principal: &SandboxProtocolPrincipal, + ) -> Result, SandboxAuthError> { + let guard = self.ownership.blocking_read(); + self.connections.require_active(principal)?; + Ok(guard) + } + + /// [`Self::owner_operation`] for an asynchronous task, which must not + /// block its runtime worker while a promotion is pending. + async fn owner_operation_async( + &self, + principal: &SandboxProtocolPrincipal, + ) -> Result, SandboxAuthError> { + let guard = self.ownership.read().await; + self.connections.require_active(principal)?; + Ok(guard) } + /// Commit an attach the registry admits, with the recovery deadline + /// its supervisor offered (`None` keeps the default) for the freeze + /// that follows the loss of this connection once it is in force. fn commit_attach( &self, principal: &SandboxProtocolPrincipal, supervisor_instance_id: openshell_sandbox_backend::boundary_protocol::SupervisorInstanceId, + takeover: Option<&openshell_core::jwt::VerifiedSupervisorTakeover>, + recovery_offer: Option, ) -> Result<(), CommitError> { - if let Some(replaced) = self + let replaced = self .connections - .attach(principal, supervisor_instance_id) - .map_err(connection_auth_error)? + .attach(principal, supervisor_instance_id, takeover) + .map_err(connection_auth_error)?; { + let mut offers = lock(&self.recovery_offers); + // An attach that raced its connection's loss leaves an entry + // behind; keep only served connections once the map grows. + if offers.len() >= MAX_CONTROL_CONNECTIONS { + let served = lock(&self.connection_shutdowns); + offers.retain(|connection, _| served.contains_key(connection)); + } + match recovery_offer { + Some(offer) => offers.insert(principal.connection_id(), offer), + None => offers.remove(&principal.connection_id()), + }; + } + if let Some(replaced) = replaced { self.close_connection(replaced); } Ok(()) } fn commit_confirm(&self, principal: &SandboxProtocolPrincipal) -> Result<(), CommitError> { - let replaced = self + let terminating = || { + Err(( + BoundaryErrorKind::Terminated, + "sandbox session is terminating".to_string(), + )) + }; + // A replay on the connection in force, such as a credential + // renewal, promotes nothing, so it neither waits for nor holds up + // owner operations. The supervisory state lock keeps the + // connection in force while it is checked. + { + let connection = lock(&self.supervisor_connection); + if self.connections.require_active(principal).is_ok() { + if matches!( + *connection, + SupervisorConnectionState::Terminating + | SupervisorConnectionState::Terminal + ) { + self.connections.mark_terminal(); + return terminating(); + } + return Ok(()); + } + } + // Wait out every authorized operation of the current owner, and + // admit none until the connection that confirms is in force. + let _promotion = self.ownership.blocking_write(); + // The registry's promotion and the supervisory state that records + // it change together: a connection lost in between is seen as the + // one in force, so its loss freezes the workload under a recovery + // deadline. + let mut connection = lock(&self.supervisor_connection); + if matches!( + *connection, + SupervisorConnectionState::Terminating | SupervisorConnectionState::Terminal + ) { + self.connections.mark_terminal(); + return terminating(); + } + let Promotion { replaced, takeover } = self .connections .confirm(principal) .map_err(connection_auth_error)?; + #[cfg(test)] + { + let hook = lock(&self.promotion_hook).take(); + if let Some(hook) = hook { + hook(); + } + } + // The member whose connection is now in force is the + // authenticated owner a pool runtime dials first. + let origin = lock(&self.connection_origins) + .get(&principal.connection_id()) + .copied(); + if let Some(origin) = origin { + *lock(&self.route_owner) = Some(origin); + } let process = { let state = lock(&self.state); match &*state { @@ -2217,41 +3064,62 @@ mod linux { | RuntimeState::Ready(_) => None, } }; + if matches!(*connection, SupervisorConnectionState::Frozen { .. }) + && let Some(process) = process { - let mut connection = lock(&self.supervisor_connection); - if matches!( - *connection, - SupervisorConnectionState::Terminating | SupervisorConnectionState::Terminal - ) { - self.connections.mark_terminal(); + if !process.boundary_runtime.resume() { return Err(( - BoundaryErrorKind::Terminated, - "sandbox session is terminating".to_string(), + BoundaryErrorKind::Process, + "frozen workload could not be resumed".to_string(), )); } - if matches!(*connection, SupervisorConnectionState::Frozen { .. }) - && let Some(process) = process - { - if !process.boundary_runtime.resume() { - return Err(( - BoundaryErrorKind::Process, - "frozen workload could not be resumed".to_string(), - )); - } - self.network_broker.set_workload_frozen(false); - tracing::info!( - connection_id = ?principal.connection_id(), - "Sandbox Protocol connection recovered; workload resumed" - ); - } - *connection = SupervisorConnectionState::Connected(principal.connection_id()); + self.network_broker.set_workload_frozen(false); + tracing::info!( + connection_id = ?principal.connection_id(), + "Sandbox Protocol connection recovered; workload resumed" + ); } - if let Some(replaced) = replaced { + *connection = SupervisorConnectionState::Connected(principal.connection_id()); + drop(connection); + if let Some(takeover) = takeover { + // Every connection of the replaced owner term is stale now, + // including ones that never attached. + let retired = lock(&self.connection_shutdowns) + .keys() + .filter(|connection| **connection != principal.connection_id()) + .copied() + .collect::>(); + for connection in retired { + self.close_connection(connection); + } + tracing::info!( + auth_epoch = takeover.auth_epoch.get(), + owner = %takeover.owner, + "Supervisor takeover confirmed; previous owner fenced" + ); + openshell_ocsf::ocsf_emit!( + openshell_ocsf::ConfigStateChangeBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(openshell_ocsf::SeverityId::Informational) + .status(openshell_ocsf::StatusId::Success) + .state(openshell_ocsf::StateId::Enabled, "transferred") + .message(format!( + "Supervisor ownership transferred [term:{} owner:{}]", + takeover.auth_epoch.get(), + takeover.owner + )) + .build() + ); + } else if let Some(replaced) = replaced { self.close_connection(replaced); } Ok(()) } + /// The supervisor pool member whose connection was confirmed last. + fn confirmed_route_owner(&self) -> Option { + *lock(&self.route_owner) + } + fn register_connection( &self, connection_id: SandboxConnectionId, @@ -2272,34 +3140,37 @@ mod linux { if let Some(shutdown) = shutdown { let _ = shutdown.send(()); } - if !self.connections.disconnect(connection_id) { - return; - } - - let recovery_id = self.next_recovery_id.fetch_add(1, Ordering::Relaxed); - let process = { - let state = lock(&self.state); - match &*state { - RuntimeState::Running(process) => Some(process.clone()), - RuntimeState::AwaitingAttach - | RuntimeState::Bound(_) - | RuntimeState::Ready(_) => None, - } - }; - { + let recovery_offer = lock(&self.recovery_offers).remove(&connection_id); + let recovery_id = { + // Taken before the registry, as promotion does, so the loss + // of a connection promoted a moment ago finds it in force. let mut connection = lock(&self.supervisor_connection); + if !self.connections.disconnect(connection_id) { + return; + } if !matches!( *connection, SupervisorConnectionState::Connected(active) if active == connection_id ) { return; } + let process = { + let state = lock(&self.state); + match &*state { + RuntimeState::Running(process) => Some(process.clone()), + RuntimeState::AwaitingAttach + | RuntimeState::Bound(_) + | RuntimeState::Ready(_) => None, + } + }; if let Some(process) = &process { self.network_broker.set_workload_frozen(true); let _ = process.boundary_runtime.freeze(); } + let recovery_id = self.next_recovery_id.fetch_add(1, Ordering::Relaxed); *connection = SupervisorConnectionState::Frozen { recovery_id }; - } + recovery_id + }; tracing::warn!( recovery_id, "Sandbox Protocol connection lost; workload frozen pending authenticated recovery" @@ -2317,9 +3188,14 @@ mod linux { .message("Sandbox Protocol connection lost; workload frozen") .build() ); + // The supervisor whose connection was in force chose how long its + // replacement may take. + let recovery_timeout = recovery_offer.unwrap_or(self.default_recovery_timeout); + #[cfg(test)] + lock(&self.scheduled_recoveries).push(recovery_timeout); let runtime = Arc::downgrade(self); self.process_runtime.spawn(async move { - tokio::time::sleep(AUTHENTICATED_RECONNECT_TIMEOUT).await; + tokio::time::sleep(recovery_timeout).await; if let Some(runtime) = runtime.upgrade() { runtime.expire_recovery(recovery_id).await; } @@ -2500,7 +3376,11 @@ mod linux { let payload_digest = envelope.payload_digest; let response = match envelope.request { Request::DiscoverPolicy => Self::discover_image_policy(), - Request::Attach { policy, .. } => self.attach(*policy), + Request::Attach { + policy, + supervisor_takeover, + .. + } => self.attach(*policy, supervisor_takeover.is_some()), Request::Confirm => self.confirm(), Request::StartAgent { sandbox_id, @@ -2782,13 +3662,20 @@ mod linux { }) } - fn attach(&self, policy: SandboxPolicyWire) -> Response { + /// Bind the launch controls at the first attach. A replacement + /// supervisor attaching under a takeover grant keeps the controls + /// already bound: they cannot change for a running generation, and + /// the snapshot returns them for its later reattaches. + fn attach(&self, policy: SandboxPolicyWire, takeover: bool) -> Response { + let mut state = lock(&self.state); + let bound = !matches!(&*state, RuntimeState::AwaitingAttach); // Reject a conflicting request before establishing the boundary // or retaining the caller's policy for later replay. - if let Err(error) = validate_vm_policy_identity(&self.config, &policy) { + if !(bound && takeover) + && let Err(error) = validate_vm_policy_identity(&self.config, &policy) + { return guest_error(BoundaryErrorKind::Denied, error); } - let mut state = lock(&self.state); let accepted = match &*state { RuntimeState::AwaitingAttach => { let prepared = match PreparedBoundary::establish(self.network_broker.clone()) { @@ -2802,7 +3689,7 @@ mod linux { RuntimeState::Bound(_) | RuntimeState::Ready(_) | RuntimeState::Running(_) => { // Idempotent retry of the same attach; a different policy // must not be silently coalesced onto the bound boundary. - lock(&self.attached_policy).as_ref() == Some(&policy) + takeover || lock(&self.attached_policy).as_ref() == Some(&policy) } }; drop(state); @@ -2863,6 +3750,7 @@ mod linux { processes.sort_by(|left, right| left.process_id.cmp(&right.process_id)); SessionSnapshotWire { generation: self.config.generation.clone(), + bound_policy: lock(&self.attached_policy).clone().map(Box::new), processes, } } @@ -3702,6 +4590,60 @@ mod linux { type CommitError = (BoundaryErrorKind, String); + /// Attach refusals that mean another supervisor process tried to take + /// the generation over, or a replaced owner tried to return. + fn attach_denial_is_takeover(error: &SandboxAuthError) -> bool { + matches!( + error, + SandboxAuthError::TakeoverRequired + | SandboxAuthError::TakeoverMismatch + | SandboxAuthError::WrongSupervisorInstance + | SandboxAuthError::StaleCredentialEpoch + ) + } + + /// Report a refused supervisor takeover or replacement attach, as a + /// configuration-state event and a detection finding. Only identifiers + /// are recorded, never a token or grant. + fn emit_takeover_denial( + reason: &str, + takeover: Option<&openshell_core::jwt::VerifiedSupervisorTakeover>, + ) { + let message = takeover.map_or_else( + || format!("Supervisor takeover denied [reason:{reason}]"), + |takeover| { + format!( + "Supervisor takeover denied [reason:{reason} term:{} owner:{} grant:{}]", + takeover.takeover.auth_epoch.get(), + takeover.takeover.owner, + takeover.token_id + ) + }, + ); + tracing::warn!(%message, "Supervisor takeover denied"); + openshell_ocsf::ocsf_emit!( + openshell_ocsf::ConfigStateChangeBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(openshell_ocsf::SeverityId::Medium) + .status(openshell_ocsf::StatusId::Failure) + .state(openshell_ocsf::StateId::Disabled, "takeover_denied") + .message(message.clone()) + .build() + ); + openshell_ocsf::ocsf_emit!( + openshell_ocsf::DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(openshell_ocsf::ActivityId::Open) + .severity(openshell_ocsf::SeverityId::Medium) + .confidence(openshell_ocsf::ConfidenceId::High) + .is_alert(true) + .finding_info(openshell_ocsf::FindingInfo::new( + "sandbox-supervisor-takeover-denied", + "Sandbox Supervisor Takeover Denied", + )) + .message(message) + .build() + ); + } + fn connection_auth_error(error: SandboxAuthError) -> CommitError { let kind = match error { SandboxAuthError::ConnectionStillActive => BoundaryErrorKind::Unavailable, @@ -4253,6 +5195,7 @@ mod linux { }; use openshell_sandbox_backend::proto::{exchange_frame::Payload, mediate_frame}; use openshell_sandbox_backend::protobuf::WireMessage as _; + use openshell_sandbox_backend::runtime_routed::{RuntimeRouter, new_route_handle}; use rcgen::PKCS_ED25519; struct FailedWaitProcess; @@ -4871,9 +5814,10 @@ mod linux { .authenticate_request(first_id, bearer_request((), &token).metadata()) .expect("first principal"); runtime - .commit_attach(&first, supervisor_instance_id) + .commit_attach(&first, supervisor_instance_id, None, None) .expect("attach first"); - runtime.commit_confirm(&first).expect("confirm first"); + // Confirmation runs on a blocking thread in production. + tokio::task::block_in_place(|| runtime.commit_confirm(&first)).expect("confirm first"); assert_eq!( *lock(&runtime.supervisor_connection), SupervisorConnectionState::Connected(first_id) @@ -4886,7 +5830,7 @@ mod linux { .expect("early replacement principal"); assert!( matches!( - runtime.commit_attach(&early, supervisor_instance_id), + runtime.commit_attach(&early, supervisor_instance_id, None, None), Err((BoundaryErrorKind::Unavailable, _)) ), "reattach before the old transport is retired must be retryable" @@ -4904,19 +5848,18 @@ mod linux { .expect("replacement principal"); assert!( runtime - .commit_attach(&replacement, test_supervisor_instance_id()) + .commit_attach(&replacement, test_supervisor_instance_id(), None, None) .is_err(), "a replacement supervisor process must not claim the generation" ); runtime - .commit_attach(&replacement, supervisor_instance_id) + .commit_attach(&replacement, supervisor_instance_id, None, None) .expect("reattach replacement"); assert_eq!( runtime.connections.require_active(&replacement), Err(SandboxAuthError::ConnectionNotAttached) ); - runtime - .commit_confirm(&replacement) + tokio::task::block_in_place(|| runtime.commit_confirm(&replacement)) .expect("reconfirm replacement"); assert_eq!( *lock(&runtime.supervisor_connection), @@ -4940,9 +5883,9 @@ mod linux { .authenticate_request(connection_id, bearer_request((), &token).metadata()) .expect("test principal"); runtime - .commit_attach(&principal, supervisor_instance_id) + .commit_attach(&principal, supervisor_instance_id, None, None) .expect("attach"); - runtime.commit_confirm(&principal).expect("confirm"); + tokio::task::block_in_place(|| runtime.commit_confirm(&principal)).expect("confirm"); runtime.transport_disconnected(connection_id); let connection_state = *lock(&runtime.supervisor_connection); let recovery_id = match connection_state { @@ -4961,7 +5904,7 @@ mod linux { .expect("replacement principal"); assert!( runtime - .commit_attach(&replacement, supervisor_instance_id) + .commit_attach(&replacement, supervisor_instance_id, None, None) .is_err() ); } @@ -5134,9 +6077,9 @@ mod linux { .authenticate_request(connection_id, bearer_request((), &token).metadata()) .expect("test principal"); runtime - .commit_attach(&principal, supervisor_instance_id) + .commit_attach(&principal, supervisor_instance_id, None, None) .expect("attach"); - runtime.commit_confirm(&principal).expect("confirm"); + tokio::task::block_in_place(|| runtime.commit_confirm(&principal)).expect("confirm"); assert_eq!( runtime.terminate_boundary().await, @@ -5162,7 +6105,7 @@ mod linux { .expect("test principal"); runtime .connections - .attach(&principal, supervisor_instance_id) + .attach(&principal, supervisor_instance_id, None) .expect("attach test connection"); runtime .connections @@ -5253,7 +6196,7 @@ mod linux { .expect("test principal"); runtime .connections - .attach(&principal, supervisor_instance_id) + .attach(&principal, supervisor_instance_id, None) .expect("attach test connection"); runtime .connections @@ -5297,7 +6240,7 @@ mod linux { .expect("test principal"); runtime .connections - .attach(&principal, test_supervisor_instance_id()) + .attach(&principal, test_supervisor_instance_id(), None) .expect("attach test connection"); runtime .connections @@ -5653,7 +6596,7 @@ mod linux { (Some(wrong_user), None, "run_as_user", wrong_user), (None, Some(wrong_group), "run_as_group", wrong_group), ] { - let response = boundary.attach(vm_identity_test_policy(user, group)); + let response = boundary.attach(vm_identity_test_policy(user, group), false); let Response::Error { kind, message } = response else { panic!("conflicting identity was attached: {response:?}"); }; @@ -5670,7 +6613,10 @@ mod linux { assert!(lock(&boundary.attached_policy).is_none()); } assert!(matches!( - boundary.attach(vm_identity_test_policy(Some("sandbox"), Some("sandbox"))), + boundary.attach( + vm_identity_test_policy(Some("sandbox"), Some("sandbox")), + false + ), Response::Attached { .. } )); } @@ -5715,6 +6661,615 @@ mod linux { } } + /// Find the member hosting `route` as a pool runtime does, then serve + /// the runtime's TLS identity, which only a supervisor expecting this + /// route completes. + async fn dial_and_serve_tls( + addresses: &[std::net::SocketAddr], + route: &str, + preferred: Option, + server_config: &Arc, + ) -> io::Result<(std::net::SocketAddr, bool)> { + let dialed = dial_route_owner(addresses, route, preferred).await?; + let reached = (dialed.member, dialed.redirected); + ControlStream::PendingTls { + stream: PlainControlStream::Tcp(dialed.stream.into_std()?), + server_config: server_config.clone(), + } + .establish_async(&tokio::runtime::Handle::current()) + .await + .map(|_| reached) + } + + #[tokio::test] + async fn runtime_follows_route_redirect_to_its_owner_and_dials_a_preferred_member_first() { + let directory = tempfile::tempdir().expect("temporary directory"); + let (server_tls, client_tls) = stage_test_tls(directory.path(), "redirect"); + let server_config = + Arc::new(load_tls_server_config(&server_tls).expect("load runtime TLS identity")); + let route = new_route_handle(); + // Pool members share the configured port on different addresses. + let member = RuntimeRouter::bind("127.0.0.1:0".parse().unwrap()).expect("bind member"); + let owner = RuntimeRouter::bind(std::net::SocketAddr::new( + "127.0.0.2".parse().unwrap(), + member.local_addr().port(), + )) + .expect("bind owner"); + let registration = owner.register(&route, &client_tls).expect("register route"); + let directory_naming = |owner_address: std::net::SocketAddr| { + std::collections::HashMap::from([(route.clone(), owner_address)]) + }; + member.set_directory( + Some(member.local_addr()), + directory_naming(owner.local_addr()), + ); + owner.set_directory( + Some(owner.local_addr()), + directory_naming(owner.local_addr()), + ); + let configured = [member.local_addr()]; + + // The configured member redirects to the owner, which names + // itself and completes the pinned handshake. + assert_eq!( + dial_and_serve_tls(&configured, &route, None, &server_config) + .await + .expect("dial through redirect"), + (owner.local_addr(), true) + ); + + // A preferred member is dialed directly, even once the + // configured member no longer knows the route. + member.set_directory(Some(member.local_addr()), std::collections::HashMap::new()); + assert_eq!( + dial_and_serve_tls( + &configured, + &route, + Some(owner.local_addr()), + &server_config + ) + .await + .expect("dial preferred member"), + (owner.local_addr(), false) + ); + + // A preferred member that no longer hosts the route is passed + // over for the configured addresses. + drop(registration); + let _registration = member.register(&route, &client_tls).expect("move route"); + member.set_directory( + Some(member.local_addr()), + directory_naming(member.local_addr()), + ); + assert_eq!( + dial_and_serve_tls( + &configured, + &route, + Some(owner.local_addr()), + &server_config + ) + .await + .expect("dial configured member"), + (member.local_addr(), false) + ); + } + + #[tokio::test] + async fn a_supervisor_from_before_route_negotiation_is_reached_with_the_first_preface() { + // A shared supervisor that only knows version 1 prefaces closes a + // version 2 one without replying, as the #4702 router did. + let legacy = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("legacy supervisor"); + let address = legacy.local_addr().expect("legacy address"); + let route = new_route_handle(); + let (seen_tx, mut seen) = tokio::sync::mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + // Connections kept for the TLS handshake that follows; they + // close when the task is aborted. + #[allow(clippy::collection_is_never_read)] + let mut kept = Vec::new(); + loop { + let (mut stream, _) = legacy.accept().await.expect("accept"); + let mut header = [0_u8; 6]; + tokio::io::AsyncReadExt::read_exact(&mut stream, &mut header) + .await + .expect("preface header"); + if header[4] != 1 { + seen_tx.send((header[4], String::new())).expect("record"); + continue; + } + let mut handle = vec![0_u8; usize::from(header[5])]; + tokio::io::AsyncReadExt::read_exact(&mut stream, &mut handle) + .await + .expect("route handle"); + seen_tx + .send((header[4], String::from_utf8(handle).expect("handle"))) + .expect("record"); + kept.push(stream); + } + }); + let (_stream, reply) = negotiate_route_at(address, &route) + .await + .expect("the legacy supervisor is reached"); + assert_eq!( + reply, + openshell_sandbox_backend::runtime_routed::RouteReply::Accept(None) + ); + assert_eq!(seen.recv().await, Some((2, String::new()))); + assert_eq!(seen.recv().await, Some((1, route))); + server.abort(); + + // A current supervisor's refusal is a reply, not a reason to fall + // back: the runtime stays on the negotiated path. + let router = RuntimeRouter::bind("127.0.0.1:0".parse().expect("router address")) + .expect("router"); + let unknown = new_route_handle(); + let (_stream, reply) = negotiate_route_at(router.local_addr(), &unknown) + .await + .expect("negotiated"); + assert_eq!( + reply, + openshell_sandbox_backend::runtime_routed::RouteReply::Unknown + ); + } + + #[tokio::test] + async fn runtime_refuses_route_redirects_off_the_supervisor_port_or_past_two_hops() { + let route = new_route_handle(); + let member = RuntimeRouter::bind("127.0.0.1:0".parse().unwrap()).expect("bind member"); + let configured = [member.local_addr()]; + + // A redirect to another port is refused before it is dialed. + let elsewhere = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind other listener"); + member.set_directory( + Some(member.local_addr()), + std::collections::HashMap::from([( + route.clone(), + elsewhere.local_addr().expect("other listener address"), + )]), + ); + let error = dial_route_owner(&configured, &route, None) + .await + .err() + .expect("an off-port redirect is refused"); + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied, "{error}"); + assert!( + tokio::time::timeout(Duration::from_millis(200), elsewhere.accept()) + .await + .is_err() + ); + + // Two members that each name the other are abandoned after two + // redirects. + let peer = RuntimeRouter::bind(std::net::SocketAddr::new( + "127.0.0.2".parse().unwrap(), + member.local_addr().port(), + )) + .expect("bind peer"); + member.set_directory( + Some(member.local_addr()), + std::collections::HashMap::from([(route.clone(), peer.local_addr())]), + ); + peer.set_directory( + Some(peer.local_addr()), + std::collections::HashMap::from([(route.clone(), member.local_addr())]), + ); + let error = dial_route_owner(&configured, &route, None) + .await + .err() + .expect("a redirect loop is abandoned"); + assert!(error.to_string().contains("hop limit"), "{error}"); + } + + /// A routed runtime of a lone shared supervisor, which never names + /// itself, keeps its one spare there and never probes the configured + /// addresses; a runtime whose supervisor named itself a pool member + /// runs discovery rounds through them. + #[test] + fn only_a_runtime_that_saw_a_pool_probes_for_its_route() { + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + for pooled in [false, true] { + let gateway = TestGateway::new(); + let boundary = takeover_test_runtime( + &gateway, + "sandbox-probe", + "generation-probe", + &process_runtime, + ); + let directory = tempfile::tempdir().expect("temporary directory"); + let (server_tls, client_tls) = stage_test_tls(directory.path(), "probe"); + let server_config = Arc::new( + load_tls_server_config(&server_tls).expect("load runtime TLS identity"), + ); + let route = new_route_handle(); + let supervisor = + RuntimeRouter::bind("127.0.0.1:0".parse().unwrap()).expect("bind supervisor"); + let _registration = supervisor + .register(&route, &client_tls) + .expect("supervisor hosts route"); + if pooled { + supervisor.set_directory( + Some(supervisor.local_addr()), + std::collections::HashMap::from([(route.clone(), supervisor.local_addr())]), + ); + } + // The configured address counts every connection the runtime + // opens before passing it on to the supervisor. + let service = process_runtime + .block_on(tokio::net::TcpListener::bind(( + "127.0.0.3", + supervisor.local_addr().port(), + ))) + .expect("bind service"); + let service_address = service.local_addr().expect("service address"); + let dialed = Arc::new(AtomicUsize::new(0)); + let target = supervisor.local_addr(); + let counter = dialed.clone(); + let forwarding = process_runtime.spawn(async move { + while let Ok((mut inbound, _)) = service.accept().await { + counter.fetch_add(1, Ordering::SeqCst); + tokio::spawn(async move { + if let Ok(mut outbound) = tokio::net::TcpStream::connect(target).await { + let _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound) + .await; + } + }); + } + }); + let control = process_runtime.spawn(run_routed_runtime_control( + SupervisorRoute { + addresses: vec![service_address], + route: route.clone(), + }, + server_config, + boundary, + Arc::new(AtomicUsize::new(0)), + Duration::from_millis(100), + )); + std::thread::sleep(Duration::from_millis(1200)); + let connections = dialed.load(Ordering::SeqCst); + control.abort(); + forwarding.abort(); + if pooled { + assert!( + connections > 1, + "a pool runtime must run discovery rounds: {connections} connections" + ); + } else { + assert_eq!( + connections, 1, + "a lone supervisor gets exactly one spare and no probes" + ); + } + } + } + + /// The gateway granted a pool runtime's route to member B while the + /// previous owner A stays reachable: A keeps its stale route and an + /// unused spare and never closes them. Discovery through the + /// configured addresses still reaches B within a few rounds without + /// touching A's connection; B authenticates its term on that + /// connection and takes over, which closes A's spare, refuses A's + /// credentials, and makes B the member the runtime dials first. When + /// a restarted process at A's address is then granted a later term, + /// it reaches the runtime and takes over from B the same way: the + /// address that held a fenced term grants nothing and blocks nothing. + #[test] + #[allow(clippy::too_many_lines)] + fn a_reachable_stale_owner_cannot_keep_a_new_owner_from_the_runtime() { + const SANDBOX: &str = "sandbox-stale-owner"; + const GENERATION: &str = "generation-stale-owner"; + const BUDGET: Duration = Duration::from_secs(10); + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let directory = tempfile::tempdir().expect("temporary directory"); + let (server_tls, client_tls) = stage_test_tls(directory.path(), "stale-owner"); + let server_config = + Arc::new(load_tls_server_config(&server_tls).expect("load runtime TLS identity")); + let route = new_route_handle(); + + // Two pool members and the pool's Service, which spreads + // connections over them in turn, all on the runtime port. + let member_a = RuntimeRouter::bind("127.0.0.1:0".parse().unwrap()).expect("bind A"); + let port = member_a.local_addr().port(); + let member_b = RuntimeRouter::bind(std::net::SocketAddr::new( + "127.0.0.2".parse().unwrap(), + port, + )) + .expect("bind B"); + let (a, b) = (member_a.local_addr(), member_b.local_addr()); + let service = process_runtime + .block_on(tokio::net::TcpListener::bind(("127.0.0.3", port))) + .expect("bind service"); + let service_address = service.local_addr().expect("service address"); + process_runtime.spawn(async move { + let backends = [a, b]; + let mut next = 0_usize; + while let Ok((mut inbound, _)) = service.accept().await { + let backend = backends[next % backends.len()]; + next += 1; + tokio::spawn(async move { + if let Ok(mut outbound) = tokio::net::TcpStream::connect(backend).await { + let _ = + tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await; + } + }); + } + }); + let naming = |owner: std::net::SocketAddr| { + std::collections::HashMap::from([(route.clone(), owner)]) + }; + member_a.set_directory(Some(a), naming(a)); + member_b.set_directory(Some(b), naming(a)); + let registration_a = member_a + .register(&route, &client_tls) + .expect("A hosts route"); + process_runtime.spawn(run_routed_runtime_control( + SupervisorRoute { + addresses: vec![service_address], + route: route.clone(), + }, + server_config, + boundary.clone(), + Arc::new(AtomicUsize::new(0)), + Duration::from_millis(200), + )); + + // A takes the runtime's spare and is confirmed on it with launch + // credentials, then never uses or closes it. + let mut stale_spare = process_runtime + .block_on(tokio::time::timeout( + BUDGET, + registration_a.accept(&client_tls), + )) + .expect("A receives a spare") + .expect("authenticated spare"); + let known_connections = || { + lock(&boundary.connection_origins) + .keys() + .copied() + .collect::>() + }; + // The connection a member opened after `known` was taken. + let connection_from = + |member: std::net::SocketAddr, known: &std::collections::HashSet<_>| { + let deadline = std::time::Instant::now() + BUDGET; + loop { + let found = lock(&boundary.connection_origins) + .iter() + .find(|(connection, origin)| { + **origin == member && !known.contains(*connection) + }) + .map(|(connection, _)| *connection); + if let Some(connection) = found { + return connection; + } + assert!( + std::time::Instant::now() < deadline, + "no new connection from {member}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + }; + let authenticate = |connection, epoch| { + boundary + .authenticate_request( + connection, + bearer_request((), &gateway.token(SANDBOX, GENERATION, epoch)).metadata(), + ) + .expect("principal") + }; + let launch_owner = + authenticate(connection_from(a, &std::collections::HashSet::default()), 1); + boundary + .commit_attach(&launch_owner, test_supervisor_instance_id(), None, None) + .expect("attach A"); + boundary.commit_confirm(&launch_owner).expect("confirm A"); + assert_eq!(boundary.confirmed_route_owner(), Some(a)); + + // The gateway grants the route to B; A, cut off from the + // gateway, still names itself. + let registration_b = member_b + .register(&route, &client_tls) + .expect("B hosts route"); + member_b.set_directory(Some(b), naming(b)); + let delivered = process_runtime + .block_on(tokio::time::timeout( + BUDGET, + registration_b.accept(&client_tls), + )) + .expect("B receives a runtime connection within the budget") + .expect("authenticated connection to B"); + // Nothing of A's was closed to find B: its unused spare holds + // only the runtime's HTTP/2 greeting. + let mut greeting = [0_u8; 256]; + let still_open = process_runtime.block_on(tokio::time::timeout( + Duration::from_millis(300), + tokio::io::AsyncReadExt::read(&mut stale_spare, &mut greeting), + )); + assert!( + matches!(still_open, Err(_) | Ok(Ok(1..))), + "A's unused spare must stay open: {still_open:?}" + ); + + // B negotiates and authenticates its term on that connection. + let mut delivered = Some(delivered); + let channel = process_runtime + .block_on( + tonic::transport::Endpoint::from_static("http://boundary.openshell.internal") + .connect_with_connector(tower::service_fn( + move |_: tonic::transport::Uri| { + let stream = delivered.take(); + async move { + stream.map(hyper_util::rt::TokioIo::new).ok_or_else(|| { + io::Error::other("the delivered connection is in use") + }) + } + }, + )), + ) + .expect("gRPC over B's connection"); + let mut control = openshell_sandbox_backend::backend_proto::isolation_backend_client::IsolationBackendClient::new(channel); + let described = process_runtime + .block_on(control.describe(DescribeIsolationBackendRequest { + supervisor: Some(openshell_sandbox_backend::isolation_protocol::metadata( + "test/supervisor", + )), + })) + .expect("describe") + .into_inner(); + assert!( + described + .backend + .expect("backend metadata") + .supported_capabilities + .iter() + .any(|capability| { + capability + == openshell_sandbox_backend::isolation_protocol::SUPERVISOR_TAKEOVER_CAPABILITY + }) + ); + process_runtime + .block_on(control.open_boundary(bearer_request( + test_open_request(&boundary.config), + &gateway.token(SANDBOX, GENERATION, 2), + ))) + .expect("B opens the boundary with its term's credential"); + let replacement = test_supervisor_instance_id(); + let grant = boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, GENERATION, 2, replacement)) + .expect("grant"); + let candidate = + authenticate(connection_from(b, &std::collections::HashSet::default()), 2); + boundary + .commit_attach(&candidate, replacement, Some(&grant), None) + .expect("stage B"); + boundary.commit_confirm(&candidate).expect("confirm B"); + + // A is fenced: its spare is closed and its credentials are + // refused, and the runtime gives its next spare to B. + assert_eq!(boundary.confirmed_route_owner(), Some(b)); + process_runtime + .block_on(tokio::time::timeout(BUDGET, async { + while matches!( + tokio::io::AsyncReadExt::read(&mut stale_spare, &mut greeting).await, + Ok(1..) + ) {} + })) + .expect("A's spare closes"); + assert!( + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, 1)).metadata(), + ) + .is_err(), + "the fenced term's credentials are refused" + ); + process_runtime + .block_on(tokio::time::timeout( + BUDGET, + registration_b.accept(&client_tls), + )) + .expect("the owner receives the next spare") + .expect("authenticated spare for B"); + + // A's process restarts at the same address and the gateway grants + // it the next term while B still holds the route. + let returned_at = std::time::Instant::now(); + drop(registration_a); + // Connections the previous process queued close with it. + let drained_by = std::time::Instant::now() + BUDGET; + while lock(&boundary.connection_origins) + .values() + .any(|origin| *origin == a) + { + assert!( + std::time::Instant::now() < drained_by, + "A's old connections linger" + ); + std::thread::sleep(Duration::from_millis(10)); + } + let before_return = known_connections(); + let registration_a = member_a + .register(&route, &client_tls) + .expect("A hosts again"); + member_a.set_directory(Some(a), naming(a)); + member_b.set_directory(Some(b), naming(a)); + let _returned = process_runtime + .block_on(tokio::time::timeout( + BUDGET, + registration_a.accept(&client_tls), + )) + .expect("A's new process receives a runtime connection within the budget") + .expect("authenticated connection to A"); + let restarted = test_supervisor_instance_id(); + let grant = boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, GENERATION, 3, restarted)) + .expect("grant"); + let returning = authenticate(connection_from(a, &before_return), 3); + boundary + .commit_attach(&returning, restarted, Some(&grant), None) + .expect("stage A's later term"); + boundary + .commit_confirm(&returning) + .expect("confirm A's later term"); + assert_eq!(boundary.confirmed_route_owner(), Some(a)); + assert!( + returned_at.elapsed() < Duration::from_secs(5), + "the later term at A's address must not wait out the old one: {:?}", + returned_at.elapsed() + ); + assert!( + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, 2)).metadata(), + ) + .is_err(), + "B's term is fenced in turn" + ); + } + + #[test] + fn route_redirects_never_lead_more_local_than_configured_addresses() { + let pod: std::net::SocketAddr = "10.42.0.7:7443".parse().unwrap(); + for allowed in ["10.42.0.8:7443", "[fd00::8]:7443"] { + assert!(route_redirect_allowed(allowed.parse().unwrap(), &[pod])); + } + for refused in [ + "10.42.0.8:7444", + "127.0.0.1:7443", + "[::1]:7443", + "[::ffff:127.0.0.1]:7443", + "0.0.0.0:7443", + ] { + assert!( + !route_redirect_allowed(refused.parse().unwrap(), &[pod]), + "{refused} must be refused" + ); + } + let loopback: std::net::SocketAddr = "127.0.0.1:7443".parse().unwrap(); + assert!(route_redirect_allowed( + "127.0.0.2:7443".parse().unwrap(), + &[loopback] + )); + } + #[test] fn runtime_resource_claim_file_must_match_admitted_claim() { let directory = tempfile::tempdir().expect("temporary directory"); @@ -6073,6 +7628,8 @@ mod linux { let request = RequestEnvelope::new(Request::Attach { supervisor_instance_id: test_supervisor_instance_id(), policy: Box::new(policy), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, }) .expect("encode attach request"); let handle = test_open_boundary(channel.clone(), opening, &token).await; @@ -6371,7 +7928,7 @@ mod linux { }; assert!(matches!( - boundary.attach(policy.clone()), + boundary.attach(policy.clone(), false), Response::Attached { .. } )); assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); @@ -6454,7 +8011,7 @@ mod linux { drop(replacement_attachment); assert!(matches!( - boundary.attach(policy.clone()), + boundary.attach(policy.clone(), false), Response::Attached { .. } )); assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); @@ -6470,7 +8027,7 @@ mod linux { let mut changed_policy = policy.clone(); changed_policy.version += 1; assert!(matches!( - boundary.attach(changed_policy), + boundary.attach(changed_policy, false), Response::Error { kind, .. } if kind == BoundaryErrorKind::Denied )); assert!(matches!( @@ -6610,7 +8167,7 @@ mod linux { "exec status must remain available after its output attachment closes" ); } - let Response::Attached { snapshot } = boundary.attach(policy) else { + let Response::Attached { snapshot } = boundary.attach(policy, false) else { panic!("reconnect attach did not return a session snapshot"); }; assert_eq!(snapshot.generation, "generation-reconnect"); @@ -6626,6 +8183,1184 @@ mod linux { assert_eq!(boundary.terminate(&process_id), Response::Terminated); } + /// A gateway signing key, the runtime's verification key for it, and + /// mints of Sandbox Protocol bearers and takeover grants. + struct TestGateway { + issuer: SessionJwtIssuer, + verification_key: GatewayVerificationKey, + } + + impl TestGateway { + fn new() -> Self { + let key = openshell_crypto::pki::generate_keypair_for(&PKCS_ED25519) + .expect("generate gateway key"); + Self { + verification_key: GatewayVerificationKey { + key_id: "test-key".to_string(), + public_key_pem: key.public_key_pem(), + }, + issuer: SessionJwtIssuer::from_ed25519_pem( + key.serialize_pem().unwrap().as_bytes(), + "test-key", + "test-gateway", + Some(DEFAULT_SESSION_TOKEN_TTL), + Arc::new(SystemJwtClock), + ) + .expect("test session issuer"), + } + } + + fn token(&self, sandbox_id: &str, generation: &str, epoch: u64) -> String { + self.issuer + .mint_pair(&SandboxRuntimeIdentity { + sandbox_id: SandboxId::parse(sandbox_id).expect("sandbox ID"), + runtime_generation: + openshell_core::sandbox_generation::SandboxGenerationId::parse( + generation, + ) + .expect("runtime generation"), + auth_epoch: CredentialEpoch::new(epoch).expect("auth epoch"), + }) + .expect("token pair") + .sandbox + .token + .expose_secret() + .to_string() + } + + fn grant( + &self, + sandbox_id: &str, + generation: &str, + epoch: u64, + instance: openshell_sandbox_backend::boundary_protocol::SupervisorInstanceId, + ) -> String { + self.issuer + .mint_supervisor_takeover( + &openshell_core::jwt::SupervisorTakeover { + sandbox_id: SandboxId::parse(sandbox_id).expect("sandbox ID"), + runtime_generation: + openshell_core::sandbox_generation::SandboxGenerationId::parse( + generation, + ) + .expect("runtime generation"), + auth_epoch: CredentialEpoch::new(epoch).expect("auth epoch"), + supervisor_instance: instance.to_string(), + owner: "supervisor-pool-1".to_string(), + }, + Duration::from_mins(5), + ) + .expect("takeover grant") + .token + .expose_secret() + .to_string() + } + } + + /// A replacement supervisor process takes over a running generation: + /// the frozen workload resumes as the same process, the replacement + /// adopts it from the snapshot instead of launching it again, and the + /// previous owner's credentials and process identity stay fenced. + #[test] + fn granted_takeover_adopts_the_running_workload_and_fences_the_previous_owner() { + const CHILD_MARKER: &str = "OPENSHELL_TEST_BOUNDARY_TAKEOVER_CHILD"; + const SANDBOX: &str = "sandbox-takeover"; + const GENERATION: &str = "generation-takeover"; + if std::env::var_os(CHILD_MARKER).is_none() { + let status = std::process::Command::new( + std::env::current_exe().expect("current test executable"), + ) + .args([ + "--exact", + "boundary_server::linux::tests::granted_takeover_adopts_the_running_workload_and_fences_the_previous_owner", + "--nocapture", + ]) + .env(CHILD_MARKER, "1") + .status() + .expect("run isolated takeover test"); + assert!(status.success(), "isolated takeover test failed"); + return; + } + + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let (network_broker, workload_launcher) = test_network_broker(); + let mut boundary = BoundaryRuntime::new( + BoundaryConfig { + boundary_id: SANDBOX.to_string(), + generation: GENERATION.to_string(), + session_id: test_session_id(), + session_rotation: openshell_core::jwt::SessionRotation::new(1) + .expect("session rotation"), + auth_epoch: CredentialEpoch::new(1).expect("auth epoch"), + gateway_id: "test-gateway".to_string(), + verification_keys: vec![gateway.verification_key.clone()], + listener: BoundaryListenerConfig::TlsTcp { + address: "127.0.0.1:5500".parse().expect("control address"), + tls: placeholder_server_tls(), + }, + resource_claims: std::collections::BTreeMap::new(), + resource_claim_files: std::collections::BTreeMap::new(), + workload_identity: test_workload_identity(), + provisioning_evidence: test_provisioning_evidence(), + child_env: std::collections::HashMap::new(), + }, + process_runtime.handle().clone(), + network_broker, + workload_launcher, + test_runtime_qualification(), + ) + .expect("test boundary runtime"); + boundary.default_recovery_timeout = Duration::from_secs(10); + let boundary = Arc::new(boundary); + let authenticate = |token: &str| { + boundary.authenticate_request( + SandboxConnectionId::new(), + bearer_request((), token).metadata(), + ) + }; + let policy = SandboxPolicyWire::from(openshell_core::policy::SandboxPolicy { + version: 1, + filesystem: openshell_core::policy::FilesystemPolicy::default(), + network: openshell_core::policy::NetworkPolicy::default(), + landlock: openshell_core::policy::LandlockPolicy::default(), + process: openshell_core::policy::ProcessPolicy::default(), + }); + + // The launch owner attaches, confirms, and starts the workload. + let launch_owner = test_supervisor_instance_id(); + let first = authenticate(&gateway.token(SANDBOX, GENERATION, 1)).expect("first"); + assert!(matches!( + boundary.attach(policy.clone(), false), + Response::Attached { .. } + )); + boundary + .commit_attach(&first, launch_owner, None, None) + .expect("attach launch owner"); + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary + .commit_confirm(&first) + .expect("confirm launch owner"); + let Response::Started { process_id, .. } = boundary.start_agent( + SANDBOX.to_string(), + AgentSpecWire { + program: "/bin/sleep".to_string(), + args: vec!["60".to_string()], + workdir: None, + timeout_secs: 120, + interactive: false, + }, + None, + None, + 0, + std::collections::HashMap::new(), + std::collections::HashMap::new(), + ) else { + panic!("launch owner could not start the workload"); + }; + let workload = boundary.running_process(&process_id).expect("running"); + let workload_pid = workload.signaler.pid(); + // The launch owner's supervision pinned the workload's tool. + let tool = |digest: char| { + Ok(openshell_isolation_interface::contract::BinaryIdentity { + executable: openshell_isolation_interface::contract::ExecutableIdentity { + path: std::path::PathBuf::from("/usr/bin/tool"), + digest: Some(digest.to_string().repeat(64).parse().expect("digest")), + }, + ancestors: Vec::new(), + cmdline_paths: Vec::new(), + }) + }; + assert!(boundary.binary_pins.admit(tool('a')).is_ok()); + + // The owner's process dies: the workload freezes, unsupervised. + boundary.transport_disconnected(first.connection_id()); + let state = *lock(&boundary.supervisor_connection); + let recovery_id = match state { + SupervisorConnectionState::Frozen { recovery_id } => recovery_id, + state => panic!("expected a frozen workload, got {state:?}"), + }; + + // Fresh credentials for a later epoch do not move ownership + // without a grant, and a grant names one process only. + let replacement = test_supervisor_instance_id(); + let ungranted = authenticate(&gateway.token(SANDBOX, GENERATION, 2)).expect("ahead"); + assert!(matches!( + boundary.commit_attach(&ungranted, replacement, None, None), + Err((BoundaryErrorKind::Denied, _)) + )); + let grant = gateway.grant(SANDBOX, GENERATION, 2, replacement); + let verified = boundary + .authenticator + .verify_takeover(&grant) + .expect("verify grant"); + assert!(matches!( + boundary.commit_attach( + &ungranted, + test_supervisor_instance_id(), + Some(&verified), + None + ), + Err((BoundaryErrorKind::Denied, _)) + )); + // A grant for another runtime generation is refused outright. + assert!( + boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, "generation-other", 2, replacement)) + .is_err() + ); + + // The granted replacement attaches with whatever launch controls + // it derived; the runtime keeps and reports the bound ones, and + // the snapshot shows the running workload to adopt. + let mut derived = policy.clone(); + derived.version += 1; + let candidate = authenticate(&gateway.token(SANDBOX, GENERATION, 2)).expect("cand"); + let Response::Attached { snapshot } = boundary.attach(derived, true) else { + panic!("takeover attach was refused"); + }; + assert_eq!(snapshot.bound_policy.as_deref(), Some(&policy)); + assert!(snapshot.processes.iter().any(|process| { + process.process_id == process_id + && process.kind == ProcessKindWire::Main + && process.status.is_none() + })); + boundary + .commit_attach(&candidate, replacement, Some(&verified), None) + .expect("stage takeover"); + // The staged replacement cannot act before it confirms, and the + // previous owner can no longer return while it is staged. + assert!(boundary.connections.require_active(&candidate).is_err()); + let returning = authenticate(&gateway.token(SANDBOX, GENERATION, 1)).expect("old"); + assert!( + boundary + .commit_attach(&returning, launch_owner, None, None) + .is_err() + ); + assert!(matches!( + *lock(&boundary.supervisor_connection), + SupervisorConnectionState::Frozen { .. } + )); + + // Confirmation promotes the replacement and resumes the same + // workload process; nothing was launched again. + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary + .commit_confirm(&candidate) + .expect("confirm takeover"); + assert_eq!( + *lock(&boundary.supervisor_connection), + SupervisorConnectionState::Connected(candidate.connection_id()) + ); + let resumed = boundary + .running_process(&process_id) + .expect("still running"); + assert!(Arc::ptr_eq(&workload, &resumed)); + assert_eq!(resumed.signaler.pid(), workload_pid); + assert!(!resumed.has_exited()); + boundary + .connections + .require_active(&candidate) + .expect("replacement is the active owner"); + // The replacement starts without supervisor-side pins, but the + // runtime never offers it a tool the workload replaced. + assert!(matches!( + boundary.binary_pins.admit(tool('b')), + Err(openshell_isolation_interface::contract::ResolveError::Failed(_)) + )); + assert!(boundary.binary_pins.admit(tool('a')).is_ok()); + + // Every credential of the replaced term is stale from now on. + assert!(authenticate(&gateway.token(SANDBOX, GENERATION, 1)).is_err()); + // The superseded recovery deadline cannot terminate the workload. + process_runtime.block_on(boundary.expire_recovery(recovery_id)); + assert_eq!( + *lock(&boundary.supervisor_connection), + SupervisorConnectionState::Connected(candidate.connection_id()) + ); + assert!(!resumed.has_exited()); + assert_eq!(boundary.terminate(&process_id), Response::Terminated); + } + + /// A runtime trusting `gateway`, launched at epoch 1, before any attach. + fn takeover_test_runtime( + gateway: &TestGateway, + sandbox: &str, + generation: &str, + process_runtime: &tokio::runtime::Runtime, + ) -> Arc { + takeover_test_runtime_with_recovery(gateway, sandbox, generation, process_runtime, None) + } + + /// [`takeover_test_runtime`] with a recovery deadline shorter than + /// configuration allows. + fn takeover_test_runtime_with_recovery( + gateway: &TestGateway, + sandbox: &str, + generation: &str, + process_runtime: &tokio::runtime::Runtime, + recovery_timeout: Option, + ) -> Arc { + let (network_broker, workload_launcher) = test_network_broker(); + let mut runtime = BoundaryRuntime::new( + BoundaryConfig { + boundary_id: sandbox.to_string(), + generation: generation.to_string(), + session_id: test_session_id(), + session_rotation: openshell_core::jwt::SessionRotation::new(1) + .expect("session rotation"), + auth_epoch: CredentialEpoch::new(1).expect("auth epoch"), + gateway_id: "test-gateway".to_string(), + verification_keys: vec![gateway.verification_key.clone()], + listener: BoundaryListenerConfig::TlsTcp { + address: "127.0.0.1:5500".parse().expect("control address"), + tls: placeholder_server_tls(), + }, + resource_claims: std::collections::BTreeMap::new(), + resource_claim_files: std::collections::BTreeMap::new(), + workload_identity: test_workload_identity(), + provisioning_evidence: test_provisioning_evidence(), + child_env: std::collections::HashMap::new(), + }, + process_runtime.handle().clone(), + network_broker, + workload_launcher, + test_runtime_qualification(), + ) + .expect("test boundary runtime"); + runtime.default_recovery_timeout = recovery_timeout.unwrap_or(Duration::from_secs(10)); + Arc::new(runtime) + } + + fn takeover_test_policy() -> SandboxPolicyWire { + SandboxPolicyWire::from(openshell_core::policy::SandboxPolicy { + version: 1, + filesystem: openshell_core::policy::FilesystemPolicy::default(), + network: openshell_core::policy::NetworkPolicy::default(), + landlock: openshell_core::policy::LandlockPolicy::default(), + process: openshell_core::policy::ProcessPolicy::default(), + }) + } + + /// Send one control request through `serve_one`, as a connection of + /// `principal`, and return the boundary's response. + fn serve_request( + runtime: &BoundaryRuntime, + principal: &SandboxProtocolPrincipal, + request: Request, + ) -> Response { + let (server, mut client) = + std::os::unix::net::UnixStream::pair().expect("control socket pair"); + write_frame( + &mut client, + &RequestEnvelope::new(request).expect("request envelope"), + ) + .expect("write request"); + serve_one(ControlStream::TestUnix(server), runtime, principal).expect("serve request"); + read_frame::(&mut client) + .expect("read response") + .response + } + + /// Candidates the registry refuses never bind launch controls or + /// establish the boundary, every refusal is reported as an OCSF + /// configuration event and detection finding, and no grant is logged. + #[test] + fn rejected_attach_candidates_never_bind_the_boundary() { + const SANDBOX: &str = "sandbox-admission"; + const GENERATION: &str = "generation-admission"; + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let capture = openshell_ocsf::test_support::OcsfCapture::default(); + let _capture = capture.install(); + let principal = |epoch| { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, epoch)).metadata(), + ) + .expect("principal") + }; + let mut other_policy = takeover_test_policy(); + other_policy.version += 7; + let attach = |takeover: Option, instance| Request::Attach { + supervisor_instance_id: instance, + policy: Box::new(other_policy.clone()), + supervisor_takeover: takeover, + supervisor_recovery_timeout_secs: None, + }; + + // An epoch-ahead bearer without a grant, a grant for another + // generation, and a grant naming another process are refused. + let candidate = test_supervisor_instance_id(); + let foreign_grant = gateway.grant(SANDBOX, "generation-other", 2, candidate); + let mismatched_grant = + gateway.grant(SANDBOX, GENERATION, 2, test_supervisor_instance_id()); + for request in [ + attach(None, candidate), + attach(Some(foreign_grant.clone()), candidate), + attach(Some(mismatched_grant.clone()), candidate), + ] { + assert!(matches!( + serve_request(&boundary, &principal(2), request), + Response::Error { + kind: BoundaryErrorKind::Denied, + .. + } + )); + assert!(lock(&boundary.attached_policy).is_none()); + assert!(matches!( + *lock(&boundary.state), + RuntimeState::AwaitingAttach + )); + } + let events = capture.json_events(); + let denials = events + .iter() + .filter(|event| { + event["message"] + .as_str() + .is_some_and(|message| message.starts_with("Supervisor takeover denied")) + }) + .count(); + assert_eq!( + denials, 6, + "a state event and a finding per refusal: {events:#?}" + ); + assert!(events.iter().any(|event| { + event["finding_info"]["uid"] == "sandbox-supervisor-takeover-denied" + })); + let exported = serde_json::to_string(&events).expect("events JSON"); + for grant in [&foreign_grant, &mismatched_grant] { + assert!( + !exported.contains(grant.as_str()), + "a grant leaked into OCSF" + ); + } + + // The launch owner then binds its own launch controls. + let launch_owner = test_supervisor_instance_id(); + assert!(matches!( + serve_request( + &boundary, + &principal(1), + Request::Attach { + supervisor_instance_id: launch_owner, + policy: Box::new(takeover_test_policy()), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: None, + } + ), + Response::Attached { .. } + )); + assert_eq!( + lock(&boundary.attached_policy).as_ref(), + Some(&takeover_test_policy()) + ); + } + + /// The supervisor whose connection is in force when supervision is + /// lost chose how long the workload stays frozen: the deadline its + /// attach offered, or the default without one. + #[test] + fn the_lost_connections_offer_governs_the_recovery_deadline() { + const SANDBOX: &str = "sandbox-recovery-offer"; + const GENERATION: &str = "generation-recovery-offer"; + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let owner = test_supervisor_instance_id(); + let connection = || { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, 1)).metadata(), + ) + .expect("principal") + }; + let attach = |offer| Request::Attach { + supervisor_instance_id: owner, + policy: Box::new(takeover_test_policy()), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: offer, + }; + let mut deadlines = Vec::new(); + // The second connection reconnects the same process without an + // offer, after the first one's loss froze the workload. + for (offer, deadline) in [ + (Some(45), Duration::from_secs(45)), + (None, boundary.default_recovery_timeout), + ] { + let principal = connection(); + assert!(matches!( + serve_request(&boundary, &principal, attach(offer)), + Response::Attached { .. } + )); + assert!(matches!( + serve_request(&boundary, &principal, Request::Confirm), + Response::Confirmed { .. } + )); + boundary.transport_disconnected(principal.connection_id()); + deadlines.push(deadline); + assert_eq!(*lock(&boundary.scheduled_recoveries), deadlines); + assert!(matches!( + *lock(&boundary.supervisor_connection), + SupervisorConnectionState::Frozen { .. } + )); + } + assert!(lock(&boundary.recovery_offers).is_empty()); + } + + /// A recovery deadline outside the accepted bounds is refused before + /// the attach binds launch controls or records anything. + #[test] + fn out_of_range_recovery_offers_are_refused_before_the_boundary_binds() { + const SANDBOX: &str = "sandbox-recovery-bounds"; + const GENERATION: &str = "generation-recovery-bounds"; + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let owner = test_supervisor_instance_id(); + let connection = || { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, 1)).metadata(), + ) + .expect("principal") + }; + let attach = |seconds| Request::Attach { + supervisor_instance_id: owner, + policy: Box::new(takeover_test_policy()), + supervisor_takeover: None, + supervisor_recovery_timeout_secs: Some(seconds), + }; + for seconds in [0, 9, 301, u32::MAX] { + assert!(matches!( + serve_request(&boundary, &connection(), attach(seconds)), + Response::Error { + kind: BoundaryErrorKind::Invalid, + .. + } + )); + assert!(lock(&boundary.attached_policy).is_none()); + assert!(matches!( + *lock(&boundary.state), + RuntimeState::AwaitingAttach + )); + assert!(lock(&boundary.recovery_offers).is_empty()); + } + for seconds in [10, 300] { + assert!(matches!( + serve_request(&boundary, &connection(), attach(seconds)), + Response::Attached { .. } + )); + } + } + + /// Promoting a replacement waits for an operation the replaced owner + /// was already authorized for, after which every later operation of + /// that owner is refused; the replacement's snapshot taken after the + /// promotion includes what the replaced owner started meanwhile. + #[test] + fn promotion_waits_for_and_then_fences_the_replaced_owner() { + const CHILD_MARKER: &str = "OPENSHELL_TEST_BOUNDARY_PROMOTION_CHILD"; + const SANDBOX: &str = "sandbox-promotion"; + const GENERATION: &str = "generation-promotion"; + if std::env::var_os(CHILD_MARKER).is_none() { + let status = std::process::Command::new( + std::env::current_exe().expect("current test executable"), + ) + .args([ + "--exact", + "boundary_server::linux::tests::promotion_waits_for_and_then_fences_the_replaced_owner", + "--nocapture", + ]) + .env(CHILD_MARKER, "1") + .status() + .expect("run isolated promotion test"); + assert!(status.success(), "isolated promotion test failed"); + return; + } + + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let principal = |epoch| { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, epoch)).metadata(), + ) + .expect("principal") + }; + let policy = takeover_test_policy(); + + // The live owner runs the workload. + let owner = test_supervisor_instance_id(); + let old = principal(1); + assert!(matches!( + boundary.attach(policy.clone(), false), + Response::Attached { .. } + )); + boundary + .commit_attach(&old, owner, None, None) + .expect("attach"); + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary.commit_confirm(&old).expect("confirm"); + let Response::Started { process_id, .. } = boundary.start_agent( + SANDBOX.to_string(), + AgentSpecWire { + program: "/bin/sleep".to_string(), + args: vec!["60".to_string()], + workdir: None, + timeout_secs: 120, + interactive: false, + }, + None, + None, + 0, + std::collections::HashMap::new(), + std::collections::HashMap::new(), + ) else { + panic!("owner could not start the workload"); + }; + + // A granted replacement stages while the owner stays connected. + let replacement = test_supervisor_instance_id(); + let candidate = principal(2); + let grant = boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, GENERATION, 2, replacement)) + .expect("grant"); + let Response::Attached { snapshot: staged } = boundary.attach(policy.clone(), true) + else { + panic!("takeover attach refused"); + }; + boundary + .commit_attach(&candidate, replacement, Some(&grant), None) + .expect("stage takeover"); + assert!( + staged + .processes + .iter() + .all(|process| process.kind == ProcessKindWire::Main) + ); + + // The owner is authorized for an operation and, still holding + // that authorization, starts an exec the replacement's staged + // snapshot does not show. + let authorized = boundary.owner_operation(&old).expect("owner in force"); + let exec_request = RequestEnvelope::new(Request::Exec { + spec: ExecSpecWire { + program: "/bin/sleep".to_string(), + args: vec!["30".to_string()], + shell: None, + runtime_helper: None, + env: Vec::new(), + workdir: None, + pty: false, + }, + }) + .expect("exec request"); + + // Promotion waits while the authorization is outstanding. + let (promoted_tx, promoted) = std::sync::mpsc::channel(); + let promoting = boundary.clone(); + let promoting_candidate = candidate.clone(); + let promotion = std::thread::spawn(move || { + assert!(matches!(promoting.confirm(), Response::Confirmed { .. })); + promoted_tx + .send(promoting.commit_confirm(&promoting_candidate)) + .expect("report promotion"); + }); + assert!( + promoted.recv_timeout(Duration::from_millis(300)).is_err(), + "promotion must wait for the authorized operation" + ); + let exec = boundary + .start_exec( + &exec_request.request_id, + &exec_request.payload_digest, + exec_request.exec_admission_deadline_unix_ms, + match exec_request.request.clone() { + Request::Exec { spec } => spec, + _ => unreachable!(), + }, + ) + .expect("authorized exec takes effect before the promotion"); + let exec_id = exec.process_id.clone(); + drop(exec); + assert!( + promoted.recv_timeout(Duration::from_millis(100)).is_err(), + "promotion still waits" + ); + drop(authorized); + promoted + .recv_timeout(Duration::from_secs(10)) + .expect("promotion completes once the operation took effect") + .expect("promotion"); + promotion.join().expect("promotion thread"); + + // Nothing of the replaced owner is admitted any more, including + // operations it had in its pipeline. + assert!(boundary.owner_operation(&old).is_err()); + for request in [ + Request::Terminate { + process_id: process_id.clone(), + }, + Request::Signal { + process_id: process_id.clone(), + signal: SignalWire::Kill, + }, + Request::UpdateProviderEnvironment { + generation: 9, + revision: 9, + provider_env: std::collections::HashMap::new(), + provider_files: std::collections::HashMap::new(), + }, + Request::TerminateBoundary, + ] { + assert!(boundary.authorize_request(&old, &request).is_err()); + } + let resumed = boundary.running_process(&process_id).expect("workload"); + assert!(!resumed.has_exited()); + + // The replacement's post-promotion snapshot shows the exec, so + // nothing the replaced owner started escapes adoption. + boundary + .connections + .check_attach(&candidate, replacement, None) + .expect("the promoted owner reattaches without its grant"); + let Response::Attached { snapshot } = boundary.attach(policy, false) else { + panic!("refresh attach refused"); + }; + assert!(snapshot.processes.iter().any(|process| { + process.process_id == exec_id + && process.kind == ProcessKindWire::Exec + && process.status.is_none() + })); + assert_eq!(boundary.terminate(&process_id), Response::Terminated); + } + + /// Run one test alone in a child process, failing it if it does not + /// finish in time instead of hanging the suite. + fn run_isolated(test: &str, marker: &str, limit: Duration) { + let mut child = std::process::Command::new( + std::env::current_exe().expect("current test executable"), + ) + .args(["--exact", test, "--nocapture"]) + .env(marker, "1") + .spawn() + .expect("run isolated test"); + let deadline = std::time::Instant::now() + limit; + loop { + if let Some(status) = child.try_wait().expect("wait for isolated test") { + assert!(status.success(), "isolated test {test} failed"); + return; + } + if std::time::Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + panic!("isolated test {test} did not finish within {limit:?}"); + } + std::thread::sleep(Duration::from_millis(50)); + } + } + + /// Whether the process is stopped, as a frozen workload is. + fn process_stopped(pid: u32) -> bool { + std::fs::read_to_string(format!("/proc/{pid}/stat")) + .ok() + .and_then(|stat| { + stat.rsplit_once(") ") + .map(|(_, fields)| fields.starts_with('T')) + }) + .unwrap_or(false) + } + + /// An owner operation holds the ownership guard across `block_on`, a + /// promotion waits for it, and a DNS answer waits behind the + /// promotion, on a runtime with a single worker. The answer waits + /// without blocking that worker, so timers keep running, the held + /// operation finishes, the promotion completes, the replaced owner's + /// answer is refused, and the new owner's answers and control + /// requests are served. A confirmation replay on the owner in force + /// never waits for owner operations. + #[test] + fn a_pending_promotion_never_stalls_a_single_worker_runtime() { + const CHILD_MARKER: &str = "OPENSHELL_TEST_BOUNDARY_BARRIER_CHILD"; + const SANDBOX: &str = "sandbox-barrier"; + const GENERATION: &str = "generation-barrier"; + if std::env::var_os(CHILD_MARKER).is_none() { + run_isolated( + "boundary_server::linux::tests::a_pending_promotion_never_stalls_a_single_worker_runtime", + CHILD_MARKER, + Duration::from_mins(1), + ); + return; + } + + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("single-worker process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime(&gateway, SANDBOX, GENERATION, &process_runtime); + let principal = |epoch| { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, epoch)).metadata(), + ) + .expect("principal") + }; + let policy = takeover_test_policy(); + let old = principal(1); + assert!(matches!( + boundary.attach(policy.clone(), false), + Response::Attached { .. } + )); + boundary + .commit_attach(&old, test_supervisor_instance_id(), None, None) + .expect("attach owner"); + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary.commit_confirm(&old).expect("confirm owner"); + let replacement = test_supervisor_instance_id(); + let candidate = principal(2); + let grant = boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, GENERATION, 2, replacement)) + .expect("grant"); + assert!(matches!( + boundary.attach(policy, true), + Response::Attached { .. } + )); + boundary + .commit_attach(&candidate, replacement, Some(&grant), None) + .expect("stage takeover"); + + // An owner operation holds the guard across `block_on`, as Exec, + // LoopbackConnect, and TerminateBoundary do; what remains of it + // needs the runtime's timer. + let (held_tx, held) = std::sync::mpsc::channel(); + let (release, released) = tokio::sync::oneshot::channel::<()>(); + let (operated_tx, operated) = std::sync::mpsc::channel(); + let operating = boundary.clone(); + let operating_owner = old.clone(); + let operation = std::thread::spawn(move || { + let guard = operating + .owner_operation(&operating_owner) + .expect("owner in force"); + held_tx.send(()).expect("report admission"); + operating.process_runtime.block_on(async { + let _ = released.await; + tokio::time::sleep(Duration::from_millis(20)).await; + }); + drop(guard); + operated_tx.send(()).expect("report operation"); + }); + held.recv_timeout(Duration::from_secs(10)) + .expect("operation admitted"); + + // A promotion queues behind it. + let (promoted_tx, promoted) = std::sync::mpsc::channel(); + let promoting = boundary.clone(); + let promoting_candidate = candidate.clone(); + let promotion = std::thread::spawn(move || { + promoted_tx + .send(promoting.commit_confirm(&promoting_candidate)) + .expect("report promotion"); + }); + let queued_by = std::time::Instant::now() + Duration::from_secs(10); + while boundary.ownership.try_read().is_ok() { + assert!( + std::time::Instant::now() < queued_by, + "promotion never queued" + ); + std::thread::sleep(Duration::from_millis(5)); + } + + // A DNS answer arrives on the persistent mediation stream of an + // owner; its task asks for the guard on the only worker. + let answer = |owner: &SandboxProtocolPrincipal, response: &[u8]| { + let (pending, answered) = + crate::network_broker::PendingDnsQuery::for_test(b"query".to_vec()); + let (inbound_tx, inbound) = tokio::sync::mpsc::channel(1); + inbound_tx + .try_send(MediationFrame { + kind: MediationFrameKind::DnsResponse, + stream_id: 1, + payload: mediation::encode_message(&DnsQueryResultWire::Response( + response.to_vec(), + )) + .expect("encode answer"), + }) + .expect("queue answer"); + let (outbound, mut queries) = tokio::sync::mpsc::channel(1); + process_runtime.spawn(run_boundary_dns_stream( + 1, + pending, + inbound, + outbound, + BoundaryMediationRoutes::default(), + (boundary.clone(), owner.clone()), + )); + let sent_by = std::time::Instant::now() + Duration::from_secs(10); + while queries.try_recv().is_err() { + assert!(std::time::Instant::now() < sent_by, "DNS query never sent"); + std::thread::sleep(Duration::from_millis(5)); + } + (answered, queries, inbound_tx) + }; + let (replaced_answer, _queries, _inbound) = answer(&old, b"stale"); + + // Timers keep running while the answer waits behind the promotion. + let (tick_tx, tick) = std::sync::mpsc::channel(); + process_runtime.spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = tick_tx.send(()); + }); + tick.recv_timeout(Duration::from_secs(5)) + .expect("the runtime keeps running while a promotion is pending"); + + // The held operation finishes, then the promotion completes. + release.send(()).expect("release the operation"); + operated + .recv_timeout(Duration::from_secs(10)) + .expect("the held operation finishes"); + promoted + .recv_timeout(Duration::from_secs(10)) + .expect("the promotion completes") + .expect("promotion"); + operation.join().expect("operation thread"); + promotion.join().expect("promotion thread"); + + // The replaced owner's answer is refused; the new owner's + // answers and control requests are served. + let refused = replaced_answer + .recv_timeout(Duration::from_secs(10)) + .expect("the replaced owner's answer completes"); + assert_eq!( + refused + .expect_err("a replaced owner's answer is refused") + .kind(), + io::ErrorKind::PermissionDenied + ); + let (current_answer, _queries, _inbound) = answer(&candidate, b"current"); + assert_eq!( + current_answer + .recv_timeout(Duration::from_secs(10)) + .expect("the new owner's answer completes") + .expect("answer"), + b"current" + ); + assert!(!matches!( + serve_request( + &boundary, + &candidate, + Request::Terminate { + process_id: "absent".to_string(), + }, + ), + Response::Error { + kind: BoundaryErrorKind::Denied, + .. + } + )); + + // A confirmation replay, such as a credential renewal, does not + // wait for an owner operation in progress. + let in_progress = boundary.owner_operation(&candidate).expect("new owner"); + let (replayed_tx, replayed) = std::sync::mpsc::channel(); + let replaying = boundary.clone(); + let replaying_owner = candidate.clone(); + std::thread::spawn(move || { + let _ = replayed_tx.send(replaying.commit_confirm(&replaying_owner)); + }); + replayed + .recv_timeout(Duration::from_secs(5)) + .expect("a replay does not wait for owner operations") + .expect("replay"); + drop(in_progress); + } + + /// A replacement whose connection is lost while it is promoted is the + /// owner in force when its loss is handled: the workload ends frozen + /// under a fresh recovery deadline, which then terminates it, whether + /// the previous owner was still connected or the workload was already + /// frozen. It never runs unsupervised. + #[test] + fn a_replacement_lost_during_promotion_leaves_the_workload_frozen_under_a_deadline() { + const CHILD_MARKER: &str = "OPENSHELL_TEST_BOUNDARY_PROMOTION_LOSS_CHILD"; + const SANDBOX: &str = "sandbox-promotion-loss"; + const GENERATION: &str = "generation-promotion-loss"; + if std::env::var_os(CHILD_MARKER).is_none() { + run_isolated( + "boundary_server::linux::tests::a_replacement_lost_during_promotion_leaves_the_workload_frozen_under_a_deadline", + CHILD_MARKER, + Duration::from_mins(2), + ); + return; + } + + for previous_owner_connected in [true, false] { + let gateway = TestGateway::new(); + let process_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("test process runtime"); + let _entered = process_runtime.enter(); + let boundary = takeover_test_runtime_with_recovery( + &gateway, + SANDBOX, + GENERATION, + &process_runtime, + Some(Duration::from_millis(500)), + ); + let principal = |epoch| { + boundary + .authenticate_request( + SandboxConnectionId::new(), + bearer_request((), &gateway.token(SANDBOX, GENERATION, epoch)) + .metadata(), + ) + .expect("principal") + }; + let policy = takeover_test_policy(); + let old = principal(1); + assert!(matches!( + boundary.attach(policy.clone(), false), + Response::Attached { .. } + )); + boundary + .commit_attach(&old, test_supervisor_instance_id(), None, None) + .expect("attach owner"); + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary.commit_confirm(&old).expect("confirm owner"); + let Response::Started { process_id, .. } = boundary.start_agent( + SANDBOX.to_string(), + AgentSpecWire { + program: "/bin/sleep".to_string(), + args: vec!["60".to_string()], + workdir: None, + timeout_secs: 120, + interactive: false, + }, + None, + None, + 0, + std::collections::HashMap::new(), + std::collections::HashMap::new(), + ) else { + panic!("owner could not start the workload"); + }; + let workload = boundary.running_process(&process_id).expect("running"); + let earlier_recovery = if previous_owner_connected { + None + } else { + boundary.transport_disconnected(old.connection_id()); + let state = *lock(&boundary.supervisor_connection); + match state { + SupervisorConnectionState::Frozen { recovery_id } => Some(recovery_id), + state => panic!("expected a frozen workload, got {state:?}"), + } + }; + + let replacement = test_supervisor_instance_id(); + let candidate = principal(2); + let grant = boundary + .authenticator + .verify_takeover(&gateway.grant(SANDBOX, GENERATION, 2, replacement)) + .expect("grant"); + assert!(matches!( + boundary.attach(policy, true), + Response::Attached { .. } + )); + boundary + .commit_attach(&candidate, replacement, Some(&grant), None) + .expect("stage takeover"); + + // The replacement's connection is lost between its promotion + // and the supervisory state recording it; the loss is given + // every chance to be handled first. + let (lost_tx, lost) = std::sync::mpsc::channel(); + let losing = boundary.clone(); + let lost_connection = candidate.connection_id(); + *lock(&boundary.promotion_hook) = Some(Box::new(move || { + let (started_tx, started) = std::sync::mpsc::channel(); + let disconnect = std::thread::spawn(move || { + started_tx.send(()).expect("report start"); + losing.transport_disconnected(lost_connection); + }); + started.recv().expect("disconnect started"); + std::thread::sleep(Duration::from_millis(200)); + lost_tx.send(disconnect).expect("hand over disconnect"); + })); + assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); + boundary.commit_confirm(&candidate).expect("promotion"); + lost.recv_timeout(Duration::from_secs(10)) + .expect("promotion ran the hook") + .join() + .expect("disconnect thread"); + + let state = *lock(&boundary.supervisor_connection); + let recovery_id = match state { + SupervisorConnectionState::Frozen { recovery_id } => recovery_id, + state => panic!( + "lost replacement (previous owner connected: \ + {previous_owner_connected}) left the workload {state:?}" + ), + }; + assert_ne!(Some(recovery_id), earlier_recovery); + assert!(boundary.connections.require_active(&candidate).is_err()); + let stopped_by = std::time::Instant::now() + Duration::from_secs(5); + while !process_stopped(workload.signaler.pid()) { + assert!( + std::time::Instant::now() < stopped_by, + "the workload must stay frozen without a supervisor" + ); + std::thread::sleep(Duration::from_millis(10)); + } + + // The fresh deadline ends the unsupervised workload. + let terminated_by = std::time::Instant::now() + Duration::from_secs(30); + while *lock(&boundary.supervisor_connection) != SupervisorConnectionState::Terminal + || !workload.has_exited() + { + assert!( + std::time::Instant::now() < terminated_by, + "the recovery deadline must terminate the workload" + ); + std::thread::sleep(Duration::from_millis(20)); + } + } + } + #[test] fn canonical_exit_preserves_pending_network_accept_and_exec() { const CHILD_MARKER: &str = "OPENSHELL_TEST_RETAINED_BOUNDARY_CHILD"; @@ -6732,7 +9467,7 @@ mod linux { // A replacement control process replays the durable lifecycle and // receives the original process rather than spawning another one. assert!(matches!( - boundary.attach(wire_policy.clone()), + boundary.attach(wire_policy.clone(), false), Response::Attached { .. } )); assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); @@ -6832,7 +9567,7 @@ mod linux { ); assert!(matches!( - boundary.attach(wire_policy), + boundary.attach(wire_policy, false), Response::Attached { .. } )); assert!(matches!(boundary.confirm(), Response::Confirmed { .. })); diff --git a/crates/openshell-sandbox/src/delegated.rs b/crates/openshell-sandbox/src/delegated.rs index 45639c4c00..144d3679a1 100644 --- a/crates/openshell-sandbox/src/delegated.rs +++ b/crates/openshell-sandbox/src/delegated.rs @@ -219,6 +219,12 @@ pub struct AgentSignaler { #[cfg(unix)] impl AgentSignaler { + /// Process ID of the agent's process group leader. + #[cfg(test)] + pub(crate) fn pid(&self) -> u32 { + self.pid + } + fn deliver(&self, signal: nix::sys::signal::Signal) -> Result<()> { let _guard = self .signal_lock diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 8ca1738dd7..3a7fd8beb4 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -3,6 +3,7 @@ //! Capability-free in-workload sandbox boundary. +pub mod binary_pins; #[cfg(target_os = "linux")] pub mod boundary_exec; #[cfg(target_os = "linux")] diff --git a/crates/openshell-sandbox/src/network_broker.rs b/crates/openshell-sandbox/src/network_broker.rs index faaff9250d..70e9b444a1 100644 --- a/crates/openshell-sandbox/src/network_broker.rs +++ b/crates/openshell-sandbox/src/network_broker.rs @@ -92,14 +92,34 @@ pub struct PendingTcpOpen { } impl PendingTcpOpen { + #[cfg(test)] pub(crate) async fn complete(self, decision: TcpOpenDecision) -> io::Result> { + self.decide(decision)?.established().await + } + + /// Hand the decision to the broker. Once this returns the decision has + /// taken effect; only the relay setup of an allowed open remains. + pub(crate) fn decide(self, decision: TcpOpenDecision) -> io::Result { self.decision .send(decision) .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "network broker stopped"))?; - if matches!(decision, TcpOpenDecision::Denied(_)) { + Ok(DecidedTcpOpen { + relay: (!matches!(decision, TcpOpenDecision::Denied(_))).then_some(self.relay), + }) + } +} + +/// A decided TCP open: the relay of an allowed open, or nothing. +pub struct DecidedTcpOpen { + relay: Option>>, +} + +impl DecidedTcpOpen { + pub(crate) async fn established(self) -> io::Result> { + let Some(relay) = self.relay else { return Ok(None); - } - self.relay + }; + relay .await .map_err(|_| { io::Error::new( @@ -127,6 +147,25 @@ impl PendingDnsQuery { .send(response) .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "DNS relay stopped")) } + + /// A UDP query as the relay queues it, and where its answer arrives. + #[cfg(test)] + pub(crate) fn for_test( + request: Vec, + ) -> (Self, std::sync::mpsc::Receiver>>) { + let (response, answers) = std::sync::mpsc::sync_channel(1); + ( + Self { + request, + transport: DnsTransport::Udp, + identity: dns_sender_identity(), + notification_to_queue: Duration::ZERO, + queued_at: Instant::now(), + response, + }, + answers, + ) + } } #[derive(Clone)] diff --git a/crates/openshell-server/README.md b/crates/openshell-server/README.md index 58f735c35c..bd0ab19332 100644 --- a/crates/openshell-server/README.md +++ b/crates/openshell-server/README.md @@ -18,6 +18,111 @@ Retries reuse that generation and session, including after gateway restart. Credential renewal retains both; credentials for an older launch cannot authorize the new one. +## Shared supervisor pool + +The `SharedSupervision` the gateway hands a compute driver that asks for it +always coordinates a pool of shared supervisor hosts, at any member count; the +driver itself is unchanged. `[openshell.gateway.supervisor_pool]` only tunes it +(`recovery_timeout_secs`), and the gateway refuses the table when no in-process +driver uses the shared supervisor. `supervisor_pool` keeps the +pool's durable state and reconciles it every two seconds on every gateway +replica, independent of which replica runs the compute driver's watch: + +- **Membership.** The replica serving a member's host stream refreshes its + `supervisor_pool_member` record from a task of its own, every two seconds and + at once when a member attaches, leaves, or detaches, each write bounded, so a + slow reconciliation pass never lets a served member expire. Only the member's + newest stream (by accept time) writes the record, so a replica still holding + a stale stream of the same process cannot overwrite it. A record older than 15 + seconds is dead; a member that announces it is leaving drops out at once; a + member whose process changed lost every instance its previous process owned. + The timings members and replicas share, and the reconnect budget that keeps a + member reconnecting past an unresponsive replica well inside the TTL, are in + `openshell_supervisor_control::pool_timing`. +- **Listings.** A pass reads each record type page by page after an immutable + key (`list_by_type_after`), so a row removed meanwhile shifts nothing out of + the listing. A listing is still not a snapshot, so before a pass drops an + owner its previous view had, or an instance that view or an owner names, it + reads the record by key; a failed read fails the pass, which then publishes + nothing. +- **Desired set.** A launch records the instance (`supervisor_instance_assignment`) + before anything else. The record holds only the handle of the launch + material, which the configured credential store keeps, as it keeps sandbox SSH + host keys. Every replica derives the desired set from these records and loads + the material of instances owned by members it serves. Recording takes single + atomic store writes and holds no lock: it reserves the route with a row named + by the route and identified by a fresh writer token, stores the material, has + a `supervisor_instance_material_candidate` row own the handle, and publishes + by advancing exactly its reservation's version. A launch publishes its + record as `Admitting`, which is not desired (nothing places, charges, or + grants it) and keeps its candidate; a driver recovery's backfill publishes + directly. Collection deletes a candidate's material only once its recording + no longer holds its reserved or admitting row, after which that recording + can never publish or be admitted; a reservation or admitting record older + than two minutes (a stalled, failed, or crashed launch) is fenced by + removing it at its version, and its material is then collected. A + withdrawal removes the material and then the record, or fences a reserved or + admitting row; an ended generation + (sandbox gone, replaced, stopped, or deleting) is withdrawn by the sweep; and + the compute driver's recovery backfills instances it runs without a record and + withdraws records published before its listing that it no longer lists. +- **Placement and ownership.** A launch admits its instance on the live member + the consistent-hash ring prefers (`gateway_ring`): one transaction creates + the owner record, charges the member's capacity, and publishes the + `Admitting` record, conditioned on that record's version, so an owner always + has its published instance and a removed record never gains one. A member at + its 256-instance capacity is skipped. With no live member with room the + launch is refused as a capacity refusal: it removes the admitting record it + created, only at the version it created (a concurrent launch of the same + generation that admitted it first wins), and then that record's candidate + and material; a removal that fails is left to fencing. One + `supervisor_instance_owner` record per instance names the member, its + process, and the ownership term (the sandbox's credential epoch). Records + change only by compare-and-swap, and a live owner keeps its instances. +- **Capacity.** One `supervisor_pool_capacity` record per member lists the + routes whose owner record, in any phase, names the member. Every owner + create, move, or removal is one `Store::write_atomically` transaction with + the capacity records it changes, conditioned on all of their versions, and a + full member is refused before anything is written. Two replicas admitting + the last slot of a member, by launch or takeover, therefore admit one, and a + stalled admission whose slot was taken meanwhile fails. A takeover prefers + the route's previous member, whose record already lists it, so a member + whose process restarted re-adopts its instances even when every member is + full. +- **Takeover.** Only an instance whose owner failed or left moves. A replica + records the intent (new owner, next term, the process the runtime will pin, + a new token lineage), advances the sandbox's persisted epoch with + `advance_identity_for_takeover`, which revokes the previous owner's + credentials on every replica, and then grants it. Any replica completes an + interrupted intent; a stopping, stopped, or deleted sandbox, a replaced + generation, or a term another lineage holds cancels it. A replica that just + started waits 30 seconds before taking over, so members reconnect after a + gateway restart first. The takeover log names its `cause`: `expired`, + `leaving`, or `restarted`. +- **Credentials.** Each replica mints a takeover term's grant and session + credentials only for members it serves, from the granted record and the + persisted identity, and mints them again after half the grant's lifetime. + Other replicas send the term without credentials; a host that already runs + it keeps it. + +- **Observations.** The replica serving an owner records what the owner reports + of its instances (`supervisor_instance_report`; no error text), one record per + instance and ownership term, and every replica judges an instance whose owner + streams elsewhere from the record of the current owner term, so launch + readiness and bootstrap completion do not depend on which replica a member + reached. A report counts for a term only once the member's stream + acknowledged the snapshot that installed the term, and only if the pass's + owner record shows that term, so a status from an earlier term is never + recorded as a later one. A record is rewritten whenever it differs from the + report, so one another writer replaced is repaired. Reports are listed before + owners, and only records of terms older than the listed owner, or of + instances without one, are removed. + +Owner records of withdrawn instances are removed once their generation ended, +or two minutes after the withdrawal, after their owner was told to stop them. +Pool members still authenticate with the one shared supervisor token, so a +member identity is asserted by its hello rather than bound to its Pod. + ## Operator credential export `GetProviderCredentials` exports selected profile-declared runtime environment diff --git a/crates/openshell-server/src/auth/principal.rs b/crates/openshell-server/src/auth/principal.rs index bc5d31381a..495be046d8 100644 --- a/crates/openshell-server/src/auth/principal.rs +++ b/crates/openshell-server/src/auth/principal.rs @@ -15,6 +15,7 @@ //! to prevent cross-sandbox access (see issue #1354). use super::identity::Identity; +use openshell_core::jwt::CredentialEpoch; /// Who is calling. /// @@ -85,6 +86,17 @@ pub struct SandboxPrincipal { pub trust_domain: Option, } +impl SandboxPrincipal { + /// The ownership term of a supervisor session this principal opens, when + /// its credential carries one. + pub fn auth_epoch(&self) -> Option { + match &self.source { + SandboxIdentitySource::BootstrapJwt { auth_epoch, .. } => *auth_epoch, + SandboxIdentitySource::ComputeDriver { .. } => None, + } + } +} + /// Shared supervisor caller. #[derive(Debug, Clone)] pub struct SupervisorPrincipal { @@ -110,7 +122,12 @@ pub struct PeerPrincipal { pub enum SandboxIdentitySource { /// Generation-bound gateway JWT validated against the persisted runtime /// identity by [`super::sandbox_jwt::SandboxSessionJwtAuthenticator`]. - BootstrapJwt { issuer: String }, + BootstrapJwt { + issuer: String, + /// Credential epoch of the bearer: the supervisor ownership term it + /// belongs to. `None` only where no bearer was verified. + auth_epoch: Option, + }, /// Driver-native credential used to bootstrap a gateway-minted JWT via /// `IssueSandboxToken`. The named compute driver authenticated only the /// sandbox identity and its concrete runtime binding; the gateway still @@ -139,7 +156,9 @@ impl Principal { /// Principals for tests. #[cfg(test)] pub mod test_support { - use super::{Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal}; + use super::{ + CredentialEpoch, Principal, SandboxIdentitySource, SandboxPrincipal, UserPrincipal, + }; use crate::auth::identity::{Identity, IdentityProvider}; /// An OIDC user with no roles or scopes. @@ -157,10 +176,23 @@ pub mod test_support { /// A sandbox holding a gateway JWT. pub fn sandbox(sandbox_id: &str) -> Principal { + sandbox_principal(sandbox_id, None) + } + + /// A sandbox holding a gateway JWT of credential epoch `auth_epoch`. + pub fn sandbox_with_epoch(sandbox_id: &str, auth_epoch: u64) -> Principal { + sandbox_principal( + sandbox_id, + Some(CredentialEpoch::new(auth_epoch).expect("credential epoch")), + ) + } + + fn sandbox_principal(sandbox_id: &str, auth_epoch: Option) -> Principal { Principal::Sandbox(SandboxPrincipal { sandbox_id: sandbox_id.to_string(), source: SandboxIdentitySource::BootstrapJwt { issuer: "openshell-gateway:test".to_string(), + auth_epoch, }, trust_domain: Some("openshell".to_string()), }) diff --git a/crates/openshell-server/src/auth/sandbox_jwt.rs b/crates/openshell-server/src/auth/sandbox_jwt.rs index 774b9f5f47..7b44633e5c 100644 --- a/crates/openshell-server/src/auth/sandbox_jwt.rs +++ b/crates/openshell-server/src/auth/sandbox_jwt.rs @@ -295,6 +295,17 @@ impl SandboxSessionJwtAuthority { self.issuer.mint_supervisor_token(ttl) } + /// Mint the takeover grant of one supervisor term of a shared pool + /// instance. The caller has already advanced the sandbox's persisted + /// credential epoch to the grant's. + pub fn mint_supervisor_takeover( + &self, + takeover: &openshell_core::jwt::SupervisorTakeover, + ttl: Duration, + ) -> Result { + self.issuer.mint_supervisor_takeover(takeover, ttl) + } + /// Verify a presented supervisor token. pub fn verify_supervisor_token( &self, @@ -366,6 +377,7 @@ impl Authenticator for SandboxSessionJwtAuthenticator { sandbox_id: authenticated.sandbox_id.to_string(), source: SandboxIdentitySource::BootstrapJwt { issuer: "launch-session".to_string(), + auth_epoch: Some(authenticated.auth_epoch), }, trust_domain: Some("openshell".to_string()), }, diff --git a/crates/openshell-server/src/auth/sandbox_session.rs b/crates/openshell-server/src/auth/sandbox_session.rs index dc9e424240..aa77b0320c 100644 --- a/crates/openshell-server/src/auth/sandbox_session.rs +++ b/crates/openshell-server/src/auth/sandbox_session.rs @@ -403,6 +403,108 @@ pub async fn rotate_gateway_token( Ok(next.clone()) } +/// Outcome of advancing a sandbox's credential epoch for a supervisor takeover. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TakeoverIdentity { + /// The epoch now equals the takeover's; earlier credentials are revoked. + Advanced, + /// The takeover's epoch was already in force (a retried transaction). + AlreadyAdvanced, + /// The takeover must not proceed: the sandbox is gone, stopping, or runs + /// another generation, or a later term already replaced this one. + Refused(&'static str), +} + +/// Phases in which a running generation may change supervisor owner. +fn takeover_allowed(phase: SandboxPhase) -> bool { + matches!( + phase, + SandboxPhase::Provisioning + | SandboxPhase::Starting + | SandboxPhase::Ready + | SandboxPhase::Error + | SandboxPhase::Completed + | SandboxPhase::Unknown + ) +} + +/// Advance the persisted credential epoch of `generation` to `auth_epoch` +/// and start a new gateway-token lineage, so every credential of earlier +/// supervisor terms is rejected by every gateway replica. Ordinary refresh +/// never changes the epoch; only an ownership takeover calls this. +/// +/// Stop and delete win: a sandbox that is stopping, stopped, deleting, or +/// gone, or that runs another generation, refuses the takeover. +pub async fn advance_identity_for_takeover( + store: &Store, + sandbox_id: &str, + generation: &str, + auth_epoch: CredentialEpoch, + gateway_token_id: Uuid, +) -> Result { + const ATTEMPTS: u32 = 5; + let mut attempt = 0; + loop { + attempt += 1; + let Some(sandbox) = store.get_message::(sandbox_id).await? else { + return Ok(TakeoverIdentity::Refused("sandbox no longer exists")); + }; + let phase = SandboxPhase::try_from(sandbox.phase()).unwrap_or(SandboxPhase::Unknown); + if !takeover_allowed(phase) { + return Ok(TakeoverIdentity::Refused( + "sandbox is stopping, stopped, restarting, or deleting", + )); + } + let Some(metadata) = sandbox.metadata.as_ref() else { + return Ok(TakeoverIdentity::Refused( + "sandbox identity metadata is missing", + )); + }; + let Ok(current) = PersistedSandboxIdentity::read(&metadata.annotations) else { + return Ok(TakeoverIdentity::Refused( + "sandbox runtime identity is invalid", + )); + }; + if current.runtime_generation.as_str() != generation { + return Ok(TakeoverIdentity::Refused( + "sandbox runs another runtime generation", + )); + } + if current.auth_epoch == auth_epoch { + // Only the intent that advanced the term may grant it; its + // credentials name its own token lineage. + return Ok(if current.gateway_token_id == gateway_token_id { + TakeoverIdentity::AlreadyAdvanced + } else { + TakeoverIdentity::Refused("another takeover already holds this term") + }); + } + if current.auth_epoch > auth_epoch { + return Ok(TakeoverIdentity::Refused( + "a later supervisor term replaced this takeover", + )); + } + let next = PersistedSandboxIdentity { + runtime_generation: current.runtime_generation.clone(), + auth_epoch, + gateway_token_id, + refresh_replay: None, + }; + match store + .update_message_cas::(sandbox_id, metadata.resource_version, |sandbox| { + if let Some(metadata) = sandbox.metadata.as_mut() { + next.write(&mut metadata.annotations); + } + }) + .await + { + Ok(_) => return Ok(TakeoverIdentity::Advanced), + Err(crate::persistence::PersistenceError::Conflict { .. }) if attempt < ATTEMPTS => {} + Err(error) => return Err(error), + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/openshell-server/src/config_file.rs b/crates/openshell-server/src/config_file.rs index 817a6a02ad..ca7ac76003 100644 --- a/crates/openshell-server/src/config_file.rs +++ b/crates/openshell-server/src/config_file.rs @@ -192,6 +192,9 @@ pub struct GatewayFileSection { pub otlp: Option, #[serde(default)] pub ocsf_log: Option, + /// Tuning of the shared supervisor pool. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub supervisor_pool: Option, // ── Disallowed-in-file fields ──────────────────────────────────────── // @@ -202,6 +205,25 @@ pub struct GatewayFileSection { pub database_url: Option, } +/// `[openshell.gateway.supervisor_pool]`: optional tuning of the shared +/// supervisor pool. +/// +/// The shared supervisor always runs as a pool that the gateway coordinates +/// for high availability: it places each sandbox instance on a live member by +/// consistent hash, records its owner durably, and moves the instances of a +/// failed or leaving member to another member, which adopts their running +/// workloads. Without this table the defaults apply. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SupervisorPoolConfig { + /// Seconds a sandbox runtime keeps its workload frozen after losing its + /// supervisor before terminating it, from 10 to 300. The pool's + /// supervisors offer it to each runtime; unset offers + /// [`openshell_supervisor_control::POOL_RECOVERY_TIMEOUT_SECS`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub recovery_timeout_secs: Option, +} + /// Gateway listener TLS fields accepted in `gateway.toml`. /// /// Client-certificate handshake policy is derived from the presence of a @@ -791,6 +813,20 @@ fn parse_and_validate(path: &Path, contents: &str) -> Result { config: ComputeDriverConfigContext<'a>, shutdown_rx: watch::Receiver, shared_supervision: &'a OnceLock, + /// Recovery deadline from `[openshell.gateway.supervisor_pool]`, when + /// configured. + supervisor_recovery_timeout_secs: Option, } impl ComputeDriverBuildContext<'_> { - /// The gateway's coordination of a shared supervisor host, for an + /// The gateway's coordination of the shared supervisor pool, for an /// in-process driver whose configuration places new generations on it. - /// The gateway owns the desired set, serves it to the host over the - /// private assignment stream, mints the supervisor token, and publishes - /// each renewal for the driver to deliver. Created on first request; every - /// call returns the same coordination. + /// The gateway owns the desired set and ownership durably, serves each + /// member over the private assignment stream, mints the supervisor + /// token, and publishes each renewal for the driver to deliver. + /// `[openshell.gateway.supervisor_pool]` only tunes it. Created on first + /// request; every call returns the same coordination. pub fn shared_supervision(&self) -> openshell_supervisor_control::SharedSupervision { self.shared_supervision - .get_or_init(openshell_supervisor_control::SharedSupervision::new) + .get_or_init(|| { + let supervision = openshell_supervisor_control::SharedSupervision::new(); + if let Some(recovery_timeout_secs) = self.supervisor_recovery_timeout_secs { + supervision.set_recovery_timeout_secs(recovery_timeout_secs); + } + supervision + }) .clone() } @@ -1826,6 +1846,9 @@ async fn build_compute_runtime( } }; info!(driver = %driver.name(), "Using compute driver"); + let supervisor_pool = driver_startup + .file + .and_then(|file| file.openshell.gateway.supervisor_pool); let runtime = match driver { ConfiguredComputeDriver::Registered(registration) => { let shared_supervision = OnceLock::new(); @@ -1839,11 +1862,20 @@ async fn build_compute_runtime( }, shutdown_rx, shared_supervision: &shared_supervision, + supervisor_recovery_timeout_secs: supervisor_pool + .and_then(|pool| pool.recovery_timeout_secs), }; let instance = registration.factory.build(build_context).await?; let supervisor_host = shared_supervision .into_inner() .map(|supervision| supervision.controller().clone()); + if supervisor_pool.is_some() && supervisor_host.is_none() { + return Err(Error::config( + "openshell.gateway.supervisor_pool tunes the shared supervisor pool, but no \ + compute driver hosts sandboxes on it; set the Kubernetes driver's \ + supervisor_topology = \"gateway\" or remove the table", + )); + } match instance { ComputeDriverInstance::InProcess(driver) => ComputeRuntime::from_driver( registration.name, @@ -1882,6 +1914,12 @@ async fn build_compute_runtime( } } } + ConfiguredComputeDriver::Remote { .. } if supervisor_pool.is_some() => { + return Err(Error::config( + "openshell.gateway.supervisor_pool tunes the shared supervisor pool, which \ + only an in-process compute driver can use; remove the table", + )); + } ConfiguredComputeDriver::Remote { name } => { let remote_config = compute::driver_config::remote_driver_config_from_context(driver_startup, &name)?; @@ -2304,6 +2342,54 @@ mod tests { #[derive(Clone, Copy)] struct TestComputeDriverFactory; + #[test] + fn shared_supervision_is_always_a_pool_tuned_by_its_table() { + let endpoint_overrides = std::collections::BTreeMap::new(); + let (_shutdown, shutdown_rx) = watch::channel(false); + for (configured, offered) in [ + ( + None, + openshell_supervisor_control::POOL_RECOVERY_TIMEOUT_SECS, + ), + (Some(75), 75), + ] { + let shared = std::sync::OnceLock::new(); + let context = super::ComputeDriverBuildContext { + config: super::ComputeDriverConfigContext { + driver_name: "kubernetes", + gateway_name: "test", + gateway_bind_address: "127.0.0.1:8080".parse().unwrap(), + gateway_log_level: "info", + driver_startup: crate::compute::driver_config::DriverStartupContext { + file: None, + guest_tls: None, + gateway_port: 8080, + gateway_tls_enabled: true, + endpoint_overrides: &endpoint_overrides, + }, + }, + shutdown_rx: shutdown_rx.clone(), + shared_supervision: &shared, + supervisor_recovery_timeout_secs: configured, + }; + // The driver asks for shared supervision as before; it is the + // pool, with the configured recovery deadline or the default. + let supervision = context.shared_supervision(); + assert_eq!( + supervision.controller().pool_recovery_timeout_secs(), + offered + ); + assert_eq!( + context + .shared_supervision() + .controller() + .pool_recovery_timeout_secs(), + offered + ); + assert!(shared.get().is_some()); + } + } + #[test] fn factory_admission_defaults_remain_enabled() { use super::ComputeDriverFactory as _; diff --git a/crates/openshell-server/src/multiplex.rs b/crates/openshell-server/src/multiplex.rs index baf9b53ef3..c1a6020473 100644 --- a/crates/openshell-server/src/multiplex.rs +++ b/crates/openshell-server/src/multiplex.rs @@ -4829,6 +4829,7 @@ mod tests { sandbox_id: "sandbox-a".to_string(), source: SandboxIdentitySource::BootstrapJwt { issuer: "openshell-gateway:test".to_string(), + auth_epoch: None, }, trust_domain: None, }), diff --git a/crates/openshell-server/src/persistence/mod.rs b/crates/openshell-server/src/persistence/mod.rs index 929e1e72e2..f14111b74d 100644 --- a/crates/openshell-server/src/persistence/mod.rs +++ b/crates/openshell-server/src/persistence/mod.rs @@ -180,6 +180,57 @@ pub struct MessagePage { const FULL_SCAN_PAGE_SIZE: u32 = 1000; +/// One conditional write of a [`Store::write_atomically`] group. +#[derive(Debug, Clone)] +pub enum AtomicWrite { + /// Insert a new object at resource version 1; the group fails when it + /// exists (or its name is taken). + Create { + object_type: String, + id: String, + name: String, + workspace: String, + payload: Vec, + }, + /// Replace an object's payload at exactly `expected_version`. + Update { + object_type: String, + id: String, + payload: Vec, + expected_version: u64, + }, + /// Delete an object at exactly `expected_version`. + Delete { + object_type: String, + id: String, + expected_version: u64, + }, +} + +impl AtomicWrite { + fn key(&self) -> (&str, &str) { + match self { + Self::Create { + object_type, id, .. + } + | Self::Update { + object_type, id, .. + } + | Self::Delete { + object_type, id, .. + } => (object_type, id), + } + } +} + +/// The order a group's writes run in: by `(object_type, id)`, so groups that +/// share rows take their row locks in one order and cannot deadlock. +fn atomic_write_order(writes: &[AtomicWrite]) -> Vec { + let mut order = (0..writes.len()).collect::>(); + order.sort_by(|left, right| writes[*left].key().cmp(&writes[*right].key())); + order +} + /// Write condition for compare-and-swap operations. #[derive(Debug, Clone, Copy)] pub enum WriteCondition { @@ -580,6 +631,25 @@ impl Store { store_dispatch_traced!(self.delete_if(object_type, id, expected_resource_version)) } + /// Apply conditional writes to several objects all together or not at + /// all, in one transaction on one connection. A create whose object or + /// name exists fails the group with `UniqueViolation`; an update or delete + /// whose object is absent or at another version fails it with `Conflict` + /// (with no current version). Returns each write's result in input order, + /// `None` for deletes. + #[tracing::instrument( + name = "store", + level = "debug", + skip_all, + fields(otel.name = "store.write_atomically", otel.status_code = tracing::field::Empty, writes = writes.len()) + )] + pub async fn write_atomically( + &self, + writes: &[AtomicWrite], + ) -> PersistenceResult>> { + store_dispatch_traced!(self.write_atomically(writes)) + } + /// Insert or update a generic named object with an application-owned scope. #[allow(clippy::too_many_arguments)] #[tracing::instrument( diff --git a/crates/openshell-server/src/persistence/postgres.rs b/crates/openshell-server/src/persistence/postgres.rs index ad5b2bec4c..a44f2399a7 100644 --- a/crates/openshell-server/src/persistence/postgres.rs +++ b/crates/openshell-server/src/persistence/postgres.rs @@ -6,9 +6,9 @@ use super::mutation_lock::{ 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, - map_migrate_error, + AtomicWrite, DraftChunkRecord, ObjectCursor, ObjectListQuery, ObjectRecord, PersistenceError, + PersistenceResult, PolicyRecord, WriteCondition, WriteResult, atomic_write_order, + current_time_ms, map_db_error, map_migrate_error, }; use crate::gateway_metrics::GaugeSlot; use crate::policy_store::{ @@ -733,6 +733,107 @@ RETURNING resource_version, created_at_ms, updated_at_ms } } + pub async fn write_atomically( + &self, + writes: &[AtomicWrite], + ) -> PersistenceResult>> { + let now_ms = current_time_ms(); + let result_of = |row: &sqlx::postgres::PgRow| { + let resource_version: i64 = row.try_get("resource_version").unwrap_or(1); + WriteResult { + resource_version: resource_version.max(1).cast_unsigned(), + created_at_ms: row.get("created_at_ms"), + updated_at_ms: row.get("updated_at_ms"), + } + }; + let mut results = vec![None; writes.len()]; + // One transaction on one data-pool connection; dropping it on an + // early return rolls every write back. Writes run in a fixed row + // order, so concurrent groups take row locks in one order. + let mut tx = self.pool.begin().await.map_err(|e| map_db_error(&e))?; + for index in atomic_write_order(writes) { + match &writes[index] { + AtomicWrite::Create { + object_type, + id, + name, + workspace, + payload, + } => { + let row = sqlx::query( + r" +INSERT INTO objects (object_type, id, name, workspace, payload, created_at_ms, updated_at_ms, labels, resource_version) +VALUES ($1, $2, $3, $4, $5, $6, $6, '{}'::jsonb, 1) +RETURNING resource_version, created_at_ms, updated_at_ms +", + ) + .bind(object_type) + .bind(id) + .bind(name) + .bind(workspace) + .bind(payload) + .bind(now_ms) + .fetch_one(&mut *tx) + .await + .map_err(|e| map_db_error(&e))?; + results[index] = Some(result_of(&row)); + } + AtomicWrite::Update { + object_type, + id, + payload, + expected_version, + } => { + let row = sqlx::query( + r" +UPDATE objects +SET payload = $4, updated_at_ms = $5, resource_version = resource_version + 1 +WHERE object_type = $1 AND id = $2 AND resource_version = $3 +RETURNING resource_version, created_at_ms, updated_at_ms +", + ) + .bind(object_type) + .bind(id) + .bind(i64::try_from(*expected_version).unwrap_or(i64::MAX)) + .bind(payload) + .bind(now_ms) + .fetch_optional(&mut *tx) + .await + .map_err(|e| map_db_error(&e))? + .ok_or(PersistenceError::Conflict { + current_resource_version: None, + })?; + results[index] = Some(result_of(&row)); + } + AtomicWrite::Delete { + object_type, + id, + expected_version, + } => { + let deleted = sqlx::query( + r" +DELETE FROM objects +WHERE object_type = $1 AND id = $2 AND resource_version = $3 +", + ) + .bind(object_type) + .bind(id) + .bind(i64::try_from(*expected_version).unwrap_or(i64::MAX)) + .execute(&mut *tx) + .await + .map_err(|e| map_db_error(&e))?; + if deleted.rows_affected() == 0 { + return Err(PersistenceError::Conflict { + current_resource_version: None, + }); + } + } + } + } + tx.commit().await.map_err(|e| map_db_error(&e))?; + Ok(results) + } + pub async fn delete_if( &self, object_type: &str, diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index cb76e3fa6f..94d52a67f2 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -2,9 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use super::{ - DraftChunkRecord, ObjectCursor, ObjectListQuery, ObjectRecord, PersistenceError, - PersistenceResult, PolicyRecord, WriteCondition, WriteResult, current_time_ms, map_db_error, - map_migrate_error, + AtomicWrite, DraftChunkRecord, ObjectCursor, ObjectListQuery, ObjectRecord, PersistenceError, + PersistenceResult, PolicyRecord, WriteCondition, WriteResult, atomic_write_order, + current_time_ms, map_db_error, map_migrate_error, }; use crate::policy_store::{ AtomicPolicyRevisionWrite, apply_draft_chunk_evaluation, draft_chunk_evaluation_inputs_match, @@ -528,6 +528,104 @@ RETURNING "resource_version", "created_at_ms", "updated_at_ms" } } + pub async fn write_atomically( + &self, + writes: &[AtomicWrite], + ) -> PersistenceResult>> { + let now_ms = current_time_ms(); + let mut results = vec![None; writes.len()]; + // Dropping the transaction on an early return rolls every write back. + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(|e| map_db_error(&e))?; + for index in atomic_write_order(writes) { + match &writes[index] { + AtomicWrite::Create { + object_type, + id, + name, + workspace, + payload, + } => { + sqlx::query( + r#" +INSERT INTO "objects" ("object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version") +VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6, '{}', 1) +"#, + ) + .bind(object_type) + .bind(id) + .bind(name) + .bind(workspace) + .bind(payload) + .bind(now_ms) + .execute(&mut *tx) + .await + .map_err(|e| map_db_error(&e))?; + results[index] = Some(WriteResult { + resource_version: 1, + created_at_ms: now_ms, + updated_at_ms: now_ms, + }); + } + AtomicWrite::Update { + object_type, + id, + payload, + expected_version, + } => { + let row = sqlx::query( + r#" +UPDATE "objects" +SET "payload" = ?4, "updated_at_ms" = ?5, "resource_version" = "resource_version" + 1 +WHERE "object_type" = ?1 AND "id" = ?2 AND "resource_version" = ?3 +RETURNING "resource_version", "created_at_ms", "updated_at_ms" +"#, + ) + .bind(object_type) + .bind(id) + .bind(i64::try_from(*expected_version).unwrap_or(i64::MAX)) + .bind(payload) + .bind(now_ms) + .fetch_optional(&mut *tx) + .await + .map_err(|e| map_db_error(&e))? + .ok_or(PersistenceError::Conflict { + current_resource_version: None, + })?; + results[index] = Some(write_result_from_row(&row)); + } + AtomicWrite::Delete { + object_type, + id, + expected_version, + } => { + let deleted = sqlx::query( + r#" +DELETE FROM "objects" +WHERE "object_type" = ?1 AND "id" = ?2 AND "resource_version" = ?3 +"#, + ) + .bind(object_type) + .bind(id) + .bind(i64::try_from(*expected_version).unwrap_or(i64::MAX)) + .execute(&mut *tx) + .await + .map_err(|e| map_db_error(&e))?; + if deleted.rows_affected() == 0 { + return Err(PersistenceError::Conflict { + current_resource_version: None, + }); + } + } + } + } + tx.commit().await.map_err(|e| map_db_error(&e))?; + Ok(results) + } + pub async fn delete_if( &self, object_type: &str, diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index be86441245..3117219abf 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -3082,3 +3082,216 @@ async fn store_operations_export_spans_with_parents() { "the store span records what it queried" ); } + +/// Exercises one atomic write group on `store`: results come back in input +/// order, and any unmet condition leaves every row as it was. Writes run in +/// `(object_type, id)` order, so each failing write below sorts after a +/// create, an update, and a delete that already succeeded inside the group. +async fn atomic_write_groups_apply_all_or_nothing(store: &Store) { + use super::AtomicWrite; + let create = |id: &str, name: &str| AtomicWrite::Create { + object_type: "atomic_test".to_string(), + id: id.to_string(), + name: name.to_string(), + workspace: String::new(), + payload: b"one".to_vec(), + }; + let update = |id: &str, expected_version| AtomicWrite::Update { + object_type: "atomic_test".to_string(), + id: id.to_string(), + payload: b"two".to_vec(), + expected_version, + }; + let delete = |id: &str, expected_version| AtomicWrite::Delete { + object_type: "atomic_test".to_string(), + id: id.to_string(), + expected_version, + }; + let results = store + .write_atomically(&[create("x", "x"), create("b", "b"), create("a", "a")]) + .await + .unwrap(); + assert_eq!(results.len(), 3); + assert!( + results + .iter() + .all(|result| result.as_ref().unwrap().resource_version == 1) + ); + let unchanged = || async { + for id in ["a", "b", "x"] { + let row = store.get("atomic_test", id).await.unwrap().unwrap(); + assert_eq!( + (row.payload.as_slice(), row.resource_version), + (b"one".as_slice(), 1), + "{id} is unchanged" + ); + } + assert!(store.get("atomic_test", "c").await.unwrap().is_none()); + }; + + // Each group updates a, deletes b, and creates c before its last write + // fails; every one of those effects is rolled back. + let failing_last = [ + // An update at a stale version. + update("x", 7), + // A create of an existing id. + create("x", "other"), + // A create of a new id under an existing name. + create("y", "x"), + // A delete at a stale version. + delete("x", 7), + ]; + for last in failing_last { + let unique = matches!(last, AtomicWrite::Create { .. }); + let outcome = store + .write_atomically(&[last, create("c", "c"), delete("b", 1), update("a", 1)]) + .await; + if unique { + assert!( + matches!(outcome, Err(PersistenceError::UniqueViolation { .. })), + "{outcome:?}" + ); + } else { + assert!( + matches!(outcome, Err(PersistenceError::Conflict { .. })), + "{outcome:?}" + ); + } + unchanged().await; + assert!(store.get("atomic_test", "y").await.unwrap().is_none()); + } + + // A group whose conditions hold applies every write. + let applied = store + .write_atomically(&[delete("b", 1), update("a", 1), create("c", "c")]) + .await + .unwrap(); + assert!(applied[0].is_none()); + assert_eq!(applied[1].as_ref().unwrap().resource_version, 2); + assert_eq!(applied[2].as_ref().unwrap().resource_version, 1); + assert!(store.get("atomic_test", "b").await.unwrap().is_none()); + assert!(store.get("atomic_test", "c").await.unwrap().is_some()); + let a = store.get("atomic_test", "a").await.unwrap().unwrap(); + assert_eq!( + (a.payload.as_slice(), a.resource_version), + (b"two".as_slice(), 2) + ); +} + +/// Concurrent groups that update the same row at the same version: exactly +/// one applies, and the rest change nothing. +async fn concurrent_atomic_write_groups_admit_one(store: &Store) { + use super::AtomicWrite; + store + .write_atomically(&[AtomicWrite::Create { + object_type: "atomic_race".to_string(), + id: "slot".to_string(), + name: "slot".to_string(), + workspace: String::new(), + payload: b"0".to_vec(), + }]) + .await + .unwrap(); + let contenders = (0..8).map(|index| { + let store = store.clone(); + async move { + store + .write_atomically(&[ + AtomicWrite::Create { + object_type: "atomic_race".to_string(), + id: format!("claim-{index}"), + name: format!("claim-{index}"), + workspace: String::new(), + payload: Vec::new(), + }, + AtomicWrite::Update { + object_type: "atomic_race".to_string(), + id: "slot".to_string(), + payload: index.to_string().into_bytes(), + expected_version: 1, + }, + ]) + .await + } + }); + let outcomes = futures::future::join_all(contenders).await; + let winners = outcomes + .iter() + .enumerate() + .filter(|(_, outcome)| outcome.is_ok()) + .map(|(index, _)| index) + .collect::>(); + assert_eq!(winners.len(), 1, "{outcomes:?}"); + for outcome in &outcomes { + if let Err(error) = outcome { + assert!( + matches!(error, PersistenceError::Conflict { .. }), + "{error:?}" + ); + } + } + let slot = store.get("atomic_race", "slot").await.unwrap().unwrap(); + assert_eq!(slot.payload, winners[0].to_string().into_bytes()); + assert_eq!(slot.resource_version, 2); + for index in 0..8 { + let claim = store + .get("atomic_race", &format!("claim-{index}")) + .await + .unwrap(); + assert_eq!(claim.is_some(), index == winners[0], "claim-{index}"); + } +} + +#[tokio::test] +async fn sqlite_atomic_write_groups_apply_all_or_nothing() { + atomic_write_groups_apply_all_or_nothing(&test_store().await).await; +} + +#[tokio::test] +async fn sqlite_concurrent_atomic_write_groups_admit_one() { + concurrent_atomic_write_groups_admit_one(&test_store().await).await; +} + +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_atomic_write_groups_apply_all_or_nothing() { + let schema = super::test_postgres::TestSchema::create("atomic_write").await; + let store = schema.connect_store().await; + atomic_write_groups_apply_all_or_nothing(&store).await; + drop(store); + 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_concurrent_atomic_write_groups_admit_one() { + let schema = super::test_postgres::TestSchema::create("atomic_race").await; + let store = schema.connect_store().await; + concurrent_atomic_write_groups_admit_one(&store).await; + drop(store); + schema.drop_schema().await; +} + +/// A group's queries all run on the connection its transaction holds, so a +/// one-connection pool neither deadlocks nor loses atomicity. +#[tokio::test] +#[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] +async fn postgres_atomic_write_groups_run_on_a_one_connection_pool() { + let schema = super::test_postgres::TestSchema::create("atomic_one").await; + // Below the operator minimum on purpose: atomic groups alone must not + // need a second connection. + let backend = super::postgres::PostgresStore::connect(schema.url(), Some(1), None) + .await + .expect("connect a one-connection store"); + backend.migrate().await.expect("migrate"); + let store = Store::Postgres(backend); + assert_eq!(store.max_connections(), 1); + tokio::time::timeout(std::time::Duration::from_mins(1), async { + atomic_write_groups_apply_all_or_nothing(&store).await; + concurrent_atomic_write_groups_admit_one(&store).await; + }) + .await + .expect("atomic write groups finish on one connection"); + drop(store); + schema.drop_schema().await; +} diff --git a/crates/openshell-server/src/supervisor_owner.rs b/crates/openshell-server/src/supervisor_owner.rs index 11a79f2a3a..f83374eb42 100644 --- a/crates/openshell-server/src/supervisor_owner.rs +++ b/crates/openshell-server/src/supervisor_owner.rs @@ -4,8 +4,10 @@ //! Shared supervisor-session ownership index for HA gateway replicas. use crate::persistence::{ObjectRecord, PersistenceError, Store, WriteCondition}; +use openshell_core::jwt::CredentialEpoch; use openshell_core::time::now_ms; use serde::{Deserialize, Serialize}; +use std::cmp::Ordering; use std::sync::Arc; use std::time::{Duration, Instant}; use thiserror::Error; @@ -44,6 +46,10 @@ struct OwnerPayload { session_id: String, supervisor_instance_id: String, connection_epoch: u64, + /// Ownership term: the credential epoch of the session's bearer. 0, as + /// in records written by older gateways, means unknown. + #[serde(default)] + auth_epoch: u64, owner_replica_id: String, owner_peer_endpoint: String, connected_at_ms: i64, @@ -64,6 +70,8 @@ pub struct OwnerRecord { pub session_id: String, pub supervisor_instance_id: String, pub connection_epoch: u64, + /// The owning session's ownership term, when known. + pub auth_epoch: Option, pub owner_replica_id: String, pub owner_peer_endpoint: String, #[allow(dead_code)] @@ -91,6 +99,7 @@ pub struct OwnerGuard { pub session_id: String, pub supervisor_instance_id: String, pub connection_epoch: u64, + pub auth_epoch: Option, pub owner_replica_id: String, pub owner_peer_endpoint: String, connected_at_ms: i64, @@ -116,6 +125,8 @@ impl SupervisorOwnerIndex { Self { store, ttl } } + /// Publish ownership for a session whose term is unknown. + #[cfg(test)] pub async fn publish( &self, sandbox_id: &str, @@ -124,6 +135,31 @@ impl SupervisorOwnerIndex { connection_epoch: u64, owner_replica_id: &str, owner_peer_endpoint: &str, + ) -> Result { + self.publish_term( + sandbox_id, + session_id, + supervisor_instance_id, + connection_epoch, + None, + owner_replica_id, + owner_peer_endpoint, + ) + .await + } + + /// Publish ownership of `sandbox_id` for a session of ownership term + /// `auth_epoch`, the credential epoch of its bearer. + #[allow(clippy::too_many_arguments)] + pub async fn publish_term( + &self, + sandbox_id: &str, + session_id: &str, + supervisor_instance_id: &str, + connection_epoch: u64, + auth_epoch: Option, + owner_replica_id: &str, + owner_peer_endpoint: &str, ) -> Result { let connected_at_ms = now_ms(); let payload = OwnerPayload { @@ -131,6 +167,7 @@ impl SupervisorOwnerIndex { session_id: session_id.to_string(), supervisor_instance_id: supervisor_instance_id.to_string(), connection_epoch, + auth_epoch: auth_epoch.map_or(0, CredentialEpoch::get), owner_replica_id: owner_replica_id.to_string(), owner_peer_endpoint: owner_peer_endpoint.to_string(), connected_at_ms, @@ -149,6 +186,7 @@ impl SupervisorOwnerIndex { &existing, supervisor_instance_id, connection_epoch, + auth_epoch, self.ttl, ) => { @@ -166,6 +204,7 @@ impl SupervisorOwnerIndex { session_id: session_id.to_string(), supervisor_instance_id: supervisor_instance_id.to_string(), connection_epoch, + auth_epoch, owner_replica_id: owner_replica_id.to_string(), owner_peer_endpoint: owner_peer_endpoint.to_string(), connected_at_ms, @@ -180,6 +219,7 @@ impl SupervisorOwnerIndex { session_id: guard.session_id.clone(), supervisor_instance_id: guard.supervisor_instance_id.clone(), connection_epoch: guard.connection_epoch, + auth_epoch: guard.auth_epoch.map_or(0, CredentialEpoch::get), owner_replica_id: guard.owner_replica_id.clone(), owner_peer_endpoint: guard.owner_peer_endpoint.clone(), connected_at_ms: guard.connected_at_ms, @@ -284,6 +324,7 @@ fn decode_owner_record(record: &ObjectRecord) -> Result<(String, OwnerRecord), O session_id: payload.session_id, supervisor_instance_id: payload.supervisor_instance_id, connection_epoch: payload.connection_epoch, + auth_epoch: CredentialEpoch::new(payload.auth_epoch).ok(), owner_replica_id: payload.owner_replica_id, owner_peer_endpoint: payload.owner_peer_endpoint, connected_at_ms: payload.connected_at_ms, @@ -293,18 +334,42 @@ fn decode_owner_record(record: &ObjectRecord) -> Result<(String, OwnerRecord), O )) } +/// How ownership term `claimed` orders against `current`. Terms are ordered +/// only when both are known. +pub fn term_order( + claimed: Option, + current: Option, +) -> Option { + claimed + .zip(current) + .map(|(claimed, current)| claimed.cmp(¤t)) +} + +/// Whether a session may replace `existing` as its sandbox's owner. +/// +/// A stale record yields to any session. A fresh record yields at once to a +/// later ownership term, whichever supervisor process holds it, and never to +/// an earlier one; a restart's new generation and a takeover both advance the +/// term. Within a term, or where either term is unknown, only a reconnect of +/// the same supervisor instance replaces it. pub fn can_supersede( existing: &OwnerRecord, supervisor_instance_id: &str, connection_epoch: u64, + auth_epoch: Option, ttl: Duration, ) -> bool { if !existing.is_fresh(ttl) { return true; } - - existing.supervisor_instance_id == supervisor_instance_id - && connection_epoch > existing.connection_epoch + match term_order(auth_epoch, existing.auth_epoch) { + Some(Ordering::Greater) => true, + Some(Ordering::Less) => false, + Some(Ordering::Equal) | None => { + existing.supervisor_instance_id == supervisor_instance_id + && connection_epoch > existing.connection_epoch + } + } } #[cfg(test)] @@ -321,6 +386,7 @@ mod tests { session_id: "s1".to_string(), supervisor_instance_id: "inst".to_string(), connection_epoch: 1, + auth_epoch: None, owner_replica_id: "gw-1".to_string(), owner_peer_endpoint: "https://gw-1".to_string(), connected_at_ms: 0, @@ -333,11 +399,49 @@ mod tests { i64::try_from(OWNER_TTL.as_millis()).unwrap() } + fn epoch(value: u64) -> CredentialEpoch { + CredentialEpoch::new(value).unwrap() + } + #[test] fn freshness_clamps_a_future_timestamp_instead_of_going_negative() { let skewed = record_updated_at(now_ms() + owner_ttl_ms() * 10); assert!(skewed.is_fresh(OWNER_TTL)); - assert!(!can_supersede(&skewed, "other-inst", 99, OWNER_TTL)); + assert!(!can_supersede(&skewed, "other-inst", 99, None, OWNER_TTL)); + } + + #[test] + fn fresh_owners_yield_only_to_later_terms_or_reconnects_within_a_term() { + let mut fresh = record_updated_at(now_ms()); + fresh.auth_epoch = Some(epoch(2)); + for (instance, connection_epoch, auth_epoch, supersedes) in [ + ("other-inst", 0, Some(epoch(3)), true), + ("inst", 9, Some(epoch(1)), false), + ("other-inst", 9, Some(epoch(2)), false), + ("inst", 1, Some(epoch(2)), false), + ("inst", 2, Some(epoch(2)), true), + ("other-inst", 9, None, false), + ("inst", 2, None, true), + ] { + assert_eq!( + can_supersede(&fresh, instance, connection_epoch, auth_epoch, OWNER_TTL), + supersedes, + "{instance} at connection epoch {connection_epoch}, term {auth_epoch:?}" + ); + } + } + + #[test] + fn a_stale_owner_yields_even_to_an_earlier_term() { + let mut stale = record_updated_at(now_ms() - owner_ttl_ms() - 1); + stale.auth_epoch = Some(epoch(3)); + assert!(can_supersede( + &stale, + "other-inst", + 0, + Some(epoch(1)), + OWNER_TTL + )); } #[test] @@ -371,7 +475,7 @@ mod tests { fn freshness_expires_past_the_ttl() { let stale = record_updated_at(now_ms() - owner_ttl_ms() - 1); assert!(!stale.is_fresh(OWNER_TTL)); - assert!(can_supersede(&stale, "other-inst", 1, OWNER_TTL)); + assert!(can_supersede(&stale, "other-inst", 1, None, OWNER_TTL)); } #[tokio::test] @@ -435,6 +539,167 @@ mod tests { assert_eq!(record.owner_replica_id, "gw-2"); } + #[tokio::test] + async fn a_later_term_supersedes_a_fresh_owner_of_another_supervisor() { + let index = test_index(OWNER_TTL).await; + index + .publish_term( + "sbx", + "s1", + "inst-a", + 7, + Some(epoch(1)), + "gw-1", + "http://gw-1", + ) + .await + .unwrap(); + // The takeover's supervisor process must not wait out the previous + // owner's record. + let guard = index + .publish_term( + "sbx", + "s2", + "inst-b", + 0, + Some(epoch(2)), + "gw-2", + "http://gw-2", + ) + .await + .unwrap(); + let record = index.read("sbx").await.unwrap().unwrap(); + assert_eq!(record.session_id, guard.session_id); + assert_eq!(record.supervisor_instance_id, "inst-b"); + assert_eq!(record.auth_epoch, Some(epoch(2))); + } + + #[tokio::test] + async fn an_earlier_term_never_supersedes_a_fresh_owner() { + let index = test_index(OWNER_TTL).await; + index + .publish_term( + "sbx", + "s2", + "inst-b", + 0, + Some(epoch(2)), + "gw-2", + "http://gw-2", + ) + .await + .unwrap(); + // Neither the replaced process nor a reconnect of the current one + // can bring the replaced term back. + for (instance, connection_epoch) in [("inst-a", 9), ("inst-b", 1)] { + let err = index + .publish_term( + "sbx", + "s3", + instance, + connection_epoch, + Some(epoch(1)), + "gw-1", + "http://gw-1", + ) + .await + .unwrap_err(); + assert!(matches!(err, OwnerError::AlreadyOwned), "{instance}"); + } + assert_eq!(index.read("sbx").await.unwrap().unwrap().session_id, "s2"); + } + + #[tokio::test] + async fn records_of_older_gateways_have_an_unknown_term() { + let index = test_index(OWNER_TTL).await; + let legacy = serde_json::json!({ + "sandbox_id": "sbx", + "session_id": "s1", + "supervisor_instance_id": "inst-a", + "connection_epoch": 1, + "owner_replica_id": "gw-1", + "owner_peer_endpoint": "http://gw-1", + "connected_at_ms": 0, + }); + index + .store + .put( + OWNER_OBJECT_TYPE, + &owner_object_id("sbx"), + "sbx", + "", + &serde_json::to_vec(&legacy).unwrap(), + None, + ) + .await + .unwrap(); + assert_eq!(index.read("sbx").await.unwrap().unwrap().auth_epoch, None); + + // An unknown term is not an earlier one: only a reconnect of the + // owning instance replaces it. + let err = index + .publish_term( + "sbx", + "s2", + "inst-b", + 5, + Some(epoch(2)), + "gw-2", + "http://gw-2", + ) + .await + .unwrap_err(); + assert!(matches!(err, OwnerError::AlreadyOwned)); + index + .publish_term( + "sbx", + "s2", + "inst-a", + 2, + Some(epoch(2)), + "gw-2", + "http://gw-2", + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn renew_keeps_the_ownership_term() { + let index = test_index(OWNER_TTL).await; + let mut guard = index + .publish_term( + "sbx", + "s1", + "inst", + 1, + Some(epoch(3)), + "gw-1", + "http://gw-1", + ) + .await + .unwrap(); + index.renew(&mut guard).await.unwrap(); + assert_eq!( + index.read("sbx").await.unwrap().unwrap().auth_epoch, + Some(epoch(3)) + ); + assert!(matches!( + index + .publish_term( + "sbx", + "s2", + "inst", + 2, + Some(epoch(2)), + "gw-1", + "http://gw-1" + ) + .await, + Err(OwnerError::AlreadyOwned) + )); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn publish_survives_the_previous_owner_releasing_concurrently() { let index = Arc::new(test_index(OWNER_TTL).await); diff --git a/crates/openshell-server/src/supervisor_pool.rs b/crates/openshell-server/src/supervisor_pool.rs new file mode 100644 index 0000000000..3f0e221719 --- /dev/null +++ b/crates/openshell-server/src/supervisor_pool.rs @@ -0,0 +1,5077 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Durable state of a highly available pool of shared supervisors. +//! +//! The pool reuses the gateway's own HA pattern for the shared supervisor's +//! sandbox instances: +//! +//! - **Membership.** The gateway replica serving a member's stream refreshes a +//! store record for it, from its own task and never behind a reconciliation +//! pass; a record older than [`MEMBER_TTL`] is dead. Only the member's +//! newest stream writes its record. A member that announces it is leaving +//! is out of the pool at once, and a member whose process changed lost +//! every instance its previous process owned. See +//! [`openshell_supervisor_control::pool_timing`]. +//! - **Placement.** A consistent-hash ring ([`GatewayRing`]) over live members +//! chooses the member a new instance launches on. +//! - **Ownership.** A store record per instance, written only with +//! compare-and-swap, names the owning member, its process, and the ownership +//! term. It is authoritative: a live owner keeps its instances whatever the +//! ring says, so membership changes move nothing that is healthy. +//! - **Takeover.** Only an instance whose owner failed or left moves, to the +//! ring's choice among live members. The takeover first records its intent +//! (new owner, next term, the process identity the runtime will pin, and a +//! new gateway-token lineage), then advances the sandbox's persisted +//! credential epoch to the new term, which revokes every credential of the +//! previous owner on every gateway replica, and only then is granted. Any +//! replica resumes an interrupted intent; a stopping, stopped, or deleted +//! sandbox, or a replaced generation, cancels it. +//! +//! - **Desired set.** Every pool instance has a store record, written on the +//! gateway replica that launches it before its owner is recorded, and +//! removed when it is withdrawn, its generation ends, or the compute +//! driver's recovery no longer lists it. Its launch material is kept in the +//! gateway's credential store, as sandbox SSH host keys are, and the record +//! holds only the handle. Every replica derives the pool's desired set from +//! these records, whether or not it runs the compute driver's watch, and +//! loads the material of the instances its own members own. +//! - **Observations.** The replica serving an owner records what the owner +//! reports of its instances, so every replica can judge an instance whose +//! owner streams elsewhere. +//! +//! Each replica serves the members connected to it the credentials of their +//! terms, minted from the granted record and the persisted identity. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use openshell_core::ObjectWorkspace as _; +use openshell_core::jwt::{SandboxId, SupervisorTakeover}; +use openshell_core::proto::datamodel::v1::{CredentialHandle, ObjectMeta, Provider}; +use openshell_core::proto::{Sandbox, SandboxPhase}; +use openshell_core::time::now_ms; +use openshell_supervisor_control::pool_timing::RECORD_WRITE_BOUND; +use openshell_supervisor_control::proto::{ + SandboxInstanceAssignment, SandboxInstanceStatus, SandboxInstanceTakeover, +}; +use openshell_supervisor_control::{ + ConnectedMember, InstanceOwner, InstancePlacement, MAX_ASSIGNED_INSTANCES, MAX_POOL_INSTANCES, + MemberReport, MemberView, PlacementFuture, PlacementRequest, PoolView, + SupervisorHostController, +}; +use prost::Message as _; +use serde::{Deserialize, Serialize}; +use tracing::{info, warn}; +use uuid::Uuid; + +use crate::auth::sandbox_jwt::SandboxSessionJwtAuthority; +use crate::auth::sandbox_session::{ + PersistedSandboxIdentity, TakeoverIdentity, advance_identity_for_takeover, +}; +use crate::credentials::CredentialRuntime; +use crate::gateway_ring::GatewayRing; +use crate::persistence::{ + AtomicWrite, ObjectCursor, ObjectRecord, PersistenceError, Store, WriteCondition, +}; + +const MEMBER_OBJECT_TYPE: &str = "supervisor_pool_member"; +const OWNER_OBJECT_TYPE: &str = "supervisor_instance_owner"; +const ASSIGNMENT_OBJECT_TYPE: &str = "supervisor_instance_assignment"; +/// Launch material stored for an instance record not yet written; the row +/// owns the material's handle until the record does, so a cancelled or +/// failed recording never loses the only handle. +const MATERIAL_CANDIDATE_OBJECT_TYPE: &str = "supervisor_instance_material_candidate"; +const REPORT_OBJECT_TYPE: &str = "supervisor_instance_report"; +/// The routes one member is charged for: exactly the routes whose owner record, +/// in any phase, names the member. Every owner change updates it in the same +/// atomic write, so a member is never charged for more than +/// [`MAX_ASSIGNED_INSTANCES`] routes across gateway replicas. +const CAPACITY_OBJECT_TYPE: &str = "supervisor_pool_capacity"; +/// Attempts of one owner change against concurrent changes of the same +/// members' capacity. +const TRANSITION_ATTEMPTS: usize = 16; +/// Credential key of an instance's launch material in the credential store. +const LAUNCH_MATERIAL_KEY: &str = "SHARED_INSTANCE_LAUNCH_MATERIAL"; + +/// How long a member's record stays live without a refresh. +pub const MEMBER_TTL: Duration = openshell_supervisor_control::pool_timing::MEMBER_TTL; +/// How often each gateway replica reconciles. +pub const RECONCILE_INTERVAL: Duration = Duration::from_secs(2); +/// How long a desired instance may lack an owner record before a replica +/// places it by takeover; a launch records its owner within this. +const UNOWNED_GRACE: Duration = Duration::from_secs(10); +/// Lifetime of the takeover grants this gateway mints. +const TAKEOVER_GRANT_TTL: Duration = Duration::from_mins(10); +/// How long minted takeover credentials are served before they are minted +/// again from the persisted identity, which bounds how stale their grant and +/// gateway-token lineage can be when a host starts the term. +const TAKEOVER_CREDENTIAL_REFRESH: Duration = Duration::from_mins(1); +/// Records read back per list page. +const LIST_LIMIT: u32 = 1024; +/// How long after this replica starts it leaves instances of members that +/// look failed alone, so members cut off by a gateway restart reconnect and +/// refresh their records before anything moves. +const TAKEOVER_STARTUP_GRACE: Duration = Duration::from_secs(30); +/// Most records of ended generations removed per pass. +const CLEANUP_PER_PASS: usize = 64; +/// Most launch materials one pass loads from the credential store. +const LOAD_PER_PASS: usize = 64; +/// How long the owner record of a withdrawn instance is kept, so its owner is +/// told to stop it, when its generation has not visibly ended. +const WITHDRAWN_OWNER_RETENTION: Duration = Duration::from_mins(2); +/// How long a recording may hold its route's reservation before any replica +/// fences it, so a recording that stalled or died cannot block its route. A +/// fenced recording can no longer publish. Liveness only: collection keeps +/// the material of a recording that may still publish by its reservation, +/// not by age. +const RESERVATION_GRACE: Duration = Duration::from_mins(2); +/// Age before which a material candidate is left to the recording that staged +/// it, which removes the candidate itself; collection is safe at any age. +const MATERIAL_CANDIDATE_GRACE: Duration = Duration::from_mins(2); + +fn member_object_id(member_id: &str) -> String { + format!("supervisor-pool-member:{member_id}") +} + +fn owner_object_id(route: &str) -> String { + format!("supervisor-instance-owner:{route}") +} + +/// One recording's row of `route`. The store's unique object name, the route, +/// keeps one row per route; the writer in the id keeps a recording from ever +/// updating another recording's row, even one that replaced its own. +fn assignment_object_id(route: &str, writer: &str) -> String { + format!("supervisor-instance-assignment:{route}:{writer}") +} + +fn capacity_object_id(member_id: &str) -> String { + format!("supervisor-pool-capacity:{member_id}") +} + +/// One ownership term's report row of `route`. A writer of an older term +/// can never overwrite a newer term's row. +fn report_object_id(route: &str, auth_epoch: u64) -> String { + format!("supervisor-instance-report:{route}:{auth_epoch}") +} + +fn report_object_name(route: &str, auth_epoch: u64) -> String { + format!("{route}:{auth_epoch}") +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct MemberPayload { + member_id: String, + member_uid: String, + process_id: String, + runtime_address: String, + #[serde(default)] + leaving: bool, + /// When the stream that wrote this record was accepted. + #[serde(default)] + connected_at_ms: i64, + gateway_replica: String, +} + +/// A pool member whose record is live. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LiveMember { + pub member_id: String, + pub process_id: String, + pub runtime_address: String, + pub leaving: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum OwnerPhase { + /// Intent recorded; the sandbox's credential epoch may not be advanced + /// yet, so no credentials of the term may be issued. + Granting, + /// The term is in force. + Granted, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct OwnerPayload { + route: String, + sandbox_id: String, + generation: String, + member_id: String, + process_id: String, + /// Ownership term: the sandbox credential epoch of the owner. + auth_epoch: u64, + /// Process identity granted for a takeover term; absent for the launch + /// term, whose process pins itself. + #[serde(default, skip_serializing_if = "Option::is_none")] + supervisor_instance_id: Option, + /// Gateway-token lineage started by a takeover term. + #[serde(default, skip_serializing_if = "Option::is_none")] + gateway_token_id: Option, + phase: OwnerPhase, + #[serde(default, skip_serializing_if = "Option::is_none")] + previous_member_id: Option, +} + +#[derive(Debug, Clone)] +struct OwnerRecord { + payload: OwnerPayload, + resource_version: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum RecordPhase { + /// A recording reserved the route; nothing is published. + Reserved, + /// A launch recorded the instance and its launch material, and is + /// admitting it onto a member. Not desired: nothing places, charges, or + /// grants it. Its owner and its publication are one atomic write, so a + /// launch that admits nothing removes it, and one that stalls or dies is + /// fenced like a reservation; its material candidate stays until then. + Admitting, + /// The pool must run the instance, with this launch material. + Published, +} + +/// A pool instance the gateway must run, or a recording's reservation of its +/// route. The launch material lives in the credential store; the record +/// holds its handle. Material candidates have the same shape. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct AssignmentPayload { + route: String, + sandbox_id: String, + generation: String, + /// Workspace of the sandbox, which scopes the material's credential. + workspace: String, + phase: RecordPhase, + /// Token of the recording that reserved the route. It publishes only by + /// advancing its own reservation. + writer: String, + /// Credential handle of the launch material (encoded `CredentialHandle`); + /// empty in a reservation. + material_handle: Vec, +} + +#[derive(Debug, Clone)] +struct AssignmentRecord { + payload: AssignmentPayload, + resource_version: u64, + /// When the route was reserved, which bounds a reservation's life. + created_at_ms: i64, + /// When the row was last written: for a published record, its + /// publication, the one change after its reservation. + updated_at_ms: i64, +} + +/// An instance record as a recording left it. +#[derive(Debug, Clone)] +struct Recording { + record: AssignmentRecord, + /// The material candidate this recording staged and still keeps: only + /// for a record it published as [`RecordPhase::Admitting`] itself. + candidate: Option, +} + +/// What an owner reported of one instance in one ownership term, for gateway +/// replicas that do not serve the owner. Error text is not recorded. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct ReportPayload { + route: String, + member_id: String, + process_id: String, + auth_epoch: u64, + sandbox_id: String, + phase: i32, + ready: bool, + exit_code: i32, + /// When the reporting stream was accepted; within one term, a newer + /// stream's report wins. + #[serde(default)] + connected_at_ms: i64, +} + +/// Listed report rows by object id, with their resource versions. +type ReportRows = BTreeMap; + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +struct CapacityPayload { + member_id: String, + routes: BTreeSet, +} + +/// A member's capacity row as read, or an empty one not yet written. +#[derive(Debug, Clone)] +struct Capacity { + payload: CapacityPayload, + resource_version: Option, +} + +impl Capacity { + fn empty(member_id: &str) -> Self { + Self { + payload: CapacityPayload { + member_id: member_id.to_string(), + routes: BTreeSet::new(), + }, + resource_version: None, + } + } + + /// Whether `route` can be charged here: it already is, or a slot is free. + fn has_room_for(&self, route: &str) -> bool { + self.payload.routes.contains(route) || self.payload.routes.len() < MAX_ASSIGNED_INSTANCES + } + + /// The write that makes this row charge exactly `routes`. + fn write(&self, routes: BTreeSet) -> Result { + let payload = serde_json::to_vec(&CapacityPayload { + member_id: self.payload.member_id.clone(), + routes, + }) + .map_err(|error| PersistenceError::Encode(error.to_string()))?; + let id = capacity_object_id(&self.payload.member_id); + Ok(match self.resource_version { + None => AtomicWrite::Create { + object_type: CAPACITY_OBJECT_TYPE.to_string(), + id, + name: self.payload.member_id.clone(), + workspace: String::new(), + payload, + }, + Some(expected_version) => AtomicWrite::Update { + object_type: CAPACITY_OBJECT_TYPE.to_string(), + id, + payload, + expected_version, + }, + }) + } +} + +fn decode_capacity(record: &ObjectRecord) -> Option { + let payload = serde_json::from_slice::(&record.payload).ok()?; + (record.id == capacity_object_id(&payload.member_id)).then_some(Capacity { + payload, + resource_version: Some(record.resource_version), + }) +} + +/// How an owner change ended. +#[derive(Debug)] +enum Transition { + /// Applied; the owner record now reads as this, or is gone. + Applied(Option>), + /// The target member has no free slot for the route. + Full, + /// The admitted instance record changed or was removed before the owner + /// was written. + InstanceChanged, + /// The owner record is not what the change was based on. + OwnerChanged, +} + +/// What reconciling one route did to its owner record. +#[derive(Debug)] +enum Reconciled { + Unchanged, + /// The record now reads as this. + Changed(Box), + /// The record is gone. + Removed, +} + +/// One gateway replica's view and maintenance of the pool. +pub struct SupervisorPool { + store: Arc, + controller: SupervisorHostController, + authority: Arc, + /// Keeps the pool's launch material. + credentials: CredentialRuntime, + replica_id: String, + member_ttl: Duration, + /// How often the records of served members are refreshed. + member_refresh: Duration, + unowned_grace: Duration, + /// No takeover before this instant; see [`TAKEOVER_STARTUP_GRACE`]. + takeovers_from: tokio::time::Instant, + /// When each desired route without an owner record was first seen. + unowned_since: Mutex>, + /// Takeover credentials minted for terms of members this replica serves. + minted: Mutex>, + /// Launch material loaded for instances owned by members this replica + /// serves, or launched here. + materials: Mutex>, + /// The last report this replica recorded of each route. + /// When each owner record of an instance that is no longer desired was + /// first seen. + withdrawn_since: Mutex>, + /// Where the next pass resumes checking records for ended generations. + sweep_from: Mutex, + /// See [`MATERIAL_CANDIDATE_GRACE`]. + candidate_grace: Duration, + /// See [`RESERVATION_GRACE`]. + reservation_grace: Duration, + /// Pauses the next recording once its material is staged. + #[cfg(test)] + staged_gate: Mutex< + Option<( + tokio::sync::oneshot::Sender<()>, + tokio::sync::oneshot::Receiver<()>, + )>, + >, + #[cfg(test)] + hooks: TestHooks, +} + +/// A pause one test sets and one task takes: it reports reaching the pause +/// and waits to be resumed. +#[cfg(test)] +type Gate = ( + tokio::sync::oneshot::Sender<()>, + tokio::sync::oneshot::Receiver<()>, +); + +/// A pause of owner changes: reports reaching it, if asked, and waits on a +/// barrier shared with competing changes. +#[cfg(test)] +type TransitionGate = ( + Option>, + Arc, +); + +/// Points where tests pause or perturb a pass. +#[cfg(test)] +#[derive(Default)] +struct TestHooks { + /// Pauses the next listing of an object type after its first page. + listing_gate: Mutex>, + /// Routes dropped from assignment and owner listings, as a listing that + /// is not a snapshot can miss rows. + listing_omits: Mutex>, + /// Pauses the next launch that admitted its instance nowhere before it + /// removes the record. + give_up_gate: Mutex>, + /// Fails the next admission write before it reaches the store. + fail_admission: Mutex>, + /// Fails the next removal of a refused launch's record. + fail_give_up: Mutex>, + /// Pauses the next reconciliation pass before it does anything. + reconcile_gate: Mutex>, + /// Replaces the members whose streams the controller serves. + served_members: Mutex>>, + /// Replaces what served members report. + served_reports: Mutex>>, + /// Pauses the next report recording after it read what members report. + report_gate: Mutex>, + /// Pauses the next owner change after it read capacity, before it + /// writes: signals arrival, then waits on the barrier. + transition_gate: Mutex>, +} + +impl std::fmt::Debug for SupervisorPool { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SupervisorPool") + .field("replica_id", &self.replica_id) + .finish_non_exhaustive() + } +} + +fn lock(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +/// The ring over live members that are not leaving. +fn ring_of(members: &[LiveMember]) -> GatewayRing { + GatewayRing::new( + members + .iter() + .filter(|member| !member.leaving) + .map(|member| member.member_id.as_str()), + ) +} + +/// Members to place `route` (of `sandbox_id`) on, best first: `same`, the +/// route's previous member, when its live process may re-adopt it; the +/// ring's choice; then the others, least charged first. Only members that +/// seem to have room for the route are listed; the owner change itself +/// re-checks room atomically. `charges` counts the routes each member is +/// charged for and lists them. +fn candidates<'a>( + ring: &GatewayRing, + members: &'a [LiveMember], + sandbox_id: &str, + route: &str, + charges: &BTreeMap>, + same: Option<&str>, +) -> Vec<&'a LiveMember> { + let routes_of = |member: &LiveMember| charges.get(&member.member_id); + let has_room = |member: &&LiveMember| { + !member.leaving + && routes_of(member).is_none_or(|routes| { + routes.contains(route) || routes.len() < MAX_ASSIGNED_INSTANCES + }) + }; + let load = |member: &LiveMember| routes_of(member).map_or(0, BTreeSet::len); + let preferred = ring.owner_for(sandbox_id); + let mut ordered = members.iter().filter(has_room).collect::>(); + ordered.sort_by_key(|member| { + ( + Some(member.member_id.as_str()) != same, + Some(member.member_id.as_str()) != preferred, + load(member), + member.member_id.clone(), + ) + }); + ordered +} + +/// The routes each member is charged for, from capacity rows. +fn charges_of(capacities: &[Capacity]) -> BTreeMap> { + capacities + .iter() + .map(|capacity| { + ( + capacity.payload.member_id.clone(), + capacity.payload.routes.clone(), + ) + }) + .collect() +} + +/// Move `route`'s charge in a local copy after an owner change. +fn recharge( + charges: &mut BTreeMap>, + route: &str, + before: Option<&str>, + after: Option<&str>, +) { + if let Some(before) = before + && let Some(routes) = charges.get_mut(before) + { + routes.remove(route); + } + if let Some(after) = after { + charges + .entry(after.to_string()) + .or_default() + .insert(route.to_string()); + } +} + +/// Whether the process that owns `owner` still serves its member. +fn owner_alive(owner: &OwnerPayload, members: &[LiveMember]) -> bool { + members.iter().any(|member| { + member.member_id == owner.member_id + && member.process_id == owner.process_id + && !member.leaving + }) +} + +/// Why `owner`, which [`owner_alive`] rejects, loses its instance: its +/// member is leaving, runs a new process, or has no live record. +fn takeover_cause(owner: &OwnerPayload, members: &[LiveMember]) -> &'static str { + match members + .iter() + .find(|member| member.member_id == owner.member_id) + { + Some(member) if member.process_id != owner.process_id => "restarted", + Some(_) => "leaving", + None => "expired", + } +} + +fn decode_owner(record: &ObjectRecord) -> Option { + let payload = serde_json::from_slice::(&record.payload).ok()?; + (record.id == owner_object_id(&payload.route)).then_some(OwnerRecord { + payload, + resource_version: record.resource_version, + }) +} + +/// The report row of each route's current, granted owner term. +fn reports_of_owners( + owners: &BTreeMap, + rows: &ReportRows, +) -> BTreeMap { + owners + .values() + .filter(|owner| owner.payload.phase == OwnerPhase::Granted) + .filter_map(|owner| { + let (report, _) = rows.get(&report_object_id( + &owner.payload.route, + owner.payload.auth_epoch, + ))?; + (report.member_id == owner.payload.member_id + && report.process_id == owner.payload.process_id) + .then(|| { + ( + report.route.clone(), + SandboxInstanceStatus { + route: report.route.clone(), + sandbox_id: report.sandbox_id.clone(), + phase: report.phase, + ready: report.ready, + exit_code: report.exit_code, + error: String::new(), + }, + ) + }) + }) + .collect() +} + +fn decode_assignment(record: &ObjectRecord) -> Option { + let payload = serde_json::from_slice::(&record.payload).ok()?; + (record.id == assignment_object_id(&payload.route, &payload.writer) + && record.name == payload.route) + .then_some(AssignmentRecord { + payload, + resource_version: record.resource_version, + created_at_ms: record.created_at_ms, + updated_at_ms: record.updated_at_ms, + }) +} + +fn decode_report(record: &ObjectRecord) -> Option { + let payload = serde_json::from_slice::(&record.payload).ok()?; + (record.id == report_object_id(&payload.route, payload.auth_epoch)).then_some(payload) +} + +/// The synthetic credential owner of one sandbox's launch material, scoped +/// like the sandbox's SSH host key. +fn material_owner(payload: &AssignmentPayload) -> (String, &str, &str) { + ( + format!("sandbox-instance-{}", payload.sandbox_id), + payload.workspace.as_str(), + payload.sandbox_id.as_str(), + ) +} + +/// Whether `generation` of `sandbox` can still run: the sandbox exists, +/// still runs that generation, and is neither stopped nor being deleted. +fn generation_runs(sandbox: Option<&Sandbox>, generation: &str) -> bool { + sandbox.is_some_and(|sandbox| { + !matches!( + SandboxPhase::try_from(sandbox.phase()).unwrap_or(SandboxPhase::Unknown), + SandboxPhase::Stopped | SandboxPhase::Deleting + ) && sandbox + .metadata + .as_ref() + .and_then(|metadata| PersistedSandboxIdentity::read(&metadata.annotations).ok()) + .is_some_and(|identity| identity.runtime_generation.as_str() == generation) + }) +} + +impl SupervisorPool { + #[must_use] + pub fn new( + store: Arc, + controller: SupervisorHostController, + authority: Arc, + credentials: CredentialRuntime, + replica_id: String, + ) -> Arc { + Arc::new(Self { + store, + controller, + authority, + credentials, + replica_id, + member_ttl: MEMBER_TTL, + member_refresh: openshell_supervisor_control::pool_timing::MEMBER_REFRESH_INTERVAL, + unowned_grace: UNOWNED_GRACE, + takeovers_from: tokio::time::Instant::now() + TAKEOVER_STARTUP_GRACE, + unowned_since: Mutex::new(HashMap::new()), + minted: Mutex::new(HashMap::new()), + materials: Mutex::new(HashMap::new()), + withdrawn_since: Mutex::new(HashMap::new()), + sweep_from: Mutex::new(0), + candidate_grace: MATERIAL_CANDIDATE_GRACE, + reservation_grace: RESERVATION_GRACE, + #[cfg(test)] + staged_gate: Mutex::new(None), + #[cfg(test)] + hooks: TestHooks::default(), + }) + } + + /// Pause the next listing of `object_type` after its first page. + #[cfg(test)] + fn pause_listing_after_first_page( + &self, + object_type: &str, + ) -> ( + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + ) { + let (paused, paused_rx) = tokio::sync::oneshot::channel(); + let (resume, resume_rx) = tokio::sync::oneshot::channel(); + *lock(&self.hooks.listing_gate) = Some((object_type.to_string(), (paused, resume_rx))); + (paused_rx, resume) + } + + /// Pause the next owner change after it read capacity, until the + /// barrier releases it. + #[cfg(test)] + fn pause_next_transition(&self, barrier: Arc) { + *lock(&self.hooks.transition_gate) = Some((None, barrier)); + } + + /// Like [`Self::pause_next_transition`]; the receiver resolves once the + /// change read capacity and waits on the barrier. + #[cfg(test)] + fn pause_next_transition_reached( + &self, + barrier: Arc, + ) -> tokio::sync::oneshot::Receiver<()> { + let (reached, arrival) = tokio::sync::oneshot::channel(); + *lock(&self.hooks.transition_gate) = Some((Some(reached), barrier)); + arrival + } + + /// The owner record of `route`. + #[cfg(test)] + async fn reload_owner(&self, route: &str) -> OwnerRecord { + self.store + .get(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await + .unwrap() + .as_ref() + .and_then(decode_owner) + .unwrap() + } + + /// Drop `route` from assignment and owner listings. + #[cfg(test)] + fn omit_from_listings(&self, route: &str) { + lock(&self.hooks.listing_omits).insert(route.to_string()); + } + + /// Pause the next recording once it has staged its material, until the + /// returned sender fires; the receiver fires when it pauses. + #[cfg(test)] + fn pause_next_recording_after_staging( + &self, + ) -> ( + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + ) { + let (staged, staged_rx) = tokio::sync::oneshot::channel(); + let (publish, publish_rx) = tokio::sync::oneshot::channel(); + *lock(&self.staged_gate) = Some((staged, publish_rx)); + (staged_rx, publish) + } + + /// Report `reports` as if served members had made them. + #[cfg(test)] + fn serve_reports(&self, reports: Vec) { + *lock(&self.hooks.served_reports) = Some(reports); + } + + /// Pause the next report recording after it read what members report. + #[cfg(test)] + fn pause_next_report_recording( + &self, + ) -> ( + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + ) { + let (paused, paused_rx) = tokio::sync::oneshot::channel(); + let (resume, resume_rx) = tokio::sync::oneshot::channel(); + *lock(&self.hooks.report_gate) = Some((paused, resume_rx)); + (paused_rx, resume) + } + + /// Pause the next reconciliation pass before it starts. + #[cfg(test)] + fn pause_next_pass( + &self, + ) -> ( + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + ) { + let (paused, paused_rx) = tokio::sync::oneshot::channel(); + let (resume, resume_rx) = tokio::sync::oneshot::channel(); + *lock(&self.hooks.reconcile_gate) = Some((paused, resume_rx)); + (paused_rx, resume) + } + + /// Serve `members` as if their streams were connected here. + #[cfg(test)] + fn serve(&self, members: Vec) { + *lock(&self.hooks.served_members) = Some(members); + } + + /// Members whose streams this replica serves. + fn served_members(&self) -> Vec { + #[cfg(test)] + { + let members = lock(&self.hooks.served_members).clone(); + if let Some(members) = members { + return members; + } + } + self.controller.connected_members() + } + + /// Install this pool as the controller's launch placement, keep the + /// records of served members fresh, and keep the pool reconciled until + /// shutdown. + pub fn start(self: &Arc, mut shutdown: tokio::sync::watch::Receiver) { + self.controller + .install_placement(Arc::new(PoolPlacement(self.clone()))); + self.start_membership_heartbeat(shutdown.clone()); + let pool = self.clone(); + tokio::spawn(async move { + let mut ticker = tokio::time::interval(RECONCILE_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = ticker.tick() => {} + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + return; + } + continue; + } + } + if let Err(error) = pool.reconcile_once().await { + warn!(%error, "supervisor pool reconciliation failed; retrying"); + } + } + }); + } + + /// Refresh served members' records every [`Self::member_refresh`], and + /// at once when a member attaches, leaves, or detaches. Independent of + /// reconciliation, so a slow pass never lets a served member expire. + fn start_membership_heartbeat( + self: &Arc, + mut shutdown: tokio::sync::watch::Receiver, + ) { + let pool = self.clone(); + tokio::spawn(async move { + let mut membership = pool.controller.membership_changes(); + let mut ticker = tokio::time::interval(pool.member_refresh); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = ticker.tick() => {} + changed = membership.changed() => { + if changed.is_err() { + return; + } + } + changed = shutdown.changed() => { + if changed.is_err() || *shutdown.borrow() { + return; + } + continue; + } + } + pool.refresh_members().await; + } + }); + } + + /// Refresh the records of members this replica serves, all at once and + /// each within [`RECORD_WRITE_BOUND`]; a write that does not finish is + /// abandoned and the next round starts over. A record a newer stream of + /// the member wrote is left alone. + async fn refresh_members(&self) { + let members = self.served_members(); + let writes = members.iter().map(|member| async move { + match tokio::time::timeout(RECORD_WRITE_BOUND, self.write_member(member)).await { + Ok(Ok(())) => {} + Ok(Err(error)) => { + warn!(member = %member.member_id, %error, "supervisor pool member refresh failed"); + } + Err(_) => { + warn!(member = %member.member_id, "supervisor pool member refresh did not finish in time"); + } + } + }); + futures::future::join_all(writes).await; + } + + async fn write_member(&self, member: &ConnectedMember) -> Result<(), PersistenceError> { + let id = member_object_id(&member.member_id); + let existing = self.store.get(MEMBER_OBJECT_TYPE, &id).await?; + // A newer stream of the member, from its process or a new one, wrote + // the record; this replica's older stream is stale. + if let Some(record) = &existing + && let Ok(current) = serde_json::from_slice::(&record.payload) + && current.connected_at_ms > member.connected_at_ms + && self.fresh(record) + { + return Ok(()); + } + let payload = MemberPayload { + member_id: member.member_id.clone(), + member_uid: member.member_uid.clone(), + process_id: member.process_id.clone(), + runtime_address: member.runtime_address.clone(), + leaving: member.leaving, + connected_at_ms: member.connected_at_ms, + gateway_replica: self.replica_id.clone(), + }; + let bytes = serde_json::to_vec(&payload) + .map_err(|error| PersistenceError::Encode(error.to_string()))?; + let condition = existing.map_or(WriteCondition::MustCreate, |record| { + WriteCondition::MatchResourceVersion(record.resource_version) + }); + match self + .store + .put_if( + MEMBER_OBJECT_TYPE, + &id, + &member.member_id, + "", + &bytes, + None, + condition, + ) + .await + { + // Another replica refreshed it at the same moment; the next + // pass decides again. + Ok(_) + | Err(PersistenceError::Conflict { .. } | PersistenceError::UniqueViolation { .. }) => { + Ok(()) + } + Err(error) => Err(error), + } + } + + fn fresh(&self, record: &ObjectRecord) -> bool { + let ttl_ms = i64::try_from(self.member_ttl.as_millis()).unwrap_or(i64::MAX); + now_ms().saturating_sub(record.updated_at_ms).max(0) < ttl_ms + } + + /// Every record of `object_type`, page by page after an immutable + /// cursor, so rows removed meanwhile cannot shift others out of the + /// listing. It is still not a snapshot: a row written meanwhile can be + /// missed, so a pass confirms an apparent disappearance before acting on + /// it (see [`Self::confirm_absences`]). + async fn list_all(&self, object_type: &str) -> Result, PersistenceError> { + let mut records = Vec::new(); + let mut after: Option = None; + loop { + let page = self + .store + .list_by_type_after(object_type, after.as_ref(), LIST_LIMIT) + .await?; + let complete = page.len() < LIST_LIMIT as usize; + after = page.last().map(ObjectCursor::from); + records.extend(page); + if complete { + break; + } + #[cfg(test)] + { + let gate = { + let mut gate = lock(&self.hooks.listing_gate); + if gate.as_ref().is_some_and(|(kind, _)| kind == object_type) { + gate.take() + } else { + None + } + }; + if let Some((_, (paused, resume))) = gate { + let _ = paused.send(()); + let _ = resume.await; + } + } + } + #[cfg(test)] + if object_type == ASSIGNMENT_OBJECT_TYPE || object_type == OWNER_OBJECT_TYPE { + let omits = lock(&self.hooks.listing_omits).clone(); + records.retain(|record| !omits.contains(&record.name)); + } + Ok(records) + } + + /// Keep what a listing seems to have lost until a point read confirms it + /// is gone: owners the published view has, and instances it or an owner + /// record names. A route that disappears is stopped or released, which a + /// missed row must never cause. A failed read fails the pass, which then + /// publishes nothing. + async fn confirm_absences( + &self, + owners: &mut BTreeMap, + desired: &mut BTreeMap, + withdrawals: &BTreeSet, + ) -> Result<(), PersistenceError> { + let Some(previous) = self.controller.pool_view() else { + return Ok(()); + }; + for route in previous.owners.keys() { + if owners.contains_key(route) { + continue; + } + if let Some(record) = self + .store + .get(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await? + .as_ref() + .and_then(decode_owner) + { + warn!( + route, + "a supervisor instance owner was missing from its listing" + ); + owners.insert(route.clone(), record); + } + } + let suspects = previous + .desired + .iter() + .chain(owners.keys()) + .filter(|route| !desired.contains_key(*route) && !withdrawals.contains(*route)) + .cloned() + .collect::>(); + for route in suspects { + if let Some(record) = self.assignment(&route).await? + && record.payload.phase == RecordPhase::Published + { + warn!( + route, + "a supervisor pool instance was missing from its listing" + ); + desired.insert(route, record); + } + } + Ok(()) + } + + /// Members whose records are live, in a stable order. + async fn live_members(&self) -> Result, PersistenceError> { + let records = self.list_all(MEMBER_OBJECT_TYPE).await?; + let ttl_ms = i64::try_from(self.member_ttl.as_millis()).unwrap_or(i64::MAX); + let mut members = Vec::new(); + for record in records { + if !self.fresh(&record) { + // Records of members that died without leaving would crowd + // live ones out of the listing limit. + if now_ms().saturating_sub(record.updated_at_ms) >= ttl_ms.saturating_mul(4) { + let _ = self + .store + .delete_if(MEMBER_OBJECT_TYPE, &record.id, record.resource_version) + .await; + } + continue; + } + let Ok(payload) = serde_json::from_slice::(&record.payload) else { + continue; + }; + members.push(LiveMember { + member_id: payload.member_id, + process_id: payload.process_id, + runtime_address: payload.runtime_address, + leaving: payload.leaving, + }); + } + members.sort_by(|left, right| left.member_id.cmp(&right.member_id)); + members.dedup_by(|left, right| left.member_id == right.member_id); + Ok(members) + } + + async fn owners(&self) -> Result, PersistenceError> { + Ok(self + .list_all(OWNER_OBJECT_TYPE) + .await? + .iter() + .filter_map(decode_owner) + .map(|record| (record.payload.route.clone(), record)) + .collect()) + } + + /// Record a new owner as a test would find it written by another + /// replica, charging its member like every owner change. + #[cfg(test)] + async fn write_owner( + &self, + payload: &OwnerPayload, + condition: WriteCondition, + ) -> Result { + assert!(matches!(condition, WriteCondition::MustCreate)); + match self.transition(&payload.route, None, Some(payload)).await? { + Transition::Applied(Some(record)) => Ok(*record), + _ => Err(PersistenceError::Conflict { + current_resource_version: None, + }), + } + } + + async fn capacity(&self, member_id: &str) -> Result { + Ok(self + .store + .get(CAPACITY_OBJECT_TYPE, &capacity_object_id(member_id)) + .await? + .as_ref() + .and_then(decode_capacity) + .unwrap_or_else(|| Capacity::empty(member_id))) + } + + async fn capacities(&self) -> Result, PersistenceError> { + Ok(self + .list_all(CAPACITY_OBJECT_TYPE) + .await? + .iter() + .filter_map(decode_capacity) + .collect()) + } + + /// Change `route`'s owner record from `before` (absent when `None`) to + /// `after` (removed when `None`), and move the route's capacity charge + /// from `before`'s member to `after`'s, in one atomic write. A target + /// member without a free slot refuses the change before anything is + /// written; a concurrent change of the same members' capacity is retried + /// on fresh rows while the owner record is still `before`. + async fn transition( + &self, + route: &str, + before: Option<&OwnerRecord>, + after: Option<&OwnerPayload>, + ) -> Result { + self.transition_admitting(route, before, after, None).await + } + + /// [`Self::transition`] that also publishes `admitting`, an + /// [`RecordPhase::Admitting`] record of the route, in the same atomic + /// write, conditioned on its version. A launch that gives the record up + /// removes exactly that version, so an owner is written only for a record + /// that is still being admitted, and a removed record never gains one. + async fn transition_admitting( + &self, + route: &str, + before: Option<&OwnerRecord>, + after: Option<&OwnerPayload>, + admitting: Option<&AssignmentRecord>, + ) -> Result { + let admission = admitting + .map(|record| { + let payload = AssignmentPayload { + phase: RecordPhase::Published, + ..record.payload.clone() + }; + serde_json::to_vec(&payload) + .map(|bytes| AtomicWrite::Update { + object_type: ASSIGNMENT_OBJECT_TYPE.to_string(), + id: assignment_object_id(&payload.route, &payload.writer), + payload: bytes, + expected_version: record.resource_version, + }) + .map_err(|error| PersistenceError::Encode(error.to_string())) + }) + .transpose()?; + let from = before.map(|record| record.payload.member_id.as_str()); + let to = after.map(|payload| payload.member_id.as_str()); + let encode = |payload: &OwnerPayload| { + serde_json::to_vec(payload).map_err(|error| PersistenceError::Encode(error.to_string())) + }; + for _ in 0..TRANSITION_ATTEMPTS { + let mut writes = Vec::new(); + if let Some(to) = to { + let capacity = self.capacity(to).await?; + if !capacity.has_room_for(route) { + return Ok(Transition::Full); + } + if !capacity.payload.routes.contains(route) { + let mut routes = capacity.payload.routes.clone(); + routes.insert(route.to_string()); + writes.push(capacity.write(routes)?); + } + } + if let Some(from) = from + && Some(from) != to + { + let capacity = self.capacity(from).await?; + if capacity.payload.routes.contains(route) { + let mut routes = capacity.payload.routes.clone(); + routes.remove(route); + writes.push(capacity.write(routes)?); + } + } + #[cfg(test)] + { + let gate = lock(&self.hooks.transition_gate).take(); + if let Some((reached, barrier)) = gate { + if let Some(reached) = reached { + let _ = reached.send(()); + } + barrier.wait().await; + } + } + if let Some(admission) = &admission { + writes.push(admission.clone()); + } + #[cfg(test)] + if admission.is_some() && lock(&self.hooks.fail_admission).take().is_some() { + return Err(PersistenceError::Database( + "test: the admission write failed".to_string(), + )); + } + let owner_index = writes.len(); + writes.push(match (before, after) { + (None, Some(payload)) => AtomicWrite::Create { + object_type: OWNER_OBJECT_TYPE.to_string(), + id: owner_object_id(route), + name: route.to_string(), + workspace: String::new(), + payload: encode(payload)?, + }, + (Some(record), Some(payload)) => AtomicWrite::Update { + object_type: OWNER_OBJECT_TYPE.to_string(), + id: owner_object_id(route), + payload: encode(payload)?, + expected_version: record.resource_version, + }, + (Some(record), None) => AtomicWrite::Delete { + object_type: OWNER_OBJECT_TYPE.to_string(), + id: owner_object_id(route), + expected_version: record.resource_version, + }, + (None, None) => return Ok(Transition::Applied(None)), + }); + match self.store.write_atomically(&writes).await { + Ok(results) => { + return Ok(Transition::Applied(after.map(|payload| { + Box::new(OwnerRecord { + payload: payload.clone(), + resource_version: results[owner_index] + .as_ref() + .map_or(0, |written| written.resource_version), + }) + }))); + } + Err( + PersistenceError::Conflict { .. } | PersistenceError::UniqueViolation { .. }, + ) => { + let current = self + .store + .get(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await? + .as_ref() + .and_then(decode_owner); + let unchanged = match (before, ¤t) { + (None, None) => true, + (Some(before), Some(current)) => { + current.resource_version == before.resource_version + } + _ => false, + }; + if !unchanged { + return Ok(Transition::OwnerChanged); + } + if let Some(record) = admitting + && self + .store + .get( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id( + &record.payload.route, + &record.payload.writer, + ), + ) + .await? + .is_none_or(|current| { + current.resource_version != record.resource_version + }) + { + return Ok(Transition::InstanceChanged); + } + } + Err(error) => return Err(error), + } + } + Err(PersistenceError::Conflict { + current_resource_version: None, + }) + } + + /// What an owner change did to `route`'s record. + async fn reconciled( + &self, + route: &str, + transition: Transition, + ) -> Result { + match transition { + Transition::Applied(Some(record)) => Ok(Reconciled::Changed(record)), + Transition::Applied(None) => Ok(Reconciled::Removed), + Transition::Full | Transition::InstanceChanged => Ok(Reconciled::Unchanged), + Transition::OwnerChanged => self.reload(route).await, + } + } + + /// Every instance record and reservation, by route. + async fn all_assignments( + &self, + ) -> Result, PersistenceError> { + Ok(self + .list_all(ASSIGNMENT_OBJECT_TYPE) + .await? + .iter() + .filter_map(decode_assignment) + .map(|record| (record.payload.route.clone(), record)) + .collect()) + } + + /// The published instance records, by route. + async fn assignments(&self) -> Result, PersistenceError> { + let mut records = self.all_assignments().await?; + records.retain(|_, record| record.payload.phase == RecordPhase::Published); + Ok(records) + } + + /// The record or reservation of `route`. + async fn assignment(&self, route: &str) -> Result, PersistenceError> { + Ok(self + .store + .get_by_name(ASSIGNMENT_OBJECT_TYPE, "", route) + .await? + .as_ref() + .and_then(decode_assignment)) + } + + /// Keep `material` in the credential store under a fresh path, so a + /// record that loses a race never overwrites or deletes the winner's. + async fn store_material( + &self, + payload: &AssignmentPayload, + material: &[u8], + ) -> Result { + let value = std::str::from_utf8(material) + .map_err(|_| "sandbox instance launch material is not UTF-8".to_string())?; + let (name, workspace, owner_id) = material_owner(payload); + let handles = self + .credentials + .store_provider_credentials_with_object_id( + &name, + workspace, + owner_id, + &Uuid::new_v4().to_string(), + &HashMap::from([(LAUNCH_MATERIAL_KEY.to_string(), value.to_string())]), + &HashMap::new(), + ) + .await + .map_err(|status| { + format!( + "store sandbox instance launch material: {}", + status.message() + ) + })?; + handles + .get(LAUNCH_MATERIAL_KEY) + .cloned() + .ok_or_else(|| "the credential store kept no launch material".to_string()) + } + + async fn load_material(&self, payload: &AssignmentPayload) -> Result, String> { + let handle = CredentialHandle::decode(payload.material_handle.as_slice()) + .map_err(|_| "invalid launch material credential handle".to_string())?; + let (name, workspace, owner_id) = material_owner(payload); + let owner = Provider { + metadata: Some(ObjectMeta { + id: owner_id.to_string(), + name, + workspace: workspace.to_string(), + ..Default::default() + }), + credential_handles: HashMap::from([(LAUNCH_MATERIAL_KEY.to_string(), handle)]), + ..Default::default() + }; + let mut resolved = self + .credentials + .resolve_provider_handles(&owner, 0) + .await + .map_err(|status| { + format!( + "load sandbox instance launch material: {}", + status.message() + ) + })?; + resolved + .values + .remove(LAUNCH_MATERIAL_KEY) + .map(String::into_bytes) + .ok_or_else(|| "sandbox instance launch material is unavailable".to_string()) + } + + async fn delete_material( + &self, + payload: &AssignmentPayload, + handle: CredentialHandle, + ) -> Result<(), String> { + let (name, workspace, owner_id) = material_owner(payload); + self.credentials + .delete_provider_credential_handles( + &name, + workspace, + owner_id, + &HashMap::from([(LAUNCH_MATERIAL_KEY.to_string(), handle)]), + ) + .await + .map_err(|status| { + format!( + "delete sandbox instance launch material: {}", + status.message() + ) + }) + } + + /// Durably record an instance the pool must run, launch material + /// included. Recording the same generation again keeps the first record. + /// + /// The recording first reserves the route with a record of its own (a + /// fresh writer token in its id), then stores the material under a fresh + /// path, has a candidate row own the handle, and publishes by advancing + /// exactly its reservation's version. Every step is one atomic store + /// write, so nothing is held across them. Collection deletes a candidate's + /// material only once that candidate's recording no longer holds its + /// reservation, after which its publication can never succeed (see + /// [`Self::remove_stale_candidates`]). A recording whose reservation is + /// fenced or withdrawn first fails and deletes its own material; a + /// failure to do so leaves the candidate for collection to retry. + async fn record_assignment( + &self, + assignment: &SandboxInstanceAssignment, + ) -> Result { + self.record_assignment_as(assignment, RecordPhase::Published) + .await + .map(|recording| recording.record) + } + + /// Record a launched instance as [`RecordPhase::Admitting`], for its + /// launch to admit onto a member. + async fn record_admitting( + &self, + assignment: &SandboxInstanceAssignment, + ) -> Result { + self.record_assignment_as(assignment, RecordPhase::Admitting) + .await + } + + /// [`Self::record_assignment`], publishing the record as `phase` + /// ([`RecordPhase::Published`] or [`RecordPhase::Admitting`]). An + /// admitting record keeps its material candidate, which this recording + /// reports; a published one releases it. + async fn record_assignment_as( + &self, + assignment: &SandboxInstanceAssignment, + phase: RecordPhase, + ) -> Result { + let read = |error: PersistenceError| format!("read supervisor pool instance: {error}"); + let existing = |record: AssignmentRecord| match record.payload.phase { + RecordPhase::Published | RecordPhase::Admitting + if record.payload.generation == assignment.generation + && record.payload.sandbox_id == assignment.sandbox_id => + { + Ok(Recording { + record, + candidate: None, + }) + } + RecordPhase::Published | RecordPhase::Admitting => { + Err("the sandbox instance route is recorded for another generation".to_string()) + } + RecordPhase::Reserved => { + Err("the sandbox instance is being recorded by another gateway task".to_string()) + } + }; + if let Some(record) = self.assignment(&assignment.route).await.map_err(read)? { + return existing(record); + } + let sandbox = self + .store + .get_message::(&assignment.sandbox_id) + .await + .map_err(read)? + .ok_or_else(|| "the sandbox of the instance does not exist".to_string())?; + let mut payload = AssignmentPayload { + route: assignment.route.clone(), + sandbox_id: assignment.sandbox_id.clone(), + generation: assignment.generation.clone(), + workspace: sandbox.object_workspace().to_string(), + phase: RecordPhase::Reserved, + writer: Uuid::new_v4().to_string(), + material_handle: Vec::new(), + }; + let id = assignment_object_id(&payload.route, &payload.writer); + let encode = |payload: &AssignmentPayload| { + serde_json::to_vec(payload) + .map_err(|error| format!("encode supervisor pool instance: {error}")) + }; + let reservation = match self + .store + .put_if( + ASSIGNMENT_OBJECT_TYPE, + &id, + &payload.route, + "", + &encode(&payload)?, + None, + WriteCondition::MustCreate, + ) + .await + { + Ok(written) => written.resource_version, + // Another recording reserved or published the route first. + Err(PersistenceError::UniqueViolation { .. }) => { + return existing( + self.assignment(&payload.route) + .await + .map_err(read)? + .ok_or_else(|| { + "the sandbox instance record was concurrently removed".to_string() + })?, + ); + } + Err(error) => return Err(format!("reserve supervisor pool instance: {error}")), + }; + let unreserve = |version| { + let store = self.store.clone(); + let id = id.clone(); + async move { + let _ = store.delete_if(ASSIGNMENT_OBJECT_TYPE, &id, version).await; + } + }; + let handle = match self + .store_material(&payload, &assignment.launch_material) + .await + { + Ok(handle) => handle, + Err(error) => { + unreserve(reservation).await; + return Err(error); + } + }; + payload.material_handle = handle.encode_to_vec(); + let candidate = Uuid::new_v4().to_string(); + if let Err(error) = self + .store + .put_if( + MATERIAL_CANDIDATE_OBJECT_TYPE, + &candidate, + &candidate, + "", + &encode(&payload)?, + None, + WriteCondition::MustCreate, + ) + .await + { + if let Err(cleanup) = self.delete_material(&payload, handle).await { + warn!(route = %payload.route, error = %cleanup, "could not remove launch material that no candidate owns"); + } + unreserve(reservation).await; + return Err(format!("stage supervisor pool instance: {error}")); + } + #[cfg(test)] + { + let gate = lock(&self.staged_gate).take(); + if let Some((staged, publish)) = gate { + let _ = staged.send(()); + let _ = publish.await; + } + } + payload.phase = phase; + let published = self + .store + .put_if( + ASSIGNMENT_OBJECT_TYPE, + &id, + &payload.route, + "", + &encode(&payload)?, + None, + WriteCondition::MatchResourceVersion(reservation), + ) + .await; + match published { + Ok(written) => { + let record = AssignmentRecord { + payload, + resource_version: written.resource_version, + created_at_ms: written.created_at_ms, + updated_at_ms: written.updated_at_ms, + }; + if phase == RecordPhase::Admitting { + // The candidate keeps the material until the record is + // admitted, or is collected once the record is gone. + return Ok(Recording { + record, + candidate: Some(candidate), + }); + } + // The record owns the material now; a candidate row that + // outlives this is released by collection. + self.release_candidate(&record.payload.route, &candidate) + .await; + lock(&self.materials).insert(record.payload.route.clone(), assignment.clone()); + Ok(Recording { + record, + candidate: None, + }) + } + Err(error) => { + // The reservation was fenced or withdrawn, so this handle can + // never be published. + payload.phase = RecordPhase::Reserved; + self.discard_candidate(&candidate, &payload).await; + Err(match error { + PersistenceError::Conflict { .. } => { + "the sandbox instance reservation was withdrawn before publication" + .to_string() + } + error => format!("publish supervisor pool instance: {error}"), + }) + } + } + } + + /// Release a candidate whose material a published record owns; a + /// failure leaves it for collection, which releases it the same way. + async fn release_candidate(&self, route: &str, candidate: &str) { + if let Err(error) = self + .store + .delete(MATERIAL_CANDIDATE_OBJECT_TYPE, candidate) + .await + { + warn!(route, %error, "could not release a supervisor pool material candidate; retrying later"); + } + } + + /// Delete a candidate's material, and then the candidate. A failure + /// keeps the candidate, so a later pass retries. + async fn discard_candidate(&self, candidate: &str, payload: &AssignmentPayload) { + let deleted = match CredentialHandle::decode(payload.material_handle.as_slice()) { + Ok(handle) => self.delete_material(payload, handle).await, + Err(_) => Err("invalid launch material credential handle".to_string()), + }; + if let Err(error) = deleted { + warn!(route = %payload.route, %error, "could not delete unrecorded launch material; retrying later"); + return; + } + if let Err(error) = self + .store + .delete(MATERIAL_CANDIDATE_OBJECT_TYPE, candidate) + .await + { + warn!(route = %payload.route, %error, "could not remove a supervisor pool material candidate; retrying later"); + } + } + + /// Fence reservations and admitting records older than + /// [`RESERVATION_GRACE`]: their recording or launch stalled or died, and + /// once the row is gone it can never publish or be admitted. A recording + /// that published first, or a launch that admitted first, keeps its + /// record. The material of a fenced record stays with its candidate, + /// which collection then deletes. + async fn fence_stale_reservations(&self, records: &BTreeMap) { + let grace_ms = i64::try_from(self.reservation_grace.as_millis()).unwrap_or(i64::MAX); + let stale = records + .values() + .filter(|record| { + matches!( + record.payload.phase, + RecordPhase::Reserved | RecordPhase::Admitting + ) && now_ms().saturating_sub(record.created_at_ms) >= grace_ms + }) + .take(CLEANUP_PER_PASS); + for record in stale { + info!( + route = %record.payload.route, + phase = ?record.payload.phase, + "fencing a supervisor pool instance record that did not publish or admit in time" + ); + if let Err(error) = self + .store + .delete_if( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id(&record.payload.route, &record.payload.writer), + record.resource_version, + ) + .await + && !matches!(error, PersistenceError::Conflict { .. }) + { + warn!(route = %record.payload.route, %error, "could not fence a supervisor pool instance reservation"); + } + } + } + + /// Settle material candidates older than the candidate grace. One whose + /// recording still holds its route's reservation is left to it; one whose + /// handle the route's published record holds is released; any other + /// belongs to a recording whose reservation is gone or was never its + /// route's, which can never publish, so its material is deleted. A + /// recording publishes only by advancing its own reservation, a row no + /// other recording ever writes and that is never recreated once removed, + /// so this check cannot race a publication. A bounded number per pass. + async fn remove_stale_candidates(&self) { + let candidates = match self + .store + .list_by_type( + MATERIAL_CANDIDATE_OBJECT_TYPE, + u32::try_from(CLEANUP_PER_PASS).unwrap_or(u32::MAX), + 0, + ) + .await + { + Ok(candidates) => candidates, + Err(error) => { + warn!(%error, "could not list supervisor pool material candidates"); + return; + } + }; + let grace_ms = i64::try_from(self.candidate_grace.as_millis()).unwrap_or(i64::MAX); + for candidate in candidates { + if now_ms().saturating_sub(candidate.created_at_ms) < grace_ms { + continue; + } + let Ok(payload) = serde_json::from_slice::(&candidate.payload) + else { + let _ = self + .store + .delete(MATERIAL_CANDIDATE_OBJECT_TYPE, &candidate.id) + .await; + continue; + }; + let record = match self.assignment(&payload.route).await { + Ok(record) => record, + Err(error) => { + warn!(route = %payload.route, %error, "could not check a supervisor pool material candidate"); + continue; + } + }; + match record { + Some(record) + if matches!( + record.payload.phase, + RecordPhase::Reserved | RecordPhase::Admitting + ) && record.payload.writer == payload.writer => {} + Some(record) + if record.payload.phase == RecordPhase::Published + && record.payload.material_handle == payload.material_handle => + { + let _ = self + .store + .delete(MATERIAL_CANDIDATE_OBJECT_TYPE, &candidate.id) + .await; + } + _ => self.discard_candidate(&candidate.id, &payload).await, + } + } + } + + /// Remove an instance record and its launch material, or a reservation or + /// admitting record, which fences its recording or launch (the material + /// of those stays with their candidate, which collection deletes). A + /// published record's material goes first, so a failure leaves the record + /// to retry from. + async fn delete_assignment(&self, record: &AssignmentRecord) -> Result<(), String> { + if record.payload.phase == RecordPhase::Published { + let handle = CredentialHandle::decode(record.payload.material_handle.as_slice()) + .map_err(|_| "invalid launch material credential handle".to_string())?; + self.delete_material(&record.payload, handle).await?; + } + match self + .store + .delete_if( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id(&record.payload.route, &record.payload.writer), + record.resource_version, + ) + .await + { + Ok(_) | Err(PersistenceError::Conflict { .. }) => { + lock(&self.materials).remove(&record.payload.route); + Ok(()) + } + Err(error) => Err(format!("remove supervisor pool instance: {error}")), + } + } + + /// Routes of the published instances, with when each was published. A + /// reservation can predate a driver recovery's listing while its + /// publication races it, so the reservation time would let recovery + /// prune an instance its listing could not have seen. + async fn recorded(&self) -> Result, String> { + Ok(self + .assignments() + .await + .map_err(|error| format!("list supervisor pool instances: {error}"))? + .into_iter() + .map(|(route, record)| (route, record.updated_at_ms)) + .collect()) + } + + /// Remove the records of instances withdrawn on this replica. + async fn remove_withdrawn(&self) { + for route in self.controller.pool_withdrawals() { + let removed = match self.assignment(&route).await { + Ok(Some(record)) => self.delete_assignment(&record).await, + Ok(None) => Ok(()), + Err(error) => Err(error.to_string()), + }; + match removed { + Ok(()) => self.controller.forget_withdrawal(&route), + Err(error) => { + warn!(route, %error, "could not remove a withdrawn supervisor pool instance; retrying"); + } + } + } + } + + /// Record a new instance and the owner of its launch term, on the member + /// the ring prefers among live members with room. + /// + /// The instance is recorded first, as [`RecordPhase::Admitting`], which + /// nothing places, charges, or grants. Its owner and its publication are + /// then one atomic write conditioned on the record's version, so every + /// owner has its instance published. A launch that admits it nowhere + /// removes the record it created at that version (never a record another + /// launch of the generation created or admitted) and then its material; + /// a removal that fails, or a launch that stalls or dies, leaves an + /// admitting record that is fenced after [`RESERVATION_GRACE`], and its + /// material is collected with its candidate. A failed write that may have + /// applied keeps the record, so a retry of the launch admits it. + async fn place_launch(&self, request: PlacementRequest) -> Result<(), String> { + let assignment = &request.assignment; + let members = self + .live_members() + .await + .map_err(|error| format!("read supervisor pool membership: {error}"))?; + let owners = self + .owners() + .await + .map_err(|error| format!("read supervisor instance owners: {error}"))?; + if owners.len() >= MAX_POOL_INSTANCES && !owners.contains_key(&assignment.route) { + return Err("the supervisor pool runs as many sandbox instances as it can".to_string()); + } + let charges = charges_of( + &self + .capacities() + .await + .map_err(|error| format!("read supervisor pool capacity: {error}"))?, + ); + let ring = ring_of(&members); + let targets = candidates( + &ring, + &members, + &assignment.sandbox_id, + &assignment.route, + &charges, + None, + ); + let no_room = || { + "no shared supervisor replica is available with room for another sandbox".to_string() + }; + if targets.is_empty() { + return Err(no_room()); + } + let recording = self.record_admitting(assignment).await?; + let admitting = + (recording.record.payload.phase == RecordPhase::Admitting).then_some(&recording.record); + let mut placed = None; + for target in targets { + let payload = OwnerPayload { + route: assignment.route.clone(), + sandbox_id: assignment.sandbox_id.clone(), + generation: assignment.generation.clone(), + member_id: target.member_id.clone(), + process_id: target.process_id.clone(), + auth_epoch: request.auth_epoch, + supervisor_instance_id: None, + gateway_token_id: None, + phase: OwnerPhase::Granted, + previous_member_id: None, + }; + match self + .transition_admitting(&assignment.route, None, Some(&payload), admitting) + .await + .map_err(|error| format!("record supervisor instance owner: {error}"))? + { + Transition::Applied(record) => { + placed = record.map(|record| *record); + break; + } + Transition::Full => {} + // Another launch of the generation admitted or recorded the + // owner first, or the record is gone. + Transition::OwnerChanged | Transition::InstanceChanged => { + placed = Some(self.launched_owner(assignment).await?); + break; + } + } + } + let Some(record) = placed else { + // Another launch of the generation may have admitted it meanwhile. + let Some(record) = self.give_up_admission(&recording, assignment).await else { + return Err(no_room()); + }; + self.placed(assignment, record); + return Ok(()); + }; + // Only the launch that recorded the instance holds its material; a + // retry that admitted an existing record leaves the first recording's + // material, which every replica loads from the record. + if let Some(candidate) = &recording.candidate { + self.release_candidate(&assignment.route, candidate).await; + lock(&self.materials).insert(assignment.route.clone(), assignment.clone()); + } + self.placed(assignment, record); + Ok(()) + } + + /// The owner another launch of `assignment`'s generation recorded, or an + /// error when the route has none or another generation's. + async fn launched_owner( + &self, + assignment: &SandboxInstanceAssignment, + ) -> Result { + match self + .reload(&assignment.route) + .await + .map_err(|error| format!("read supervisor instance owner: {error}"))? + { + Reconciled::Changed(existing) + if existing.payload.generation == assignment.generation => + { + Ok(*existing) + } + Reconciled::Changed(_) => { + Err("the sandbox instance route is owned by another generation".to_string()) + } + _ => Err("the sandbox instance record was withdrawn before admission".to_string()), + } + } + + /// A launch admitted its instance nowhere: remove the admitting record it + /// created, at the version it created, and then its material. A record + /// another launch created, or one admitted meanwhile, is left alone; the + /// owner such a launch recorded is returned. A removal that fails leaves + /// the record for fencing and the material for collection. + async fn give_up_admission( + &self, + recording: &Recording, + assignment: &SandboxInstanceAssignment, + ) -> Option { + let record = &recording.record; + let Some(candidate) = &recording.candidate else { + return self.launched_owner(assignment).await.ok(); + }; + #[cfg(test)] + { + let gate = lock(&self.hooks.give_up_gate).take(); + if let Some((paused, resume)) = gate { + let _ = paused.send(()); + let _ = resume.await; + } + if lock(&self.hooks.fail_give_up).take().is_some() { + warn!(route = %assignment.route, "test: removing a refused admitting record failed"); + return None; + } + } + match self + .store + .delete_if( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id(&record.payload.route, &record.payload.writer), + record.resource_version, + ) + .await + { + Ok(true) => { + info!( + sandbox_id = %assignment.sandbox_id, + route = %assignment.route, + "removed the record of a sandbox instance no supervisor pool member had room for" + ); + self.discard_candidate(candidate, &record.payload).await; + None + } + // Admitted by another launch of the generation, or fenced. + Ok(false) | Err(PersistenceError::Conflict { .. }) => { + self.launched_owner(assignment).await.ok() + } + Err(error) => { + warn!(route = %assignment.route, %error, "could not remove a refused sandbox instance record; it is fenced later"); + None + } + } + } + + /// Serve a newly placed owner at once. + fn placed(&self, assignment: &SandboxInstanceAssignment, record: OwnerRecord) { + info!( + sandbox_id = %assignment.sandbox_id, + route = %assignment.route, + member = %record.payload.member_id, + "placed shared sandbox instance on supervisor pool member" + ); + if record.payload.phase == OwnerPhase::Granted + && record.payload.supervisor_instance_id.is_none() + { + self.controller.place_owner( + &assignment.route, + InstanceOwner { + member_id: record.payload.member_id, + process_id: record.payload.process_id, + auth_epoch: record.payload.auth_epoch, + supervisor_instance_id: None, + takeover: None, + }, + ); + } + } + + /// One reconciliation pass: remove withdrawn instances; then place, take + /// over, complete, and clean up from the records; record what served + /// owners report; load the material served members need; and publish the + /// resulting view. + pub async fn reconcile_once(&self) -> Result<(), PersistenceError> { + #[cfg(test)] + { + let gate = lock(&self.hooks.reconcile_gate).take(); + if let Some((paused, resume)) = gate { + let _ = paused.send(()); + let _ = resume.await; + } + } + self.remove_withdrawn().await; + let members = self.live_members().await?; + let ring = ring_of(&members); + let listed_at = tokio::time::Instant::now(); + // Listed before the owners, so a report row of a term newer than the + // route's listed owner cannot exist among them, and cleanup, which + // removes only rows of older terms or of routes without an owner, + // never removes a row a new term wrote meanwhile. + let report_rows = self.report_rows().await?; + let mut owners = self.owners().await?; + // Listed after the owners: an instance is published before its owner + // is recorded, so every owner listed has its instance listed unless + // it was withdrawn. + let withdrawals = self + .controller + .pool_withdrawals() + .into_iter() + .collect::>(); + let records = self.all_assignments().await?; + self.fence_stale_reservations(&records).await; + self.remove_stale_candidates().await; + let mut desired = records + .into_iter() + .filter(|(route, record)| { + record.payload.phase == RecordPhase::Published && !withdrawals.contains(route) + }) + .collect::>(); + self.confirm_absences(&mut owners, &mut desired, &withdrawals) + .await?; + let mut charges = charges_of(&self.capacities().await?); + let stopping = self.controller.stopping_routes(); + self.forget_owned(&desired, &owners); + for (route, instance) in &desired { + if stopping.contains(route) { + continue; + } + let takeovers_allowed = tokio::time::Instant::now() >= self.takeovers_from; + let outcome = match owners.get(route) { + None if takeovers_allowed => { + self.place_unowned(route, &instance.payload, &ring, &members, &charges) + .await + } + // An intent recorded by any replica, including one that + // crashed before granting it. + Some(record) if record.payload.phase == OwnerPhase::Granting => { + self.complete_takeover(record.clone()).await + } + Some(record) if takeovers_allowed && !owner_alive(&record.payload, &members) => { + self.take_over(record, &ring, &members, &charges).await + } + // A live owner keeps its instance; nothing is placed or + // taken over during the startup grace. + _ => Ok(Reconciled::Unchanged), + }; + match outcome { + Ok(Reconciled::Changed(record)) => { + recharge( + &mut charges, + route, + owners + .get(route) + .map(|previous| previous.payload.member_id.as_str()), + Some(&record.payload.member_id), + ); + owners.insert(route.clone(), *record); + } + Ok(Reconciled::Removed) => { + let previous = owners.remove(route); + recharge( + &mut charges, + route, + previous + .as_ref() + .map(|previous| previous.payload.member_id.as_str()), + None, + ); + } + Ok(Reconciled::Unchanged) => {} + Err(error) => { + warn!(route, %error, "supervisor pool could not reconcile an instance owner"); + } + } + } + self.remove_withdrawn_owners(&desired, &mut owners).await; + self.sweep_ended_instances(&desired).await; + self.record_reports(&owners, &report_rows).await; + self.remove_stale_reports(&owners, &report_rows).await; + let reports = reports_of_owners(&owners, &report_rows); + let materials = self.load_materials(&owners, &desired).await; + // Only granted terms are served; pending intents stay recorded for + // the next pass. + if let Some(view) = self + .view(&members, &owners, &desired, materials, reports) + .await + { + self.controller.update_pool(view, listed_at); + } + Ok(()) + } + + /// Remove the owner records of instances that are no longer desired: at + /// once when their generation ended, else once their owner had time to be + /// told to stop them. + async fn remove_withdrawn_owners( + &self, + desired: &BTreeMap, + owners: &mut BTreeMap, + ) { + let now = tokio::time::Instant::now(); + let withdrawn = { + let mut since = lock(&self.withdrawn_since); + since.retain(|route, _| owners.contains_key(route) && !desired.contains_key(route)); + owners + .values() + .filter(|record| !desired.contains_key(&record.payload.route)) + .take(CLEANUP_PER_PASS) + .map(|record| { + let first_seen = *since.entry(record.payload.route.clone()).or_insert(now); + (record.clone(), now.duration_since(first_seen)) + }) + .collect::>() + }; + for (record, withdrawn_for) in withdrawn { + let removed = if withdrawn_for >= WITHDRAWN_OWNER_RETENTION { + self.remove_owner(&record).await + } else { + self.remove_if_ended(&record).await + }; + match removed { + Ok(Reconciled::Changed(current)) => { + owners.insert(record.payload.route.clone(), *current); + } + // Its report rows are removed with the other rows of routes + // without an owner. + Ok(Reconciled::Removed) => { + owners.remove(&record.payload.route); + } + Ok(Reconciled::Unchanged) => {} + Err(error) => { + warn!(route = %record.payload.route, %error, "could not remove the owner of a withdrawn supervisor pool instance"); + } + } + } + } + + /// Withdraw instances whose generation visibly ended, a bounded number + /// of records per pass, so an instance whose withdrawal was lost with a + /// gateway replica does not stay desired. + async fn sweep_ended_instances(&self, desired: &BTreeMap) { + if desired.is_empty() { + return; + } + let start = { + let mut from = lock(&self.sweep_from); + let start = *from % desired.len(); + *from = start + CLEANUP_PER_PASS; + start + }; + let batch = desired + .values() + .cycle() + .skip(start) + .take(CLEANUP_PER_PASS.min(desired.len())) + .cloned() + .collect::>(); + for record in batch { + let sandbox = match self + .store + .get_message::(&record.payload.sandbox_id) + .await + { + Ok(sandbox) => sandbox, + Err(error) => { + warn!(route = %record.payload.route, %error, "could not check a supervisor pool instance's generation"); + continue; + } + }; + if generation_runs(sandbox.as_ref(), &record.payload.generation) { + continue; + } + info!( + sandbox_id = %record.payload.sandbox_id, + route = %record.payload.route, + "withdrawing a supervisor pool instance whose generation ended" + ); + self.controller.unassign(&record.payload.route); + } + } + + /// Every report row, by object id. + async fn report_rows(&self) -> Result { + Ok(self + .list_all(REPORT_OBJECT_TYPE) + .await? + .iter() + .filter_map(|record| { + decode_report(record) + .map(|payload| (record.id.clone(), (payload, record.resource_version))) + }) + .collect()) + } + + /// Reports of the members this replica serves. + fn served_reports(&self) -> Vec { + #[cfg(test)] + { + let reports = lock(&self.hooks.served_reports).clone(); + if let Some(reports) = reports { + return reports; + } + } + self.controller.member_reports() + } + + /// Record what the owners this replica serves report of their instances, + /// each only if the term its stream acknowledged is the granted term this + /// pass's owner records show, so a status from an earlier term is never + /// relabelled as a later one, even from the same process. A row is + /// rewritten whenever it differs from the report, conditioned on the row + /// as listed, so a row another writer replaced or removed is repaired by + /// the next pass; a row a newer stream of the same term wrote is kept. + async fn record_reports(&self, owners: &BTreeMap, rows: &ReportRows) { + let reports = self.served_reports(); + #[cfg(test)] + { + let gate = lock(&self.hooks.report_gate).take(); + if let Some((paused, resume)) = gate { + let _ = paused.send(()); + let _ = resume.await; + } + } + for report in reports { + let route = report.status.route.clone(); + if !owners.get(&route).is_some_and(|owner| { + owner.payload.phase == OwnerPhase::Granted + && owner.payload.member_id == report.member_id + && owner.payload.process_id == report.process_id + && owner.payload.auth_epoch == report.auth_epoch + }) { + continue; + } + let auth_epoch = report.auth_epoch; + let payload = ReportPayload { + route: route.clone(), + member_id: report.member_id, + process_id: report.process_id, + auth_epoch, + sandbox_id: report.status.sandbox_id, + phase: report.status.phase, + ready: report.status.ready, + exit_code: report.status.exit_code, + connected_at_ms: report.connected_at_ms, + }; + let id = report_object_id(&route, auth_epoch); + let listed = rows.get(&id); + if listed.is_some_and(|(current, _)| { + current == &payload || current.connected_at_ms > payload.connected_at_ms + }) { + continue; + } + let Ok(bytes) = serde_json::to_vec(&payload) else { + continue; + }; + let condition = listed.map_or(WriteCondition::MustCreate, |(_, version)| { + WriteCondition::MatchResourceVersion(*version) + }); + match self + .store + .put_if( + REPORT_OBJECT_TYPE, + &id, + &report_object_name(&route, auth_epoch), + "", + &bytes, + None, + condition, + ) + .await + { + // Another writer changed the row since it was listed; the + // next pass compares against what it wrote. + Ok(_) + | Err( + PersistenceError::Conflict { .. } | PersistenceError::UniqueViolation { .. }, + ) => {} + Err(error) => { + warn!(route, %error, "could not record a supervisor pool instance report"); + } + } + } + } + + /// Remove listed report rows of terms older than their route's owner, or + /// of routes without one; a bounded number per pass, each only as listed. + async fn remove_stale_reports( + &self, + owners: &BTreeMap, + rows: &ReportRows, + ) { + let stale = rows + .iter() + .filter(|(_, (report, _))| { + owners + .get(&report.route) + .is_none_or(|owner| owner.payload.auth_epoch > report.auth_epoch) + }) + .take(CLEANUP_PER_PASS); + for (id, (_, version)) in stale { + let _ = self.store.delete_if(REPORT_OBJECT_TYPE, id, *version).await; + } + } + + /// Launch material of the desired instances owned, or about to be taken + /// over, by member processes this replica serves; at most + /// [`LOAD_PER_PASS`] new loads per pass. + async fn load_materials( + &self, + owners: &BTreeMap, + desired: &BTreeMap, + ) -> BTreeMap { + let served = self + .controller + .connected_members() + .into_iter() + .map(|member| (member.member_id, member.process_id)) + .collect::>(); + let wanted = desired + .iter() + .filter(|(route, _)| { + owners.get(*route).is_some_and(|owner| { + served.contains(&( + owner.payload.member_id.clone(), + owner.payload.process_id.clone(), + )) + }) + }) + .collect::>(); + let missing = { + let mut materials = lock(&self.materials); + materials.retain(|route, assignment| { + wanted + .get(route) + .is_some_and(|record| record.payload.generation == assignment.generation) + }); + wanted + .iter() + .filter(|(route, _)| !materials.contains_key(**route)) + .take(LOAD_PER_PASS) + .map(|(_, record)| (*record).clone()) + .collect::>() + }; + for record in missing { + match self.load_material(&record.payload).await { + Ok(launch_material) => { + lock(&self.materials).insert( + record.payload.route.clone(), + SandboxInstanceAssignment { + route: record.payload.route, + sandbox_id: record.payload.sandbox_id, + generation: record.payload.generation, + launch_material, + takeover: None, + }, + ); + } + Err(error) => { + warn!(route = %record.payload.route, %error, "could not load a supervisor pool instance's launch material; retrying"); + } + } + } + lock(&self.materials) + .iter() + .filter(|(route, _)| wanted.contains_key(*route)) + .map(|(route, assignment)| (route.clone(), assignment.clone())) + .collect() + } + + /// The stored owner record of `route` after a write lost a race. + async fn reload(&self, route: &str) -> Result { + Ok(self + .store + .get(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await? + .as_ref() + .and_then(decode_owner) + .map_or(Reconciled::Removed, |record| { + Reconciled::Changed(Box::new(record)) + })) + } + + /// Track desired routes without an owner record; forget the rest. + fn forget_owned( + &self, + desired: &BTreeMap, + owners: &BTreeMap, + ) { + let now = tokio::time::Instant::now(); + let mut unowned = lock(&self.unowned_since); + unowned.retain(|route, _| desired.contains_key(route) && !owners.contains_key(route)); + for route in desired.keys() { + if !owners.contains_key(route) { + unowned.entry(route.clone()).or_insert(now); + } + } + } + + /// Place a desired instance that has no owner record, such as one whose + /// launch predates the pool. Whichever process may already supervise + /// it is fenced, as for any takeover. + async fn place_unowned( + &self, + route: &str, + assignment: &AssignmentPayload, + ring: &GatewayRing, + members: &[LiveMember], + charges: &BTreeMap>, + ) -> Result { + let unowned_for = lock(&self.unowned_since) + .get(route) + .map(tokio::time::Instant::elapsed); + if unowned_for.is_none_or(|elapsed| elapsed < self.unowned_grace) { + return Ok(Reconciled::Unchanged); + } + let targets = candidates(ring, members, &assignment.sandbox_id, route, charges, None); + if targets.is_empty() { + return Ok(Reconciled::Unchanged); + } + let Some(sandbox) = self + .store + .get_message::(&assignment.sandbox_id) + .await? + else { + return Ok(Reconciled::Unchanged); + }; + let Some(identity) = sandbox + .metadata + .as_ref() + .and_then(|metadata| PersistedSandboxIdentity::read(&metadata.annotations).ok()) + .filter(|identity| identity.runtime_generation.as_str() == assignment.generation) + else { + return Ok(Reconciled::Unchanged); + }; + for target in targets { + let intent = OwnerPayload { + route: route.to_string(), + sandbox_id: assignment.sandbox_id.clone(), + generation: assignment.generation.clone(), + member_id: target.member_id.clone(), + process_id: target.process_id.clone(), + auth_epoch: identity.auth_epoch.get().saturating_add(1), + supervisor_instance_id: Some(Uuid::new_v4().to_string()), + gateway_token_id: Some(Uuid::new_v4().to_string()), + phase: OwnerPhase::Granting, + previous_member_id: None, + }; + match self.transition(route, None, Some(&intent)).await? { + Transition::Applied(Some(record)) => return self.complete_takeover(*record).await, + Transition::Full => {} + // A launch or another replica recorded it first. + transition => return self.reconciled(route, transition).await, + } + } + Ok(Reconciled::Unchanged) + } + + /// Move an instance whose owner failed or left to the ring's choice. + async fn take_over( + &self, + record: &OwnerRecord, + ring: &GatewayRing, + members: &[LiveMember], + charges: &BTreeMap>, + ) -> Result { + // A member whose process restarted re-adopts its routes first: their + // slots are already charged to it, so this needs no free capacity. + let targets = candidates( + ring, + members, + &record.payload.sandbox_id, + &record.payload.route, + charges, + Some(&record.payload.member_id), + ); + for target in targets { + let intent = OwnerPayload { + member_id: target.member_id.clone(), + process_id: target.process_id.clone(), + auth_epoch: record.payload.auth_epoch.saturating_add(1), + supervisor_instance_id: Some(Uuid::new_v4().to_string()), + gateway_token_id: Some(Uuid::new_v4().to_string()), + phase: OwnerPhase::Granting, + previous_member_id: Some(record.payload.member_id.clone()), + ..record.payload.clone() + }; + match self + .transition(&record.payload.route, Some(record), Some(&intent)) + .await? + { + Transition::Applied(Some(intent)) => { + info!( + sandbox_id = %intent.payload.sandbox_id, + route = %intent.payload.route, + from = %record.payload.member_id, + to = %intent.payload.member_id, + term = intent.payload.auth_epoch, + cause = %takeover_cause(&record.payload, members), + "taking over a shared sandbox instance from a failed supervisor" + ); + return self.complete_takeover(*intent).await; + } + Transition::Full => {} + // Another replica moved it first. + transition => return self.reconciled(&record.payload.route, transition).await, + } + } + Ok(Reconciled::Unchanged) + } + + /// Advance the sandbox's credential epoch to the intent's term and grant + /// it, or cancel the intent when the sandbox refuses the takeover. + async fn complete_takeover(&self, record: OwnerRecord) -> Result { + let (Some(epoch), Some(gateway_token_id)) = ( + openshell_core::jwt::CredentialEpoch::new(record.payload.auth_epoch).ok(), + record + .payload + .gateway_token_id + .as_deref() + .and_then(|id| Uuid::parse_str(id).ok()), + ) else { + return Ok(Reconciled::Unchanged); + }; + match advance_identity_for_takeover( + &self.store, + &record.payload.sandbox_id, + &record.payload.generation, + epoch, + gateway_token_id, + ) + .await? + { + TakeoverIdentity::Advanced | TakeoverIdentity::AlreadyAdvanced => { + let granted = OwnerPayload { + phase: OwnerPhase::Granted, + ..record.payload.clone() + }; + let transition = self + .transition(&record.payload.route, Some(&record), Some(&granted)) + .await?; + self.reconciled(&record.payload.route, transition).await + } + TakeoverIdentity::Refused(reason) => { + info!( + sandbox_id = %record.payload.sandbox_id, + route = %record.payload.route, + reason, + "cancelled a shared sandbox instance takeover" + ); + let transition = self + .transition(&record.payload.route, Some(&record), None) + .await?; + self.reconciled(&record.payload.route, transition).await + } + } + } + + /// Remove the owner record of an instance that is no longer desired once + /// its generation has ended: the sandbox is gone, runs another one, or + /// was stopped or is being deleted, so the generation never runs again. + async fn remove_if_ended(&self, record: &OwnerRecord) -> Result { + let sandbox = self + .store + .get_message::(&record.payload.sandbox_id) + .await?; + if generation_runs(sandbox.as_ref(), &record.payload.generation) { + return Ok(Reconciled::Unchanged); + } + self.remove_owner(record).await + } + + async fn remove_owner(&self, record: &OwnerRecord) -> Result { + let transition = self + .transition(&record.payload.route, Some(record), None) + .await?; + self.reconciled(&record.payload.route, transition).await + } + + /// The pool view to serve: live members and the granted owners, with + /// takeover credentials for the members this replica serves. `None` keeps + /// the published view, when credentials a served owner needs cannot be + /// minted yet. + async fn view( + &self, + members: &[LiveMember], + owners: &BTreeMap, + desired: &BTreeMap, + assignments: BTreeMap, + reports: BTreeMap, + ) -> Option { + let served = self + .controller + .connected_members() + .into_iter() + .map(|member| (member.member_id, member.process_id)) + .collect::>(); + let mut view = PoolView { + members: members + .iter() + .map(|member| MemberView { + member_id: member.member_id.clone(), + process_id: member.process_id.clone(), + runtime_address: member.runtime_address.clone(), + leaving: member.leaving, + }) + .collect(), + owners: BTreeMap::new(), + desired: desired.keys().cloned().collect(), + assignments, + reports, + }; + let mut current_terms = std::collections::HashSet::new(); + for record in owners.values() { + let owner = &record.payload; + if owner.phase != OwnerPhase::Granted { + continue; + } + current_terms.insert((owner.route.clone(), owner.auth_epoch)); + let takeover = if owner.supervisor_instance_id.is_some() + && served.contains(&(owner.member_id.clone(), owner.process_id.clone())) + { + match self.takeover_credentials(owner).await { + Ok(Some(credentials)) => Some(credentials), + // The sandbox already moved past this record's term; the + // next pass reads the record that replaced it. + Ok(None) => continue, + Err(error) => { + warn!(route = %owner.route, %error, "could not mint supervisor takeover credentials; keeping the published pool view"); + return None; + } + } + } else { + None + }; + view.owners.insert( + owner.route.clone(), + InstanceOwner { + member_id: owner.member_id.clone(), + process_id: owner.process_id.clone(), + auth_epoch: owner.auth_epoch, + supervisor_instance_id: owner.supervisor_instance_id.clone(), + takeover, + }, + ); + } + lock(&self.minted).retain(|term, _| current_terms.contains(term)); + Some(view) + } + + /// Credentials of a granted takeover term, minted again every + /// [`TAKEOVER_CREDENTIAL_REFRESH`]. A host already supervising the term + /// keeps its own credentials, refreshed since; only one that starts the + /// term uses these. + async fn takeover_credentials( + &self, + owner: &OwnerPayload, + ) -> Result, String> { + let term = (owner.route.clone(), owner.auth_epoch); + if let Some((minted, at)) = lock(&self.minted).get(&term) + && at.elapsed() < TAKEOVER_CREDENTIAL_REFRESH + { + return Ok(Some(minted.clone())); + } + let minted = self.mint_takeover(owner).await?; + if let Some(minted) = &minted { + lock(&self.minted).insert(term, (minted.clone(), tokio::time::Instant::now())); + } + Ok(minted) + } + + /// Mint a takeover term's grant and session credentials from the granted + /// record and the persisted identity, which must be at the same term. + async fn mint_takeover( + &self, + owner: &OwnerPayload, + ) -> Result, String> { + let Some(instance) = &owner.supervisor_instance_id else { + return Ok(None); + }; + let Some(sandbox) = self + .store + .get_message::(&owner.sandbox_id) + .await + .map_err(|error| error.to_string())? + else { + return Ok(None); + }; + let Some(identity) = sandbox + .metadata + .as_ref() + .and_then(|metadata| PersistedSandboxIdentity::read(&metadata.annotations).ok()) + .filter(|identity| { + identity.runtime_generation.as_str() == owner.generation + && identity.auth_epoch.get() == owner.auth_epoch + }) + else { + return Ok(None); + }; + let launch = self + .authority + .mint_persisted_launch(&owner.sandbox_id, &identity) + .map_err(|status| status.message().to_string())?; + let grant = self + .authority + .mint_supervisor_takeover( + &SupervisorTakeover { + sandbox_id: SandboxId::parse(owner.sandbox_id.clone()) + .map_err(|error| error.to_string())?, + runtime_generation: identity.runtime_generation.clone(), + auth_epoch: identity.auth_epoch, + supervisor_instance: instance.clone(), + owner: owner.member_id.clone(), + }, + TAKEOVER_GRANT_TTL, + ) + .map_err(|error| error.to_string())?; + let auth = launch.supervisor; + Ok(Some(SandboxInstanceTakeover { + auth_epoch: owner.auth_epoch, + supervisor_instance_id: instance.clone(), + grant: grant.token.expose_secret().to_string(), + gateway_token: auth.gateway_token.expose_secret().to_string(), + gateway_expires_at: auth.gateway_expires_at, + sandbox_token: auth.sandbox_token.expose_secret().to_string(), + sandbox_expires_at: auth.sandbox_expires_at, + })) + } +} + +/// The pool's launch placement, installed into the controller. +#[derive(Debug)] +struct PoolPlacement(Arc); + +/// Run a durable pool write on its own task, so a caller that is cancelled, +/// such as a launch whose request ended, cannot stop it halfway. +async fn detached( + write: impl Future> + Send + 'static, +) -> Result { + tokio::spawn(write) + .await + .map_err(|_| "supervisor pool write task failed".to_string())? +} + +impl InstancePlacement for PoolPlacement { + fn place_launch(&self, request: PlacementRequest) -> PlacementFuture { + let pool = self.0.clone(); + Box::pin(detached(async move { pool.place_launch(request).await })) + } + + fn recorded(&self) -> PlacementFuture> { + let pool = self.0.clone(); + Box::pin(async move { pool.recorded().await }) + } + + fn record_recovered(&self, assignment: SandboxInstanceAssignment) -> PlacementFuture { + let pool = self.0.clone(); + Box::pin(detached(async move { + pool.record_assignment(&assignment).await.map(drop) + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::sandbox_session::PersistedSandboxIdentity; + use crate::persistence::test_postgres::TestSchema; + use openshell_core::jwt::{AuthenticatedSandboxSession, CredentialEpoch}; + use openshell_core::proto::datamodel::v1::ObjectMeta; + use openshell_core::proto::{SandboxPhase, SandboxStatus}; + use openshell_core::sandbox_generation::SandboxGenerationId; + use openshell_supervisor_control::proto::SandboxInstanceAssignment; + + struct Fixture { + store: Arc, + authority: Arc, + credentials: CredentialRuntime, + } + + impl Fixture { + async fn new() -> Self { + Self::with_store(crate::persistence::test_store().await) + } + + /// A fixture over `store`, such as one connected to `PostgreSQL`. + fn with_store(store: Store) -> Self { + let material = openshell_bootstrap::jwt::generate_jwt_key().expect("jwt key"); + let store = Arc::new(store); + let config = crate::Config::new(None) + .with_credential_drivers(["test-static"]) + .with_default_credential_driver(Some("test-static")); + Self { + credentials: CredentialRuntime::from_config_with_store(&config, store.clone()) + .expect("credential runtime"), + store, + authority: Arc::new( + SandboxSessionJwtAuthority::from_pem( + material.signing_key_pem.as_bytes(), + material.public_key_pem.as_bytes(), + material.kid, + "test-gateway", + Some(Duration::from_mins(15)), + ) + .expect("authority"), + ), + } + } + + /// A gateway replica's pool over this store, with its own controller. + fn replica(&self, replica_id: &str) -> Arc { + let controller = SupervisorHostController::new(); + let mut pool = SupervisorPool::new( + self.store.clone(), + controller, + self.authority.clone(), + self.credentials.clone(), + replica_id.to_string(), + ); + let unique = Arc::get_mut(&mut pool).expect("unique"); + unique.unowned_grace = Duration::ZERO; + unique.candidate_grace = Duration::ZERO; + unique.takeovers_from = tokio::time::Instant::now(); + pool + } + + /// A replica whose member records expire after `ttl` and are + /// refreshed every `refresh`. + fn timed_replica( + &self, + replica_id: &str, + ttl: Duration, + refresh: Duration, + ) -> Arc { + let mut pool = self.replica(replica_id); + let unique = Arc::get_mut(&mut pool).expect("unique"); + unique.member_ttl = ttl; + unique.member_refresh = refresh; + pool + } + + /// A replica that just started, within its takeover grace. + fn starting_replica(&self, replica_id: &str) -> Arc { + let controller = SupervisorHostController::new(); + SupervisorPool::new( + self.store.clone(), + controller, + self.authority.clone(), + self.credentials.clone(), + replica_id.to_string(), + ) + } + + async fn sandbox(&self, sandbox_id: &str, generation: &str, phase: SandboxPhase) { + let identity = PersistedSandboxIdentity { + runtime_generation: SandboxGenerationId::parse(generation).unwrap(), + auth_epoch: CredentialEpoch::new(1).unwrap(), + gateway_token_id: Uuid::new_v4(), + refresh_replay: None, + }; + let mut metadata = ObjectMeta { + id: sandbox_id.to_string(), + name: sandbox_id.to_string(), + workspace: "default".to_string(), + ..Default::default() + }; + identity.write(&mut metadata.annotations); + self.store + .put_message(&Sandbox { + metadata: Some(metadata), + status: Some(SandboxStatus { + phase: phase as i32, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect("persist sandbox"); + } + + async fn identity(&self, sandbox_id: &str) -> PersistedSandboxIdentity { + let sandbox = self + .store + .get_message::(sandbox_id) + .await + .unwrap() + .unwrap(); + PersistedSandboxIdentity::read(&sandbox.metadata.unwrap().annotations).unwrap() + } + + async fn set_phase(&self, sandbox_id: &str, phase: SandboxPhase) { + self.store + .update_message_cas::(sandbox_id, 0, |sandbox| { + sandbox.set_phase(phase as i32); + }) + .await + .unwrap(); + } + + /// The persisted credential epoch of `sandbox_id`. + async fn auth_epoch(&self, sandbox_id: &str) -> u64 { + let sandbox = self + .store + .get_message::(sandbox_id) + .await + .unwrap() + .unwrap(); + PersistedSandboxIdentity::read(&sandbox.metadata.unwrap().annotations) + .unwrap() + .auth_epoch + .get() + } + + /// A launch of `route` that read a free slot on `m1`, the pool's only + /// member, and found it taken by another sandbox once it tried to + /// admit: it is refused. + async fn refused_launch(&self, pool: &Arc, route: &str, sandbox: &str) { + self.member("m1", "p1", false).await; + self.fill("m1", "p1", MAX_ASSIGNED_INSTANCES - 1).await; + self.sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + let release = Arc::new(tokio::sync::Barrier::new(2)); + let reached = pool.pause_next_transition_reached(release.clone()); + let launching = { + let pool = pool.clone(); + let request = launch(route, sandbox, "generation-1"); + tokio::spawn(async move { pool.place_launch(request).await }) + }; + reached.await.unwrap(); + self.sandbox("sandbox-last", "generation-1", SandboxPhase::Ready) + .await; + self.publish("rt-last-slot-0123", "sandbox-last", "generation-1") + .await; + self.own("rt-last-slot-0123", "sandbox-last", "m1", "p1") + .await; + release.wait().await; + assert!(launching.await.unwrap().is_err(), "{route} is refused"); + } + + async fn member(&self, member_id: &str, process_id: &str, leaving: bool) { + let pool = self.replica("membership-writer"); + pool.write_member(&ConnectedMember { + member_id: member_id.to_string(), + member_uid: format!("uid-{member_id}"), + process_id: process_id.to_string(), + runtime_address: "10.0.0.1:5500".to_string(), + leaving, + connected_at_ms: now_ms(), + }) + .await + .unwrap(); + } + + async fn remove_owner_record(&self, route: &str) { + self.store + .delete(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await + .unwrap(); + } + + async fn remove_member(&self, member_id: &str) { + self.store + .delete(MEMBER_OBJECT_TYPE, &member_object_id(member_id)) + .await + .unwrap(); + } + + async fn owner(&self, route: &str) -> Option { + self.store + .get(OWNER_OBJECT_TYPE, &owner_object_id(route)) + .await + .unwrap() + .and_then(|record| decode_owner(&record)) + .map(|record| record.payload) + } + + async fn candidates(&self) -> usize { + self.store + .list_by_type(MATERIAL_CANDIDATE_OBJECT_TYPE, 100, 0) + .await + .unwrap() + .len() + } + + /// A published instance record written directly, as another replica + /// would. + async fn publish(&self, route: &str, sandbox_id: &str, generation: &str) { + let payload = AssignmentPayload { + route: route.to_string(), + sandbox_id: sandbox_id.to_string(), + generation: generation.to_string(), + workspace: "default".to_string(), + phase: RecordPhase::Published, + writer: Uuid::new_v4().to_string(), + material_handle: b"handle".to_vec(), + }; + self.store + .put_if( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id(route, &payload.writer), + route, + "", + &serde_json::to_vec(&payload).unwrap(), + None, + WriteCondition::MustCreate, + ) + .await + .unwrap(); + } + + /// The launch-term owner of `route`, written directly. + async fn own(&self, route: &str, sandbox_id: &str, member: &str, process: &str) { + self.replica("owner-writer") + .write_owner( + &OwnerPayload { + route: route.to_string(), + sandbox_id: sandbox_id.to_string(), + generation: "generation-1".to_string(), + member_id: member.to_string(), + process_id: process.to_string(), + auth_epoch: 1, + supervisor_instance_id: None, + gateway_token_id: None, + phase: OwnerPhase::Granted, + previous_member_id: None, + }, + WriteCondition::MustCreate, + ) + .await + .unwrap(); + } + + /// `count` running instances owned by `member`'s `process`, each + /// with its sandbox and published record. + async fn fill(&self, member: &str, process: &str, count: usize) -> Vec { + let mut routes = Vec::with_capacity(count); + for index in 0..count { + let route = format!("rt-{member}-{index:05}"); + let sandbox = format!("sandbox-{member}-{index}"); + self.sandbox(&sandbox, "generation-1", SandboxPhase::Ready) + .await; + self.publish(&route, &sandbox, "generation-1").await; + self.own(&route, &sandbox, member, process).await; + routes.push(route); + } + routes + } + + /// Write `payload` as its term's report row, as any writer could. + async fn write_report(&self, payload: ReportPayload) { + self.store + .put_if( + REPORT_OBJECT_TYPE, + &report_object_id(&payload.route, payload.auth_epoch), + &report_object_name(&payload.route, payload.auth_epoch), + "", + &serde_json::to_vec(&payload).unwrap(), + None, + WriteCondition::Unconditional, + ) + .await + .unwrap(); + } + + /// Every report row of `route`, by term. + async fn report_rows(&self, route: &str) -> BTreeMap { + self.replica("report-reader") + .report_rows() + .await + .unwrap() + .into_values() + .map(|(payload, _)| (payload.auth_epoch, payload)) + .filter(|(_, payload)| payload.route == route) + .collect() + } + + /// Routes `member` is charged for. + async fn charged(&self, member: &str) -> BTreeSet { + self.replica("capacity-reader") + .capacity(member) + .await + .unwrap() + .payload + .routes + } + + async fn withdraw_record(&self, route: &str) { + let record = self.instance(route).await.unwrap(); + self.store + .delete( + ASSIGNMENT_OBJECT_TYPE, + &assignment_object_id(route, &record.payload.writer), + ) + .await + .unwrap(); + } + + async fn instance(&self, route: &str) -> Option { + self.store + .get_by_name(ASSIGNMENT_OBJECT_TYPE, "", route) + .await + .unwrap() + .and_then(|record| decode_assignment(&record)) + } + } + + fn assignment(route: &str, sandbox_id: &str, generation: &str) -> SandboxInstanceAssignment { + SandboxInstanceAssignment { + route: route.to_string(), + sandbox_id: sandbox_id.to_string(), + generation: generation.to_string(), + launch_material: format!("{{\"launch\":\"{route}\"}}").into_bytes(), + takeover: None, + } + } + + fn launch(route: &str, sandbox_id: &str, generation: &str) -> PlacementRequest { + PlacementRequest { + assignment: assignment(route, sandbox_id, generation), + auth_epoch: 1, + } + } + + #[tokio::test] + async fn launches_follow_the_ring_and_healthy_owners_never_move() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let pool = fixture.replica("gw-a"); + let mut routes = Vec::new(); + for index in 0..24 { + let (route, sandbox) = (format!("rt-route-{index:04}"), format!("sandbox-{index}")); + fixture + .sandbox(&sandbox, "generation-1", SandboxPhase::Ready) + .await; + pool.place_launch(launch(&route, &sandbox, "generation-1")) + .await + .unwrap(); + routes.push((route, sandbox)); + } + let ring = GatewayRing::new(["m1", "m2"]); + let mut before = BTreeMap::new(); + for (route, sandbox) in &routes { + let owner = fixture.owner(route).await.unwrap(); + assert_eq!(Some(owner.member_id.as_str()), ring.owner_for(sandbox)); + assert_eq!(owner.auth_epoch, 1); + assert!(owner.supervisor_instance_id.is_none()); + before.insert(route.clone(), owner); + } + assert!(before.values().any(|owner| owner.member_id == "m1")); + assert!(before.values().any(|owner| owner.member_id == "m2")); + + // A new member changes where new launches go, not who owns what. + fixture.member("m3", "p3", false).await; + pool.reconcile_once().await.unwrap(); + for (route, owner) in &before { + assert_eq!(fixture.owner(route).await.as_ref(), Some(owner)); + } + let view = pool.controller.pool_view().unwrap(); + assert_eq!(view.owners.len(), routes.len()); + assert_eq!(view.members.len(), 3); + // With no live member a launch is refused before its workload runs. + let empty = Fixture::new().await; + assert!( + empty + .replica("gw-b") + .place_launch(launch("rt-route-none", "sandbox-x", "generation-1")) + .await + .is_err() + ); + } + + #[tokio::test] + async fn a_failed_owner_is_taken_over_once_and_fenced_everywhere() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let replica_a = fixture.replica("gw-a"); + let replica_b = fixture.replica("gw-b"); + replica_a + .place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + let launch_identity = fixture.identity("sandbox-a").await; + + // The owner's process dies; another member is live. Two replicas + // reconcile at once and grant exactly one new term. + fixture.member("m2", "p2", false).await; + fixture.remove_member("m1").await; + let (first, second) = tokio::join!(replica_a.reconcile_once(), replica_b.reconcile_once()); + first.unwrap(); + second.unwrap(); + replica_a.reconcile_once().await.unwrap(); + let owner = fixture.owner("rt-route-a000").await.unwrap(); + assert_eq!(owner.member_id, "m2"); + assert_eq!(owner.process_id, "p2"); + assert_eq!(owner.auth_epoch, 2); + assert_eq!(owner.phase, OwnerPhase::Granted); + assert_eq!(owner.previous_member_id.as_deref(), Some("m1")); + let identity = fixture.identity("sandbox-a").await; + assert_eq!(identity.auth_epoch.get(), 2); + assert_eq!( + Some(identity.gateway_token_id.to_string()), + owner.gateway_token_id + ); + + // Every credential of the replaced term is refused by the gateway, + // and refresh cannot carry it into the new term. + let stale = AuthenticatedSandboxSession { + sandbox_id: SandboxId::parse("sandbox-a").unwrap(), + runtime_generation: launch_identity.runtime_generation.clone(), + auth_epoch: launch_identity.auth_epoch, + token_id: launch_identity.gateway_token_id, + issued_at: 1, + expires_at: 2, + }; + assert!( + crate::auth::sandbox_session::authorize_persisted(&fixture.store, &stale) + .await + .is_err() + ); + let request_hash = + crate::auth::sandbox_session::RefreshRequestHash::from_extension_services(&[]).unwrap(); + assert!( + crate::auth::sandbox_session::authorize_refresh( + &fixture.store, + &stale, + &request_hash, + now_ms() / 1000 + ) + .await + .is_err() + ); + + // The new term's credentials name its process and the new lineage. + let minted = replica_a.mint_takeover(&owner).await.unwrap().unwrap(); + assert_eq!(minted.auth_epoch, 2); + assert_eq!( + Some(minted.supervisor_instance_id.as_str()), + owner.supervisor_instance_id.as_deref() + ); + let gateway = replica_a + .authority + .verify_gateway_token(&minted.gateway_token) + .unwrap(); + assert_eq!(gateway.auth_epoch.get(), 2); + crate::auth::sandbox_session::authorize_persisted(&fixture.store, &gateway) + .await + .expect("new term's gateway credential is current"); + + // A healthy new owner stays put on later passes. + replica_b.reconcile_once().await.unwrap(); + assert_eq!(fixture.owner("rt-route-a000").await.unwrap(), owner); + } + + #[tokio::test] + async fn restarted_and_leaving_members_lose_their_instances() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + for sandbox in ["sandbox-a", "sandbox-b"] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + } + let pool = fixture.replica("gw-a"); + pool.place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + pool.place_launch(launch("rt-route-b000", "sandbox-b", "generation-1")) + .await + .unwrap(); + + // The member's container restarts: its new process owns nothing until + // the instances are taken over, here by the same member. + fixture.member("m1", "p1-restarted", false).await; + pool.reconcile_once().await.unwrap(); + for route in ["rt-route-a000", "rt-route-b000"] { + let owner = fixture.owner(route).await.unwrap(); + assert_eq!( + (owner.member_id.as_str(), owner.process_id.as_str()), + ("m1", "p1-restarted") + ); + assert_eq!(owner.auth_epoch, 2); + } + + // A member that announces it is leaving is drained at once. + fixture.member("m2", "p2", false).await; + fixture.member("m1", "p1-restarted", true).await; + pool.reconcile_once().await.unwrap(); + for route in ["rt-route-a000", "rt-route-b000"] { + let owner = fixture.owner(route).await.unwrap(); + assert_eq!(owner.member_id, "m2"); + assert_eq!(owner.auth_epoch, 3); + } + } + + #[tokio::test] + async fn stop_delete_and_replaced_generations_cancel_takeover() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + let pool = fixture.replica("gw-a"); + for (sandbox, route) in [ + ("sandbox-stopping", "rt-route-s000"), + ("sandbox-deleted", "rt-route-d000"), + ("sandbox-restarted", "rt-route-r000"), + ] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + pool.place_launch(launch(route, sandbox, "generation-1")) + .await + .unwrap(); + } + fixture + .set_phase("sandbox-stopping", SandboxPhase::Stopping) + .await; + fixture + .store + .delete("sandbox", "sandbox-deleted") + .await + .unwrap(); + fixture + .sandbox("sandbox-restarted", "generation-2", SandboxPhase::Ready) + .await; + fixture.member("m2", "p2", false).await; + fixture.remove_member("m1").await; + pool.reconcile_once().await.unwrap(); + for route in ["rt-route-s000", "rt-route-d000", "rt-route-r000"] { + assert!(fixture.owner(route).await.is_none(), "{route}"); + } + assert_eq!( + fixture.identity("sandbox-stopping").await.auth_epoch.get(), + 1 + ); + assert_eq!( + fixture.identity("sandbox-restarted").await.auth_epoch.get(), + 1 + ); + } + + #[tokio::test] + async fn interrupted_takeovers_and_unowned_instances_are_completed() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + fixture + .sandbox("sandbox-b", "generation-1", SandboxPhase::Ready) + .await; + let pool = fixture.replica("gw-a"); + // A replica crashed after recording intent, before advancing the + // sandbox identity. + pool.write_owner( + &OwnerPayload { + route: "rt-route-a000".to_string(), + sandbox_id: "sandbox-a".to_string(), + generation: "generation-1".to_string(), + member_id: "m1".to_string(), + process_id: "p1".to_string(), + auth_epoch: 2, + supervisor_instance_id: Some(Uuid::new_v4().to_string()), + gateway_token_id: Some(Uuid::new_v4().to_string()), + phase: OwnerPhase::Granting, + previous_member_id: Some("m0".to_string()), + }, + WriteCondition::MustCreate, + ) + .await + .unwrap(); + for (route, sandbox) in [ + ("rt-route-a000", "sandbox-a"), + // Recorded by a driver recovery: launched before the pool + // recorded owners. + ("rt-route-b000", "sandbox-b"), + ] { + pool.record_assignment(&assignment(route, sandbox, "generation-1")) + .await + .unwrap(); + } + pool.reconcile_once().await.unwrap(); + pool.reconcile_once().await.unwrap(); + let resumed = fixture.owner("rt-route-a000").await.unwrap(); + assert_eq!(resumed.phase, OwnerPhase::Granted); + assert_eq!(resumed.auth_epoch, 2); + assert_eq!(fixture.identity("sandbox-a").await.auth_epoch.get(), 2); + // An unowned instance is placed by a takeover, which fences + // whichever process may already supervise it. + let placed = fixture.owner("rt-route-b000").await.unwrap(); + assert_eq!(placed.phase, OwnerPhase::Granted); + assert_eq!(placed.auth_epoch, 2); + assert!(placed.supervisor_instance_id.is_some()); + assert_eq!(fixture.identity("sandbox-b").await.auth_epoch.get(), 2); + // Only granted owners are served. + let view = pool.controller.pool_view().unwrap(); + assert_eq!(view.owners.len(), 2); + } + + #[tokio::test] + async fn one_pass_completes_every_pending_intent_behind_a_healthy_route() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let pool = fixture.replica("gw-a"); + for sandbox in ["sandbox-a", "sandbox-b", "sandbox-c", "sandbox-d"] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + } + // The first route in reconciliation order is healthy. + pool.place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + let healthy = fixture.owner("rt-route-a000").await.unwrap(); + let intent = |route: &str, sandbox: &str, member: &str, process: &str| OwnerPayload { + route: route.to_string(), + sandbox_id: sandbox.to_string(), + generation: "generation-1".to_string(), + member_id: member.to_string(), + process_id: process.to_string(), + auth_epoch: 2, + supervisor_instance_id: Some(Uuid::new_v4().to_string()), + gateway_token_id: Some(Uuid::new_v4().to_string()), + phase: OwnerPhase::Granting, + previous_member_id: Some("m0".to_string()), + }; + // A gateway replica crashed after recording intent (b), after also + // advancing the sandbox identity (c), and after recording intent for + // a term another takeover already holds (d). + let b = intent("rt-route-b000", "sandbox-b", "m1", "p1"); + let c = intent("rt-route-c000", "sandbox-c", "m2", "p2"); + let d = intent("rt-route-d000", "sandbox-d", "m2", "p2"); + for pending in [&b, &c, &d] { + pool.write_owner(pending, WriteCondition::MustCreate) + .await + .unwrap(); + } + let advance = |sandbox: &'static str, token: Uuid| { + let store = fixture.store.clone(); + async move { + advance_identity_for_takeover( + &store, + sandbox, + "generation-1", + CredentialEpoch::new(2).unwrap(), + token, + ) + .await + .unwrap() + } + }; + let c_token = Uuid::parse_str(c.gateway_token_id.as_deref().unwrap()).unwrap(); + assert_eq!( + advance("sandbox-c", c_token).await, + TakeoverIdentity::Advanced + ); + assert_eq!( + advance("sandbox-d", Uuid::new_v4()).await, + TakeoverIdentity::Advanced + ); + for (route, sandbox) in [ + ("rt-route-b000", "sandbox-b"), + ("rt-route-c000", "sandbox-c"), + ("rt-route-d000", "sandbox-d"), + ] { + pool.record_assignment(&assignment(route, sandbox, "generation-1")) + .await + .unwrap(); + } + + pool.reconcile_once().await.unwrap(); + assert_eq!(fixture.owner("rt-route-a000").await.unwrap(), healthy); + for (route, sandbox, pending) in [ + ("rt-route-b000", "sandbox-b", &b), + ("rt-route-c000", "sandbox-c", &c), + ] { + let granted = fixture.owner(route).await.unwrap(); + assert_eq!( + granted, + OwnerPayload { + phase: OwnerPhase::Granted, + ..pending.clone() + }, + "{route}" + ); + let identity = fixture.identity(sandbox).await; + assert_eq!(identity.auth_epoch.get(), 2); + assert_eq!( + Some(identity.gateway_token_id.to_string()), + granted.gateway_token_id + ); + } + // The intent that lost its term is cancelled, never granted with + // credentials of a lineage the gateway does not recognize, and the + // instance is placed again at a later term. + assert!(fixture.owner("rt-route-d000").await.is_none()); + let view = pool.controller.pool_view().unwrap(); + assert_eq!(view.owners.len(), 3); + assert!(!view.owners.contains_key("rt-route-d000")); + pool.reconcile_once().await.unwrap(); + let replaced = fixture.owner("rt-route-d000").await.unwrap(); + assert_eq!( + (replaced.phase, replaced.auth_epoch), + (OwnerPhase::Granted, 3) + ); + assert_eq!(fixture.identity("sandbox-d").await.auth_epoch.get(), 3); + } + + #[tokio::test] + async fn records_of_ended_generations_are_removed() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let pool = fixture.replica("gw-a"); + pool.place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + // Withdrawn while still the sandbox's generation: the record and its + // material go at once; the owner record stays for its owner to be + // told to stop it. + pool.controller.unassign("rt-route-a000"); + pool.reconcile_once().await.unwrap(); + assert!(fixture.instance("rt-route-a000").await.is_none()); + assert!(pool.controller.pool_withdrawals().is_empty()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + assert!(fixture.owner("rt-route-a000").await.is_some()); + fixture + .sandbox("sandbox-a", "generation-2", SandboxPhase::Ready) + .await; + pool.reconcile_once().await.unwrap(); + assert!(fixture.owner("rt-route-a000").await.is_none()); + + // A stopped sandbox never runs its generation again, even when its + // withdrawal was lost with a gateway replica. + fixture + .sandbox("sandbox-b", "generation-1", SandboxPhase::Ready) + .await; + pool.place_launch(launch("rt-route-b000", "sandbox-b", "generation-1")) + .await + .unwrap(); + pool.reconcile_once().await.unwrap(); + assert!(fixture.owner("rt-route-b000").await.is_some()); + fixture.set_phase("sandbox-b", SandboxPhase::Stopped).await; + pool.reconcile_once().await.unwrap(); + pool.reconcile_once().await.unwrap(); + assert!(fixture.instance("rt-route-b000").await.is_none()); + assert!(fixture.owner("rt-route-b000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + } + + #[tokio::test] + async fn a_starting_sandbox_is_taken_over_but_not_by_a_replica_that_just_started() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Starting) + .await; + let starting = fixture.starting_replica("gw-new"); + starting + .place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + fixture.member("m2", "p2", false).await; + fixture.remove_member("m1").await; + // Members cut off by a gateway restart get time to reconnect. + starting.reconcile_once().await.unwrap(); + assert_eq!(fixture.owner("rt-route-a000").await.unwrap().auth_epoch, 1); + // A replica past its grace takes the starting instance over. + let settled = fixture.replica("gw-settled"); + settled.reconcile_once().await.unwrap(); + let owner = fixture.owner("rt-route-a000").await.unwrap(); + assert_eq!((owner.member_id.as_str(), owner.auth_epoch), ("m2", 2)); + assert_eq!(owner.phase, OwnerPhase::Granted); + } + + #[tokio::test] + async fn every_replica_derives_the_desired_set_and_material_from_records() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let launching = fixture.replica("gw-launch"); + let other = fixture.replica("gw-other"); + let launched = assignment("rt-route-a000", "sandbox-a", "generation-1"); + launching + .place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + // The record keeps only the material's credential handle. + let record = fixture.instance("rt-route-a000").await.unwrap(); + let stored = fixture + .store + .get_by_name(ASSIGNMENT_OBJECT_TYPE, "", "rt-route-a000") + .await + .unwrap() + .unwrap(); + assert!(!String::from_utf8_lossy(&stored.payload).contains("{\"launch\"")); + assert_eq!(record.payload.workspace, "default"); + + // A replica that never saw the launch, and runs no compute watch, + // derives the instance, its owner, and its material from the records. + other.reconcile_once().await.unwrap(); + let view = other.controller.pool_view().unwrap(); + assert!(view.desired.contains("rt-route-a000")); + assert_eq!(view.owners["rt-route-a000"].member_id, "m1"); + assert_eq!( + other.load_material(&record.payload).await.unwrap(), + launched.launch_material + ); + assert!( + other + .recorded() + .await + .unwrap() + .contains_key("rt-route-a000") + ); + + // Recording the generation again keeps the first record and material; + // another generation on its route is refused. + other.record_assignment(&launched).await.unwrap(); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert!( + other + .record_assignment(&assignment("rt-route-a000", "sandbox-a", "generation-2")) + .await + .is_err() + ); + + // A withdrawal through either replica removes the record and its + // material, and leaves the desired set of both. + other.controller.unassign("rt-route-a000"); + other.reconcile_once().await.unwrap(); + assert!(fixture.instance("rt-route-a000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + launching.reconcile_once().await.unwrap(); + assert!( + !launching + .controller + .pool_view() + .unwrap() + .desired + .contains("rt-route-a000") + ); + } + + #[tokio::test] + async fn racing_records_of_one_instance_keep_one_material() { + let fixture = Fixture::new().await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let (first, second) = (fixture.replica("gw-a"), fixture.replica("gw-b")); + let instance = assignment("rt-route-a000", "sandbox-a", "generation-1"); + let results: [_; 2] = tokio::join!( + first.record_assignment(&instance), + second.record_assignment(&instance) + ) + .into(); + // One recording reserves the route and publishes; the other finds the + // reservation or the record, and stores no material. + let record = fixture.instance("rt-route-a000").await.unwrap(); + assert_eq!(record.payload.phase, RecordPhase::Published); + for result in results { + match result { + Ok(recorded) => assert_eq!( + recorded.payload.material_handle, + record.payload.material_handle + ), + Err(error) => assert!(error.contains("being recorded"), "{error}"), + } + } + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert_eq!(fixture.candidates().await, 0); + assert_eq!( + first.load_material(&record.payload).await.unwrap(), + instance.launch_material + ); + } + + #[tokio::test] + async fn a_fenced_recording_whose_cleanup_fails_keeps_its_material_owned() { + let fixture = Fixture::new().await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let pool = fixture.replica("gw-a"); + let instance = assignment("rt-route-a000", "sandbox-a", "generation-1"); + let (staged, publish) = pool.pause_next_recording_after_staging(); + let recording = { + let (pool, instance) = (pool.clone(), instance.clone()); + tokio::spawn(async move { pool.record_assignment(&instance).await }) + }; + staged.await.unwrap(); + // The instance is withdrawn while it is being recorded: removing its + // reservation fences the recording. + pool.controller.unassign("rt-route-a000"); + pool.remove_withdrawn().await; + assert!(fixture.instance("rt-route-a000").await.is_none()); + // The recording cannot publish, and its cleanup fails; its material + // stays owned by its candidate. + fixture.credentials.fail_next_delete(); + publish.send(()).unwrap(); + let error = recording.await.unwrap().unwrap_err(); + assert!(error.contains("withdrawn before publication"), "{error}"); + assert!(fixture.instance("rt-route-a000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert_eq!(fixture.candidates().await, 1); + // Collection deletes it. + pool.remove_stale_candidates().await; + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + assert_eq!(fixture.candidates().await, 0); + } + + #[tokio::test] + async fn collection_never_deletes_the_material_of_a_recording_in_progress() { + let fixture = Fixture::new().await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + // Candidates are collectable at once (grace zero). + let pool = fixture.replica("gw-a"); + let instance = assignment("rt-route-a000", "sandbox-a", "generation-1"); + let (staged, publish) = pool.pause_next_recording_after_staging(); + let recording = { + let (pool, instance) = (pool.clone(), instance.clone()); + tokio::spawn(async move { pool.record_assignment(&instance).await }) + }; + staged.await.unwrap(); + // Staged and unpublished: material stored, candidate written, the + // route only reserved. Collection runs now, alone and as part of a + // whole pass, and must not delete the material. + assert_eq!(fixture.candidates().await, 1); + assert_eq!( + fixture + .instance("rt-route-a000") + .await + .unwrap() + .payload + .phase, + RecordPhase::Reserved + ); + pool.remove_stale_candidates().await; + pool.reconcile_once().await.unwrap(); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert_eq!(fixture.candidates().await, 1); + // A driver recovery begins listing after the reservation but before + // the publication, which its listing may then miss. + let reserved_at = fixture + .instance("rt-route-a000") + .await + .unwrap() + .created_at_ms; + while now_ms() <= reserved_at { + tokio::time::sleep(Duration::from_millis(1)).await; + } + let listing_began = now_ms(); + while now_ms() <= listing_began { + tokio::time::sleep(Duration::from_millis(1)).await; + } + // The recording publishes, and its record's material resolves. + publish.send(()).unwrap(); + let record = recording.await.unwrap().unwrap(); + // Recovery sees when it was published, not reserved, so it keeps it. + assert!(pool.recorded().await.unwrap()["rt-route-a000"] > listing_began); + pool.remove_stale_candidates().await; + assert_eq!( + pool.load_material(&record.payload).await.unwrap(), + instance.launch_material + ); + assert_eq!(fixture.candidates().await, 0); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + } + + #[tokio::test] + async fn a_stalled_recording_is_fenced_and_never_publishes_its_material() { + let fixture = Fixture::new().await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let mut pool = fixture.replica("gw-a"); + Arc::get_mut(&mut pool).expect("unique").reservation_grace = Duration::ZERO; + let instance = assignment("rt-route-a000", "sandbox-a", "generation-1"); + let (staged, publish) = pool.pause_next_recording_after_staging(); + let recording = { + let (pool, instance) = (pool.clone(), instance.clone()); + tokio::spawn(async move { pool.record_assignment(&instance).await }) + }; + staged.await.unwrap(); + // The recording outlives its reservation's grace: a pass fences it and + // collects its material. + pool.reconcile_once().await.unwrap(); + assert!(fixture.instance("rt-route-a000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + assert_eq!(fixture.candidates().await, 0); + // When it resumes, its publication fails: no record ever names the + // deleted material. + publish.send(()).unwrap(); + assert!(recording.await.unwrap().is_err()); + assert!(fixture.instance("rt-route-a000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + // The route can be recorded again. + let record = pool.record_assignment(&instance).await.unwrap(); + assert_eq!( + pool.load_material(&record.payload).await.unwrap(), + instance.launch_material + ); + } + + #[tokio::test] + async fn a_fenced_recording_cannot_publish_over_the_next_reservation() { + let fixture = Fixture::new().await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let instance = assignment("rt-route-a000", "sandbox-a", "generation-1"); + let record_paused = |pool: Arc| { + let (staged, publish) = pool.pause_next_recording_after_staging(); + let instance = instance.clone(); + let recording = tokio::spawn(async move { pool.record_assignment(&instance).await }); + (staged, publish, recording) + }; + // The first recording stalls after staging and is fenced; its + // material is collected. + let (staged, publish_first, first) = record_paused(fixture.replica("gw-a")); + staged.await.unwrap(); + let mut fencing = fixture.replica("gw-fence"); + Arc::get_mut(&mut fencing) + .expect("unique") + .reservation_grace = Duration::ZERO; + fencing.reconcile_once().await.unwrap(); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + // A second recording reserves the route anew, at the same store + // version the first one reserved at, and stalls after staging. + let second_pool = fixture.replica("gw-b"); + let (staged, publish_second, second) = record_paused(second_pool.clone()); + staged.await.unwrap(); + // The first recording resumes: it cannot publish over the second's + // reservation. + publish_first.send(()).unwrap(); + assert!(first.await.unwrap().is_err()); + assert_eq!( + fixture + .instance("rt-route-a000") + .await + .unwrap() + .payload + .phase, + RecordPhase::Reserved + ); + // The second publishes, and its material resolves. + publish_second.send(()).unwrap(); + let record = second.await.unwrap().unwrap(); + assert_eq!( + second_pool.load_material(&record.payload).await.unwrap(), + instance.launch_material + ); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + } + + #[tokio::test] + async fn a_cancelled_launch_still_completes_its_durable_records() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let placement = PoolPlacement(fixture.replica("gw-a")); + let (entered, release) = fixture.credentials.gate_next_store(); + let launching = tokio::spawn(placement.place_launch(launch( + "rt-route-a000", + "sandbox-a", + "generation-1", + ))); + entered.await.unwrap(); + // The launch's caller goes away while its material is being stored. + launching.abort(); + release.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while fixture.owner("rt-route-a000").await.is_none() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the detached launch completes"); + assert!(fixture.instance("rt-route-a000").await.is_some()); + assert_eq!(fixture.candidates().await, 0); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + } + + #[tokio::test] + async fn recorded_reports_of_the_current_owner_term_reach_every_replica() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-a", "generation-1", SandboxPhase::Ready) + .await; + let pool = fixture.replica("gw-a"); + pool.place_launch(launch("rt-route-a000", "sandbox-a", "generation-1")) + .await + .unwrap(); + let report = |process: &str, ready| ReportPayload { + route: "rt-route-a000".to_string(), + member_id: "m1".to_string(), + process_id: process.to_string(), + auth_epoch: 1, + sandbox_id: "sandbox-a".to_string(), + phase: 2, + ready, + exit_code: 0, + connected_at_ms: 0, + }; + let record = |payload: ReportPayload| fixture.write_report(payload); + // The replica serving the owner recorded that it is ready; this + // replica, which does not serve it, reads the verdict. + record(report("p1", true)).await; + pool.reconcile_once().await.unwrap(); + assert!(matches!( + pool.controller.observed("rt-route-a000"), + openshell_supervisor_control::ObservedInstance::Reported(status) if status.ready + )); + // A report of another process is not the current owner's. + record(report("p1-restarted", true)).await; + pool.reconcile_once().await.unwrap(); + assert_eq!( + pool.controller.observed("rt-route-a000"), + openshell_supervisor_control::ObservedInstance::NoHost + ); + } + + /// One more instance than a listing page, the live one listed last and + /// owned by `m1`; returns the routes and a replica that places nothing, + /// whose first pass has loaded the view. + async fn two_page_pool(fixture: &Fixture) -> (Vec, Arc) { + fixture.member("m1", "p1", false).await; + let count = LIST_LIMIT as usize + 1; + let routes = (0..count) + .map(|index| format!("rt-route-{index:05}")) + .collect::>(); + for (index, route) in routes.iter().enumerate() { + let sandbox = format!("sandbox-{index}"); + fixture + .sandbox(&sandbox, "generation-1", SandboxPhase::Ready) + .await; + fixture.publish(route, &sandbox, "generation-1").await; + } + fixture + .own( + &routes[count - 1], + &format!("sandbox-{}", count - 1), + "m1", + "p1", + ) + .await; + let pool = fixture.starting_replica("gw-a"); + pool.reconcile_once().await.unwrap(); + assert!( + pool.controller + .pool_view() + .unwrap() + .desired + .contains(&routes[count - 1]) + ); + (routes, pool) + } + + #[tokio::test] + async fn a_row_removed_between_listing_pages_shifts_nothing_out() { + let fixture = Fixture::new().await; + let (routes, pool) = two_page_pool(&fixture).await; + let live = routes.last().unwrap(); + // Another replica withdraws an instance page one already returned. + let (paused, resume) = pool.pause_listing_after_first_page(ASSIGNMENT_OBJECT_TYPE); + let listing = { + let pool = pool.clone(); + tokio::spawn(async move { pool.list_all(ASSIGNMENT_OBJECT_TYPE).await }) + }; + paused.await.unwrap(); + fixture.withdraw_record(&routes[0]).await; + resume.send(()).unwrap(); + // Page two starts after page one's last row instead of skipping one. + let listed = listing.await.unwrap().unwrap(); + assert_eq!(listed.len(), routes.len()); + assert!(listed.iter().any(|record| record.name == *live)); + } + + #[tokio::test] + async fn a_row_removed_mid_pass_never_stops_a_live_route() { + let fixture = Fixture::new().await; + let (routes, pool) = two_page_pool(&fixture).await; + let live = routes.last().unwrap(); + let (paused, resume) = pool.pause_listing_after_first_page(ASSIGNMENT_OBJECT_TYPE); + let pass = { + let pool = pool.clone(); + tokio::spawn(async move { pool.reconcile_once().await }) + }; + paused.await.unwrap(); + fixture.withdraw_record(&routes[0]).await; + resume.send(()).unwrap(); + pass.await.unwrap().unwrap(); + // The live route stays desired and owned, so its owner is never told + // to stop it; the withdrawn one is no longer desired. + let view = pool.controller.pool_view().unwrap(); + assert!(view.desired.contains(live)); + assert_eq!(view.owners[live.as_str()].member_id, "m1"); + assert!(!pool.controller.stopping_routes().contains(live)); + pool.reconcile_once().await.unwrap(); + assert!( + !pool + .controller + .pool_view() + .unwrap() + .desired + .contains(&routes[0]) + ); + } + + #[tokio::test] + async fn a_route_a_listing_misses_is_confirmed_before_it_is_stopped() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + for (route, sandbox) in [ + ("rt-route-live", "sandbox-a"), + ("rt-route-gone", "sandbox-b"), + ] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + fixture.publish(route, sandbox, "generation-1").await; + fixture.own(route, sandbox, "m1", "p1").await; + } + let pool = fixture.starting_replica("gw-a"); + pool.reconcile_once().await.unwrap(); + // A listing misses a live route's instance and owner rows; a real + // withdrawal removes the other's instance record. + pool.omit_from_listings("rt-route-live"); + fixture.withdraw_record("rt-route-gone").await; + pool.reconcile_once().await.unwrap(); + let view = pool.controller.pool_view().unwrap(); + assert!(view.desired.contains("rt-route-live")); + assert!(view.owners.contains_key("rt-route-live")); + assert!(!pool.controller.stopping_routes().contains("rt-route-live")); + // The confirmed withdrawal still stops its instance. + assert!(!view.desired.contains("rt-route-gone")); + assert!(pool.controller.stopping_routes().contains("rt-route-gone")); + } + + #[tokio::test] + async fn the_last_slot_of_a_member_is_admitted_once_across_replicas() { + the_last_slot_of_a_member_is_admitted_once_across_replicas_on(&Fixture::new().await).await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_the_last_slot_of_a_member_is_admitted_once_across_replicas() { + let schema = TestSchema::create("supervisor_pool").await; + the_last_slot_of_a_member_is_admitted_once_across_replicas_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn the_last_slot_of_a_member_is_admitted_once_across_replicas_on(fixture: &Fixture) { + fixture.member("m1", "p1", false).await; + fixture.fill("m1", "p1", MAX_ASSIGNED_INSTANCES - 1).await; + for sandbox in ["sandbox-x", "sandbox-y"] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + } + // Both launches read 255 charged routes before either writes. + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + let (first, second) = (fixture.replica("gw-a"), fixture.replica("gw-b")); + first.pause_next_transition(barrier.clone()); + second.pause_next_transition(barrier); + let (x, y) = tokio::join!( + first.place_launch(launch("rt-route-x000", "sandbox-x", "generation-1")), + second.place_launch(launch("rt-route-y000", "sandbox-y", "generation-1")), + ); + assert_eq!( + usize::from(x.is_ok()) + usize::from(y.is_ok()), + 1, + "{x:?} {y:?}" + ); + assert_eq!(fixture.charged("m1").await.len(), MAX_ASSIGNED_INSTANCES); + let owned = [ + fixture.owner("rt-route-x000").await, + fixture.owner("rt-route-y000").await, + ]; + assert_eq!(owned.iter().flatten().count(), 1); + } + + #[tokio::test] + async fn a_takeover_and_a_launch_compete_for_the_last_slot() { + a_takeover_and_a_launch_compete_for_the_last_slot_on(&Fixture::new().await).await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_takeover_and_a_launch_compete_for_the_last_slot() { + let schema = TestSchema::create("supervisor_pool").await; + a_takeover_and_a_launch_compete_for_the_last_slot_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_takeover_and_a_launch_compete_for_the_last_slot_on(fixture: &Fixture) { + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let orphaned = fixture.fill("m1", "p1", 1).await.remove(0); + fixture.fill("m2", "p2", MAX_ASSIGNED_INSTANCES - 1).await; + fixture + .sandbox("sandbox-new", "generation-1", SandboxPhase::Ready) + .await; + // m1 fails; its route and a new launch both need m2's last slot. + fixture.remove_member("m1").await; + let (taking, launching) = (fixture.replica("gw-a"), fixture.replica("gw-b")); + let members = taking.live_members().await.unwrap(); + let charges = charges_of(&taking.capacities().await.unwrap()); + let record = taking.reload_owner(&orphaned).await; + let ring = ring_of(&members); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + taking.pause_next_transition(barrier.clone()); + launching.pause_next_transition(barrier); + let (taken, launched) = tokio::join!( + taking.take_over(&record, &ring, &members, &charges), + launching.place_launch(launch("rt-route-new0", "sandbox-new", "generation-1")), + ); + taken.unwrap(); + let moved = fixture.owner(&orphaned).await.unwrap().member_id == "m2"; + assert_ne!(moved, launched.is_ok(), "exactly one gets the slot"); + assert_eq!(fixture.charged("m2").await.len(), MAX_ASSIGNED_INSTANCES); + // A takeover that lost the slot leaves the route charged to m1. + assert_eq!(fixture.charged("m1").await.contains(&orphaned), !moved); + } + + #[tokio::test] + async fn a_stalled_admission_cannot_take_a_slot_consumed_meanwhile() { + a_stalled_admission_cannot_take_a_slot_consumed_meanwhile_on(&Fixture::new().await).await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_stalled_admission_cannot_take_a_slot_consumed_meanwhile() { + let schema = TestSchema::create("supervisor_pool").await; + a_stalled_admission_cannot_take_a_slot_consumed_meanwhile_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_stalled_admission_cannot_take_a_slot_consumed_meanwhile_on(fixture: &Fixture) { + fixture.member("m1", "p1", false).await; + fixture.fill("m1", "p1", MAX_ASSIGNED_INSTANCES - 1).await; + for sandbox in ["sandbox-x", "sandbox-y"] { + fixture + .sandbox(sandbox, "generation-1", SandboxPhase::Ready) + .await; + } + // One launch stalls after reading 255 charged routes; another + // admits the last slot meanwhile. + let stalled = fixture.replica("gw-a"); + let release = Arc::new(tokio::sync::Barrier::new(2)); + let reached = stalled.pause_next_transition_reached(release.clone()); + let stalled_launch = { + let stalled = stalled.clone(); + tokio::spawn(async move { + stalled + .place_launch(launch("rt-route-x000", "sandbox-x", "generation-1")) + .await + }) + }; + reached.await.unwrap(); + fixture + .replica("gw-b") + .place_launch(launch("rt-route-y000", "sandbox-y", "generation-1")) + .await + .unwrap(); + release.wait().await; + assert!(stalled_launch.await.unwrap().is_err()); + assert!(fixture.owner("rt-route-x000").await.is_none()); + assert_eq!(fixture.charged("m1").await.len(), MAX_ASSIGNED_INSTANCES); + // The refused launch left nothing: no record, no material (only the + // admitted launch's), no candidate, nothing cached. + assert!(fixture.instance("rt-route-x000").await.is_none()); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert_eq!(fixture.candidates().await, 0); + assert!(!lock(&stalled.materials).contains_key("rt-route-x000")); + // Room appears later; no pass places, charges, or grants it. + fixture.member("m2", "p2", false).await; + let pass = fixture.replica("gw-c"); + pass.reconcile_once().await.unwrap(); + pass.reconcile_once().await.unwrap(); + assert!(fixture.owner("rt-route-x000").await.is_none()); + assert!(fixture.charged("m2").await.is_empty()); + assert_eq!(fixture.auth_epoch("sandbox-x").await, 1); + assert!(fixture.instance("rt-route-x000").await.is_none()); + } + + #[tokio::test] + async fn a_refused_launch_whose_cleanup_fails_is_fenced_and_never_placed() { + a_refused_launch_whose_cleanup_fails_is_fenced_and_never_placed_on(&Fixture::new().await) + .await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_refused_launch_whose_cleanup_fails_is_fenced_and_never_placed() { + let schema = TestSchema::create("supervisor_pool").await; + a_refused_launch_whose_cleanup_fails_is_fenced_and_never_placed_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_refused_launch_whose_cleanup_fails_is_fenced_and_never_placed_on(fixture: &Fixture) { + let pool = fixture.replica("gw-a"); + *lock(&pool.hooks.fail_give_up) = Some(()); + fixture + .refused_launch(&pool, "rt-route-x000", "sandbox-x") + .await; + // The record stays, but only as an admitting record with its + // candidate: not desired, so room appearing places nothing. + let kept = fixture.instance("rt-route-x000").await.unwrap(); + assert_eq!(kept.payload.phase, RecordPhase::Admitting); + assert_eq!(fixture.candidates().await, 1); + fixture.member("m2", "p2", false).await; + let pass = fixture.replica("gw-b"); + pass.reconcile_once().await.unwrap(); + assert!(fixture.owner("rt-route-x000").await.is_none()); + assert!(fixture.charged("m2").await.is_empty()); + assert_eq!(fixture.auth_epoch("sandbox-x").await, 1); + // Once its grace passes, any replica fences it and collects its + // material. + let mut fencing = fixture.replica("gw-c"); + Arc::get_mut(&mut fencing).unwrap().reservation_grace = Duration::ZERO; + fencing.reconcile_once().await.unwrap(); + fencing.reconcile_once().await.unwrap(); + assert!(fixture.instance("rt-route-x000").await.is_none()); + assert_eq!(fixture.candidates().await, 0); + assert_eq!(fixture.credentials.stored_credential_count(), Some(0)); + assert!(fixture.owner("rt-route-x000").await.is_none()); + assert!(fixture.charged("m2").await.is_empty()); + } + + #[tokio::test] + async fn a_retried_launch_admits_the_record_a_refused_attempt_kept() { + a_retried_launch_admits_the_record_a_refused_attempt_kept_on(&Fixture::new().await).await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_retried_launch_admits_the_record_a_refused_attempt_kept() { + let schema = TestSchema::create("supervisor_pool").await; + a_retried_launch_admits_the_record_a_refused_attempt_kept_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_retried_launch_admits_the_record_a_refused_attempt_kept_on(fixture: &Fixture) { + let pool = fixture.replica("gw-a"); + *lock(&pool.hooks.fail_give_up) = Some(()); + fixture + .refused_launch(&pool, "rt-route-x000", "sandbox-x") + .await; + // The launch is retried once there is room, on another replica. + fixture.member("m2", "p2", false).await; + fixture + .replica("gw-b") + .place_launch(launch("rt-route-x000", "sandbox-x", "generation-1")) + .await + .unwrap(); + let owner = fixture.owner("rt-route-x000").await.unwrap(); + assert_eq!((owner.member_id.as_str(), owner.auth_epoch), ("m2", 1)); + let record = fixture.instance("rt-route-x000").await.unwrap(); + assert_eq!(record.payload.phase, RecordPhase::Published); + assert!(fixture.charged("m2").await.contains("rt-route-x000")); + // One material, now the record's; its candidate is released. + let mut collecting = fixture.replica("gw-c"); + Arc::get_mut(&mut collecting).unwrap().reservation_grace = Duration::ZERO; + collecting.reconcile_once().await.unwrap(); + assert_eq!(fixture.candidates().await, 0); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert!(fixture.instance("rt-route-x000").await.is_some()); + assert!(fixture.owner("rt-route-x000").await.is_some()); + } + + #[tokio::test] + async fn a_failed_admission_keeps_the_record_for_a_retry_with_its_first_material() { + a_failed_admission_keeps_the_record_for_a_retry_with_its_first_material_on( + &Fixture::new().await, + ) + .await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_failed_admission_keeps_the_record_for_a_retry_with_its_first_material() { + let schema = TestSchema::create("supervisor_pool").await; + a_failed_admission_keeps_the_record_for_a_retry_with_its_first_material_on( + &Fixture::with_store(schema.connect_store().await), + ) + .await; + schema.drop_schema().await; + } + + async fn a_failed_admission_keeps_the_record_for_a_retry_with_its_first_material_on( + fixture: &Fixture, + ) { + fixture.member("m1", "p1", false).await; + fixture + .sandbox("sandbox-x", "generation-1", SandboxPhase::Ready) + .await; + let first = launch("rt-route-x000", "sandbox-x", "generation-1"); + let failing = fixture.replica("gw-a"); + *lock(&failing.hooks.fail_admission) = Some(()); + // The admission write fails without a known outcome: the launch + // fails and keeps its admitting record, which nothing places. + assert!(failing.place_launch(first.clone()).await.is_err()); + let kept = fixture.instance("rt-route-x000").await.unwrap(); + assert_eq!(kept.payload.phase, RecordPhase::Admitting); + assert!(fixture.owner("rt-route-x000").await.is_none()); + fixture.replica("gw-b").reconcile_once().await.unwrap(); + assert!(fixture.owner("rt-route-x000").await.is_none()); + assert!(fixture.charged("m1").await.is_empty()); + // A retry with other launch material admits the kept record; the + // first recording's material stays the instance's everywhere. + let mut retry = first.clone(); + retry.assignment.launch_material = b"{\"launch\":\"retried\"}".to_vec(); + let retrying = fixture.replica("gw-c"); + retrying.place_launch(retry).await.unwrap(); + assert!(!lock(&retrying.materials).contains_key("rt-route-x000")); + let record = fixture.instance("rt-route-x000").await.unwrap(); + assert_eq!(record.payload.phase, RecordPhase::Published); + assert_eq!( + retrying.load_material(&record.payload).await.unwrap(), + first.assignment.launch_material + ); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert_eq!( + fixture.owner("rt-route-x000").await.unwrap().member_id, + "m1" + ); + } + + #[tokio::test] + async fn a_refused_launch_never_removes_a_record_another_launch_admitted() { + a_refused_launch_never_removes_a_record_another_launch_admitted_on(&Fixture::new().await) + .await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_refused_launch_never_removes_a_record_another_launch_admitted() { + let schema = TestSchema::create("supervisor_pool").await; + a_refused_launch_never_removes_a_record_another_launch_admitted_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_refused_launch_never_removes_a_record_another_launch_admitted_on(fixture: &Fixture) { + // The refused launch pauses before removing the record it created; + // meanwhile a retry of the same launch admits that record elsewhere. + let pool = fixture.replica("gw-a"); + let (paused, resume) = { + let (paused, paused_rx) = tokio::sync::oneshot::channel(); + let (resume, resume_rx) = tokio::sync::oneshot::channel(); + *lock(&pool.hooks.give_up_gate) = Some((paused, resume_rx)); + (paused_rx, resume) + }; + fixture.member("m1", "p1", false).await; + fixture.fill("m1", "p1", MAX_ASSIGNED_INSTANCES - 1).await; + fixture + .sandbox("sandbox-x", "generation-1", SandboxPhase::Ready) + .await; + let release = Arc::new(tokio::sync::Barrier::new(2)); + let reached = pool.pause_next_transition_reached(release.clone()); + let launching = { + let pool = pool.clone(); + tokio::spawn(async move { + pool.place_launch(launch("rt-route-x000", "sandbox-x", "generation-1")) + .await + }) + }; + reached.await.unwrap(); + fixture + .sandbox("sandbox-last", "generation-1", SandboxPhase::Ready) + .await; + fixture + .publish("rt-last-slot-0123", "sandbox-last", "generation-1") + .await; + fixture + .own("rt-last-slot-0123", "sandbox-last", "m1", "p1") + .await; + release.wait().await; + tokio::time::timeout(Duration::from_secs(30), paused) + .await + .expect("the refused launch reaches its removal") + .unwrap(); + fixture.member("m2", "p2", false).await; + fixture + .replica("gw-b") + .place_launch(launch("rt-route-x000", "sandbox-x", "generation-1")) + .await + .unwrap(); + let _ = resume.send(()); + // The first launch finds the record admitted by the retry and + // succeeds with it; nothing of the instance is removed. + launching.await.unwrap().unwrap(); + let record = fixture.instance("rt-route-x000").await.unwrap(); + assert_eq!(record.payload.phase, RecordPhase::Published); + assert_eq!( + fixture.owner("rt-route-x000").await.unwrap().member_id, + "m2" + ); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + let mut collecting = fixture.replica("gw-c"); + Arc::get_mut(&mut collecting).unwrap().reservation_grace = Duration::ZERO; + collecting.reconcile_once().await.unwrap(); + assert_eq!(fixture.candidates().await, 0); + assert_eq!(fixture.credentials.stored_credential_count(), Some(1)); + assert!(fixture.instance("rt-route-x000").await.is_some()); + } + + #[tokio::test] + async fn a_failed_contender_for_an_owned_route_charges_nothing() { + a_failed_contender_for_an_owned_route_charges_nothing_on(&Fixture::new().await).await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_failed_contender_for_an_owned_route_charges_nothing() { + let schema = TestSchema::create("supervisor_pool").await; + a_failed_contender_for_an_owned_route_charges_nothing_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_failed_contender_for_an_owned_route_charges_nothing_on(fixture: &Fixture) { + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let contender = fixture.replica("gw-b"); + let mut payload = fixture.owner(&route).await.unwrap(); + payload.member_id = "m2".to_string(); + payload.process_id = "p2".to_string(); + assert!(matches!( + contender + .transition(&route, None, Some(&payload)) + .await + .unwrap(), + Transition::OwnerChanged + )); + assert!(fixture.charged("m2").await.is_empty()); + assert!(fixture.charged("m1").await.contains(&route)); + assert_eq!(fixture.owner(&route).await.unwrap().member_id, "m1"); + } + + #[tokio::test] + async fn a_restarted_member_of_a_full_pool_re_adopts_its_routes_in_place() { + a_restarted_member_of_a_full_pool_re_adopts_its_routes_in_place_on(&Fixture::new().await) + .await; + } + + #[tokio::test] + #[ignore = "requires a disposable PostgreSQL database in OPENSHELL_TEST_POSTGRES_URL; run mise run test:rust:postgres"] + async fn postgres_a_restarted_member_of_a_full_pool_re_adopts_its_routes_in_place() { + let schema = TestSchema::create("supervisor_pool").await; + a_restarted_member_of_a_full_pool_re_adopts_its_routes_in_place_on(&Fixture::with_store( + schema.connect_store().await, + )) + .await; + schema.drop_schema().await; + } + + async fn a_restarted_member_of_a_full_pool_re_adopts_its_routes_in_place_on(fixture: &Fixture) { + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let routes = fixture.fill("m1", "p1", MAX_ASSIGNED_INSTANCES).await; + fixture.fill("m2", "p2", MAX_ASSIGNED_INSTANCES).await; + fixture + .sandbox("sandbox-new", "generation-1", SandboxPhase::Ready) + .await; + // m1's supervisor restarts inside its Pod: same member, new process. + fixture.member("m1", "p1-restarted", false).await; + let pool = fixture.replica("gw-a"); + let launching = fixture.replica("gw-b"); + let (pass, launched) = tokio::join!( + pool.reconcile_once(), + launching.place_launch(launch("rt-route-new0", "sandbox-new", "generation-1")), + ); + pass.unwrap(); + // Every route is re-adopted by the new process at the next term, + // without needing a free slot, and the launch found none. + for route in &routes { + let owner = fixture.owner(route).await.unwrap(); + assert_eq!( + ( + owner.member_id.as_str(), + owner.process_id.as_str(), + owner.auth_epoch, + owner.phase + ), + ("m1", "p1-restarted", 2, OwnerPhase::Granted), + "{route}" + ); + } + assert!(launched.is_err()); + assert!(fixture.owner("rt-route-new0").await.is_none()); + assert_eq!(fixture.charged("m1").await.len(), MAX_ASSIGNED_INSTANCES); + assert_eq!(fixture.charged("m2").await.len(), MAX_ASSIGNED_INSTANCES); + } + + /// A member stream accepted at `connected_at_ms`. + fn stream(member_id: &str, process_id: &str, connected_at_ms: i64) -> ConnectedMember { + ConnectedMember { + member_id: member_id.to_string(), + member_uid: format!("uid-{member_id}"), + process_id: process_id.to_string(), + runtime_address: "10.0.0.1:5500".to_string(), + leaving: false, + connected_at_ms, + } + } + + /// Short membership timings, so expiry happens within a test. + const TEST_TTL: Duration = Duration::from_millis(600); + const TEST_REFRESH: Duration = Duration::from_millis(100); + + /// Start `pool`; the returned sender stops it. + fn run(pool: &Arc) -> tokio::sync::watch::Sender { + let (stop, stopped) = tokio::sync::watch::channel(false); + pool.start(stopped); + stop + } + + async fn member_record(fixture: &Fixture, member_id: &str) -> MemberPayload { + let record = fixture + .store + .get(MEMBER_OBJECT_TYPE, &member_object_id(member_id)) + .await + .unwrap() + .unwrap(); + serde_json::from_slice(&record.payload).unwrap() + } + + #[tokio::test] + async fn a_stalled_pass_never_lets_a_served_member_expire() { + let fixture = Fixture::new().await; + let serving = fixture.timed_replica("gw-a", TEST_TTL, TEST_REFRESH); + serving.serve(vec![stream("m1", "p1", now_ms())]); + let (paused, resume) = serving.pause_next_pass(); + let stop = run(&serving); + paused.await.unwrap(); + // The pass stays stalled for several TTLs. + tokio::time::sleep(TEST_TTL * 4).await; + let observer = fixture.timed_replica("gw-b", TEST_TTL, TEST_REFRESH); + let live = observer.live_members().await.unwrap(); + assert_eq!( + live.iter() + .map(|member| member.member_id.as_str()) + .collect::>(), + ["m1"] + ); + let _ = resume.send(()); + stop.send_replace(true); + } + + #[tokio::test] + async fn a_member_that_reconnects_past_a_frozen_replica_keeps_its_instances() { + let fixture = Fixture::new().await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let first_stream = now_ms(); + let frozen = fixture.timed_replica("gw-a", TEST_TTL, TEST_REFRESH); + frozen.serve(vec![stream("m1", "p1", first_stream)]); + let peer = fixture.timed_replica("gw-b", TEST_TTL, TEST_REFRESH); + peer.serve(vec![stream("m2", "p2", first_stream)]); + let stop_frozen = run(&frozen); + let stop_peer = run(&peer); + tokio::time::sleep(TEST_REFRESH * 2).await; + // gw-a freezes; m1 notices and reconnects through gw-b well within + // the TTL, keeping its process. + stop_frozen.send_replace(true); + tokio::time::sleep(TEST_REFRESH * 2).await; + peer.serve(vec![ + stream("m1", "p1", first_stream + 1_000), + stream("m2", "p2", first_stream), + ]); + // Any replica reconciling meanwhile leaves m1's instance alone. + let judge = fixture.timed_replica("gw-c", TEST_TTL, TEST_REFRESH); + let until = tokio::time::Instant::now() + TEST_TTL * 4; + while tokio::time::Instant::now() < until { + judge.reconcile_once().await.unwrap(); + let owner = fixture.owner(&route).await.unwrap(); + assert_eq!( + (owner.member_id.as_str(), owner.auth_epoch), + ("m1", 1), + "m1 kept its instance" + ); + tokio::time::sleep(TEST_REFRESH).await; + } + // gw-a thaws with the member's stale stream still open: its refresh + // leaves the newer stream's record alone. + frozen.refresh_members().await; + let record = member_record(&fixture, "m1").await; + assert_eq!( + (record.gateway_replica.as_str(), record.connected_at_ms), + ("gw-b", first_stream + 1_000) + ); + stop_peer.send_replace(true); + } + + #[tokio::test] + async fn a_member_no_replica_serves_expires_and_its_instances_move() { + let fixture = Fixture::new().await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let lost = fixture.timed_replica("gw-a", TEST_TTL, TEST_REFRESH); + lost.serve(vec![stream("m1", "p1", now_ms())]); + let peer = fixture.timed_replica("gw-b", TEST_TTL, TEST_REFRESH); + peer.serve(vec![stream("m2", "p2", now_ms())]); + let stop_lost = run(&lost); + let stop_peer = run(&peer); + tokio::time::sleep(TEST_REFRESH * 2).await; + // m1 is cut off from every replica. + lost.serve(Vec::new()); + let judge = fixture.timed_replica("gw-c", TEST_TTL, TEST_REFRESH); + judge.reconcile_once().await.unwrap(); + assert_eq!(fixture.owner(&route).await.unwrap().member_id, "m1"); + tokio::time::sleep(TEST_TTL + TEST_REFRESH * 2).await; + judge.reconcile_once().await.unwrap(); + let owner = fixture.owner(&route).await.unwrap(); + assert_eq!((owner.member_id.as_str(), owner.auth_epoch), ("m2", 2)); + stop_lost.send_replace(true); + stop_peer.send_replace(true); + } + + /// What `member`'s `process` reported of `route` under `auth_epoch`. + fn member_report( + (member_id, process_id, auth_epoch): (&str, &str, u64), + route: &str, + sandbox_id: &str, + (phase, ready): (i32, bool), + connected_at_ms: i64, + ) -> MemberReport { + MemberReport { + member_id: member_id.to_string(), + process_id: process_id.to_string(), + auth_epoch, + connected_at_ms, + status: SandboxInstanceStatus { + route: route.to_string(), + sandbox_id: sandbox_id.to_string(), + phase, + ready, + exit_code: 0, + error: String::new(), + }, + } + } + + /// Grant `route` to `member`'s `process` at the next term, as a + /// completed takeover would. + async fn grant_next_term(fixture: &Fixture, route: &str, member_id: &str, process_id: &str) { + let mover = fixture.replica("mover"); + let record = mover.reload_owner(route).await; + let next = OwnerPayload { + member_id: member_id.to_string(), + process_id: process_id.to_string(), + auth_epoch: record.payload.auth_epoch + 1, + supervisor_instance_id: Some(Uuid::new_v4().to_string()), + gateway_token_id: Some(Uuid::new_v4().to_string()), + phase: OwnerPhase::Granted, + previous_member_id: Some(record.payload.member_id.clone()), + ..record.payload.clone() + }; + assert!(matches!( + mover + .transition(route, Some(&record), Some(&next)) + .await + .unwrap(), + Transition::Applied(Some(_)) + )); + } + + /// The status a replica's last pass published for `route`. + fn published_report(pool: &SupervisorPool, route: &str) -> Option<(i32, bool)> { + pool.controller + .pool_view() + .unwrap() + .reports + .get(route) + .map(|status| (status.phase, status.ready)) + } + + #[tokio::test] + async fn a_paused_report_of_an_earlier_term_never_hides_the_new_terms() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + fixture.member("m2", "p2", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let sandbox = fixture.owner(&route).await.unwrap().sandbox_id; + // The replica serving m1 read m1's term-1 report and stalls before + // writing it. + let old = fixture.replica("gw-a"); + old.serve_reports(vec![member_report( + ("m1", "p1", 1), + &route, + &sandbox, + (2, true), + 1, + )]); + let (paused, resume) = old.pause_next_report_recording(); + let stalled = { + let old = old.clone(); + tokio::spawn(async move { old.reconcile_once().await }) + }; + paused.await.unwrap(); + // m2 takes the route over at term 2 and reports it starting. + grant_next_term(&fixture, &route, "m2", "p2").await; + let new = fixture.replica("gw-b"); + new.serve_reports(vec![member_report( + ("m2", "p2", 2), + &route, + &sandbox, + (1, false), + 1, + )]); + new.reconcile_once().await.unwrap(); + // The stalled write lands afterwards. + let _ = resume.send(()); + stalled.await.unwrap().unwrap(); + let judge = fixture.replica("gw-c"); + judge.reconcile_once().await.unwrap(); + assert_eq!(published_report(&judge, &route), Some((1, false))); + // The earlier term's row is gone; term 2's row is untouched, and an + // unchanged report leaves it so. + let rows = fixture.report_rows(&route).await; + assert_eq!(rows.keys().copied().collect::>(), [2]); + new.reconcile_once().await.unwrap(); + judge.reconcile_once().await.unwrap(); + assert_eq!(published_report(&judge, &route), Some((1, false))); + } + + #[tokio::test] + async fn an_unchanged_report_whose_row_was_replaced_or_removed_is_rewritten() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let sandbox = fixture.owner(&route).await.unwrap().sandbox_id; + let serving = fixture.replica("gw-a"); + serving.serve_reports(vec![member_report( + ("m1", "p1", 1), + &route, + &sandbox, + (2, true), + 5, + )]); + serving.reconcile_once().await.unwrap(); + let written = fixture.report_rows(&route).await.remove(&1).unwrap(); + assert!(written.ready); + // Another writer replaced the row; the same report repairs it. + fixture + .write_report(ReportPayload { + ready: false, + connected_at_ms: 0, + ..written.clone() + }) + .await; + serving.reconcile_once().await.unwrap(); + assert_eq!(fixture.report_rows(&route).await.get(&1), Some(&written)); + // So it does after the row was removed. + fixture + .store + .delete(REPORT_OBJECT_TYPE, &report_object_id(&route, 1)) + .await + .unwrap(); + serving.reconcile_once().await.unwrap(); + assert_eq!(fixture.report_rows(&route).await.get(&1), Some(&written)); + } + + #[tokio::test] + async fn the_newest_stream_of_a_term_keeps_its_report() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let sandbox = fixture.owner(&route).await.unwrap().sandbox_id; + let current = fixture.replica("gw-b"); + current.serve_reports(vec![member_report( + ("m1", "p1", 1), + &route, + &sandbox, + (2, true), + 2_000, + )]); + // A replica still holding the process's older stream. + let stale = fixture.replica("gw-a"); + stale.serve_reports(vec![member_report( + ("m1", "p1", 1), + &route, + &sandbox, + (1, false), + 1_000, + )]); + current.reconcile_once().await.unwrap(); + stale.reconcile_once().await.unwrap(); + let row = fixture.report_rows(&route).await.remove(&1).unwrap(); + assert_eq!((row.ready, row.connected_at_ms), (true, 2_000)); + } + + #[tokio::test] + async fn reports_count_only_for_their_owners_term() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let sandbox = fixture.owner(&route).await.unwrap().sandbox_id; + // The route returns to the same process at term 2; its stream has + // only acknowledged term 1. + grant_next_term(&fixture, &route, "m1", "p1").await; + let serving = fixture.replica("gw-a"); + serving.serve_reports(vec![member_report( + ("m1", "p1", 1), + &route, + &sandbox, + (2, true), + 1, + )]); + serving.reconcile_once().await.unwrap(); + assert!(fixture.report_rows(&route).await.is_empty()); + assert_eq!(published_report(&serving, &route), None); + } + + #[tokio::test] + async fn cleanup_removes_only_rows_older_than_the_listed_owner() { + let fixture = Fixture::new().await; + fixture.member("m1", "p1", false).await; + let route = fixture.fill("m1", "p1", 1).await.remove(0); + let pool = fixture.replica("gw-a"); + let listed = pool.owners().await.unwrap(); + for auth_epoch in [1, 2] { + fixture + .write_report(ReportPayload { + route: route.clone(), + member_id: "m1".to_string(), + process_id: "p1".to_string(), + auth_epoch, + sandbox_id: "sandbox".to_string(), + phase: 2, + ready: true, + exit_code: 0, + connected_at_ms: 0, + }) + .await; + } + // Term 2's row is newer than the listed owner (term 1): kept, as is + // the listed owner's own row. + let rows = pool.report_rows().await.unwrap(); + pool.remove_stale_reports(&listed, &rows).await; + assert_eq!( + fixture + .report_rows(&route) + .await + .keys() + .copied() + .collect::>(), + [1, 2] + ); + // Once the owner is listed at term 2, term 1's row goes. + grant_next_term(&fixture, &route, "m1", "p1").await; + let listed = pool.owners().await.unwrap(); + pool.remove_stale_reports(&listed, &rows).await; + assert_eq!( + fixture + .report_rows(&route) + .await + .keys() + .copied() + .collect::>(), + [2] + ); + // A row changed since it was listed is left for the next pass. + fixture + .write_report(ReportPayload { + route: route.clone(), + member_id: "m1".to_string(), + process_id: "p1".to_string(), + auth_epoch: 2, + sandbox_id: "sandbox".to_string(), + phase: 2, + ready: false, + exit_code: 0, + connected_at_ms: 0, + }) + .await; + let rows = pool.report_rows().await.unwrap(); + fixture.remove_owner_record(&route).await; + fixture + .write_report(ReportPayload { + route: route.clone(), + member_id: "m1".to_string(), + process_id: "p1".to_string(), + auth_epoch: 2, + sandbox_id: "sandbox".to_string(), + phase: 3, + ready: false, + exit_code: 0, + connected_at_ms: 0, + }) + .await; + pool.remove_stale_reports(&BTreeMap::new(), &rows).await; + assert_eq!( + fixture + .report_rows(&route) + .await + .keys() + .copied() + .collect::>(), + [2] + ); + // Without an owner, a row as listed goes. + let rows = pool.report_rows().await.unwrap(); + pool.remove_stale_reports(&BTreeMap::new(), &rows).await; + assert!(fixture.report_rows(&route).await.is_empty()); + } + + #[test] + fn a_takeover_names_why_the_owner_lost_its_instance() { + let owner = OwnerPayload { + route: "rt-a".to_string(), + sandbox_id: "sandbox-a".to_string(), + generation: "generation-1".to_string(), + member_id: "m1".to_string(), + process_id: "p1".to_string(), + auth_epoch: 1, + supervisor_instance_id: None, + gateway_token_id: None, + phase: OwnerPhase::Granted, + previous_member_id: None, + }; + let member = |process_id: &str, leaving| LiveMember { + member_id: "m1".to_string(), + process_id: process_id.to_string(), + runtime_address: "10.0.0.1:5500".to_string(), + leaving, + }; + assert_eq!(takeover_cause(&owner, &[]), "expired"); + assert_eq!(takeover_cause(&owner, &[member("p1", true)]), "leaving"); + assert_eq!(takeover_cause(&owner, &[member("p2", false)]), "restarted"); + } + + #[test] + fn placement_prefers_the_ring_and_moves_on_only_when_a_member_is_full() { + let member = |id: &str, leaving| LiveMember { + member_id: id.to_string(), + process_id: format!("p-{id}"), + runtime_address: "10.0.0.1:5500".to_string(), + leaving, + }; + let members = vec![member("m1", false), member("m2", false), member("m3", true)]; + let ring = ring_of(&members); + let sandbox = (0..64) + .map(|index| format!("sandbox-{index}")) + .find(|sandbox| ring.owner_for(sandbox) == Some("m1")) + .expect("a sandbox the ring places on m1"); + let routes_for = |member: &str, count: usize| { + ( + member.to_string(), + (0..count) + .map(|index| format!("rt-{member}-{index}")) + .collect::>(), + ) + }; + let order = |charges: BTreeMap>, route: &str, same| { + candidates(&ring, &members, &sandbox, route, &charges, same) + .into_iter() + .map(|member| member.member_id.clone()) + .collect::>() + }; + let charges = |m1, m2| BTreeMap::from([routes_for("m1", m1), routes_for("m2", m2)]); + assert_eq!( + order(charges(MAX_ASSIGNED_INSTANCES - 1, 0), "rt-new", None), + ["m1", "m2"] + ); + // A full member is skipped for the least charged one with room; a + // leaving member never receives instances. + assert_eq!( + order(charges(MAX_ASSIGNED_INSTANCES, 3), "rt-new", None), + ["m2"] + ); + assert!( + order( + charges(MAX_ASSIGNED_INSTANCES, MAX_ASSIGNED_INSTANCES), + "rt-new", + None + ) + .is_empty() + ); + // A full member still has room for a route it is charged for, and + // its restarted process re-adopts it before any other member. + assert_eq!( + order( + charges(MAX_ASSIGNED_INSTANCES, MAX_ASSIGNED_INSTANCES), + "rt-m2-0", + Some("m2") + ), + ["m2"] + ); + } +} diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index a0fb2654be..248f5fab7b 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -17,6 +17,7 @@ use tonic::{Request, Response, Status}; use tracing::{debug, info, warn}; use uuid::Uuid; +use openshell_core::jwt::CredentialEpoch; use openshell_core::proto::{ GatewayMessage, GetSandboxProviderStatusRequest, GetSandboxProviderStatusResponse, PeerNotifyConfigUpdateRequest, PeerNotifyConfigUpdateResponse, PeerRelayFrame, PeerRelayInit, @@ -39,7 +40,7 @@ use crate::gateway_metrics::{ use crate::grpc::provider_readiness::ProviderReadinessEvidence; use crate::persistence::{ObjectId, ObjectWorkspace}; use crate::supervisor_owner::{ - OWNER_TTL, OwnerError, OwnerGuard, SupervisorOwnerIndex, can_supersede, + OWNER_TTL, OwnerError, OwnerGuard, SupervisorOwnerIndex, can_supersede, term_order, }; const HEARTBEAT_INTERVAL_SECS: u32 = 15; @@ -306,6 +307,8 @@ struct LiveSession { /// Uniquely identifies this session instance. Used by cleanup to avoid /// removing a session that has since been superseded by a reconnect. session_id: String, + /// Ownership term of the session, when known. + auth_epoch: Option, tx: mpsc::Sender, /// Present only for sessions that apply streamed configuration. config: Option, @@ -353,6 +356,17 @@ pub(crate) struct EndpointReportCursor { pub(crate) report_digest: [u8; 32], } +/// Outcome of registering a supervisor session. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Registration { + /// No session was registered for the sandbox. + New, + /// The previous session was replaced and told to close. + Superseded, + /// A session of a later ownership term holds the sandbox. + Refused, +} + /// Holds a oneshot sender that will deliver the upgraded relay stream or a /// target-open failure reported by the supervisor. type RelayStreamSender = oneshot::Sender>; @@ -461,17 +475,42 @@ impl SupervisorSessionRegistry { shutdown: oneshot::Sender<()>, mode: SessionMode, ) -> bool { + self.register_term(sandbox_id, session_id, None, tx, shutdown, mode) + == Registration::Superseded + } + + /// Register a session of ownership term `auth_epoch`. + /// + /// The owner record admits sessions in term order, but a session can + /// register after a later term replaced its record. A registered session + /// of a later term therefore keeps the sandbox and this one is refused. + /// Otherwise any previous session is replaced and told to close. + pub(crate) fn register_term( + &self, + sandbox_id: String, + session_id: String, + auth_epoch: Option, + tx: mpsc::Sender, + shutdown: oneshot::Sender<()>, + mode: SessionMode, + ) -> Registration { let config = match mode { SessionMode::Push => Some(ConfigSessionState::new(tx.clone())), SessionMode::Poll => None, }; let mut sessions = self.sessions.lock().unwrap(); + if sessions.get(&sandbox_id).is_some_and(|current| { + term_order(auth_epoch, current.auth_epoch) == Some(std::cmp::Ordering::Less) + }) { + return Registration::Refused; + } let previous = sessions.remove(&sandbox_id); sessions.insert( sandbox_id.clone(), LiveSession { sandbox_id, session_id, + auth_epoch, tx, config, shutdown: Some(shutdown), @@ -489,9 +528,9 @@ impl SupervisorSessionRegistry { if let Some(shutdown) = prev.shutdown { let _ = shutdown.send(()); } - true + Registration::Superseded } - None => false, + None => Registration::New, } } @@ -2170,6 +2209,8 @@ struct SessionSetup { sandbox_id: String, instance_id: String, connection_epoch: u64, + /// Ownership term: the credential epoch of the session's bearer. + auth_epoch: Option, supports_session_redirect: bool, mode: SessionMode, bootstrap: Option, @@ -2186,6 +2227,7 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> { sandbox_id, instance_id, connection_epoch, + auth_epoch, supports_session_redirect, mode, bootstrap, @@ -2244,29 +2286,43 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> { ); let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); let owner_guard = owner_index - .publish( + .publish_term( &sandbox_id, &session_id, &instance_id, connection_epoch, + auth_epoch, &state.replica_id, &owner_peer_endpoint, ) .await .map_err(owner_error_to_status)?; - let superseded = state.supervisor_sessions.register_with_mode( + match state.supervisor_sessions.register_term( sandbox_id.clone(), session_id.clone(), + auth_epoch, session_tx.clone(), shutdown_tx, mode, - ); - if superseded { - info!( + ) { + Registration::New => {} + Registration::Superseded => info!( sandbox_id = %sandbox_id, session_id = %session_id, "supervisor session: superseded previous session" - ); + ), + Registration::Refused => { + info!( + sandbox_id = %sandbox_id, + session_id = %session_id, + auth_epoch = auth_epoch.map_or(0, CredentialEpoch::get), + "supervisor session: refused; a later ownership term holds the sandbox" + ); + abandon_session_setup(&state, &owner_index, &owner_guard, "registration").await; + return Err(Status::unavailable( + "a later supervisor ownership term holds this sandbox", + )); + } } if let Some(bootstrap) = bootstrap_in_flight { state.supervisor_sessions.begin_config_bootstrap( @@ -2386,6 +2442,7 @@ async fn accept_supervisor_session(setup: SessionSetup) -> Result<(), Status> { session_id = %session_id, instance_id = %instance_id, connection_epoch, + auth_epoch = auth_epoch.map_or(0, CredentialEpoch::get), replica_id = %state.replica_id, config_push = pushes_config, "supervisor session: accepted" @@ -2474,9 +2531,13 @@ pub async fn handle_connect_supervisor( if sandbox_id.is_empty() { return Err(Status::invalid_argument("sandbox_id is required")); } - if let Some(principal) = principal.as_ref() { - crate::auth::guard::ensure_sandbox_principal_scope(principal, &sandbox_id)?; - } + // The bearer's credential epoch is the session's ownership term. + let auth_epoch = match principal.as_ref() { + Some(principal) => { + crate::auth::guard::ensure_sandbox_principal_scope(principal, &sandbox_id)?.auth_epoch() + } + None => None, + }; // Configuration is pushed only in push mode and only to supervisors that // apply it; the rest keep polling. let pushes_config = state.config.config_delivery_mode @@ -2514,6 +2575,7 @@ pub async fn handle_connect_supervisor( &owner, &hello.instance_id, hello.connection_epoch, + auth_epoch, OWNER_TTL, ), Ok(None) => true, @@ -2565,6 +2627,7 @@ pub async fn handle_connect_supervisor( sandbox_id, instance_id: hello.instance_id, connection_epoch: hello.connection_epoch, + auth_epoch, supports_session_redirect: hello.supports_session_redirect, mode, bootstrap, @@ -2875,6 +2938,13 @@ async fn run_session_loop( msg = tokio_stream::StreamExt::next(inbound) => { match msg.transpose() { Ok(Some(msg)) => { + // The shutdown signal can lose this select, so check + // that no reconnect or later ownership term replaced + // this session before acting on its message. + if !state.supervisor_sessions.is_current_session(sandbox_id, session_id) { + info!(sandbox_id = %sandbox_id, session_id = %session_id, "supervisor session: replaced; dropping its message and closing"); + break; + } if matches!(msg.payload.as_ref(), Some(supervisor_message::Payload::Heartbeat(_))) && !renew_supervisor_owner(state, sandbox_id, session_id, owner_guard).await { @@ -3657,6 +3727,7 @@ mod tests { sandbox_id: sandbox_id.to_string(), source: SandboxIdentitySource::BootstrapJwt { issuer: "openshell-gateway:test".to_string(), + auth_epoch: None, }, trust_domain: Some("openshell".to_string()), }) @@ -4713,6 +4784,7 @@ mod tests { session_id: "session-a".to_string(), supervisor_instance_id: "instance-a".to_string(), connection_epoch: 1, + auth_epoch: None, owner_replica_id: replica.to_string(), owner_peer_endpoint: format!("http://{replica}:8080"), connected_at_ms: 0, @@ -5511,4 +5583,565 @@ mod tests { assert!(ring_redirect("gw-0", "gw-1", Some(String::new())).is_none()); assert!(ring_redirect("gw-0", "gw-1", Some("local://gw-1".to_string())).is_none()); } + + // ---- ownership terms ---- + + fn term(value: u64) -> CredentialEpoch { + CredentialEpoch::new(value).unwrap() + } + + #[test] + fn registry_refuses_an_earlier_term_while_a_later_term_is_registered() { + let registry = SupervisorSessionRegistry::new(); + let (later_tx, _later_rx) = mpsc::channel(1); + let (earlier_tx, _earlier_rx) = mpsc::channel(1); + let (later_shutdown, mut later_closed) = oneshot::channel(); + assert_eq!( + registry.register_term( + "sbx".into(), + "s-later".into(), + Some(term(2)), + later_tx, + later_shutdown, + SessionMode::Poll, + ), + Registration::New + ); + + assert_eq!( + registry.register_term( + "sbx".into(), + "s-earlier".into(), + Some(term(1)), + earlier_tx, + make_shutdown(), + SessionMode::Poll, + ), + Registration::Refused + ); + assert!(registry.is_current_session("sbx", "s-later")); + assert!( + later_closed.try_recv().is_err(), + "the later term must stay open" + ); + } + + #[test] + fn registry_later_term_replaces_and_closes_an_earlier_term() { + let registry = SupervisorSessionRegistry::new(); + let (earlier_tx, _earlier_rx) = mpsc::channel(1); + let (later_tx, _later_rx) = mpsc::channel(1); + let (earlier_shutdown, mut earlier_closed) = oneshot::channel(); + registry.register_term( + "sbx".into(), + "s-earlier".into(), + Some(term(1)), + earlier_tx, + earlier_shutdown, + SessionMode::Poll, + ); + + assert_eq!( + registry.register_term( + "sbx".into(), + "s-later".into(), + Some(term(2)), + later_tx, + make_shutdown(), + SessionMode::Poll, + ), + Registration::Superseded + ); + assert!(registry.is_current_session("sbx", "s-later")); + assert!( + earlier_closed.try_recv().is_ok(), + "the earlier term must be told to close" + ); + } + + #[test] + fn registry_replaces_sessions_of_equal_or_unknown_terms_as_before() { + for (first, second) in [ + (Some(term(1)), Some(term(1))), + (None, Some(term(1))), + (Some(term(2)), None), + (None, None), + ] { + let registry = SupervisorSessionRegistry::new(); + let (tx, _rx) = mpsc::channel(1); + registry.register_term( + "sbx".into(), + "s1".into(), + first, + tx.clone(), + make_shutdown(), + SessionMode::Poll, + ); + assert_eq!( + registry.register_term( + "sbx".into(), + "s2".into(), + second, + tx, + make_shutdown(), + SessionMode::Poll, + ), + Registration::Superseded, + "{first:?} then {second:?}" + ); + assert!(registry.is_current_session("sbx", "s2")); + } + } + + #[tokio::test] + async fn a_session_replaced_by_a_later_term_acts_on_none_of_its_messages() { + use futures::StreamExt as _; + + let state = crate::grpc::test_support::test_server_state().await; + let mut owner_guard = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .publish_term( + "sb-replaced", + "session-a", + "instance-a", + 1, + Some(term(1)), + &state.replica_id, + "local://this", + ) + .await + .unwrap(); + let (tx, mut rx) = mpsc::channel(4); + state.supervisor_sessions.register_term( + "sb-replaced".into(), + "session-a".into(), + Some(term(1)), + tx.clone(), + make_shutdown(), + SessionMode::Poll, + ); + let (channel_id, _relay_rx) = state + .supervisor_sessions + .open_relay("sb-replaced", Duration::from_secs(1)) + .await + .unwrap(); + rx.recv().await.expect("RelayOpen reaches the first term"); + let (later_tx, _later_rx) = mpsc::channel(4); + state.supervisor_sessions.register_term( + "sb-replaced".into(), + "session-b".into(), + Some(term(2)), + later_tx, + make_shutdown(), + SessionMode::Poll, + ); + + // The replaced supervisor answers before it sees its stream close. + let failure = SupervisorMessage { + payload: Some(supervisor_message::Payload::RelayOpenResult( + openshell_core::proto::RelayOpenResult { + channel_id: channel_id.clone(), + success: false, + error: "target refused".into(), + }, + )), + }; + let mut inbound = futures::stream::iter([Ok(failure)]).chain(futures::stream::pending()); + let session = AcceptedSession { + state: Arc::clone(&state), + sandbox_id: "sb-replaced".into(), + session_id: "session-a".into(), + instance_id: "instance-a".into(), + tx, + supports_session_redirect: false, + pushes_config: false, + }; + let exit = tokio::time::timeout( + Duration::from_secs(5), + run_session_loop( + &session, + &mut inbound, + oneshot::channel().1, + &mut owner_guard, + ), + ) + .await + .expect("a replaced session closes on its next message"); + + assert_eq!(exit, SessionLoopExit::Other); + assert!( + state + .supervisor_sessions + .pending_relays + .lock() + .unwrap() + .contains_key(&channel_id), + "the replaced session must not fail a relay its replacement serves" + ); + assert!( + state + .supervisor_sessions + .is_current_session("sb-replaced", "session-b") + ); + } + + #[tokio::test] + async fn renewal_closes_a_session_whose_record_a_later_term_took_on_another_replica() { + use futures::StreamExt as _; + + let state = crate::grpc::test_support::test_server_state().await; + let mut sandbox = sandbox_record("sb-moved", "sb-moved"); + sandbox.set_phase(SandboxPhase::Ready as i32); + state.store.put_message(&sandbox).await.unwrap(); + let owners = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); + let mut owner_guard = owners + .publish_term( + "sb-moved", + "session-replaced", + "instance-a", + 1, + Some(term(1)), + &state.replica_id, + "local://this", + ) + .await + .unwrap(); + let (tx, _rx) = mpsc::channel(4); + state.supervisor_sessions.register_term( + "sb-moved".into(), + "session-replaced".into(), + Some(term(1)), + tx.clone(), + make_shutdown(), + SessionMode::Poll, + ); + + // The takeover's supervisor connects to another replica while this + // one still renews the previous term's fresh record. + owners + .publish_term( + "sb-moved", + "session-takeover", + "instance-b", + 0, + Some(term(2)), + "gateway-b", + "http://gateway-b:8080", + ) + .await + .expect("a later term replaces a fresh record of another replica"); + + let heartbeat = SupervisorMessage { + payload: Some(supervisor_message::Payload::Heartbeat( + openshell_core::proto::SupervisorHeartbeat {}, + )), + }; + let mut inbound = futures::stream::iter([Ok(heartbeat)]).chain(futures::stream::pending()); + let session = AcceptedSession { + state: Arc::clone(&state), + sandbox_id: "sb-moved".into(), + session_id: "session-replaced".into(), + instance_id: "instance-a".into(), + tx, + supports_session_redirect: false, + pushes_config: false, + }; + let exit = tokio::time::timeout( + Duration::from_secs(5), + run_session_loop( + &session, + &mut inbound, + oneshot::channel().1, + &mut owner_guard, + ), + ) + .await + .expect("the next renewal finds the record taken"); + assert_eq!(exit, SessionLoopExit::Other); + drop(session); + finish_supervisor_session(&state, "sb-moved", "session-replaced", exit, &owner_guard).await; + + assert!(!state.supervisor_sessions.has_session("sb-moved")); + let owner = owners.read("sb-moved").await.unwrap().unwrap(); + assert_eq!(owner.session_id, "session-takeover"); + assert_eq!(owner.auth_epoch, Some(term(2))); + // The takeover's live record keeps the sandbox Ready. + assert_eq!(sandbox_phase(&state, "sb-moved").await, SandboxPhase::Ready); + } + + /// A `ConnectSupervisor` stream of one supervisor process. + struct TermStream { + server: tokio::task::JoinHandle>, + _outbound: mpsc::Sender, + inbound: tonic::Streaming, + } + + impl Drop for TermStream { + fn drop(&mut self) { + self.server.abort(); + } + } + + /// Connect as supervisor `instance_id` holding a gateway credential of + /// epoch `auth_epoch`, and return the stream with its accepted session ID. + async fn connect_term( + state: &Arc, + sandbox_id: &str, + instance_id: &str, + connection_epoch: u64, + auth_epoch: u64, + ) -> Result<(TermStream, String), Status> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let principal = + crate::auth::principal::test_support::sandbox_with_epoch(sandbox_id, auth_epoch); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service( + openshell_core::proto::open_shell_server::OpenShellServer::with_interceptor( + crate::grpc::OpenShellService::new(Arc::clone(state)), + move |mut request: Request<()>| { + request.extensions_mut().insert(principal.clone()); + Ok(request) + }, + ), + ) + .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener)), + ); + let mut client = open_shell_client::OpenShellClient::connect(format!("http://{address}")) + .await + .unwrap(); + let (outbound, rx) = mpsc::channel(4); + outbound + .send(SupervisorMessage { + payload: Some(supervisor_message::Payload::Hello(SupervisorHello { + sandbox_id: sandbox_id.into(), + instance_id: instance_id.into(), + connection_epoch, + ..Default::default() + })), + }) + .await + .unwrap(); + let inbound = match client.connect_supervisor(ReceiverStream::new(rx)).await { + Ok(response) => response.into_inner(), + Err(status) => { + server.abort(); + return Err(status); + } + }; + let mut stream = TermStream { + server, + _outbound: outbound, + inbound, + }; + let accepted = tokio::time::timeout(Duration::from_secs(5), stream.inbound.message()) + .await + .expect("gateway response before timeout")? + .expect("stream open"); + let Some(gateway_message::Payload::SessionAccepted(accepted)) = accepted.payload else { + panic!("expected SessionAccepted"); + }; + Ok((stream, accepted.session_id)) + } + + #[tokio::test] + async fn a_later_term_takes_over_a_live_session_of_another_supervisor_at_once() { + let state = state_with_sandbox("sb-takeover").await; + let (mut replaced, _) = connect_term(&state, "sb-takeover", "instance-a", 3, 1) + .await + .expect("the first term connects"); + + let (_current, current_session) = connect_term(&state, "sb-takeover", "instance-b", 0, 2) + .await + .expect("a later term must not wait out the previous owner's record"); + + assert!( + state + .supervisor_sessions + .is_current_session("sb-takeover", ¤t_session) + ); + let owner = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL) + .read("sb-takeover") + .await + .unwrap() + .unwrap(); + assert_eq!(owner.session_id, current_session); + assert_eq!(owner.supervisor_instance_id, "instance-b"); + assert_eq!(owner.auth_epoch, Some(term(2))); + // The gateway ends the replaced supervisor's stream. + let closed = tokio::time::timeout(Duration::from_secs(5), async { + while let Ok(Some(_)) = replaced.inbound.message().await {} + }) + .await; + assert!(closed.is_ok(), "the replaced term's stream must close"); + } + + #[tokio::test] + async fn an_earlier_term_is_refused_while_a_later_term_holds_the_sandbox() { + let state = state_with_sandbox("sb-fenced").await; + let (_current, current_session) = connect_term(&state, "sb-fenced", "instance-b", 0, 2) + .await + .expect("the later term connects"); + + // Neither the replaced process nor an earlier-term connection that + // would otherwise count as a reconnect may reclaim the sandbox. + for instance in ["instance-a", "instance-b"] { + let Err(status) = connect_term(&state, "sb-fenced", instance, 1, 1).await else { + panic!("{instance} at an earlier term must be refused"); + }; + assert_eq!(status.code(), tonic::Code::Unavailable, "{instance}"); + } + assert!( + state + .supervisor_sessions + .is_current_session("sb-fenced", ¤t_session) + ); + } + + #[tokio::test] + async fn an_equal_term_still_refuses_another_supervisor_instance() { + let state = state_with_sandbox("sb-same-term").await; + let (_owner, owner_session) = connect_term(&state, "sb-same-term", "instance-a", 0, 1) + .await + .expect("the owner connects"); + + let Err(status) = connect_term(&state, "sb-same-term", "instance-b", 1, 1).await else { + panic!("another instance of the same term must wait for the owner's record"); + }; + assert_eq!(status.code(), tonic::Code::Unavailable); + assert!( + state + .supervisor_sessions + .is_current_session("sb-same-term", &owner_session) + ); + } + + /// Authenticate `token` for `ReportMainProcessExit` as the gateway router + /// does, then run the handler as the resulting principal. + async fn report_main_process_exit_as( + state: &Arc, + authenticator: &crate::auth::sandbox_jwt::SandboxSessionJwtAuthenticator, + token: &str, + sandbox_id: &str, + ) -> Result { + use crate::auth::authenticator::{Authentication, Authenticator as _}; + + let mut headers = http::HeaderMap::new(); + headers.insert( + http::header::AUTHORIZATION, + format!("Bearer {token}").parse().unwrap(), + ); + let principal = match authenticator + .authenticate(&headers, "/openshell.v1.OpenShell/ReportMainProcessExit") + .await + { + Ok(Authentication::Accepted(principal)) => principal, + Ok(_) => return Err(Status::unauthenticated("bearer was not accepted")), + Err(rejection) => return Err(rejection.status), + }; + let mut request = Request::new(ReportMainProcessExitRequest { + sandbox_id: sandbox_id.into(), + instance_id: "instance-a".into(), + exit_code: 1, + }); + request.extensions_mut().insert(principal.clone()); + handle_report_main_process_exit(state, request).await?; + Ok(principal) + } + + #[tokio::test] + async fn a_replaced_term_cannot_report_the_main_process_exit() { + use crate::auth::sandbox_session::{ + PersistedSandboxIdentity, TakeoverIdentity, advance_identity_for_takeover, + }; + + let state = crate::grpc::test_support::test_server_state().await; + let launch = PersistedSandboxIdentity::new().unwrap(); + let mut sandbox = sandbox_record("sb-exit", "sb-exit"); + launch.write(&mut sandbox.metadata.as_mut().unwrap().annotations); + sandbox.set_phase(SandboxPhase::Ready as i32); + state.store.put_message(&sandbox).await.unwrap(); + let material = openshell_bootstrap::jwt::generate_jwt_key().unwrap(); + let authority = Arc::new( + crate::auth::sandbox_jwt::SandboxSessionJwtAuthority::from_pem( + material.signing_key_pem.as_bytes(), + material.public_key_pem.as_bytes(), + material.kid, + "test-gateway", + Some(Duration::from_mins(15)), + ) + .unwrap(), + ); + let authenticator = crate::auth::sandbox_jwt::SandboxSessionJwtAuthenticator::new( + Arc::clone(&authority), + state.store.clone(), + ); + let replaced = authority + .mint_persisted_launch("sb-exit", &launch) + .unwrap() + .supervisor + .gateway_token; + + // A takeover moves the instance to another supervisor process. + let takeover_epoch = CredentialEpoch::new(launch.auth_epoch.get() + 1).unwrap(); + assert_eq!( + advance_identity_for_takeover( + &state.store, + "sb-exit", + launch.runtime_generation.as_str(), + takeover_epoch, + Uuid::new_v4(), + ) + .await + .unwrap(), + TakeoverIdentity::Advanced + ); + let before = state + .store + .get_message::("sb-exit") + .await + .unwrap() + .unwrap(); + + let error = report_main_process_exit_as( + &state, + &authenticator, + replaced.expose_secret(), + "sb-exit", + ) + .await + .expect_err("the replaced term's credential is revoked"); + assert_eq!(error.code(), tonic::Code::Unauthenticated); + let unchanged = state + .store + .get_message::("sb-exit") + .await + .unwrap() + .unwrap(); + assert_eq!(unchanged, before); + + // The current term's credential carries its epoch and is accepted. + let current = + PersistedSandboxIdentity::read(&before.metadata.unwrap().annotations).unwrap(); + let token = authority + .mint_persisted_launch("sb-exit", ¤t) + .unwrap() + .supervisor + .gateway_token; + let principal = + report_main_process_exit_as(&state, &authenticator, token.expose_secret(), "sb-exit") + .await + .expect("the current term reports the exit"); + let Principal::Sandbox(principal) = principal else { + panic!("expected a sandbox principal"); + }; + assert_eq!(principal.auth_epoch(), Some(takeover_epoch)); + let reported = state + .store + .get_message::("sb-exit") + .await + .unwrap() + .unwrap(); + assert_eq!(reported.status.unwrap().exit_code, Some(1)); + } } diff --git a/crates/openshell-supervisor-control/README.md b/crates/openshell-supervisor-control/README.md index 124caa5080..bbbb47cbdf 100644 --- a/crates/openshell-supervisor-control/README.md +++ b/crates/openshell-supervisor-control/README.md @@ -12,18 +12,46 @@ re-implementing the control contracts. ## Modules -- `controller`: `SupervisorHostController` holds the desired set of - `SandboxInstanceAssignment`s and serves it to the host. Desired state is - uninitialized until the complete set has been derived from durable records - once, so a restarted gateway never presents an empty set as authoritative; - an explicit empty set after initialization does stop instances. Full - replacements are compare-and-replace against the revision read before the - records were listed. One host session at a time: a newer connection epoch - of the same host process or a new process supersedes the current stream, a - delayed older hello is refused, and reports from superseded streams, acks - of revisions never sent, regressing acks, and oversized reports are - dropped. A stream ends when the token that opened it expires; hosted - instances keep running. At most 256 instances are assigned. +- `controller`: `SupervisorHostController` serves the shared supervisor, + which is always a pool (high availability): several hosts (pool members) + are served at once and each instance has one durable owner: a member, its + process, and an ownership term. Every member names itself in its hello; a + hello without a member, or with an invalid one, is refused with the + `OPENSHELL_SUPERVISOR_POOL_*` setting to fix. One session per member: a + newer connection epoch of the same member process or a new process of the + member supersedes its stream, a delayed older hello is refused, and reports + from superseded streams, acks of revisions never sent, regressing acks, + and oversized reports are dropped. A stream ends when the token that opened + it expires; hosted instances keep running. The pool's recovery deadline + defaults to `POOL_RECOVERY_TIMEOUT_SECS` (`set_recovery_timeout_secs` + overrides it). The gateway's pool reconciler replaces the `PoolView` as a whole + (`update_pool`): the desired set from the gateway's durable records, the + owners, the launch material this replica loaded, and the latest recorded + reports of owners served elsewhere. Owners placed, instances launched, and + instances withdrawn here since the records were listed are kept as they are. + Each member's snapshot carries the instances its process owns (a takeover + term with its credentials, or with the term alone when another gateway + replica minted them), `stop_routes` for withdrawn instances it owns, + `keep_routes` for owned desired instances whose material this replica has + not loaded yet, the route directory for runtime redirects, and the recovery + deadline its hosts offer each runtime (`POOL_RECOVERY_TIMEOUT_SECS`, 60 s, + unless configured). Omission means release, never stop. Nothing is sent + before the first view. An instance is judged by its owning process's + session, or by that owner term's recorded report on other replicas. + Placement keeps each member at 256 instances. Each stream remembers the term + of every route its sent snapshots assigned and the revision that installed + it; `member_reports` attributes a reported status to that term only once the + stream acknowledged that revision, so the gateway never records an earlier + term's status under a later one. `membership_changes` changes when a member + stream attaches, ends, or announces leaving, so the gateway refreshes + membership records at once. The gateway closes a member's stream after + `pool_timing::GATEWAY_SILENCE_TIMEOUT` of silence, and drops a member that + stops reading as soon. +- `pool_timing`: the timing contract between pool hosts and the gateway: the + 15 s membership TTL, the 2 s record refresh, 1 s heartbeats, host silence + detection and bounded reconnect attempts, and `reconnect_budget`, the worst + case for a member to reconnect past an unresponsive replica and be refreshed + elsewhere, which a test keeps well inside the TTL. - `registration`: `SandboxInstanceRegistration` is the JSON launch material one assignment carries and the driver persists with the generation. It grants the supervisor exactly that generation's gateway and Sandbox Protocol @@ -36,7 +64,12 @@ re-implementing the control contracts. to the compute driver. `launch` validates the driver's registration, persists the encoded launch material through the driver's operation, and only then assigns the instance; a host at capacity is refused before the - workload runs. `begin_recovery` captures + workload runs. For a pool, the gateway's installed `InstancePlacement` + durably records the instance and its launch term's owner before the + assignment, and a launch no member can take is refused like a full host + (`LaunchError::Refused`). A pool's withdrawals are removed from the durable + records by the gateway, and a recovery reconciles those records with the + driver's listing instead of replacing the desired set. `begin_recovery` captures the revision before the driver lists its durable records and `complete_recovery` reuses current assignments, reads the rest through the driver's operation, binds each to its recorded route and sandbox, and keeps diff --git a/crates/openshell-supervisor-control/src/controller.rs b/crates/openshell-supervisor-control/src/controller.rs index 2d0aa885d9..5037f90799 100644 --- a/crates/openshell-supervisor-control/src/controller.rs +++ b/crates/openshell-supervisor-control/src/controller.rs @@ -4,28 +4,32 @@ //! Gateway side of the private supervisor host protocol //! (`openshell.supervisor.host.v1`). //! -//! The gateway owns the desired set of sandbox instances a shared supervisor -//! runs. The [`SharedSupervision`](crate::SharedSupervision) coordinator -//! publishes that set here on behalf of a compute driver; the gateway's gRPC -//! server hands each host stream authenticated by a supervisor token to -//! [`SupervisorHostController::attach`], which sends complete snapshots and -//! records what the host reports. The gateway also mints that token and -//! publishes each renewal here for the driver to deliver to the host. +//! Shared supervisors always run as a pool (high availability) that the +//! gateway coordinates. The [`SharedSupervision`](crate::SharedSupervision) +//! coordinator publishes launches and withdrawals here on behalf of a compute +//! driver; the gateway's gRPC server hands each member stream authenticated +//! by a supervisor token to [`SupervisorHostController::attach`], which sends +//! complete snapshots and records what the member reports. The gateway also +//! mints that token and publishes each renewal here for the driver to deliver. //! -//! Desired state is uninitialized until the complete set has been derived -//! from durable records once. Nothing is sent before that, so a restarted -//! gateway can never present an empty map as authoritative and stop running -//! instances. An explicit empty snapshot, after initialization, does stop -//! them. Full resyncs are compare-and-replace against the revision read -//! before the records were listed, so a list that raced a lifecycle edit -//! never resurrects a withdrawn instance or discards a new one. +//! The pool's desired set, the durable owner of every instance (a member's +//! process and an ownership term), and the latest reports of owners served by +//! other gateway replicas come from the gateway's durable records through +//! [`PoolView`], on every gateway replica alike; each member is sent only the +//! instances its current process owns, with the credentials of their term. A +//! member releases, without stopping, an instance a snapshot omits; only +//! withdrawn instances are listed for stopping. Nothing is sent to members +//! until the first view has been loaded, so a restarted gateway never +//! releases running instances. //! -//! Exactly one host serves a gateway: a new attach from a newer connection of -//! the same host process, or from a new host process, supersedes the current -//! session; reports and acks from a superseded stream are dropped, and a -//! delayed hello from an older connection of the current process is refused. +//! Every member names itself in its hello. A new attach from a newer +//! connection of the same member process, or from a new process of the +//! member, supersedes the member's current session; reports and acks from a +//! superseded stream are dropped, and a delayed hello from an older +//! connection of the current process is refused. Several members are served +//! at once. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, Mutex, MutexGuard}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -36,25 +40,38 @@ use tonic::{Status, Streaming}; use crate::proto::{ AssignmentSnapshot, HostAccepted, HostGatewayHeartbeat, HostGatewayMessage, HostHello, - HostMessage, HostStatus, SandboxInstanceAssignment, SandboxInstanceStatus, - host_gateway_message, host_message, + HostMessage, HostStatus, PoolMember, RouteOwner, SandboxInstanceAssignment, + SandboxInstanceStatus, SandboxInstanceTakeover, host_gateway_message, host_message, }; -/// Interval at which both ends send heartbeats on an idle stream. -pub const HOST_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); -/// Silence after which the gateway closes a host stream. -const HOST_SILENCE_TIMEOUT: Duration = Duration::from_mins(1); +/// Recovery deadline a supervisor pool offers each sandbox runtime unless the +/// gateway configures another. +/// +/// Membership expiry, one reconciliation pass, and a replacement instance's +/// startup fit with room to spare. +pub const POOL_RECOVERY_TIMEOUT_SECS: u32 = 60; + +/// Heartbeat interval of pool members, whose liveness drives failover. +pub const POOL_HEARTBEAT_INTERVAL: Duration = crate::pool_timing::HEARTBEAT_INTERVAL; +/// Silence after which the gateway closes a pool member's stream. +const POOL_SILENCE_TIMEOUT: Duration = crate::pool_timing::GATEWAY_SILENCE_TIMEOUT; /// Bound on the host sending its hello. const HELLO_TIMEOUT: Duration = Duration::from_secs(10); -/// Bound on queueing one gateway-to-host message; a host that cannot drain -/// its stream this long is disconnected rather than stalling the session. -const SEND_TIMEOUT: Duration = Duration::from_secs(30); +/// Bound on queueing one gateway-to-host message. A member's served stream +/// keeps its membership record alive, so a member that stops reading is +/// dropped as soon as one that stops writing. +const SEND_TIMEOUT: Duration = crate::pool_timing::GATEWAY_SILENCE_TIMEOUT; /// Gateway-to-host messages queued before the stream applies backpressure. const OUTBOUND_QUEUE: usize = 16; +/// How long a withdrawn pool instance stays on its owner's stop list. +const POOL_STOP_RETENTION: Duration = Duration::from_mins(2); /// Most sandbox instances one shared supervisor host runs. The gateway /// refuses to assign beyond it, and the host refuses snapshots beyond it, so /// a sandbox is never created that its host could not accept. pub const MAX_ASSIGNED_INSTANCES: usize = 256; +/// Most sandbox instances a supervisor pool runs across its members; the +/// gateway places at most [`MAX_ASSIGNED_INSTANCES`] on each member. +pub const MAX_POOL_INSTANCES: usize = 16 * MAX_ASSIGNED_INSTANCES; /// Most instances one host status report may describe: every assignment plus /// instances still being removed. pub const MAX_REPORTED_INSTANCES: usize = 4 * MAX_ASSIGNED_INSTANCES; @@ -86,7 +103,8 @@ pub struct SupervisorCredential { /// What the gateway knows about one assigned route on the host. #[derive(Clone, Debug, PartialEq, Eq)] pub enum ObservedInstance { - /// No host session is connected. + /// No host session is connected (for a pool: the owner's process is not + /// connected to this gateway, or the route has no owner yet). NoHost, /// A host is connected but has not reconciled to the snapshot that /// assigned this route. @@ -98,6 +116,118 @@ pub enum ObservedInstance { Reported(SandboxInstanceStatus), } +/// The process and ownership term that supervise one pool instance. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InstanceOwner { + /// Pool member that owns the instance. + pub member_id: String, + /// Host process of that member that owns the term. + pub process_id: String, + /// Ownership term: the sandbox credential epoch of the owner. + pub auth_epoch: u64, + /// Supervisor process identity a term after the launch term was granted + /// for; absent for the launch term. + pub supervisor_instance_id: Option, + /// Credentials of a term after the launch term, for the owner's host. + /// Absent for the launch term, whose credentials the launch material + /// carries. Minted by this gateway only for members it serves. + pub takeover: Option, +} + +impl InstanceOwner { + /// The takeover term an assignment carries: with its credentials when + /// this gateway minted them, else the term alone, so a host already + /// running the instance for that term keeps it. Credentials minted for + /// one term differ between gateway replicas; the term does not. + fn takeover_term(&self) -> Option { + let supervisor_instance_id = self.supervisor_instance_id.as_ref()?; + Some( + self.takeover + .clone() + .unwrap_or_else(|| SandboxInstanceTakeover { + auth_epoch: self.auth_epoch, + supervisor_instance_id: supervisor_instance_id.clone(), + ..SandboxInstanceTakeover::default() + }), + ) + } +} + +/// A live pool member as the gateway's durable membership records show it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MemberView { + pub member_id: String, + /// Host process currently serving the member. + pub process_id: String, + /// Routed runtime listener address of the member. + pub runtime_address: String, + /// The member announced it is leaving the pool. + pub leaving: bool, +} + +/// The gateway's durable view of a supervisor pool, replaced as a whole by +/// the gateway's pool reconciler on every replica. +#[derive(Clone, Default, PartialEq, Eq)] +pub struct PoolView { + pub members: Vec, + /// Owner of each placed route. + pub owners: BTreeMap, + /// Every route the pool must run, from the gateway's durable records. An + /// owned route that is not desired was withdrawn. + pub desired: BTreeSet, + /// Assignments, launch material included, this gateway replica loaded + /// for routes its members own. An owned desired route missing here is + /// kept by its owner as it is. + pub assignments: BTreeMap, + /// Latest report of each route's current owner term, recorded by the + /// gateway replica that serves the owner, for judging instances whose + /// owner streams to another replica. + pub reports: BTreeMap, +} + +impl std::fmt::Debug for PoolView { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PoolView") + .field("members", &self.members) + .field("owners", &self.owners) + .field("desired", &self.desired) + .field("assignments", &self.assignments.keys().collect::>()) + .field("reports", &self.reports.keys().collect::>()) + .finish() + } +} + +/// What a pool member's stream served here last reported of one instance, +/// for the gateway to record durably under the ownership term it belongs to, +/// so other replicas can judge the instance. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MemberReport { + pub member_id: String, + pub process_id: String, + /// Ownership term the reporting stream had applied for the route when it + /// reported: the term of the snapshot that last installed the route on + /// the stream, once the stream acknowledged it. + pub auth_epoch: u64, + /// When the reporting stream was accepted (Unix milliseconds). + pub connected_at_ms: i64, + pub status: SandboxInstanceStatus, +} + +/// A connected pool member's stream, for the gateway's membership records. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ConnectedMember { + pub member_id: String, + pub member_uid: String, + pub process_id: String, + pub runtime_address: String, + /// The member reported that it is leaving. + pub leaving: bool, + /// When this stream was accepted (Unix milliseconds); a newer stream of a + /// member supersedes an older one on another gateway replica. + pub connected_at_ms: i64, +} + struct Assigned { assignment: SandboxInstanceAssignment, /// Snapshot revision that first carried this exact assignment. @@ -109,22 +239,12 @@ struct Desired { initialized: bool, revision: u64, assignments: BTreeMap, - /// Edits made before initialization, applied over the first full list. - pending: BTreeMap>, + /// Withdrawn instances to stop, with the owner and the time of + /// withdrawal. + stopping: BTreeMap, } impl Desired { - fn snapshot(&self) -> Option { - self.initialized.then(|| AssignmentSnapshot { - revision: self.revision, - instances: self - .assignments - .values() - .map(|assigned| assigned.assignment.clone()) - .collect(), - }) - } - /// Insert or replace one assignment; returns whether anything changed. fn put(&mut self, assignment: SandboxInstanceAssignment, revision: u64) -> bool { match self.assignments.get(&assignment.route) { @@ -143,36 +263,100 @@ impl Desired { } } +/// Pool state: the durable view and when each owner term was first recorded +/// here, for telling a reconciling host from one that dropped an instance. +struct Pool { + loaded: bool, + view: PoolView, + owned_since: BTreeMap, + /// Owners placed by launches on this gateway, with when, so a view whose + /// records were listed before a placement does not drop it. + placed: BTreeMap, + /// Instances launched through this gateway, with when, kept over views + /// listed before them. + launched: BTreeMap, + /// Instances withdrawn here whose durable records the gateway has not + /// removed yet; they are not desired whatever a view says. + withdrawals: BTreeSet, + /// Recovery deadline offered to every instance's runtime. + recovery_timeout_secs: u32, +} + +impl Pool { + fn new(recovery_timeout_secs: u32) -> Self { + Self { + loaded: false, + view: PoolView::default(), + owned_since: BTreeMap::new(), + placed: BTreeMap::new(), + launched: BTreeMap::new(), + withdrawals: BTreeSet::new(), + recovery_timeout_secs, + } + } +} + struct HostSession { id: u64, identity: HostIdentity, instance_id: String, connection_epoch: u64, + member: PoolMember, superseded: watch::Sender, /// Highest snapshot revision sent on this stream; acks beyond it are /// impossible and rejected. sent_revision: u64, + /// Latest report of this session. + status: Option, + connected_at_ms: i64, + /// Ownership term of each route the snapshots sent on this stream assign + /// or keep, with the revision that first installed that term. + terms: BTreeMap, + /// Terms of the last snapshot built for this stream, committed to + /// `terms` once that snapshot is sent. + pending_terms: Option, +} + +/// Terms of one snapshot: the revision, the routes it assigns, and the routes +/// it keeps as they are. +struct PendingTerms { + revision: u64, + assigned: BTreeMap, + kept: BTreeMap, +} + +impl HostSession { + fn member_id(&self) -> &str { + &self.member.member_id + } } #[derive(Default)] struct Observed { - session: Option, - /// Latest report of the current session. - status: Option, + sessions: BTreeMap, next_session: u64, } struct Inner { desired: Mutex, + /// The supervisor pool: its durable view and what this gateway placed, + /// launched, and withdrawn since. + pool: Mutex, + /// The gateway's durable placement of new pool instances. + placement: Mutex>>, /// Latest supervisor token the gateway minted; `None` until the first. credential: watch::Sender>, - /// Current desired revision; 0 while uninitialized. + /// Current revision of everything snapshots are derived from; 0 while + /// uninitialized. revision: watch::Sender, + /// Counts changes to the pool members this gateway serves: a member + /// stream attached or ended, or a member announced leaving. + membership: watch::Sender, observed: Mutex, } -/// Gateway-owned desired assignments and observed status of one shared -/// supervisor host. +/// Gateway-owned desired assignments and observed status of the shared +/// supervisor pool. #[derive(Clone)] pub struct SupervisorHostController { inner: Arc, @@ -195,13 +379,9 @@ impl std::fmt::Debug for SupervisorHostController { /// Why a change to the desired set was refused. #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] pub enum AssignmentError { - /// A full resync was derived from a desired set that has since changed: - /// `baseline` is the revision the caller read before deriving its list and - /// `current` the one now in force; derive the list again from it. - #[error("desired assignments changed from revision {baseline} to {current} during the resync")] - Conflict { baseline: u64, current: u64 }, - /// The change would assign more instances than one host runs. - #[error("the shared supervisor hosts at most {MAX_ASSIGNED_INSTANCES} sandbox instances")] + /// The change would assign more instances than the pool runs, or no + /// member has room. + #[error("the shared supervisor has no capacity for another sandbox instance")] Capacity, } @@ -217,13 +397,33 @@ impl SupervisorHostController { Self { inner: Arc::new(Inner { desired: Mutex::new(Desired::default()), + pool: Mutex::new(Pool::new(POOL_RECOVERY_TIMEOUT_SECS)), + placement: Mutex::new(None), credential: watch::Sender::new(None), revision: watch::Sender::new(0), + membership: watch::Sender::new(0), observed: Mutex::new(Observed::default()), }), } } + /// Offer every instance's runtime `recovery_timeout_secs` instead of + /// [`POOL_RECOVERY_TIMEOUT_SECS`]. Call it before any host attaches. + pub fn set_recovery_timeout_secs(&self, recovery_timeout_secs: u32) { + lock(&self.inner.pool).recovery_timeout_secs = recovery_timeout_secs; + } + + /// Install the gateway's durable placement of new pool instances. + pub fn install_placement(&self, placement: Arc) { + *lock(&self.inner.placement) = Some(placement); + } + + /// The installed placement of new pool instances. + #[must_use] + pub fn placement(&self) -> Option> { + lock(&self.inner.placement).clone() + } + /// Publish a freshly minted supervisor token for delivery to the host. pub fn publish_credential(&self, credential: SupervisorCredential) { self.inner.credential.send_replace(Some(credential)); @@ -242,9 +442,8 @@ impl SupervisorHostController { lock(&self.inner.desired).initialized } - /// Current desired revision: 0 while uninitialized. Read it before - /// deriving a full list for [`Self::replace_assignments`]; the - /// coordinator does so through + /// Current desired revision: 0 while uninitialized. The coordinator reads + /// it through /// [`SharedSupervision::begin_recovery`](crate::SharedSupervision::begin_recovery). #[must_use] pub fn revision(&self) -> u64 { @@ -260,91 +459,42 @@ impl SupervisorHostController { .map(|assigned| assigned.assignment.clone()) } - /// Replace the whole desired set with `assignments`, derived from - /// durable records while the revision was `baseline`. - /// - /// Initializes the controller; edits made before initialization are - /// applied on top of the list. After initialization the replacement is - /// compare-and-replace: when the revision moved past `baseline` while - /// the list was being derived, nothing changes and the conflict is - /// returned so the caller derives the list again. - pub fn replace_assignments( - &self, - baseline: u64, - assignments: impl IntoIterator, - ) -> Result { - let mut desired = lock(&self.inner.desired); - if desired.initialized && desired.revision != baseline { - return Err(AssignmentError::Conflict { - baseline, - current: desired.revision, - }); - } - // Resolve the list and the pending overlay first, so a refused - // replacement leaves the current set untouched. - let mut resolved = BTreeMap::new(); - for assignment in assignments { - resolved.insert(assignment.route.clone(), assignment); - } - for (route, edit) in &desired.pending { - match edit { - Some(assignment) => { - resolved.insert(route.clone(), assignment.clone()); - } - None => { - resolved.remove(route); - } - } - } - if resolved.len() > MAX_ASSIGNED_INSTANCES { - return Err(AssignmentError::Capacity); - } - desired.pending.clear(); - let next = desired.revision + 1; - let mut changed = !desired.initialized; - for assignment in resolved.values() { - changed |= desired.put(assignment.clone(), next); - } - let before = desired.assignments.len(); - desired + /// Every current assignment. + #[must_use] + pub fn assignments(&self) -> Vec { + lock(&self.inner.desired) .assignments - .retain(|route, _| resolved.contains_key(route)); - changed |= desired.assignments.len() != before; - desired.initialized = true; - if changed { - desired.revision = next; - self.inner.revision.send_replace(next); - } - Ok(desired.revision) + .values() + .map(|assigned| assigned.assignment.clone()) + .collect() + } + + fn bump(&self, desired: &mut Desired) -> u64 { + desired.revision += 1; + self.inner.revision.send_replace(desired.revision); + desired.revision } - /// Assign one instance. Before initialization the edit is kept until the - /// first full list arrives. Refused when it would exceed the host's - /// capacity, so the caller fails its create before anything runs. + /// Assign one instance, served at once and kept over views listed + /// before it; its durable record makes it desired on every gateway + /// replica. Refused beyond the pool's capacity, so the caller fails its + /// create before anything runs. pub fn assign(&self, assignment: SandboxInstanceAssignment) -> Result<(), AssignmentError> { let mut desired = lock(&self.inner.desired); - if !desired.initialized { - let pending_additions = desired - .pending - .values() - .filter(|edit| edit.is_some()) - .count(); - if !desired.pending.contains_key(&assignment.route) - && pending_additions >= MAX_ASSIGNED_INSTANCES - { - return Err(AssignmentError::Capacity); - } - desired - .pending - .insert(assignment.route.clone(), Some(assignment)); - return Ok(()); - } if !desired.assignments.contains_key(&assignment.route) - && desired.assignments.len() >= MAX_ASSIGNED_INSTANCES + && desired.assignments.len() >= MAX_POOL_INSTANCES { return Err(AssignmentError::Capacity); } + { + let mut pool = lock(&self.inner.pool); + pool.launched + .insert(assignment.route.clone(), tokio::time::Instant::now()); + pool.withdrawals.remove(&assignment.route); + pool.view.desired.insert(assignment.route.clone()); + } let next = desired.revision + 1; + desired.stopping.remove(&assignment.route); if desired.put(assignment, next) { desired.revision = next; self.inner.revision.send_replace(next); @@ -352,35 +502,281 @@ impl SupervisorHostController { Ok(()) } - /// Withdraw one instance so the host stops it. + /// Withdraw one instance so its owner stops it. The withdrawal is also + /// listed for the gateway to remove the instance's durable record (see + /// [`Self::pool_withdrawals`]). pub fn unassign(&self, route: &str) { let mut desired = lock(&self.inner.desired); - if !desired.initialized { - desired.pending.insert(route.to_string(), None); + let owner = { + let mut pool = lock(&self.inner.pool); + pool.launched.remove(route); + pool.withdrawals.insert(route.to_string()); + pool.view.desired.remove(route); + pool.view.owners.get(route).cloned() + }; + let removed = desired.assignments.remove(route).is_some(); + if let Some(owner) = &owner { + desired.stopping.insert( + route.to_string(), + (owner.clone(), tokio::time::Instant::now()), + ); + } + if removed || owner.is_some() { + self.bump(&mut desired); + } + } + + /// Replace the durable pool view, derived from records listed at + /// `listed_at`; the view's loaded assignments become the desired + /// assignments this gateway serves. Owners placed and instances launched + /// here since then are kept until a view includes them, and instances + /// withdrawn here stay withdrawn until their records are gone. An owned + /// instance that leaves the desired set is listed for its owner to stop. + /// Owners whose term is new to this gateway are recorded as assigned at + /// the next revision, so their host is considered reconciling until it + /// acknowledges it. + pub fn update_pool(&self, mut view: PoolView, listed_at: tokio::time::Instant) { + let mut desired = lock(&self.inner.desired); + let mut pool_guard = lock(&self.inner.pool); + let pool = &mut *pool_guard; + pool.placed + .retain(|_, (_, placed_at)| *placed_at >= listed_at); + for (route, (owner, _)) in &pool.placed { + view.owners + .entry(route.clone()) + .or_insert_with(|| owner.clone()); + } + pool.launched + .retain(|_, launched_at| *launched_at >= listed_at); + for route in pool.launched.keys() { + view.desired.insert(route.clone()); + if let Some(assigned) = desired.assignments.get(route) { + view.assignments + .entry(route.clone()) + .or_insert_with(|| assigned.assignment.clone()); + } + } + for route in &pool.withdrawals { + view.desired.remove(route); + view.assignments.remove(route); + } + let now = tokio::time::Instant::now(); + let stopping_before = desired.stopping.len(); + if pool.loaded { + for route in pool.view.desired.difference(&view.desired) { + if let Some(owner) = pool.view.owners.get(route) { + desired + .stopping + .entry(route.clone()) + .or_insert_with(|| (owner.clone(), now)); + } + } + } + // A desired route is not being stopped, whatever an earlier + // withdrawal recorded. + desired + .stopping + .retain(|route, _| !view.desired.contains(route)); + let stops_changed = desired.stopping.len() != stopping_before; + if pool.loaded && desired.initialized && pool.view == view && !stops_changed { return; } - if desired.assignments.remove(route).is_some() { - desired.revision += 1; - let revision = desired.revision; - self.inner.revision.send_replace(revision); + let next = desired.revision + 1; + let previous = std::mem::take(&mut pool.owned_since); + for (route, owner) in &view.owners { + let unchanged = pool.view.owners.get(route).is_some_and(|current| { + current.member_id == owner.member_id + && current.process_id == owner.process_id + && current.auth_epoch == owner.auth_epoch + }); + let since = previous + .get(route) + .copied() + .filter(|_| unchanged) + .unwrap_or(next); + pool.owned_since.insert(route.clone(), since); + } + let assignments = view + .assignments + .iter() + .map(|(route, assignment)| { + let revision = desired + .assignments + .get(route) + .filter(|assigned| assigned.assignment == *assignment) + .map_or(next, |assigned| assigned.revision); + ( + route.clone(), + Assigned { + assignment: assignment.clone(), + revision, + }, + ) + }) + .collect(); + desired.assignments = assignments; + desired.initialized = true; + pool.view = view; + pool.loaded = true; + drop(pool_guard); + desired.revision = next; + self.inner.revision.send_replace(next); + } + + /// Routes withdrawn here whose durable records the gateway still has to + /// remove; see [`Self::forget_withdrawal`]. + #[must_use] + pub fn pool_withdrawals(&self) -> Vec { + lock(&self.inner.pool).withdrawals.iter().cloned().collect() + } + + /// The durable record of a route withdrawn here is gone. + pub fn forget_withdrawal(&self, route: &str) { + lock(&self.inner.pool).withdrawals.remove(route); + } + + /// Recovery deadline the pool offers each instance's runtime. + #[must_use] + pub fn pool_recovery_timeout_secs(&self) -> u32 { + lock(&self.inner.pool).recovery_timeout_secs + } + + /// What the pool members served here last reported of their instances. + /// The caller attributes each report to the owner term its own records + /// show, and drops reports of routes the reporting process does not own. + #[must_use] + pub fn member_reports(&self) -> Vec { + let observed = lock(&self.inner.observed); + let mut reports = Vec::new(); + for session in observed.sessions.values() { + let member_id = session.member_id(); + let Some(status) = &session.status else { + continue; + }; + // A report counts only for a term the stream acknowledged; a + // status from before the term's snapshot describes an earlier + // term, even from the same process. + reports.extend(status.instances.iter().filter_map(|instance| { + let (auth_epoch, since) = session.terms.get(&instance.route)?; + (status.applied_revision >= *since).then(|| MemberReport { + member_id: member_id.to_string(), + process_id: session.instance_id.clone(), + auth_epoch: *auth_epoch, + connected_at_ms: session.connected_at_ms, + status: instance.clone(), + }) + })); + } + reports + } + + /// Pool routes withdrawn for stopping; they must not be placed or taken + /// over again. + #[must_use] + pub fn stopping_routes(&self) -> BTreeSet { + let mut desired = lock(&self.inner.desired); + let now = tokio::time::Instant::now(); + desired + .stopping + .retain(|_, (_, since)| now.duration_since(*since) < POOL_STOP_RETENTION); + desired.stopping.keys().cloned().collect() + } + + /// Record a newly placed owner without waiting for the next pool + /// reconciliation, so a launched instance is served at once. + pub fn place_owner(&self, route: &str, owner: InstanceOwner) { + let mut desired = lock(&self.inner.desired); + let mut pool_guard = lock(&self.inner.pool); + let pool = &mut *pool_guard; + pool.placed.insert( + route.to_string(), + (owner.clone(), tokio::time::Instant::now()), + ); + // Until the first load, the reconciler's view will include it. + if !pool.loaded || pool.view.owners.get(route) == Some(&owner) { + return; } + let next = desired.revision + 1; + pool.view.owners.insert(route.to_string(), owner); + pool.owned_since.insert(route.to_string(), next); + drop(pool_guard); + desired.revision = next; + self.inner.revision.send_replace(next); + } + + /// The durable pool view this gateway serves, once loaded. + #[must_use] + pub fn pool_view(&self) -> Option { + let pool = lock(&self.inner.pool); + pool.loaded.then(|| pool.view.clone()) + } + + /// Pool members whose streams this gateway serves, for its durable + /// membership records. + #[must_use] + pub fn connected_members(&self) -> Vec { + lock(&self.inner.observed) + .sessions + .values() + .map(|session| ConnectedMember { + member_id: session.member.member_id.clone(), + member_uid: session.member.member_uid.clone(), + process_id: session.instance_id.clone(), + runtime_address: session.member.runtime_address.clone(), + leaving: session.status.as_ref().is_some_and(|status| status.leaving), + connected_at_ms: session.connected_at_ms, + }) + .collect() + } + + /// Changes whenever [`Self::connected_members`] gains or loses a stream + /// or a member announces leaving, so membership records follow at once. + #[must_use] + pub fn membership_changes(&self) -> watch::Receiver { + self.inner.membership.subscribe() + } + + fn membership_changed(&self) { + self.inner + .membership + .send_modify(|changes| *changes = changes.wrapping_add(1)); } /// Whether a host session is currently connected. #[must_use] pub fn host_connected(&self) -> bool { - lock(&self.inner.observed).session.is_some() + !lock(&self.inner.observed).sessions.is_empty() } - /// What the connected host reports for `route`. + /// What the pool reports for `route`: the session of the route's owning + /// process, or the owner's recorded report when it streams elsewhere. #[must_use] pub fn observed(&self, route: &str) -> ObservedInstance { let desired = lock(&self.inner.desired); + let pool = lock(&self.inner.pool); let observed = lock(&self.inner.observed); - if observed.session.is_none() { + let Some(owner) = pool.view.owners.get(route) else { return ObservedInstance::NoHost; - } - let Some(status) = &observed.status else { + }; + let Some(session) = observed.sessions.values().find(|session| { + session.member_id() == owner.member_id && session.instance_id == owner.process_id + }) else { + // The owner streams to another gateway replica, which records its + // reports for every replica. + return pool + .view + .reports + .get(route) + .map_or(ObservedInstance::NoHost, |status| { + ObservedInstance::Reported(status.clone()) + }); + }; + let assigned = desired + .assignments + .get(route) + .map(|assigned| assigned.revision); + let assigned_revision = assigned.max(pool.owned_since.get(route).copied()); + let Some(status) = &session.status else { return ObservedInstance::Pending; }; if let Some(instance) = status @@ -390,16 +786,14 @@ impl SupervisorHostController { { return ObservedInstance::Reported(instance.clone()); } - match desired.assignments.get(route) { - Some(assigned) if assigned.revision > status.applied_revision => { - ObservedInstance::Pending - } + match assigned_revision { + Some(revision) if revision > status.applied_revision => ObservedInstance::Pending, _ => ObservedInstance::Absent, } } /// Serve one authenticated host stream. Reads the hello, supersedes any - /// current session, and returns the gateway-to-host stream. + /// session it replaces, and returns the gateway-to-host stream. pub async fn attach( &self, identity: HostIdentity, @@ -431,6 +825,11 @@ impl SupervisorHostController { "supervisor host hello names no instance", )); } + let Some(member) = &hello.member else { + return Err(Status::failed_precondition(MISSING_MEMBER)); + }; + validate_member(member).map_err(Status::invalid_argument)?; + let heartbeat = POOL_HEARTBEAT_INTERVAL; let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); let expires_at = identity.expires_at; let (session_id, superseded) = self.begin_session(identity, &hello)?; @@ -438,7 +837,7 @@ impl SupervisorHostController { payload: Some(host_gateway_message::Payload::Accepted(HostAccepted { session_id: session_id.to_string(), heartbeat_interval: Some(prost_types::Duration { - seconds: i64::try_from(HOST_HEARTBEAT_INTERVAL.as_secs()).unwrap_or(15), + seconds: i64::try_from(heartbeat.as_secs()).unwrap_or(15), nanos: 0, }), })), @@ -455,17 +854,31 @@ impl SupervisorHostController { Ok(ReceiverStream::new(rx)) } - /// Start a session for `hello`, superseding the current one. A hello from - /// the current host process with a connection epoch that is not newer is - /// a delayed or replayed connection and is refused. + /// Start a session for `hello`, superseding the same member's session. A + /// hello that names no member is refused. A hello from the current host + /// process with a connection epoch that is not newer is a delayed or + /// replayed connection and is refused. fn begin_session( &self, identity: HostIdentity, hello: &HostHello, ) -> Result<(u64, watch::Receiver), Status> { + let member = hello + .member + .clone() + .ok_or_else(|| Status::failed_precondition(MISSING_MEMBER))?; let mut observed = lock(&self.inner.observed); - if let Some(current) = &observed.session - && current.instance_id == hello.instance_id + let member_id = member.member_id.clone(); + let replaced = observed + .sessions + .values() + .filter(|session| session.member_id() == member_id) + .map(|session| session.id) + .collect::>(); + if let Some(current) = replaced + .iter() + .filter_map(|id| observed.sessions.get(id)) + .find(|session| session.instance_id == hello.instance_id) && hello.connection_epoch <= current.connection_epoch { return Err(Status::failed_precondition( @@ -474,60 +887,97 @@ impl SupervisorHostController { } observed.next_session += 1; let id = observed.next_session; - if let Some(previous) = observed.session.take() { - tracing::info!( - subject = %previous.identity.subject, - instance_id = %previous.instance_id, - connection_epoch = previous.connection_epoch, - "superseding supervisor host session" - ); - previous.superseded.send_replace(true); + for previous in replaced { + if let Some(previous) = observed.sessions.remove(&previous) { + tracing::info!( + subject = %previous.identity.subject, + instance_id = %previous.instance_id, + member = previous.member_id(), + connection_epoch = previous.connection_epoch, + "superseding supervisor host session" + ); + previous.superseded.send_replace(true); + } } tracing::info!( subject = %identity.subject, instance_id = %hello.instance_id, + member = %member_id, connection_epoch = hello.connection_epoch, session = id, "supervisor host attached" ); let (superseded, watch) = watch::channel(false); - observed.session = Some(HostSession { + observed.sessions.insert( id, - identity, - instance_id: hello.instance_id.clone(), - connection_epoch: hello.connection_epoch, - superseded, - sent_revision: 0, - }); - observed.status = None; + HostSession { + id, + identity, + instance_id: hello.instance_id.clone(), + connection_epoch: hello.connection_epoch, + member, + superseded, + sent_revision: 0, + status: None, + connected_at_ms: unix_millis(), + terms: BTreeMap::new(), + pending_terms: None, + }, + ); + drop(observed); + self.membership_changed(); Ok((id, watch)) } fn end_session(&self, id: u64) { let mut observed = lock(&self.inner.observed); - if observed - .session - .as_ref() - .is_some_and(|session| session.id == id) - { - let session = observed.session.take(); - observed.status = None; - if let Some(session) = session { - tracing::info!( - subject = %session.identity.subject, - instance_id = %session.instance_id, - session = id, - "supervisor host detached; hosted instances keep running" - ); - } - } + let Some(session) = observed.sessions.remove(&id) else { + return; + }; + drop(observed); + tracing::info!( + subject = %session.identity.subject, + instance_id = %session.instance_id, + member = session.member_id(), + session = id, + "supervisor host detached; hosted instances keep running" + ); + self.membership_changed(); } - /// Note that `revision` was sent on session `id`. + /// Note that `revision` was sent on session `id`. A route keeps the + /// revision that installed its term while snapshots repeat the term; a + /// new term is installed by `revision`; a kept route keeps only the term + /// it already had, and a route the snapshot omits has none. fn record_sent(&self, id: u64, revision: u64) { let mut observed = lock(&self.inner.observed); - if let Some(session) = observed.session.as_mut().filter(|session| session.id == id) { + if let Some(session) = observed.sessions.get_mut(&id) { session.sent_revision = session.sent_revision.max(revision); + if let Some(pending) = session + .pending_terms + .take() + .filter(|pending| pending.revision == revision) + { + let previous = std::mem::take(&mut session.terms); + let installed = |route: &String, epoch: u64| { + previous + .get(route) + .filter(|(installed, _)| *installed == epoch) + .map(|(_, since)| *since) + }; + session.terms = pending + .assigned + .iter() + .map(|(route, epoch)| { + let since = installed(route, *epoch).unwrap_or(revision); + (route.clone(), (*epoch, since)) + }) + .chain(pending.kept.iter().filter_map(|(route, epoch)| { + let since = installed(route, *epoch)?; + Some((route.clone(), (*epoch, since))) + })) + .collect(); + } } } @@ -536,7 +986,7 @@ impl SupervisorHostController { /// oversized reports are dropped. fn record_status(&self, id: u64, status: HostStatus) -> Result<(), &'static str> { let mut observed = lock(&self.inner.observed); - let Some(session) = observed.session.as_ref().filter(|session| session.id == id) else { + let Some(session) = observed.sessions.get_mut(&id) else { return Err("report from a superseded supervisor host session"); }; if status.applied_revision > session.sent_revision { @@ -545,19 +995,135 @@ impl SupervisorHostController { if status.instances.len() > MAX_REPORTED_INSTANCES { return Err("supervisor host status reports too many instances"); } - if observed + if session .status .as_ref() .is_some_and(|previous| status.applied_revision < previous.applied_revision) { return Err("supervisor host acknowledged an older revision than before"); } - observed.status = Some(status); + let leaving = status.leaving + && !session + .status + .as_ref() + .is_some_and(|previous| previous.leaving); + session.status = Some(status); + drop(observed); + if leaving { + // Membership records follow the leaving flag at once. + self.membership_changed(); + let mut desired = lock(&self.inner.desired); + self.bump(&mut desired); + } Ok(()) } - fn snapshot(&self) -> Option { - lock(&self.inner.desired).snapshot() + /// The snapshot for session `id`'s member: the instances its process + /// owns, the stop list, and the route directory. `None` until it can be + /// authoritative. + fn snapshot(&self, id: u64) -> Option { + let mut desired = lock(&self.inner.desired); + if !desired.initialized { + return None; + } + let pool = lock(&self.inner.pool); + if !pool.loaded { + return None; + } + let mut observed = lock(&self.inner.observed); + let session = observed.sessions.get_mut(&id)?; + let member_id = session.member_id().to_string(); + let instance_id = session.instance_id.clone(); + let owned_by_session = + |owner: &InstanceOwner| owner.member_id == member_id && owner.process_id == instance_id; + let mut assigned_terms = BTreeMap::new(); + let instances = desired + .assignments + .values() + .filter_map(|assigned| { + let owner = pool.view.owners.get(&assigned.assignment.route)?; + owned_by_session(owner).then(|| { + assigned_terms.insert(assigned.assignment.route.clone(), owner.auth_epoch); + SandboxInstanceAssignment { + takeover: owner.takeover_term(), + ..assigned.assignment.clone() + } + }) + }) + .collect(); + let now = tokio::time::Instant::now(); + desired + .stopping + .retain(|_, (_, since)| now.duration_since(*since) < POOL_STOP_RETENTION); + // Withdrawn here, or owned yet no longer desired anywhere. + let stop_routes = desired + .stopping + .iter() + .filter(|(_, (owner, _))| owner.member_id == *member_id) + .map(|(route, _)| route.clone()) + .chain( + pool.view + .owners + .iter() + .filter(|(route, owner)| { + owner.member_id == *member_id + && !pool.view.desired.contains(*route) + && !desired.stopping.contains_key(*route) + }) + .map(|(route, _)| route.clone()), + ) + .collect::>() + .into_iter() + .collect(); + let addresses = pool + .view + .members + .iter() + .map(|member| (member.member_id.as_str(), member.runtime_address.as_str())) + .collect::>(); + let route_directory = pool + .view + .owners + .iter() + .filter_map(|(route, owner)| { + Some(RouteOwner { + route: route.clone(), + runtime_address: (*addresses.get(owner.member_id.as_str())?).to_string(), + }) + }) + .collect(); + // Owned desired instances whose launch material this gateway has not + // loaded yet are kept as they are rather than released. + let kept_terms = pool + .view + .owners + .iter() + .filter(|(route, owner)| { + owned_by_session(owner) + && pool.view.desired.contains(*route) + && !desired.assignments.contains_key(*route) + && !desired.stopping.contains_key(*route) + }) + .map(|(route, owner)| (route.clone(), owner.auth_epoch)) + .collect::>(); + let keep_routes = kept_terms.keys().cloned().collect(); + session.pending_terms = Some(PendingTerms { + revision: desired.revision, + assigned: assigned_terms, + kept: kept_terms, + }); + Some(AssignmentSnapshot { + revision: desired.revision, + instances, + pool: true, + stop_routes, + route_directory, + keep_routes, + recovery_timeout: Some(prost_types::Duration { + seconds: i64::from(pool.recovery_timeout_secs), + nanos: 0, + }), + }) } async fn run_session( @@ -573,10 +1139,13 @@ impl SupervisorHostController { // reconnects with the renewed token it has been delivered. let credential_expiry = tokio::time::sleep(credential_remaining(expires_at)); tokio::pin!(credential_expiry); + let (heartbeat_interval, silence_timeout, send_timeout) = + (POOL_HEARTBEAT_INTERVAL, POOL_SILENCE_TIMEOUT, SEND_TIMEOUT); let mut revision = self.inner.revision.subscribe(); + let mut last_snapshot: Option = None; let mut sent_revision = 0; let mut last_heard = tokio::time::Instant::now(); - let mut heartbeat = tokio::time::interval(HOST_HEARTBEAT_INTERVAL); + let mut heartbeat = tokio::time::interval(heartbeat_interval); heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); // A send that cannot complete within the bound ends the session, so a // stalled host can neither block supersession nor delay heartbeats. @@ -587,7 +1156,7 @@ impl SupervisorHostController { Ok(HostGatewayMessage { payload: Some(message), }), - SEND_TIMEOUT, + send_timeout, ) .await .is_ok() @@ -595,21 +1164,31 @@ impl SupervisorHostController { }; loop { let current = *revision.borrow_and_update(); - if current > sent_revision { - let Some(snapshot) = self.snapshot() else { - continue; - }; - tracing::debug!( - session = session_id, - revision = snapshot.revision, - instances = snapshot.instances.len(), - "sending supervisor host assignment snapshot" - ); - let revision_sent = snapshot.revision; - if !send(host_gateway_message::Payload::Assignments(snapshot)).await { - return; + if current > sent_revision + && let Some(snapshot) = self.snapshot(session_id) + { + // Members are only sent snapshots that change something they + // act on. + let unchanged = last_snapshot.as_ref().is_some_and(|previous| { + previous.instances == snapshot.instances + && previous.stop_routes == snapshot.stop_routes + && previous.route_directory == snapshot.route_directory + && previous.keep_routes == snapshot.keep_routes + }); + if !unchanged { + tracing::debug!( + session = session_id, + revision = snapshot.revision, + instances = snapshot.instances.len(), + "sending supervisor host assignment snapshot" + ); + let revision_sent = snapshot.revision; + if !send(host_gateway_message::Payload::Assignments(snapshot.clone())).await { + return; + } + self.record_sent(session_id, revision_sent); + last_snapshot = Some(snapshot); } - self.record_sent(session_id, revision_sent); sent_revision = current; } tokio::select! { @@ -634,7 +1213,7 @@ impl SupervisorHostController { Err(Status::invalid_argument( "unexpected message on supervisor host stream", )), - SEND_TIMEOUT, + send_timeout, ) .await; return; @@ -646,7 +1225,7 @@ impl SupervisorHostController { None => return, }, _ = heartbeat.tick() => { - if last_heard.elapsed() > HOST_SILENCE_TIMEOUT { + if last_heard.elapsed() > silence_timeout { tracing::warn!(session = session_id, "supervisor host stream went silent; closing it"); return; } @@ -664,6 +1243,50 @@ impl SupervisorHostController { } } +/// Why a hello without a pool member is refused. +const MISSING_MEMBER: &str = "the shared supervisor did not name its pool member; set \ +OPENSHELL_SUPERVISOR_POOL_MEMBER_ID, OPENSHELL_SUPERVISOR_POOL_MEMBER_UID, and \ +OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP (the OpenShell Helm chart sets them from the Pod name, \ +UID, and IP)"; + +/// Check a pool member's identity, naming the setting that is wrong. +fn validate_member(member: &PoolMember) -> Result<(), &'static str> { + let valid = |value: &str| { + !value.is_empty() && value.len() <= 253 && !value.chars().any(char::is_whitespace) + }; + if !valid(&member.member_id) { + return Err( + "the shared supervisor's pool member id is empty or invalid; set \ +OPENSHELL_SUPERVISOR_POOL_MEMBER_ID to its Pod name", + ); + } + if !valid(&member.member_uid) { + return Err( + "the shared supervisor's pool member UID is empty or invalid; set \ +OPENSHELL_SUPERVISOR_POOL_MEMBER_UID to its Pod UID", + ); + } + if member + .runtime_address + .parse::() + .is_err() + { + return Err( + "the shared supervisor's pool member runtime address is invalid; set \ +OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP to its Pod IP", + ); + } + Ok(()) +} + +fn unix_millis() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| { + i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX) + }) +} + /// Time until `expires_at` (Unix seconds); zero once it has passed. fn credential_remaining(expires_at: i64) -> Duration { let now = SystemTime::now() @@ -684,11 +1307,13 @@ mod tests { sandbox_id: format!("sandbox-{route}"), generation: "gen".to_string(), launch_material: material.to_vec(), + takeover: None, } } fn status(applied_revision: u64, routes: &[&str]) -> HostStatus { HostStatus { + leaving: false, applied_revision, instances: routes .iter() @@ -705,106 +1330,35 @@ mod tests { } #[test] - fn full_resyncs_initialize_once_and_then_compare_and_replace() { + fn assignments_are_refused_beyond_the_pools_capacity() { let controller = SupervisorHostController::new(); - let revision = controller.inner.revision.subscribe(); - // Edits before the first full list wait for it. - controller - .assign(assignment("rt-created-before-list", b"a")) - .unwrap(); - controller.unassign("rt-stopped-before-list"); - assert!(!controller.is_initialized()); - assert!(controller.snapshot().is_none()); - assert!(!revision.has_changed().unwrap()); - - let baseline = controller.revision(); - assert_eq!( - controller.replace_assignments( - baseline, - [ - assignment("rt-stopped-before-list", b"old"), - assignment("rt-listed", b"b"), - ], - ), - Ok(1) - ); - let snapshot = controller.snapshot().unwrap(); - assert_eq!(snapshot.revision, 1); - assert_eq!( - snapshot - .instances - .iter() - .map(|instance| instance.route.as_str()) - .collect::>(), - ["rt-created-before-list", "rt-listed"] - ); - assert!(revision.has_changed().unwrap()); - - // An identical list changes nothing. - assert_eq!( - controller.replace_assignments( - 1, - [ - assignment("rt-created-before-list", b"a"), - assignment("rt-listed", b"b"), - ], - ), - Ok(1) - ); - controller.assign(assignment("rt-listed", b"b")).unwrap(); - assert_eq!(controller.revision(), 1); - - // A list derived at revision 1 that races a stop cannot resurrect the - // stopped instance; the caller derives it again. - controller.unassign("rt-listed"); - assert_eq!(controller.revision(), 2); - assert_eq!( - controller.replace_assignments( - 1, - [ - assignment("rt-created-before-list", b"a"), - assignment("rt-listed", b"b"), - ], - ), - Err(AssignmentError::Conflict { - baseline: 1, - current: 2 - }) - ); - assert!(controller.assignment("rt-listed").is_none()); - assert_eq!(controller.replace_assignments(2, []), Ok(3)); - assert!(controller.snapshot().unwrap().instances.is_empty()); - - // The gateway refuses the assignment its host could not accept. - let full = - (0..MAX_ASSIGNED_INSTANCES).map(|index| assignment(&format!("rt-{index}"), b"x")); - assert_eq!(controller.replace_assignments(3, full), Ok(4)); + for index in 0..MAX_POOL_INSTANCES { + controller + .assign(assignment(&format!("rt-{index}"), b"x")) + .unwrap(); + } assert_eq!( controller.assign(assignment("rt-one-too-many", b"x")), Err(AssignmentError::Capacity) ); + // Assigning a route again is not a new instance. assert_eq!(controller.assign(assignment("rt-0", b"x")), Ok(())); } #[test] fn sessions_fence_stale_streams_epochs_and_acks() { let controller = SupervisorHostController::new(); - controller - .replace_assignments(0, [assignment("rt-a", b"a")]) - .unwrap(); + controller.update_pool( + loaded_view( + vec![member_view("m1", "boot-1")], + [("rt-a", owner("m1", "boot-1", 1))], + ), + tokio::time::Instant::now(), + ); assert_eq!(controller.observed("rt-a"), ObservedInstance::NoHost); - let identity = HostIdentity { - subject: "openshell-supervisor:test".to_string(), - token_id: "token".to_string(), - expires_at: i64::MAX, - }; - let hello = |instance_id: &str, connection_epoch| HostHello { - instance_id: instance_id.to_string(), - connection_epoch, - }; let (first, superseded) = controller - .begin_session(identity.clone(), &hello("boot-1", 1)) + .begin_session(identity(), &member_hello("m1", "boot-1", 1)) .unwrap(); assert_eq!(controller.observed("rt-a"), ObservedInstance::Pending); @@ -812,20 +1366,25 @@ mod tests { // refused; a newer one supersedes the stream. assert!( controller - .begin_session(identity.clone(), &hello("boot-1", 1)) + .begin_session(identity(), &member_hello("m1", "boot-1", 1)) .is_err() ); let (current, _) = controller - .begin_session(identity.clone(), &hello("boot-1", 2)) + .begin_session(identity(), &member_hello("m1", "boot-1", 2)) .unwrap(); assert!(*superseded.borrow()); assert!(controller.record_status(first, status(0, &[])).is_err()); // Acks are bounded by what this stream was sent and never regress. - assert!(controller.record_status(current, status(1, &[])).is_err()); - controller.record_sent(current, 1); + let revision = controller.revision(); + assert!( + controller + .record_status(current, status(revision, &[])) + .is_err() + ); + controller.record_sent(current, revision); controller - .record_status(current, status(1, &["rt-a"])) + .record_status(current, status(revision, &["rt-a"])) .unwrap(); assert!(controller.record_status(current, status(0, &[])).is_err()); assert!(matches!( @@ -833,21 +1392,545 @@ mod tests { ObservedInstance::Reported(status) if status.ready )); - // A newer assignment is pending until the host reconciles past it. + // A newer instance of the process is pending until its member + // reconciles past it; a route without an owner has no verdict. controller.assign(assignment("rt-b", b"b")).unwrap(); + controller.place_owner("rt-b", owner("m1", "boot-1", 1)); assert_eq!(controller.observed("rt-b"), ObservedInstance::Pending); - assert_eq!(controller.observed("rt-unknown"), ObservedInstance::Absent); - controller.record_sent(current, 2); - controller.record_status(current, status(2, &[])).unwrap(); + assert_eq!(controller.observed("rt-unknown"), ObservedInstance::NoHost); + let next = controller.revision(); + controller.record_sent(current, next); + controller + .record_status(current, status(next, &[])) + .unwrap(); assert_eq!(controller.observed("rt-b"), ObservedInstance::Absent); - // A new host process always takes over; ending a stale session is a no-op. + // A new process of the member always takes over; ending a stale + // session is a no-op. let (replacement, _) = controller - .begin_session(identity, &hello("boot-2", 1)) + .begin_session(identity(), &member_hello("m1", "boot-2", 1)) .unwrap(); controller.end_session(current); assert!(controller.host_connected()); controller.end_session(replacement); assert_eq!(controller.observed("rt-a"), ObservedInstance::NoHost); } + + fn member_hello(member: &str, process: &str, connection_epoch: u64) -> HostHello { + HostHello { + instance_id: process.to_string(), + connection_epoch, + member: Some(PoolMember { + member_id: member.to_string(), + member_uid: format!("uid-{member}"), + runtime_address: format!("10.0.0.{}:5500", member.len()), + }), + } + } + + fn owner(member: &str, process: &str, auth_epoch: u64) -> InstanceOwner { + InstanceOwner { + member_id: member.to_string(), + process_id: process.to_string(), + auth_epoch, + supervisor_instance_id: (auth_epoch > 1).then(|| format!("instance-{auth_epoch}")), + takeover: (auth_epoch > 1).then(|| SandboxInstanceTakeover { + auth_epoch, + supervisor_instance_id: format!("instance-{auth_epoch}"), + grant: format!("grant-{auth_epoch}"), + ..SandboxInstanceTakeover::default() + }), + } + } + + fn member_view(member: &str, process: &str) -> MemberView { + MemberView { + member_id: member.to_string(), + process_id: process.to_string(), + runtime_address: format!("10.0.0.{}:5500", member.len()), + leaving: false, + } + } + + fn identity() -> HostIdentity { + HostIdentity { + subject: "openshell-supervisor:test".to_string(), + token_id: "token".to_string(), + expires_at: i64::MAX, + } + } + + /// A pool view in which every owned route is desired and loaded. + fn loaded_view( + members: Vec, + owners: impl IntoIterator, + ) -> PoolView { + let owners = owners + .into_iter() + .map(|(route, owner)| (route.to_string(), owner)) + .collect::>(); + PoolView { + members, + desired: owners.keys().cloned().collect(), + assignments: owners + .keys() + .map(|route| (route.clone(), assignment(route, route.as_bytes()))) + .collect(), + owners, + reports: BTreeMap::new(), + } + } + + fn routes(snapshot: &AssignmentSnapshot) -> Vec<&str> { + snapshot + .instances + .iter() + .map(|instance| instance.route.as_str()) + .collect() + } + + #[test] + fn membership_changes_follow_member_streams_and_leaving() { + let controller = SupervisorHostController::new(); + let mut changes = controller.membership_changes(); + changes.mark_unchanged(); + let (first, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 1)) + .unwrap(); + assert!(changes.has_changed().unwrap(), "a member attached"); + changes.mark_unchanged(); + // A newer stream of the same member supersedes the first. + let (second, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 2)) + .unwrap(); + assert!(changes.has_changed().unwrap(), "a member reconnected"); + changes.mark_unchanged(); + // Ending the superseded stream changes nothing anyone serves. + controller.end_session(first); + assert!(!changes.has_changed().unwrap()); + let status = |leaving| HostStatus { + applied_revision: 0, + instances: Vec::new(), + leaving, + }; + controller.record_status(second, status(false)).unwrap(); + assert!(!changes.has_changed().unwrap(), "an ordinary report"); + controller.record_status(second, status(true)).unwrap(); + assert!(changes.has_changed().unwrap(), "the member is leaving"); + changes.mark_unchanged(); + controller.end_session(second); + assert!(changes.has_changed().unwrap(), "the member's stream ended"); + } + + #[tokio::test] + async fn pool_members_receive_only_what_their_process_owns() { + let controller = SupervisorHostController::new(); + controller.set_recovery_timeout_secs(75); + let (one, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 1)) + .unwrap(); + let (two, _) = controller + .begin_session(identity(), &member_hello("m-2", "p2", 1)) + .unwrap(); + // Members coexist; nothing is authoritative before ownership loads, + // so a restarted gateway never makes a member release its instances. + assert!(controller.snapshot(one).is_none()); + assert_eq!(controller.connected_members().len(), 2); + + controller.update_pool( + loaded_view( + vec![member_view("m1", "p1"), member_view("m-2", "p2")], + [ + ("rt-a", owner("m1", "p1", 1)), + ("rt-b", owner("m-2", "p2", 1)), + ], + ), + tokio::time::Instant::now(), + ); + let first = controller.snapshot(one).unwrap(); + assert!(first.pool); + // Every member offers its instances' runtimes the pool's deadline. + assert_eq!( + first.recovery_timeout, + Some(prost_types::Duration { + seconds: 75, + nanos: 0 + }) + ); + assert_eq!(routes(&first), ["rt-a"]); + assert!(first.instances[0].takeover.is_none()); + assert_eq!(routes(&controller.snapshot(two).unwrap()), ["rt-b"]); + // Every member can redirect runtimes to any instance's owner. + assert_eq!( + first + .route_directory + .iter() + .map(|entry| (entry.route.as_str(), entry.runtime_address.as_str())) + .collect::>(), + [("rt-a", "10.0.0.2:5500"), ("rt-b", "10.0.0.3:5500")] + ); + + // A takeover moves rt-b to m1 with the new term's credentials; m2's + // snapshot omits it, which releases it without stopping it. + controller.update_pool( + loaded_view( + vec![member_view("m1", "p1"), member_view("m-2", "p2")], + [ + ("rt-a", owner("m1", "p1", 1)), + ("rt-b", owner("m1", "p1", 2)), + ], + ), + tokio::time::Instant::now(), + ); + let moved = controller.snapshot(one).unwrap(); + assert_eq!(routes(&moved), ["rt-a", "rt-b"]); + assert_eq!( + moved.instances[1] + .takeover + .as_ref() + .map(|takeover| takeover.auth_epoch), + Some(2) + ); + let released = controller.snapshot(two).unwrap(); + assert!(released.instances.is_empty()); + assert!(released.stop_routes.is_empty()); + + // A term whose credentials another gateway replica minted still names + // its term, so a host that runs the instance for it keeps it. + let mut unminted = owner("m1", "p1", 2); + unminted.takeover = None; + controller.update_pool( + loaded_view( + vec![member_view("m1", "p1"), member_view("m-2", "p2")], + [("rt-a", owner("m1", "p1", 1)), ("rt-b", unminted)], + ), + tokio::time::Instant::now(), + ); + let term = controller.snapshot(one).unwrap().instances[1] + .takeover + .clone() + .expect("takeover term"); + assert_eq!( + (term.auth_epoch, term.supervisor_instance_id.as_str()), + (2, "instance-2") + ); + assert!(term.grant.is_empty() && term.gateway_token.is_empty()); + + // Withdrawing an instance lists it for its owner, and only its owner, + // to stop, and for the gateway to remove its durable record. + controller.unassign("rt-a"); + assert_eq!(controller.snapshot(one).unwrap().stop_routes, ["rt-a"]); + assert!(controller.snapshot(two).unwrap().stop_routes.is_empty()); + assert!(controller.stopping_routes().contains("rt-a")); + assert_eq!(controller.pool_withdrawals(), ["rt-a"]); + + // A restarted process of a member supersedes that member's stream + // only; it owns nothing until the gateway grants it a term. + let (restarted, _) = controller + .begin_session(identity(), &member_hello("m1", "p1-restarted", 1)) + .unwrap(); + assert!(controller.snapshot(one).is_none()); + assert!(controller.snapshot(restarted).unwrap().instances.is_empty()); + assert_eq!( + routes(&controller.snapshot(two).unwrap()), + Vec::<&str>::new() + ); + // A delayed hello from an older connection of a member's current + // process is refused. + assert!( + controller + .begin_session(identity(), &member_hello("m-2", "p2", 1)) + .is_err() + ); + } + + #[test] + fn reports_count_only_for_the_term_their_stream_acknowledged() { + let controller = SupervisorHostController::new(); + let (one, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 1)) + .unwrap(); + let members = || vec![member_view("m1", "p1"), member_view("m-2", "p2")]; + let send = |owner: InstanceOwner| { + controller.update_pool( + loaded_view(members(), [("rt-a", owner)]), + tokio::time::Instant::now(), + ); + let revision = controller.snapshot(one).unwrap().revision; + controller.record_sent(one, revision); + revision + }; + let terms = || { + controller + .member_reports() + .into_iter() + .map(|report| (report.status.route, report.auth_epoch)) + .collect::>() + }; + let first = send(owner("m1", "p1", 1)); + controller + .record_status(one, status(first, &["rt-a"])) + .unwrap(); + assert_eq!(terms(), [("rt-a".to_string(), 1)]); + // Another member takes rt-a over, then it returns to the same, + // still running process at a later term, before the process applied + // either snapshot. + send(owner("m-2", "p2", 2)); + assert!(terms().is_empty()); + let third = send(owner("m1", "p1", 3)); + // Its status from term 1 still lists rt-a as ready; that is not term + // 3's readiness. + controller + .record_status(one, status(first, &["rt-a"])) + .unwrap(); + assert!(terms().is_empty()); + controller + .record_status(one, status(third, &["rt-a"])) + .unwrap(); + assert_eq!(terms(), [("rt-a".to_string(), 3)]); + // A later snapshot that repeats the term keeps it installed. + let later = send(owner("m1", "p1", 3)); + assert!(later >= third); + assert_eq!(terms(), [("rt-a".to_string(), 3)]); + } + + #[test] + fn pool_observations_follow_the_owning_process_through_a_takeover() { + let controller = SupervisorHostController::new(); + let (two, _) = controller + .begin_session(identity(), &member_hello("m-2", "p2", 1)) + .unwrap(); + let (one, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 1)) + .unwrap(); + // An unplaced instance has no verdict. + assert_eq!(controller.observed("rt-b"), ObservedInstance::NoHost); + controller.update_pool( + loaded_view( + vec![member_view("m1", "p1"), member_view("m-2", "p2")], + [("rt-b", owner("m-2", "p2", 1))], + ), + tokio::time::Instant::now(), + ); + let revision = controller.snapshot(two).unwrap().revision; + controller.record_sent(two, revision); + controller + .record_status(two, status(revision, &["rt-b"])) + .unwrap(); + assert!(matches!( + controller.observed("rt-b"), + ObservedInstance::Reported(status) if status.ready + )); + // m1 reconciled before the takeover; it does not host rt-b, but that + // is not evidence against the instance. + controller.record_sent(one, revision); + controller + .record_status(one, status(revision, &[])) + .unwrap(); + + controller.update_pool( + loaded_view( + vec![member_view("m1", "p1"), member_view("m-2", "p2")], + [("rt-b", owner("m1", "p1", 2))], + ), + tokio::time::Instant::now(), + ); + // The new owner has not reconciled to the takeover yet. + assert_eq!(controller.observed("rt-b"), ObservedInstance::Pending); + let takeover = controller.snapshot(one).unwrap().revision; + controller.record_sent(one, takeover); + controller + .record_status(one, status(takeover, &["rt-b"])) + .unwrap(); + assert!(matches!( + controller.observed("rt-b"), + ObservedInstance::Reported(_) + )); + // Each stream's reports are offered for recording under the term it + // acknowledged: the new owner's under term 2, and the previous + // owner's, which has not applied its release yet, under term 1. + let mut reports = controller + .member_reports() + .into_iter() + .map(|report| { + ( + report.status.route, + report.member_id, + report.process_id, + report.auth_epoch, + ) + }) + .collect::>(); + reports.sort(); + assert_eq!( + reports, + [ + ("rt-b".to_string(), "m-2".to_string(), "p2".to_string(), 1), + ("rt-b".to_string(), "m1".to_string(), "p1".to_string(), 2), + ] + ); + + // An owner process this gateway does not serve gives no verdict of + // its own; another replica's recorded report of the owner gives one. + let mut elsewhere = loaded_view( + vec![member_view("m1", "p1")], + [("rt-b", owner("m-3", "p3", 3))], + ); + controller.update_pool(elsewhere.clone(), tokio::time::Instant::now()); + assert_eq!(controller.observed("rt-b"), ObservedInstance::NoHost); + let mut recorded = status(0, &["rt-b"]).instances.remove(0); + recorded.ready = false; + elsewhere + .reports + .insert("rt-b".to_string(), recorded.clone()); + controller.update_pool(elsewhere, tokio::time::Instant::now()); + assert_eq!( + controller.observed("rt-b"), + ObservedInstance::Reported(recorded) + ); + } + + #[tokio::test] + async fn pool_views_keep_fresh_placements_and_owned_instances() { + let controller = SupervisorHostController::new(); + let (one, _) = controller + .begin_session(identity(), &member_hello("m1", "p1", 1)) + .unwrap(); + // `loaded` routes carry material this replica read; the rest are + // desired and owned only. + let view = |owned: &[&'static str], loaded: &[&str]| { + let mut view = loaded_view( + vec![member_view("m1", "p1")], + owned.iter().map(|route| (*route, owner("m1", "p1", 1))), + ); + view.assignments + .retain(|route, _| loaded.contains(&route.as_str())); + view + }; + controller.update_pool(view(&["rt-a"], &["rt-a"]), tokio::time::Instant::now()); + + // A launch through this gateway places and assigns rt-b while a + // reconciliation pass that listed records before it runs; the pass's + // view must not release the new instance. + let listed_before_launch = tokio::time::Instant::now(); + controller.place_owner("rt-b", owner("m1", "p1", 1)); + controller.assign(assignment("rt-b", b"b")).unwrap(); + controller.update_pool(view(&["rt-a"], &["rt-a"]), listed_before_launch); + assert_eq!(routes(&controller.snapshot(one).unwrap()), ["rt-a", "rt-b"]); + controller.update_pool( + view(&["rt-a", "rt-b"], &["rt-a", "rt-b"]), + tokio::time::Instant::now(), + ); + assert_eq!(routes(&controller.snapshot(one).unwrap()), ["rt-a", "rt-b"]); + + // An owned instance whose launch material this gateway has not loaded + // yet, such as one launched through another gateway, is kept as is. + controller.update_pool( + view(&["rt-a", "rt-b", "rt-c"], &["rt-a", "rt-b"]), + tokio::time::Instant::now(), + ); + let snapshot = controller.snapshot(one).unwrap(); + assert_eq!(routes(&snapshot), ["rt-a", "rt-b"]); + assert_eq!(snapshot.keep_routes, ["rt-c"]); + assert!( + snapshot + .route_directory + .iter() + .any(|entry| entry.route == "rt-c") + ); + // Once loaded, it is assigned. + controller.update_pool( + view(&["rt-a", "rt-b", "rt-c"], &["rt-a", "rt-b", "rt-c"]), + tokio::time::Instant::now(), + ); + let snapshot = controller.snapshot(one).unwrap(); + assert_eq!(routes(&snapshot), ["rt-a", "rt-b", "rt-c"]); + assert!(snapshot.keep_routes.is_empty()); + + // A withdrawal here is stopped, not kept, and stays withdrawn over + // views that still list its record, until the record is gone. + controller.unassign("rt-b"); + let snapshot = controller.snapshot(one).unwrap(); + assert_eq!(snapshot.stop_routes, ["rt-b"]); + controller.update_pool( + view(&["rt-a", "rt-b", "rt-c"], &["rt-a", "rt-b", "rt-c"]), + tokio::time::Instant::now(), + ); + let snapshot = controller.snapshot(one).unwrap(); + assert_eq!(routes(&snapshot), ["rt-a", "rt-c"]); + assert_eq!(snapshot.stop_routes, ["rt-b"]); + controller.forget_withdrawal("rt-b"); + assert!(controller.pool_withdrawals().is_empty()); + + // An instance withdrawn through another replica leaves the desired + // set while its owner record remains: its owner stops it. + let mut withdrawn_elsewhere = view(&["rt-a", "rt-c"], &["rt-a", "rt-c"]); + withdrawn_elsewhere.desired.remove("rt-c"); + withdrawn_elsewhere.assignments.remove("rt-c"); + controller.update_pool(withdrawn_elsewhere, tokio::time::Instant::now()); + let snapshot = controller.snapshot(one).unwrap(); + assert_eq!(routes(&snapshot), ["rt-a"]); + assert_eq!(snapshot.stop_routes, ["rt-b", "rt-c"]); + assert!(snapshot.keep_routes.is_empty()); + + // An instance desired again is no longer stopped. + controller.update_pool( + view(&["rt-a", "rt-b", "rt-c"], &["rt-a", "rt-b", "rt-c"]), + tokio::time::Instant::now(), + ); + assert!(controller.snapshot(one).unwrap().stop_routes.is_empty()); + assert!(controller.stopping_routes().is_empty()); + } + + #[test] + fn pool_member_identities_name_the_setting_to_fix() { + let member = member_hello("m1", "p1", 1).member.unwrap(); + assert_eq!(validate_member(&member), Ok(())); + for (broken, setting) in [ + ( + PoolMember { + member_id: String::new(), + ..member.clone() + }, + "OPENSHELL_SUPERVISOR_POOL_MEMBER_ID", + ), + ( + PoolMember { + member_uid: "uid with space".to_string(), + ..member.clone() + }, + "OPENSHELL_SUPERVISOR_POOL_MEMBER_UID", + ), + ( + PoolMember { + runtime_address: "not-an-address".to_string(), + ..member + }, + "OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP", + ), + ] { + let error = validate_member(&broken).unwrap_err(); + assert!(error.contains(setting), "{error}"); + } + // A host that names no member is told every setting. + let controller = SupervisorHostController::new(); + let error = controller + .begin_session( + identity(), + &HostHello { + instance_id: "p1".to_string(), + connection_epoch: 1, + member: None, + }, + ) + .unwrap_err(); + for setting in [ + "OPENSHELL_SUPERVISOR_POOL_MEMBER_ID", + "OPENSHELL_SUPERVISOR_POOL_MEMBER_UID", + "OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP", + ] { + assert!(error.message().contains(setting), "{}", error.message()); + } + assert!(!error.message().contains(" "), "{}", error.message()); + } } diff --git a/crates/openshell-supervisor-control/src/coordinator.rs b/crates/openshell-supervisor-control/src/coordinator.rs index cedbb42bc2..b0e6551c74 100644 --- a/crates/openshell-supervisor-control/src/coordinator.rs +++ b/crates/openshell-supervisor-control/src/coordinator.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Gateway-owned coordination of one shared supervisor host. +//! Gateway-owned coordination of the shared supervisor pool. //! //! The gateway constructs a [`SharedSupervision`] and hands it to the compute //! driver that provisions the host's sandbox instances. The driver supplies @@ -15,9 +15,11 @@ //! host, and which of its own objects must exist for an instance to count as //! supervised. +use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Display; use std::future::Future; -use std::time::Duration; +use std::pin::Pin; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::watch; @@ -30,6 +32,12 @@ use crate::status::SandboxInstanceStatusReport; /// Interval at which [`SharedSupervision::wait_for_ready`] re-reads the /// host's observations. const READINESS_POLL: Duration = Duration::from_secs(1); +/// Age, before a driver recovery began listing, past which a pool record the +/// listing omits is withdrawn. +/// +/// Younger records may come from a launch the listing raced on another +/// gateway replica, whose clock may also differ. +const RECOVERY_PRUNE_MARGIN: Duration = Duration::from_mins(1); /// Why a launch was not assigned. #[derive(Debug, thiserror::Error)] @@ -40,18 +48,71 @@ pub enum LaunchError { /// The driver could not persist the launch material; nothing was assigned. #[error("{0}")] Persist(E), - /// The host is at capacity. The material was persisted; the caller rolls - /// the generation back before its workload runs. + /// The host, or every member of a pool, is at capacity or unavailable. + /// The material was persisted; the caller rolls the generation back + /// before its workload runs. #[error("{0}")] Refused(AssignmentError), } -/// Revision of the desired set read before the driver listed its durable -/// records. Completing a recovery with it makes the replacement -/// compare-and-replace against that revision. +/// A new pool instance to record durably before it is assigned. +#[derive(Clone, PartialEq, Eq)] +pub struct PlacementRequest { + /// The instance's assignment, launch material included. + pub assignment: crate::proto::SandboxInstanceAssignment, + /// Credential epoch of the launch credentials. + pub auth_epoch: u64, +} + +impl std::fmt::Debug for PlacementRequest { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PlacementRequest") + .field("route", &self.assignment.route) + .field("sandbox_id", &self.assignment.sandbox_id) + .field("generation", &self.assignment.generation) + .field("auth_epoch", &self.auth_epoch) + .finish_non_exhaustive() + } +} + +/// Outcome of one durable pool operation. +pub type PlacementFuture = Pin> + Send>>; + +/// The gateway's durable records of a supervisor pool's instances. +/// +/// Installed by the gateway before instances launch, they make the pool's +/// desired set, ownership, and launch material available to every gateway +/// replica, whichever replica launched an instance or runs the compute +/// driver's watch. +pub trait InstancePlacement: Send + Sync + std::fmt::Debug { + /// Durably record a new instance, launch material included, and the + /// owner of its launch term: the pool member the consistent-hash ring + /// prefers among live members with room, and its current process. Fails + /// when no live member has room. + fn place_launch(&self, request: PlacementRequest) -> PlacementFuture; + + /// Routes of the durably recorded instances, with when each was + /// published (Unix milliseconds). + fn recorded(&self) -> PlacementFuture>; + + /// Durably record an instance a driver recovery found without a record, + /// such as one launched before the host became a pool; the pool places + /// it, fencing whichever process may supervise it. + fn record_recovered( + &self, + assignment: crate::proto::SandboxInstanceAssignment, + ) -> PlacementFuture; +} + +/// When a recovery began, read before the driver listed its durable records. +/// +/// A pool record published after it may come from a launch the listing raced, +/// so it is never withdrawn for missing from the listing. #[must_use = "complete the recovery with the records listed after this baseline was read"] pub struct RecoveryBaseline { - revision: u64, + /// When the baseline was read (Unix milliseconds). + read_at_ms: i64, } /// A generation that may run on the shared supervisor, as the driver's @@ -66,13 +127,11 @@ pub struct RecoverableInstance { /// Result of a recovery pass. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RecoveryOutcome { - /// The desired set matches the records; carries the revision in force. + /// The records match the listing; carries the desired revision in force. Published(u64), - /// The desired set changed while the records were listed; the next pass - /// derives them again. - Conflict, - /// Launch material could not be read or does not match its record, or the - /// set exceeds the host's capacity; the previous set stays published. + /// The gateway has not installed its records yet, launch material could + /// not be read or does not match its record, or a record could not be + /// written; the next pass tries again. Kept, } @@ -91,7 +150,7 @@ pub enum InstanceSupervision { Unknown, } -/// The gateway's coordination of one shared supervisor host: the desired set +/// The gateway's coordination of the shared supervisor pool: the desired set /// it publishes, the token it delivers, and the observations it reads back. /// /// Clones share one controller. The gateway keeps one for serving the host's @@ -107,6 +166,16 @@ impl SharedSupervision { Self::default() } + /// Offer each instance's runtime `recovery_timeout_secs` instead of + /// [`POOL_RECOVERY_TIMEOUT_SECS`](crate::controller::POOL_RECOVERY_TIMEOUT_SECS). + /// The gateway also installs its durable records with + /// [`SupervisorHostController::install_placement`] before instances + /// launch. + pub fn set_recovery_timeout_secs(&self, recovery_timeout_secs: u32) { + self.controller + .set_recovery_timeout_secs(recovery_timeout_secs); + } + /// The controller the gateway serves to the host and publishes the /// supervisor token into. #[must_use] @@ -126,9 +195,11 @@ impl SharedSupervision { /// persist its encoded launch material through `persist`, and only then /// assign the instance so the host starts it. The launch material is /// persisted first so a restarted gateway recovers the assignment from the - /// generation's own records. An identical assignment changes nothing. - /// Capacity is refused as [`LaunchError::Refused`], so the caller fails - /// its create before the workload runs. + /// generation's own records. The pool records the instance durably for + /// every gateway replica, on a live member with room; without one, or + /// beyond the pool's capacity, the launch is refused as + /// [`LaunchError::Refused`], so the caller fails its create before the + /// workload runs. An identical assignment changes nothing. pub async fn launch( &self, registration: &SandboxInstanceRegistration, @@ -143,13 +214,41 @@ impl SharedSupervision { persist(assignment.launch_material.clone()) .await .map_err(LaunchError::Persist)?; - self.controller - .assign(assignment) - .map_err(LaunchError::Refused) + // The instance and the owner of its launch term are durable before + // any member is assigned it, so every gateway replica serves it and a + // later failure of that member is recovered by takeover. + let Some(placement) = self.controller.placement() else { + tracing::warn!( + route = %registration.route, + "the supervisor pool is not ready to place sandbox instances" + ); + return Err(LaunchError::Refused(AssignmentError::Capacity)); + }; + if let Err(error) = placement + .place_launch(PlacementRequest { + assignment: assignment.clone(), + auth_epoch: registration.auth.auth_epoch.get(), + }) + .await + { + tracing::warn!( + route = %registration.route, + sandbox_id = %registration.sandbox_id, + %error, + "could not place a shared sandbox instance in the supervisor pool" + ); + return Err(LaunchError::Refused(AssignmentError::Capacity)); + } + let route = assignment.route.clone(); + self.controller.assign(assignment).map_err(|error| { + // A refused instance must not stay desired durably. + self.controller.unassign(&route); + LaunchError::Refused(error) + }) } - /// Withdraw a generation's sandbox instance so the host stops it. Nothing - /// happens for a route that is not assigned. + /// Withdraw a generation's sandbox instance so its owner stops it. The + /// gateway removes the instance's durable record right after. pub fn withdraw(&self, route: &str) { self.controller.unassign(route); } @@ -160,23 +259,20 @@ impl SharedSupervision { /// being undone. pub fn begin_recovery(&self) -> RecoveryBaseline { RecoveryBaseline { - revision: self.controller.revision(), + read_at_ms: unix_millis(), } } - /// Replace the desired set with `candidates`, the generations that may - /// run according to the records listed after `baseline` was read. Each - /// candidate's launch material comes from its current assignment or, for - /// a generation this gateway has not assigned since it started, from - /// `read` on the driver's record, which returns the persisted bytes or - /// `None` when they are gone. Material is decoded without quoting values, - /// bound to the recorded route and sandbox, and validated. - /// - /// The first completed pass initializes the desired set; until then - /// nothing is sent to the host. A failed or mismatched read keeps the - /// previous set, because an incomplete set would stop running instances. - /// A pass that raced a lifecycle edit is refused and derived again next - /// time. + /// Reconcile the pool's durable records, which are the desired set every + /// gateway replica serves, with `candidates`, the generations that may run + /// according to the records the driver listed after `baseline` was read. + /// A candidate without a record is recorded from its launch material, + /// read by `read` from the driver's record (`None` when it is gone); + /// material is decoded without quoting values, bound to the recorded route + /// and sandbox, and validated. A record the listing omits is withdrawn + /// when it was published before the listing. Until the gateway installs + /// its records, or when a read fails or material is unusable, the pass + /// reports [`RecoveryOutcome::Kept`]. pub async fn complete_recovery( &self, baseline: RecoveryBaseline, @@ -188,62 +284,91 @@ impl SharedSupervision { E: Display, F: Future>, E>> + Send, { - let mut assignments = Vec::with_capacity(candidates.len()); - for candidate in &candidates { - if let Some(assignment) = self.controller.assignment(&candidate.route) { - assignments.push(assignment); + self.reconcile_pool(&baseline, &candidates, read).await + } + + /// Reconcile a pool's durable records with the generations a driver + /// recovery listed after `baseline`. + async fn reconcile_pool( + &self, + baseline: &RecoveryBaseline, + candidates: &[RecoverableInstance], + read: impl Fn(&R) -> F + Send + Sync, + ) -> RecoveryOutcome + where + R: Send + Sync, + E: Display, + F: Future>, E>> + Send, + { + let Some(placement) = self.controller.placement() else { + return RecoveryOutcome::Kept; + }; + let recorded = match placement.recorded().await { + Ok(recorded) => recorded, + Err(error) => { + tracing::warn!(%error, "could not list the supervisor pool's instances; recovery skipped"); + return RecoveryOutcome::Kept; + } + }; + let mut outcome = RecoveryOutcome::Published(self.controller.revision()); + for candidate in candidates { + if recorded.contains_key(&candidate.route) { continue; } - let material = match read(&candidate.record).await { - Ok(Some(material)) => material, - Ok(None) => { - tracing::warn!( - sandbox_id = %candidate.sandbox_id, - route = %candidate.route, - "shared sandbox instance has no launch material; it stays unassigned" - ); - continue; - } + let assignment = match read(&candidate.record).await { + Ok(Some(material)) => recovered_assignment(candidate, &material), + Ok(None) => continue, Err(error) => { tracing::warn!( sandbox_id = %candidate.sandbox_id, route = %candidate.route, %error, - "could not read shared sandbox instance launch material; keeping the previous assignments" + "could not read shared sandbox instance launch material" ); - return RecoveryOutcome::Kept; + outcome = RecoveryOutcome::Kept; + continue; } }; - match recovered_assignment(candidate, &material) { - Ok(assignment) => assignments.push(assignment), - Err(error) => { - tracing::warn!( - sandbox_id = %candidate.sandbox_id, - route = %candidate.route, - %error, - "shared sandbox instance launch material is unusable; keeping the previous assignments" - ); - return RecoveryOutcome::Kept; - } + let recorded = match assignment { + Ok(assignment) => placement.record_recovered(assignment).await, + Err(error) => Err(error.to_string()), + }; + if let Err(error) = recorded { + tracing::warn!( + sandbox_id = %candidate.sandbox_id, + route = %candidate.route, + %error, + "could not record a recovered shared sandbox instance in the supervisor pool" + ); + outcome = RecoveryOutcome::Kept; } } - match self - .controller - .replace_assignments(baseline.revision, assignments) - { - Ok(revision) => RecoveryOutcome::Published(revision), - Err(AssignmentError::Conflict { .. }) => RecoveryOutcome::Conflict, - Err(error) => { - tracing::warn!(%error, "could not publish the shared sandbox instances; keeping the previous assignments"); - RecoveryOutcome::Kept + let listed = candidates + .iter() + .map(|candidate| candidate.route.as_str()) + .collect::>(); + let margin = i64::try_from(RECOVERY_PRUNE_MARGIN.as_millis()).unwrap_or(i64::MAX); + for (route, published_at_ms) in &recorded { + if !listed.contains(route.as_str()) + && *published_at_ms < baseline.read_at_ms.saturating_sub(margin) + { + tracing::info!( + route, + "withdrawing a supervisor pool instance its driver no longer runs" + ); + self.withdraw(route); } } + outcome } /// What the host's observations say about the instance of `sandbox_id` /// on `route`. Only a host that has reconciled to the assigning snapshot /// renders a verdict; a disconnected or still-reconciling host is not - /// evidence that a running instance is gone. + /// evidence that a running instance is gone. For a pool only the owning + /// process's host renders it, so an instance being taken over by another + /// member reads as degraded or unknown, never unavailable, until its new + /// owner has reconciled. #[must_use] pub fn supervision(&self, route: &str, sandbox_id: &str) -> InstanceSupervision { match self.controller.observed(route) { @@ -281,6 +406,14 @@ impl SharedSupervision { } } +fn unix_millis() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| { + i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX) + }) +} + /// Decode persisted launch material for `candidate`, binding it to the /// recorded route and sandbox before trusting it. fn recovered_assignment( @@ -301,90 +434,104 @@ fn recovered_assignment( mod tests { use super::*; use crate::proto::SandboxInstanceAssignment; + use std::sync::Mutex; + + /// Durable pool records kept in memory. + #[derive(Debug, Default)] + struct RecordedPool { + created_at_ms: Mutex>, + recovered: Mutex>, + } + + impl InstancePlacement for std::sync::Arc { + fn place_launch(&self, request: PlacementRequest) -> PlacementFuture { + self.created_at_ms + .lock() + .unwrap() + .insert(request.assignment.route, unix_millis()); + Box::pin(async { Ok(()) }) + } + + fn recorded(&self) -> PlacementFuture> { + let recorded = self.created_at_ms.lock().unwrap().clone(); + Box::pin(async move { Ok(recorded) }) + } + + fn record_recovered(&self, assignment: SandboxInstanceAssignment) -> PlacementFuture { + self.recovered.lock().unwrap().push(assignment.route); + Box::pin(async { Ok(()) }) + } + } fn candidate(route: &str, sandbox_id: &str) -> RecoverableInstance<&'static str> { RecoverableInstance { route: route.to_string(), sandbox_id: sandbox_id.to_string(), - record: "record", + record: if route.starts_with("rt-unreadable") { + "unreadable" + } else { + "record" + }, } } #[tokio::test] - async fn recovery_keeps_the_previous_set_unless_every_record_is_usable() { + async fn pool_recovery_reconciles_durable_records_with_the_drivers_listing() { let supervision = SharedSupervision::new(); - let assigned = SandboxInstanceAssignment { - route: "rt-assigned-0123456789".to_string(), - sandbox_id: "sandbox-a".to_string(), - generation: "gen".to_string(), - launch_material: b"opaque".to_vec(), - }; - // An assignment made before the first recovery is carried by it. - supervision.controller().assign(assigned.clone()).unwrap(); - - // A record whose material cannot be read keeps the set uninitialized. - let baseline = supervision.begin_recovery(); - let outcome = supervision - .complete_recovery( - baseline, - vec![ - candidate(&assigned.route, &assigned.sandbox_id), - candidate("rt-unreadable-0123456789", "sandbox-b"), - ], - |_record| async { Err::>, _>("boom") }, - ) - .await; - assert_eq!(outcome, RecoveryOutcome::Kept); - assert!(!supervision.controller().is_initialized()); - - // Material that is not a registration is refused without quoting it. - let baseline = supervision.begin_recovery(); + let records = std::sync::Arc::new(RecordedPool::default()); + // A recovery before the gateway installed its records changes nothing. let outcome = supervision .complete_recovery( - baseline, - vec![candidate("rt-garbage-0123456789", "sandbox-c")], - |_record| async { Ok::<_, &str>(Some(b"{\"secret\":1}".to_vec())) }, + supervision.begin_recovery(), + vec![candidate("rt-listed-0123456789", "sandbox-a")], + |_record| async { Ok::>, &str>(None) }, ) .await; assert_eq!(outcome, RecoveryOutcome::Kept); - assert!(!supervision.controller().is_initialized()); + supervision + .controller() + .install_placement(std::sync::Arc::new(records.clone())); - // A record without material is skipped; the current assignment is - // reused without a read; the pass initializes the set. + let old = unix_millis() - 2 * 60 * 1000; + records.created_at_ms.lock().unwrap().extend([ + ("rt-listed-0123456789".to_string(), old), + ("rt-stale-0123456789".to_string(), old), + ]); let baseline = supervision.begin_recovery(); + // A launch the listing raced on another replica is younger than it. + records + .created_at_ms + .lock() + .unwrap() + .insert("rt-raced-0123456789".to_string(), unix_millis()); let outcome = supervision .complete_recovery( baseline, vec![ - candidate(&assigned.route, &assigned.sandbox_id), - candidate("rt-gone-0123456789", "sandbox-d"), + candidate("rt-listed-0123456789", "sandbox-a"), + candidate("rt-unreadable-0123456789", "sandbox-b"), + candidate("rt-gone-0123456789", "sandbox-c"), ], - |_record| async { Ok::>, &str>(None) }, + |record: &&str| { + let unreadable = *record == "unreadable"; + async move { + if unreadable { + Ok::<_, &str>(Some(b"{\"secret\":1}".to_vec())) + } else { + Ok(None) + } + } + }, ) .await; - assert_eq!(outcome, RecoveryOutcome::Published(1)); - assert!(supervision.controller().is_initialized()); + // Unusable material keeps the pass from reporting success and is + // never recorded; a listed record is not read again. + assert_eq!(outcome, RecoveryOutcome::Kept); + assert!(records.recovered.lock().unwrap().is_empty()); + // Only the old record the listing omits is withdrawn. assert_eq!( - supervision.controller().assignment(&assigned.route), - Some(assigned.clone()) - ); - - // A baseline read before a lifecycle edit is refused, and the edit stands. - let baseline = supervision.begin_recovery(); - supervision.withdraw(&assigned.route); - let outcome = supervision - .complete_recovery( - baseline, - vec![candidate(&assigned.route, &assigned.sandbox_id)], - |_record| async { Ok::>, &str>(None) }, - ) - .await; - assert_eq!(outcome, RecoveryOutcome::Conflict); - assert!( - supervision - .controller() - .assignment(&assigned.route) - .is_none() + supervision.controller().pool_withdrawals(), + ["rt-stale-0123456789"] ); } } diff --git a/crates/openshell-supervisor-control/src/lib.rs b/crates/openshell-supervisor-control/src/lib.rs index 4401ff3d84..ffb743f8e9 100644 --- a/crates/openshell-supervisor-control/src/lib.rs +++ b/crates/openshell-supervisor-control/src/lib.rs @@ -27,6 +27,7 @@ pub mod controller; pub mod coordinator; pub mod credential; +pub mod pool_timing; pub mod registration; pub mod status; @@ -37,14 +38,18 @@ pub mod proto { } pub use controller::{ - AssignmentError, HOST_HEARTBEAT_INTERVAL, HostIdentity, MAX_ASSIGNED_INSTANCES, - MAX_REPORTED_INSTANCES, ObservedInstance, SupervisorCredential, SupervisorHostController, + AssignmentError, ConnectedMember, HostIdentity, InstanceOwner, MAX_ASSIGNED_INSTANCES, + MAX_POOL_INSTANCES, MAX_REPORTED_INSTANCES, MemberReport, MemberView, ObservedInstance, + POOL_HEARTBEAT_INTERVAL, POOL_RECOVERY_TIMEOUT_SECS, PoolView, SupervisorCredential, + SupervisorHostController, }; pub use coordinator::{ - InstanceSupervision, LaunchError, RecoverableInstance, RecoveryBaseline, RecoveryOutcome, - SharedSupervision, + InstancePlacement, InstanceSupervision, LaunchError, PlacementFuture, PlacementRequest, + RecoverableInstance, RecoveryBaseline, RecoveryOutcome, SharedSupervision, }; pub use credential::deliver_credentials; +/// Recovery deadlines, in seconds, a supervisor may offer a sandbox runtime. +pub use openshell_sandbox_backend::boundary_protocol::SUPERVISOR_RECOVERY_TIMEOUT_RANGE; pub use registration::{ INSTANCE_REMOVAL_GRACE, MAX_LAUNCH_MATERIAL_BYTES, ProxyCaMaterial, RegistrationError, RegistrationTraceContext, SandboxInstanceRegistration, SecretPem, diff --git a/crates/openshell-supervisor-control/src/pool_timing.rs b/crates/openshell-supervisor-control/src/pool_timing.rs new file mode 100644 index 0000000000..6041b87394 --- /dev/null +++ b/crates/openshell-supervisor-control/src/pool_timing.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Timing contract between a supervisor pool's hosts and the gateway. +//! +//! The gateway replica that serves a member's stream refreshes the member's +//! durable record; a record older than [`MEMBER_TTL`] is dead, and any replica +//! then takes the member's instances over. A healthy member whose serving +//! replica stops answering without resetting the connection (frozen, or +//! partitioned from the member only) must therefore notice the silence, +//! reconnect through the gateway Service to a replica that answers, and be +//! refreshed there before its record expires. [`reconnect_budget`] is that +//! worst case; it stays well inside [`MEMBER_TTL`] even when the first +//! [`COVERED_UNANSWERED_ATTEMPTS`] attempts reach the unresponsive replica +//! again. A member cut off from every replica still expires after +//! [`MEMBER_TTL`], so a true partition fails over as before. +//! +//! Conversely, a replica stops refreshing a member that died without closing +//! its stream once the stream has been silent for +//! [`GATEWAY_SILENCE_TIMEOUT`], so the member's record expires and its +//! instances move well inside the runtime's recovery deadline. + +use std::time::Duration; + +/// Age after which a member's durable record means the member is dead. +pub const MEMBER_TTL: Duration = Duration::from_secs(15); +/// Interval at which a gateway replica refreshes the records of the members +/// it serves, independently of its reconciliation passes. A member that +/// attaches or announces leaving is refreshed at once. +pub const MEMBER_REFRESH_INTERVAL: Duration = Duration::from_secs(2); +/// Interval of heartbeats in both directions on a pool member's stream. +pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1); +/// Silence after which a gateway replica closes a pool member's stream and +/// stops refreshing its record. Also bounds queueing one message to the +/// member, so a member that stops reading is dropped as soon. +pub const GATEWAY_SILENCE_TIMEOUT: Duration = Duration::from_secs(5); +/// Silence after which a pool host abandons its stream and reconnects. +pub const HOST_SILENCE_TIMEOUT: Duration = Duration::from_secs(3); +/// How often a pool host checks its stream for silence. +pub const HOST_SILENCE_CHECK: Duration = Duration::from_millis(250); +/// Bound on one pool host attempt to connect and be accepted. +/// +/// It applies to the first attempts after a stream is lost. Later attempts +/// get the host's ordinary bound, so a slow but responsive gateway is still +/// reached. +pub const HOST_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(1); +/// First pool host reconnect backoff; it doubles up to +/// [`HOST_RECONNECT_BACKOFF_MAX`]. +pub const HOST_RECONNECT_BACKOFF_MIN: Duration = Duration::from_millis(100); +/// Longest pool host reconnect backoff, before jitter. +pub const HOST_RECONNECT_BACKOFF_MAX: Duration = Duration::from_millis(400); +/// Jitter added to a pool host backoff, at most this fraction of it. +pub const HOST_RECONNECT_JITTER: f64 = 0.5; +/// Attempts after a lost stream that use [`HOST_ATTEMPT_TIMEOUT`] and that +/// [`reconnect_budget`] allows to reach an unresponsive replica again. +pub const COVERED_UNANSWERED_ATTEMPTS: u32 = 3; +/// Bound on the gateway writing one member's record. A replica writes the +/// records of all the members it serves at once, so one refresh round takes +/// no longer either. +pub const RECORD_WRITE_BOUND: Duration = Duration::from_millis(500); + +/// Longest pool host backoff before attempt `attempt` (1-based) after a lost +/// stream, jitter included. +#[must_use] +pub fn host_backoff_bound(attempt: u32) -> Duration { + let doubled = HOST_RECONNECT_BACKOFF_MIN + .saturating_mul(1_u32 << attempt.saturating_sub(1).min(16)) + .min(HOST_RECONNECT_BACKOFF_MAX); + doubled.mul_f64(1.0 + HOST_RECONNECT_JITTER) +} + +/// Worst case from a member's serving replica going unresponsive to another +/// replica refreshing the member's record. +/// +/// The first `unanswered` attempts reach the unresponsive replica again. The +/// sum is the record's age when the replica stopped, silence detection, each +/// unanswered attempt and its backoff, the answered attempt, and two refresh +/// rounds on the answering replica: one already in flight when the member +/// attaches, and the one the attach starts. +#[must_use] +pub fn reconnect_budget(unanswered: u32) -> Duration { + let mut budget = MEMBER_REFRESH_INTERVAL + HOST_SILENCE_TIMEOUT + HOST_SILENCE_CHECK; + for attempt in 1..=unanswered { + budget += HOST_ATTEMPT_TIMEOUT + host_backoff_bound(attempt); + } + budget + HOST_ATTEMPT_TIMEOUT + RECORD_WRITE_BOUND * 2 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_member_reconnects_well_before_its_record_expires() { + let budget = reconnect_budget(COVERED_UNANSWERED_ATTEMPTS); + // At least a fifth of the TTL is left after the covered worst case. + assert!( + budget + MEMBER_TTL / 5 <= MEMBER_TTL, + "reconnect budget {budget:?} leaves too little of the {MEMBER_TTL:?} TTL" + ); + // Heartbeats are frequent enough that silence means several were + // missed, and the runtime's recovery deadline outlasts expiry, + // takeover, and a replacement's start. + assert!(HOST_SILENCE_TIMEOUT >= HEARTBEAT_INTERVAL * 3); + assert!(GATEWAY_SILENCE_TIMEOUT >= HEARTBEAT_INTERVAL * 3); + // A member that died without closing its stream: the stream goes + // silent, its record's last refresh ages out, and a pass takes its + // instances over, all within half the runtime's recovery deadline. + assert!( + (GATEWAY_SILENCE_TIMEOUT + + MEMBER_REFRESH_INTERVAL + + MEMBER_TTL + + MEMBER_REFRESH_INTERVAL) + * 2 + < Duration::from_secs(u64::from(crate::controller::POOL_RECOVERY_TIMEOUT_SECS)) + ); + assert!( + MEMBER_TTL * 2 + < Duration::from_secs(u64::from(crate::controller::POOL_RECOVERY_TIMEOUT_SECS)) + ); + assert_eq!(host_backoff_bound(1), Duration::from_millis(150)); + assert_eq!(host_backoff_bound(10), Duration::from_millis(600)); + } +} diff --git a/crates/openshell-supervisor-control/src/registration.rs b/crates/openshell-supervisor-control/src/registration.rs index 508227554f..dfd80910c3 100644 --- a/crates/openshell-supervisor-control/src/registration.rs +++ b/crates/openshell-supervisor-control/src/registration.rs @@ -193,6 +193,7 @@ impl SandboxInstanceRegistration { sandbox_id: self.sandbox_id.clone(), generation: self.auth.runtime_generation.to_string(), launch_material, + takeover: None, }) } @@ -240,6 +241,7 @@ mod tests { sandbox_id: "sandbox".to_string(), generation: "gen".to_string(), launch_material: b"{}".to_vec(), + takeover: None, }; let error = SandboxInstanceRegistration::from_assignment(&mismatched).unwrap_err(); assert!(error.to_string().contains("not a valid registration")); diff --git a/crates/openshell-supervisor/README.md b/crates/openshell-supervisor/README.md index 10964b8542..184bf22c78 100644 --- a/crates/openshell-supervisor/README.md +++ b/crates/openshell-supervisor/README.md @@ -152,15 +152,19 @@ control listener: it dials the gateway and receives its assignments from it. validates the snapshot as a whole (at most 256 instances, no route twice), decodes each assignment's launch material (the JSON `SandboxInstanceRegistration` of `openshell-supervisor-control`), admits it into - the `SandboxInstanceHost` registry keyed by route with the registration - digest as its fingerprint, and removes every hosted route the snapshot - omits. An identical assignment is a no-op; one that cannot be admitted + the `SandboxInstanceHost` registry keyed by route with its pool fingerprint + (see **Pool members**), relinquishes every hosted route the snapshot omits, + and stops the routes it lists for stopping. A snapshot that is not a pool's + is refused. An identical assignment is a no-op; one that cannot be admitted (invalid material, or a different fingerprint for a live route) is reported as a failed instance and never touches the running one. Snapshots are - reconciled one at a time on a separate task, so heartbeats and status keep - flowing while instances stop; only the newest unapplied snapshot is kept, - and a connection finishes the reconciliation it started before the next - connection begins. Private files live under `--session-root` (proxy CA, + reconciled one at a time by one worker that outlives connections, so + heartbeats and status keep flowing while instances start or stop, and a + lost connection is replaced without waiting for a slow reconciliation; only + the newest unapplied snapshot is kept, and work queued by an older + connection is dropped when a new one begins. On shutdown the worker drops + queued work and finishes the reconciliation in flight before the host + relinquishes its instances. Private files live under `--session-root` (proxy CA, trust directory, SSH socket); the registry owns removal and keeps the route reserved until the instance has ended and its files are deleted. - **Status.** After each reconciliation and whenever an instance is admitted, @@ -169,15 +173,46 @@ control listener: it dials the gateway and receives its assignments from it. once the gateway has accepted the instance's own supervisor session. - **Disconnects.** Losing the gateway stream never stops an instance. The host reconnects with backoff, reports what it runs, and applies the first - snapshot of the new session; only a snapshot that omits an instance, or - process shutdown, stops it. + snapshot of the new session; only a snapshot that lists an instance for + stopping stops it. The host uses the timings of + `openshell_supervisor_control::pool_timing`: it abandons a stream after 3 s + without hearing from the gateway, bounds sends by the same silence, gives + each early reconnect attempt one 1 s deadline across connecting, attaching, + and being accepted, and backs off 100 to 400 ms with jitter, so it reaches a + responsive gateway replica well before its membership record expires. - **Runtime connections.** One routed listener (`--runtime-listen`) accepts every instance's runtime-initiated Sandbox Protocol connections; see `openshell-sandbox-backend` for the routing preface. - **Health.** `--address` is a plain TCP listener that accepts and closes kubelet `tcpSocket` probe connections; it carries no control traffic and - reports only that the process runs. SIGTERM and SIGINT stop every instance, - which ends their workloads' supervision. + reports only that the process runs. On SIGTERM or SIGINT the member + announces that it is leaving, waits for the gateway to move its instances, + and relinquishes whatever it still hosts; their workloads keep running. +- **Pool members.** `--pool-member-id`, `--pool-member-uid`, and + `--pool-advertise-ip` (or `OPENSHELL_SUPERVISOR_POOL_MEMBER_ID`, + `_MEMBER_UID`, `_ADVERTISE_IP`, which the chart fills from the Pod's name, + UID, and IP) are required with `--mode=shared`: every shared supervisor is a + member of the gateway's high-availability pool, and its hello names the + member and its runtime address. Missing values fail startup with the flag + and variable to set. A pool snapshot lists only + the instances this process owns. An instance it omits moved to another + member and is **relinquished**: the instance is cancelled without its stop + path, so the runtime keeps the workload frozen until the new owner takes + it over. Only `stop_routes` are stopped, and `keep_routes` are left as + they are. A pool instance's admission fingerprint is its launch material + and ownership term, not its credentials, so the same term re-minted by + another gateway replica keeps the running instance. A takeover term + replaces the launch credentials with the term's and attaches with its + grant; the replacement adopts the running workload instead of starting it + and keeps the launch material's proxy CA and SSH host key. A term that + arrives without credentials, or whose route is still being released, is + deferred, reported as starting, and retried at the heartbeat cadence. The + host follows the gateway's heartbeat interval, offers each instance's + runtime the recovery deadline of the latest snapshot + (`recovery_timeout_secs`) when it attaches, feeds the route directory to the + runtime router for redirects, and on SIGTERM announces that it is + leaving, waits up to 40 seconds for the gateway to move its instances, and + relinquishes what remains. Sandbox instances get no upstream proxy, extra CA bundle, or provider workload identity: the Kubernetes driver refuses the topology when an upstream proxy or provider diff --git a/crates/openshell-supervisor/src/backend_setup.rs b/crates/openshell-supervisor/src/backend_setup.rs index 91d4894c09..a0cd52ea87 100644 --- a/crates/openshell-supervisor/src/backend_setup.rs +++ b/crates/openshell-supervisor/src/backend_setup.rs @@ -38,6 +38,10 @@ pub struct BackendServices { pub ca_file_paths: Arc>>, pub provider_credentials: ProviderCredentialState, pub sandbox_bearer: SessionBearerTokenSlot, + /// Gateway-granted authority to replace the generation's current owner. + pub takeover: Option, + /// Recovery deadline (seconds) offered to the runtime at every attach. + pub supervisor_recovery_timeout_secs: Option, } /// Startup result; native payload decoding never reaches shared runtime @@ -324,17 +328,20 @@ impl PreparedBackend for OpenShellLaunch { services: BackendServices, ) -> std::result::Result { let payload = self.descriptor.backend_descriptor()?.payload; + let mut backend = openshell_sandbox_backend::OpenShellRuntimeBackend::new( + services.ca_file_paths, + services.provider_credentials, + services.sandbox_bearer, + ) + .with_transport(self.transport) + .with_recovery_timeout_secs(services.supervisor_recovery_timeout_secs); + if let Some(takeover) = services.takeover { + backend = backend.with_takeover(takeover); + } Ok(BuiltBackend { payload, proxy_listener: None, - backend: Arc::new( - openshell_sandbox_backend::OpenShellRuntimeBackend::new( - services.ca_file_paths, - services.provider_credentials, - services.sandbox_bearer, - ) - .with_transport(self.transport), - ), + backend: Arc::new(backend), }) } } diff --git a/crates/openshell-supervisor/src/backend_setup/registration.rs b/crates/openshell-supervisor/src/backend_setup/registration.rs index b27ec6cf19..df337cc7f1 100644 --- a/crates/openshell-supervisor/src/backend_setup/registration.rs +++ b/crates/openshell-supervisor/src/backend_setup/registration.rs @@ -113,6 +113,11 @@ impl PreparedBackend for RegisteredLaunch { self: Box, services: BackendServices, ) -> std::result::Result { + if services.takeover.is_some() { + return Err(BackendError::Unsupported( + "registered isolation backends do not support supervisor takeover".to_string(), + )); + } let payload = self.payload.clone(); let backend = Arc::new(OpenShellRuntimeBackend::registered( self.name, diff --git a/crates/openshell-supervisor/src/backend_setup/tests.rs b/crates/openshell-supervisor/src/backend_setup/tests.rs index b4704bf623..887b138860 100644 --- a/crates/openshell-supervisor/src/backend_setup/tests.rs +++ b/crates/openshell-supervisor/src/backend_setup/tests.rs @@ -186,6 +186,8 @@ impl TestSetup { HashMap::from([("TEST_CONFIG".into(), "initial".into())]), ), sandbox_bearer: self.auth.sandbox_bearer_slot().unwrap(), + takeover: None, + supervisor_recovery_timeout_secs: None, } } } @@ -297,6 +299,7 @@ impl PreparedBackend for TestLaunch { snapshot: SessionSnapshotWire { generation: generation.clone(), processes: vec![], + bound_policy: None, }, } } @@ -704,6 +707,8 @@ async fn shared_startup_uses_selected_backend_through_readiness_and_shutdown() { provider_workload_api_socket: None, // The binary loads the proxy trust directory from its environment. proxy_tls: openshell_supervisor_network::run::ProxyTlsFiles::from_environment(), + takeover: None, + supervisor_recovery_timeout_secs: None, }; let bindings = crate::SandboxInstanceBindings::dedicated( &config, diff --git a/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs b/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs index 699880f1ad..8e6def6552 100644 --- a/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs +++ b/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs @@ -304,6 +304,8 @@ async fn config_push_regression_legacy_gateway_never_observes_ready_before_workl trace_context: None, provider_workload_api_socket: None, proxy_tls: openshell_supervisor_network::run::ProxyTlsFiles::default(), + takeover: None, + supervisor_recovery_timeout_secs: None, }; let bindings = SandboxInstanceBindings::dedicated( &config, diff --git a/crates/openshell-supervisor/src/lib.rs b/crates/openshell-supervisor/src/lib.rs index 46eb9e8839..ae8e160c23 100644 --- a/crates/openshell-supervisor/src/lib.rs +++ b/crates/openshell-supervisor/src/lib.rs @@ -791,6 +791,13 @@ pub struct SandboxInstanceConfig { /// SPIFFE Workload API socket for provider token grants; absent fails closed. pub provider_workload_api_socket: Option, pub proxy_tls: openshell_supervisor_network::run::ProxyTlsFiles, + /// Gateway-granted authority to replace the generation's current owner. + /// The instance then adopts the running workload instead of starting it. + pub takeover: Option, + /// Recovery deadline (seconds) offered to the sandbox runtime at every + /// attach: how long it keeps the workload frozen after losing this + /// supervisor. `None` leaves the runtime default. + pub supervisor_recovery_timeout_secs: Option, } impl SandboxInstanceConfig { @@ -939,6 +946,8 @@ pub async fn run_sandbox_instance( trace_context, provider_workload_api_socket, proxy_tls, + takeover, + supervisor_recovery_timeout_secs, } = config; // The owner binds the event context and log sink to the instance's // threads; the instance itself uses the rest. @@ -1164,6 +1173,8 @@ pub async fn run_sandbox_instance( ca_file_paths: ca_file_paths.clone(), provider_credentials: provider_credentials.clone(), sandbox_bearer, + takeover, + supervisor_recovery_timeout_secs, }, policy.clone(), openshell_isolation_interface::AgentSpec { @@ -1249,18 +1260,20 @@ pub async fn run_sandbox_instance( .is_none() .then(|| remote_boundary.0.network_mediation_source()); let remote_host_gateway_ip = remote_boundary.0.host_gateway_ip(); - let (remote_ready, backend_name, ca_file_paths) = { - let (bound, backend_name, ca_file_paths) = remote_boundary; + let takeover_attach = remote_boundary.0.is_takeover(); + let adopts_running_agent = remote_boundary.0.adopts_running_agent(); + let (bound, backend_name, ca_file_paths) = remote_boundary; + let confirm = Box::pin(async { let ready = bound .confirm() .instrument(tracing::info_span!(parent: &startup, "supervisor.boundary.confirm")) .await .map_err(|error| miette::miette!(error.to_string()))?; info!(backend = %backend_name, "Isolation boundary enforcement confirmed"); - (ready, backend_name, ca_file_paths) - }; - - let mut networking = Some( + Ok::<_, miette::Report>(ready) + }); + // Both startup futures are boxed: the order below holds them at once. + let start_networking = Box::pin(async { openshell_supervisor_network::run::run_networking( &policy, None, @@ -1294,8 +1307,37 @@ pub async fn run_sandbox_instance( remote_network_source, proxy_listener, ) - .await?, - ); + .await + }); + // A replacement supervisor confirms only once its own enforcement is up: + // confirmation resumes a frozen workload, whose network opens then wait + // for this supervisor's mediation rather than for nothing. Mediation + // accepts start once the boundary is confirmed. + let (remote_ready, networking) = if takeover_attach { + let networking = start_networking.await?; + (confirm.await?, networking) + } else { + let ready = confirm.await?; + (ready, start_networking.await?) + }; + let mut networking = Some(networking); + if takeover_attach { + info!( + adopts_running_agent, + "Supervisor takeover confirmed; previous owner fenced" + ); + ocsf_emit!( + // Lifecycle activity 3 is Start. + AppLifecycleBuilder::new(ocsf_ctx()) + .activity(ActivityId::Reset) + .severity(SeverityId::Informational) + .status(StatusId::Success) + .message(format!( + "Supervisor took over the sandbox runtime [adopted_workload:{adopts_running_agent}]" + )) + .build() + ); + } ca_file_paths .lock() @@ -1490,7 +1532,17 @@ pub async fn run_sandbox_instance( .await .map_err(|error| start_agent_error(error.to_string()))?; workload_started_tx.send_replace(true); - info!(backend = %backend_name, "Isolation boundary agent started"); + if running.agent_adopted() { + info!( + backend = %backend_name, + process_id = %running.agent().process_id(), + adopted_execs = running.adopted_exec_processes().len(), + "Adopted the running agent from the previous supervisor" + ); + observe_adopted_execs(running.adopted_exec_processes()); + } else { + info!(backend = %backend_name, "Isolation boundary agent started"); + } let agent = running.agent(); let boundary_access = openshell_supervisor_process::delegated::start_boundary_access( sandbox_id.as_deref(), @@ -1640,6 +1692,28 @@ pub async fn run_sandbox_instance( Ok(exit_code) } +/// Exec processes a replacement supervisor adopted outlive their original +/// clients; record when each exits so their end is not silent. +fn observe_adopted_execs(processes: &[Arc]) { + for process in processes { + let process = process.clone(); + tokio::spawn(async move { + match process.wait().await { + Ok(status) => info!( + process_id = %process.process_id(), + ?status, + "Adopted exec process exited" + ), + Err(error) => debug!( + process_id = %process.process_id(), + %error, + "Stopped observing adopted exec process" + ), + } + }); + } +} + fn persist_main_exit_marker(path: &std::path::Path, exit_code: i32) -> std::io::Result<()> { let parent = path .parent() @@ -4367,6 +4441,8 @@ mod tests { trace_context: None, provider_workload_api_socket: None, proxy_tls: openshell_supervisor_network::run::ProxyTlsFiles::default(), + takeover: None, + supervisor_recovery_timeout_secs: None, }; let bindings = SandboxInstanceBindings::dedicated( &config, diff --git a/crates/openshell-supervisor/src/main.rs b/crates/openshell-supervisor/src/main.rs index de29e9509a..7f183d8559 100644 --- a/crates/openshell-supervisor/src/main.rs +++ b/crates/openshell-supervisor/src/main.rs @@ -83,6 +83,24 @@ struct Args { #[arg(long, required_if_eq("mode", "shared"))] session_root: Option, + /// Stable name of this shared supervisor as a member of the gateway's + /// supervisor pool (its Pod name), so the gateway can place, keep, and + /// fail over sandbox instances among members. Required for + /// `--mode=shared`. + #[arg(long, env = "OPENSHELL_SUPERVISOR_POOL_MEMBER_ID")] + pool_member_id: Option, + + /// Identity of this incarnation of the pool member (its Pod UID). + /// Required for `--mode=shared`. + #[arg(long, env = "OPENSHELL_SUPERVISOR_POOL_MEMBER_UID")] + pool_member_uid: Option, + + /// Address other pool members redirect this member's sandbox runtimes + /// to (its Pod IP); the runtime listener's port is appended. Required for + /// `--mode=shared`. + #[arg(long, env = "OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP")] + pool_advertise_ip: Option, + /// Supervisor responsibility to run. #[arg(long, value_enum, default_value_t)] role: SupervisorRole, @@ -596,6 +614,8 @@ fn main() -> Result<()> { .filter(|path| !path.is_empty()) .map(PathBuf::from), proxy_tls: openshell_supervisor_network::run::ProxyTlsFiles::from_environment(), + takeover: None, + supervisor_recovery_timeout_secs: None, }; // Log push already shares these credentials and the logging // layers these OCSF JSON settings. diff --git a/crates/openshell-supervisor/src/sandbox_instance_host.rs b/crates/openshell-supervisor/src/sandbox_instance_host.rs index 66aa30e57e..88be0e13a9 100644 --- a/crates/openshell-supervisor/src/sandbox_instance_host.rs +++ b/crates/openshell-supervisor/src/sandbox_instance_host.rs @@ -136,6 +136,17 @@ pub enum SandboxInstanceHostError { /// Final status of one removal, shared by every remover. type RemovalReport = watch::Receiver>; +/// How a sandbox instance leaves the host. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Departure { + /// Stop the instance, which terminates its workload. + Stop, + /// Stop supervising without terminating the workload: the instance is + /// cancelled, so its runtime loses this supervisor and keeps the workload + /// frozen until another supervisor takes the generation over. + Relinquish, +} + struct HostedSandboxInstance { /// Distinguishes this instance from a later one that reuses its key. serial: u64, @@ -344,7 +355,29 @@ impl SandboxInstanceHost { /// removal tries again. Returns the final status, or `None` when `key` is /// not hosted. pub async fn remove(&self, key: &str, grace: Duration) -> Option { - let (mut removed, sandbox_id) = self.begin_removal(key, grace)?; + self.depart(key, grace, Departure::Stop).await + } + + /// Stop supervising one sandbox instance without stopping its workload, + /// delete its private files, and forget it, as when its ownership moved + /// to another supervisor. The instance is cancelled rather than asked to + /// stop, so it never terminates the workload; the runtime freezes the + /// workload until a supervisor granted the generation takes it over, and + /// terminates it if none does in time. A removal already in progress is + /// joined instead. Returns the final status, or `None` when `key` is not + /// hosted. + pub async fn relinquish(&self, key: &str) -> Option { + self.depart(key, RUNTIME_SHUTDOWN_TIMEOUT, Departure::Relinquish) + .await + } + + async fn depart( + &self, + key: &str, + grace: Duration, + departure: Departure, + ) -> Option { + let (mut removed, sandbox_id) = self.begin_removal(key, grace, departure)?; // The removal task ends without a report only with its runtime. let reported = removed .wait_for(Option::is_some) @@ -361,13 +394,19 @@ impl SandboxInstanceHost { })) } - /// Stop every sandbox instance, for process shutdown. - pub async fn remove_all(&self, grace: Duration) { + /// Relinquish every sandbox instance, for the shutdown of a supervisor + /// whose instances other supervisors take over. + pub async fn relinquish_all(&self) { let keys = lock(&self.instances).keys().cloned().collect::>(); - futures::future::join_all(keys.iter().map(|key| self.remove(key, grace))).await; + futures::future::join_all(keys.iter().map(|key| self.relinquish(key))).await; } - fn begin_removal(&self, key: &str, grace: Duration) -> Option<(RemovalReport, String)> { + fn begin_removal( + &self, + key: &str, + grace: Duration, + departure: Departure, + ) -> Option<(RemovalReport, String)> { let mut instances = lock(&self.instances); let instance = instances.get_mut(key)?; // Join a removal in progress; retry one that kept the key. @@ -376,9 +415,16 @@ impl SandboxInstanceHost { { return Some((removed.clone(), instance.sandbox_id.clone())); } - instance.stop.send_replace(true); - if *instance.phase.borrow() == SandboxInstancePhaseReport::Starting { - instance.abort.send_replace(true); + match departure { + Departure::Stop => { + instance.stop.send_replace(true); + if *instance.phase.borrow() == SandboxInstancePhaseReport::Starting { + instance.abort.send_replace(true); + } + } + Departure::Relinquish => { + instance.abort.send_replace(true); + } } let (report, removed) = watch::channel(None); instance.removal = Some(removed.clone()); @@ -1074,7 +1120,72 @@ mod tests { Some(SandboxInstancePhaseReport::Exited { exit_code: 4 }) ); assert!(!dir.exists()); - host.remove_all(Duration::from_secs(10)).await; + assert!(host.list().is_empty()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn relinquishing_ends_supervision_without_stopping_the_workload() { + let (root, host) = private_host(); + let dir = root.path().join("instances").join("relinquished"); + let credentials = GatewayCredentials::from_gateway_token("relinquished-token").unwrap(); + // Only the stop path terminates the workload. + let (terminated_tx, mut terminated) = mpsc::channel::<()>(1); + let body: SandboxInstanceBody = Box::new(move |controls| { + Box::pin(async move { + controls.readiness.send_replace(true); + controls.shutdown.await; + let _ = terminated_tx.send(()).await; + Ok(0) + }) + }); + admit( + &host, + "relinquished", + "sandbox-relinquished", + 1, + SandboxInstanceBindings { + gateway_credentials: credentials.clone(), + ..unbound() + }, + body, + ) + .unwrap(); + wait_supervising(&host, "relinquished").await; + + let status = host + .relinquish("relinquished") + .await + .expect("instance was hosted"); + assert!(!status.ready); + assert!(status.phase.is_terminal(), "{status:?}"); + assert_eq!( + terminated.try_recv(), + Err(mpsc::error::TryRecvError::Disconnected), + "relinquishing must never run the workload's stop path" + ); + assert!(host.status("relinquished").is_none()); + assert!(!dir.exists()); + let revoked = GatewayEndpoint::new("http://127.0.0.1:9", credentials); + let error = openshell_core::grpc_client::fetch_settings_snapshot(&revoked, "relinquished") + .await + .unwrap_err(); + assert!(error.to_string().contains("revoked"), "{error}"); + // The key is free for a later term of the same instance. + assert_eq!( + admit( + &host, + "relinquished", + "sandbox-relinquished", + 2, + unbound(), + Box::new(|controls| Box::pin(async move { + controls.shutdown.await; + Ok(0) + })), + ), + Ok(SandboxInstanceAdmitted::Created) + ); + host.relinquish_all().await; assert!(host.list().is_empty()); } diff --git a/crates/openshell-supervisor/src/shared.rs b/crates/openshell-supervisor/src/shared.rs index 4dc3e521ef..849ccee4a2 100644 --- a/crates/openshell-supervisor/src/shared.rs +++ b/crates/openshell-supervisor/src/shared.rs @@ -3,6 +3,12 @@ //! Supervisor host shared by the sandboxes of one gateway. +/// How long a leaving pool member waits for the gateway to move its sandbox +/// instances to other members before it relinquishes the rest. Shorter than +/// the Pod's termination grace period. +#[cfg(unix)] +const POOL_LEAVE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(40); + use miette::Result; #[cfg(unix)] use std::future::Future; @@ -34,6 +40,39 @@ fn run_platform(args: crate::Args) -> Result<()> { let instance_root = args .session_root .ok_or_else(|| miette::miette!("--session-root is required for --mode=shared"))?; + // Every shared supervisor is a member of the gateway's supervisor pool. + let missing = [ + ( + args.pool_member_id.is_none(), + "--pool-member-id (OPENSHELL_SUPERVISOR_POOL_MEMBER_ID, the Pod name)", + ), + ( + args.pool_member_uid.is_none(), + "--pool-member-uid (OPENSHELL_SUPERVISOR_POOL_MEMBER_UID, the Pod UID)", + ), + ( + args.pool_advertise_ip.is_none(), + "--pool-advertise-ip (OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP, the Pod IP)", + ), + ] + .into_iter() + .filter_map(|(missing, setting)| missing.then_some(setting)) + .collect::>(); + let (Some(name), Some(incarnation), Some(advertise_ip)) = ( + args.pool_member_id, + args.pool_member_uid, + args.pool_advertise_ip, + ) else { + return Err(miette::miette!( + "--mode=shared runs a supervisor pool member and needs {}", + missing.join(", ") + )); + }; + let pool = openshell_supervisor::shared_host::PoolIdentity { + member_id: name, + member_uid: incarnation, + runtime_address: std::net::SocketAddr::new(advertise_ip, runtime_listen.port()), + }; // The supervisor token travels only over TLS; the host verifies the // gateway the way every sandbox instance does, with `OPENSHELL_TLS_CA` // and `OPENSHELL_GATEWAY_TLS_SERVER_NAME`. @@ -95,6 +134,7 @@ fn run_platform(args: crate::Args) -> Result<()> { instance_root, runtime_listen, removal_grace: openshell_supervisor_control::INSTANCE_REMOVAL_GRACE, + pool, }, )?; let health = tokio::net::TcpListener::bind(address) @@ -104,6 +144,7 @@ fn run_platform(args: crate::Args) -> Result<()> { tracing::info!( runtime = %host.runtime_addr(), health = %address, + pool_member = %host.pool().member_id, "Shared supervisor ready for sandbox instance assignments" ); let (stop, stopped) = tokio::sync::watch::channel(false); @@ -120,11 +161,15 @@ fn run_platform(args: crate::Args) -> Result<()> { _ = terminate.recv() => {} _ = interrupt.recv() => {} } + // A pool member hands its instances to other members while it still + // hears from the gateway. + host.leave(POOL_LEAVE_TIMEOUT).await; stop.send_replace(true); let _ = probes.await; let _ = attachment.await; - // Sandboxes lose supervision with this process and fail closed. - host.remove_all().await; + // The member relinquishes what it still hosts to the members the + // gateway grants it to; the workloads keep running. + host.shut_down().await; Ok(()) }) } diff --git a/crates/openshell-supervisor/src/shared_host.rs b/crates/openshell-supervisor/src/shared_host.rs index b78ce7469c..78cf9b8af8 100644 --- a/crates/openshell-supervisor/src/shared_host.rs +++ b/crates/openshell-supervisor/src/shared_host.rs @@ -1,21 +1,28 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Supervisor shared by the sandboxes of one gateway. +//! Supervisor shared by the sandboxes of one gateway, as a member of the +//! gateway's supervisor pool (high availability). //! //! The gateway owns which sandbox instances, one running generation of a -//! stable Sandbox resource each, this process hosts. The host dials the -//! gateway presenting the supervisor token the gateway minted for it, -//! receives complete assignment snapshots over the private -//! `openshell.supervisor.host.v1` stream, reconciles the +//! stable Sandbox resource each, this process hosts: the instances its +//! process owns. The host dials the gateway presenting the supervisor token +//! the gateway minted for it, receives complete assignment snapshots over the +//! private `openshell.supervisor.host.v1` stream, reconciles the //! [`SandboxInstanceHost`] registry to each one, and reports what it //! observes. An assignment carries exactly the gateway-issued credentials of //! its generation; the host never acts with any other identity toward -//! sandbox APIs. Losing the gateway stream never stops an instance; -//! only a later snapshot that omits it does. Snapshots are reconciled one at -//! a time, and a connection finishes the reconciliation it started before the -//! next connection begins, so a superseded stream can never apply after its -//! replacement. Runtime connections arrive through one routed listener. +//! sandbox APIs. Runtime connections arrive through one routed listener. +//! +//! Losing the gateway stream never stops an instance. An instance a snapshot +//! omits has moved to another member: the host relinquishes it, ending +//! supervision without stopping the workload, which its runtime keeps frozen +//! until the new owner takes it over. Only instances a snapshot lists for +//! stopping are stopped. Snapshots are reconciled one at a time by one worker +//! that outlives connections, and work an older connection queued is dropped +//! when a new one begins. A member that shuts down announces that it is +//! leaving, so the gateway moves its instances first, and relinquishes +//! whatever it still hosts. use std::collections::BTreeMap; use std::future::Future; @@ -27,10 +34,12 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use miette::Result; use openshell_core::proto::supervisor_host::v1::supervisor_host_client::SupervisorHostClient; use openshell_core::proto::supervisor_host::v1::{ - AssignmentSnapshot, HostHeartbeat, HostHello, HostMessage, HostStatus, host_gateway_message, + AssignmentSnapshot, HostGatewayMessage, HostHeartbeat, HostHello, HostMessage, HostStatus, + PoolMember, SandboxInstanceAssignment, SandboxInstanceTakeover, host_gateway_message, host_message, }; use openshell_extension_core::BearerTokenSlot; +use openshell_sandbox_backend::BoundaryTakeover; use openshell_sandbox_backend::runtime_routed::RuntimeRouter; use openshell_supervisor_control::{ SandboxInstancePhaseReport, SandboxInstanceRegistration, SandboxInstanceStatusReport, @@ -43,18 +52,29 @@ use crate::sandbox_instance_host::{ SandboxInstanceHostError, }; -/// Interval at which the host sends heartbeats on an idle stream. +/// Assignment lifecycle messages. Each names the route in its text, after +/// the phrase: the console's shorthand format keeps only an event's message, +/// and operators (and the supervisor pool E2E tests) follow which member +/// supervises a route by these lines. The structured fields stay as well. +const ADMITTED: &str = "admitted assigned sandbox instance"; +const REMOVING_WITHDRAWN: &str = "removing withdrawn sandbox instance"; +const RELINQUISHED_TO_OTHER_MEMBER: &str = + "relinquishing a sandbox instance owned by another pool member"; +const RELINQUISHED_FOR_LATER_TERM: &str = + "relinquishing a sandbox instance for its later ownership term"; + +/// Interval at which the host sends heartbeats on an idle stream, unless the +/// gateway asks for a shorter one. const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); -/// Silence after which the host abandons a stream and reconnects. -const GATEWAY_SILENCE_TIMEOUT: Duration = Duration::from_mins(1); -/// Bound on dialing the gateway, on the gateway answering the attach call, -/// and on it accepting the host. +/// Shortest heartbeat interval the host accepts from the gateway. +const MIN_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1); +/// Bound on connecting, attaching, and being accepted, once more than the +/// covered attempts since the last accepted stream went unanswered, so a slow +/// but responsive gateway is still reached. const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); -/// Bound on queueing one host-to-gateway message; a gateway that cannot -/// drain the stream this long is treated as lost. +/// Bound on queueing one host-to-gateway message, short of the stream's +/// silence budget; a gateway that cannot drain the stream is treated as lost. const SEND_TIMEOUT: Duration = Duration::from_secs(30); -const RECONNECT_BACKOFF_MIN: Duration = Duration::from_secs(1); -const RECONNECT_BACKOFF_MAX: Duration = Duration::from_secs(30); /// Host-to-gateway messages queued before the stream applies backpressure. const OUTBOUND_QUEUE: usize = 16; /// Most sandbox instances one snapshot may assign to a host; the gateway @@ -67,6 +87,83 @@ const MAX_GATEWAY_MESSAGE_BYTES: usize = /// Largest host message: a status report is a few hundred bytes per instance. const MAX_HOST_MESSAGE_BYTES: usize = 1024 * 1024; +/// How a host notices a silent gateway and paces its reconnects: it must +/// reach a responsive gateway replica well before its membership record +/// expires (see [`openshell_supervisor_control::pool_timing`]). +#[derive(Clone, Copy, Debug, PartialEq)] +struct StreamTimings { + /// Silence after which a stream is abandoned. + silence_timeout: Duration, + /// How often silence is checked. + silence_check: Duration, + /// First reconnect backoff; it doubles up to `backoff_max`. + backoff_min: Duration, + backoff_max: Duration, + /// Most jitter added to a backoff, as a fraction of it. + jitter: f64, + /// Bound on connecting, attaching, and being accepted while no more than + /// `covered_attempts` attempts failed since the last accepted stream; + /// later attempts get [`CONNECT_TIMEOUT`]. + attempt_timeout: Duration, + covered_attempts: u32, +} + +impl StreamTimings { + fn pool() -> Self { + use openshell_supervisor_control::pool_timing; + Self { + silence_timeout: pool_timing::HOST_SILENCE_TIMEOUT, + silence_check: pool_timing::HOST_SILENCE_CHECK, + backoff_min: pool_timing::HOST_RECONNECT_BACKOFF_MIN, + backoff_max: pool_timing::HOST_RECONNECT_BACKOFF_MAX, + jitter: pool_timing::HOST_RECONNECT_JITTER, + attempt_timeout: pool_timing::HOST_ATTEMPT_TIMEOUT, + covered_attempts: pool_timing::COVERED_UNANSWERED_ATTEMPTS, + } + } + + /// Bound on an attempt after `unanswered` attempts since the last + /// accepted stream that were never accepted. + fn attempt_timeout(&self, unanswered: u32) -> Duration { + if unanswered <= self.covered_attempts { + self.attempt_timeout + } else { + CONNECT_TIMEOUT + } + } + + /// Backoff before the `failed`-th (1-based) attempt since the last + /// accepted stream; `sample` in `[0, 1)` scales the jitter. + fn backoff(&self, failed: u32, sample: f64) -> Duration { + let doubled = self + .backoff_min + .saturating_mul(1_u32 << failed.saturating_sub(1).min(16)) + .min(self.backoff_max); + doubled.mul_f64(self.jitter.mul_add(sample.clamp(0.0, 1.0), 1.0)) + } +} + +/// A scheduling jitter sample in `[0, 1)`; not security-relevant. +fn jitter_sample() -> f64 { + use std::hash::BuildHasher as _; + let bits = std::collections::hash_map::RandomState::new().hash_one(SystemTime::now()); + #[allow(clippy::cast_precision_loss)] + let sample = (bits >> 11) as f64 / (1_u64 << 53) as f64; + sample +} + +/// Identity of a shared supervisor process in a supervisor pool. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PoolIdentity { + /// Stable name of the member (its Pod name). + pub member_id: String, + /// Identity of this incarnation of the member (its Pod UID). + pub member_uid: String, + /// Address of this member's routed runtime listener that other members + /// redirect its instances' runtimes to. + pub runtime_address: SocketAddr, +} + /// Process-level inputs of a shared supervisor. pub struct SharedHostConfig { /// Gateway gRPC endpoint every sandbox instance uses with its own @@ -78,6 +175,8 @@ pub struct SharedHostConfig { pub runtime_listen: SocketAddr, /// Time a withdrawn sandbox instance may take to stop its workload. pub removal_grace: Duration, + /// This process's identity as a member of the supervisor pool. + pub pool: PoolIdentity, } /// The sandbox instance registry and the routed runtime listener its @@ -93,10 +192,71 @@ pub struct SharedHost { /// Routes of the latest applied snapshot, so a withdrawal whose cleanup /// failed can be retried without the gateway resending the snapshot. desired_routes: Mutex>>, + /// This process's identity as a member of the supervisor pool. + pool: PoolIdentity, + /// Routes of the latest snapshot listed for stopping. Every other route + /// it omits is relinquished. + stop_routes: Mutex>, + /// The latest snapshot while some of its assignments wait for + /// an earlier instance of their route to leave, or for the credentials of + /// their term. + deferred: Mutex>, + /// Deferred assignments, reported as starting. The gateway + /// reads an assigned route the host does not report as gone, so a + /// deferral must never look like one. + deferred_reports: Mutex>, + /// The pool member announced that it is leaving. + leaving: watch::Sender, + /// A gateway currently accepts this host's stream. + gateway_connected: watch::Sender, + /// Recovery deadline (seconds) the latest snapshot offers the runtimes + /// of instances this host starts or takes over; `None` keeps the runtime + /// default. + recovery_timeout_secs: Mutex>, + /// Lets tests stall reconciliations and see how many run at once. + #[cfg(test)] + reconcile_probe: ReconcileProbe, +} + +/// Stalls the next reconciliation until released, and records how many ran +/// and how many run at once. +#[cfg(test)] +#[derive(Default)] +struct ReconcileProbe { + gate: Mutex>>, + entered: std::sync::atomic::AtomicUsize, + running: std::sync::atomic::AtomicUsize, + peak: std::sync::atomic::AtomicUsize, +} + +#[cfg(test)] +impl ReconcileProbe { + async fn enter(&self) -> impl Drop + '_ { + use std::sync::atomic::Ordering; + struct Running<'a>(&'a std::sync::atomic::AtomicUsize); + impl Drop for Running<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + self.entered.fetch_add(1, Ordering::SeqCst); + let now = self.running.fetch_add(1, Ordering::SeqCst) + 1; + self.peak.fetch_max(now, Ordering::SeqCst); + let gate = self + .gate + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + let running = Running(&self.running); + if let Some(gate) = gate { + let _ = gate.await; + } + running + } } /// A rejected sandbox instance assignment. -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, thiserror::Error)] pub enum ControlError { #[error("{0}")] Invalid(String), @@ -106,6 +266,9 @@ pub enum ControlError { Removing, #[error("{0}")] Unavailable(String), + /// The assignment can be admitted later without the gateway resending it. + #[error("{0}")] + Deferred(String), } impl From for ControlError { @@ -134,31 +297,132 @@ impl SharedHost { removal_grace: config.removal_grace, rejected: Mutex::new(BTreeMap::new()), desired_routes: Mutex::new(None), + pool: config.pool, + stop_routes: Mutex::new(std::collections::BTreeSet::new()), + deferred: Mutex::new(None), + deferred_reports: Mutex::new(BTreeMap::new()), + leaving: watch::Sender::new(false), + gateway_connected: watch::Sender::new(false), + recovery_timeout_secs: Mutex::new(None), + #[cfg(test)] + reconcile_probe: ReconcileProbe::default(), })) } + /// Pool identity of this process. + #[must_use] + pub fn pool(&self) -> &PoolIdentity { + &self.pool + } + + /// Whether this pool member announced that it is leaving. + #[must_use] + pub fn is_leaving(&self) -> bool { + *self.leaving.borrow() + } + + /// Announce that this pool member is leaving, and wait up to `limit` for + /// the gateway to move its instances to other members; each is + /// relinquished as a snapshot omits it. + pub async fn leave(&self, limit: Duration) { + tracing::info!( + hosted = self.host.list().len(), + "pool member is leaving; asking the gateway to move its sandbox instances" + ); + self.leaving.send_replace(true); + let mut changes = self.host.changes(); + let mut connected = self.gateway_connected.subscribe(); + // Only a gateway that hears the announcement moves the instances; one + // that stops hearing from this member takes them over after the + // membership TTL instead. + let drained = tokio::time::timeout(limit, async { + while !self.host.list().is_empty() && *connected.borrow_and_update() { + tokio::select! { + hosted = changes.changed() => if hosted.is_err() { return }, + gateway = connected.changed() => if gateway.is_err() { return }, + } + } + }) + .await; + if drained.is_err() { + tracing::warn!( + remaining = self.host.list().len(), + "pool member is leaving before the gateway moved all of its sandbox instances" + ); + } + } + + /// End supervision of every hosted instance at process shutdown: they are + /// relinquished to the members that take them over, and their workloads + /// keep running. + pub async fn shut_down(&self) { + self.host.relinquish_all().await; + } + /// Address of the routed runtime listener. #[must_use] pub fn runtime_addr(&self) -> SocketAddr { self.router.local_addr() } - /// Start a sandbox instance for `registration`, or confirm an identical - /// one. - pub fn register( + /// Admit `registration` under `fingerprint`. A new instance takes over + /// its generation with `takeover`; one whose term's credentials have not + /// arrived yet is deferred, while an instance already hosted for the same + /// fingerprint is kept. + fn admit( &self, registration: SandboxInstanceRegistration, + fingerprint: [u8; 32], + takeover: Result, ControlError>, ) -> Result { - registration - .validate() - .map_err(|error| ControlError::Invalid(error.to_string()))?; let admission = SandboxInstanceAdmission { key: registration.route.clone(), - fingerprint: registration_digest(®istration)?, + fingerprint, sandbox_id: registration.sandbox_id.clone(), }; - self.host - .admit(admission, |dir| self.instance_launch(registration, dir)) + self.host.admit(admission, |dir| { + self.instance_launch(registration, takeover?, dir) + }) + } + + /// Admit one assignment of a pool snapshot for its ownership term. + async fn admit_pool_assignment( + &self, + assignment: &SandboxInstanceAssignment, + ) -> Result { + let mut registration = SandboxInstanceRegistration::from_assignment(assignment) + .map_err(|error| ControlError::Invalid(error.to_string()))?; + let fingerprint = pool_fingerprint( + &assignment.launch_material, + registration.auth.auth_epoch.get(), + assignment.takeover.as_ref(), + )?; + let takeover = match &assignment.takeover { + None => Ok(None), + Some(takeover) => takeover_term(registration.auth.auth_epoch, takeover).map(|term| { + registration.auth.auth_epoch = term.auth_epoch; + registration.auth.gateway_token = term.gateway_token; + registration.auth.gateway_expires_at = takeover.gateway_expires_at; + registration.auth.sandbox_token = term.sandbox_token; + registration.auth.sandbox_expires_at = takeover.sandbox_expires_at; + Some(term.boundary) + }), + }; + match self.admit(registration.clone(), fingerprint, takeover.clone()) { + // This process was granted a later term of an instance it still + // supervises for an earlier one; hand the earlier term over first. + Err(ControlError::Conflict) => { + let route = &assignment.route; + tracing::info!( + route = %route, + sandbox_id = %assignment.sandbox_id, + "{RELINQUISHED_FOR_LATER_TERM} {route}" + ); + self.host.relinquish(&assignment.route).await; + self.admit(registration, fingerprint, takeover) + } + outcome => outcome, + } } /// Status of one hosted sandbox instance. @@ -170,20 +434,29 @@ impl SharedHost { /// Status of every hosted sandbox instance. An assignment of the latest /// snapshot that could not be admitted is reported as failed, even when /// an earlier instance still runs under its route, so the gateway sees - /// the rejection rather than the instance it could not replace. + /// the rejection rather than the instance it could not replace. A + /// deferred pool assignment is reported as starting, even while an + /// earlier instance of its route is still leaving. #[must_use] pub fn list(&self) -> Vec { let rejected = self .rejected .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); + let deferred = self + .deferred_reports + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); let mut reports = self .host .list() .into_iter() - .filter(|hosted| !rejected.contains_key(&hosted.route)) + .filter(|hosted| { + !rejected.contains_key(&hosted.route) && !deferred.contains_key(&hosted.route) + }) .collect::>(); reports.extend(rejected.values().cloned()); + reports.extend(deferred.values().cloned()); reports } @@ -199,26 +472,24 @@ impl SharedHost { self.host.remove(route, self.removal_grace).await } - /// Stop every sandbox instance, for process shutdown. - pub async fn remove_all(&self) { - self.host.remove_all(self.removal_grace).await; - } - /// Make the registry match `snapshot`: admit every assignment and remove /// every hosted instance it omits. Assignments that cannot be admitted are /// reported as failed instances; they never affect a running instance. /// + /// A snapshot that would run more than [`MAX_ASSIGNED_INSTANCES`] + /// instances still applies its stops, directory, and kept routes, and + /// keeps every hosted instance it assigns; only the new instances beyond + /// capacity are left unstarted, as kept routes are, until a later + /// snapshot fits. + /// /// # Errors /// - /// A malformed snapshot (too many instances, or one route assigned twice) - /// is refused as a whole before anything is admitted or removed. + /// A malformed snapshot (one route assigned twice, or the wrong kind for + /// this host) is refused as a whole before anything is admitted or + /// removed. pub async fn reconcile(&self, snapshot: &AssignmentSnapshot) -> Result<(), ControlError> { - if snapshot.instances.len() > MAX_ASSIGNED_INSTANCES { - return Err(ControlError::Invalid(format!( - "snapshot assigns {} instances; at most {MAX_ASSIGNED_INSTANCES} are supported", - snapshot.instances.len() - ))); - } + #[cfg(test)] + let _running = self.reconcile_probe.enter().await; let mut routes = std::collections::BTreeSet::new(); for assignment in &snapshot.instances { if !routes.insert(assignment.route.clone()) { @@ -227,24 +498,89 @@ impl SharedHost { )); } } + self.offer_recovery_timeout(snapshot.recovery_timeout.as_ref()); + // Only a pool snapshot says what to stop and where other members + // are; a snapshot without it would release everything it omits. + if !snapshot.pool { + return Err(ControlError::Invalid( + "the gateway sent a snapshot that is not a supervisor pool's; run a gateway \ + that supports the shared supervisor pool" + .to_string(), + )); + } + *self + .stop_routes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + snapshot.stop_routes.iter().cloned().collect(); + // Runtimes that reach this member for an instance another member + // supervises are redirected there. + let owners = snapshot + .route_directory + .iter() + .filter_map(|entry| { + let address = entry.runtime_address.parse().ok(); + if address.is_none() { + tracing::warn!(route = %entry.route, "ignoring a malformed route directory address"); + } + Some((entry.route.clone(), address?)) + }) + .collect(); + self.router + .set_directory(Some(self.pool.runtime_address), owners); + // Owned instances the gateway could not describe yet are kept as they + // are: neither started, released, nor stopped. + routes.extend(snapshot.keep_routes.iter().cloned()); + let hosted = self + .host + .list() + .into_iter() + .map(|status| status.route) + .collect::>(); + let (admitted, unstarted) = within_capacity(snapshot, &hosted); + if !unstarted.is_empty() { + tracing::warn!( + assigned = snapshot.instances.len(), + unstarted = unstarted.len(), + "snapshot assigns more sandbox instances than this host runs; leaving the excess unstarted" + ); + } *self .desired_routes .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(routes); let mut rejected = BTreeMap::new(); - for assignment in &snapshot.instances { - let outcome = SandboxInstanceRegistration::from_assignment(assignment) - .map_err(|error| ControlError::Invalid(error.to_string())) - .and_then(|registration| self.register(registration)); + let mut deferred = BTreeMap::new(); + for assignment in admitted { + let outcome = self.admit_pool_assignment(assignment).await; match outcome { Ok(SandboxInstanceAdmitted::Created) => { + let route = &assignment.route; tracing::info!( - route = %assignment.route, + route = %route, sandbox_id = %assignment.sandbox_id, - "admitted assigned sandbox instance" + "{ADMITTED} {route}" ); } Ok(SandboxInstanceAdmitted::Unchanged) => {} + // An earlier instance of the route is still leaving, or the + // term's credentials are on their way; retried shortly. + Err(error @ (ControlError::Removing | ControlError::Deferred(_))) => { + tracing::debug!( + route = %assignment.route, + %error, + "deferred assigned sandbox instance" + ); + deferred.insert( + assignment.route.clone(), + SandboxInstanceStatusReport { + route: assignment.route.clone(), + sandbox_id: assignment.sandbox_id.clone(), + phase: SandboxInstancePhaseReport::Starting, + ready: false, + }, + ); + } Err(error) => { tracing::warn!( route = %assignment.route, @@ -270,10 +606,47 @@ impl SharedHost { .rejected .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = rejected; + *self + .deferred + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + (!deferred.is_empty()).then(|| snapshot.clone()); + *self + .deferred_reports + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = deferred; self.remove_withdrawn().await; Ok(()) } + /// Retry what the latest snapshot left undone without the gateway + /// resending it: deferred admissions and withdrawals whose cleanup failed. + pub async fn retry_pending(&self) { + let deferred = self + .deferred + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + match deferred { + Some(snapshot) => { + if let Err(error) = self.reconcile(&snapshot).await { + tracing::warn!(%error, "could not retry a deferred assignment snapshot"); + } + } + None => self.remove_withdrawn().await, + } + } + + /// Whether something of the latest snapshot is still to be retried. + #[must_use] + pub fn has_pending(&self) -> bool { + self.deferred + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_some() + || self.has_withdrawn_leftovers() + } + /// Remove every hosted instance the latest applied snapshot omits. Also /// run between snapshots: a removal whose file cleanup failed keeps its /// key reserved and reports a failure, and the gateway does not resend an @@ -294,13 +667,31 @@ impl SharedHost { .into_iter() .filter(|hosted| !desired.contains(&hosted.route)) .collect::>(); - futures::future::join_all(withdrawn.iter().map(|hosted| async move { - tracing::info!( - route = %hosted.route, - sandbox_id = %hosted.sandbox_id, - "removing withdrawn sandbox instance" - ); - self.remove(&hosted.route).await; + let stop_routes = self + .stop_routes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + futures::future::join_all(withdrawn.iter().map(|hosted| { + let stop = stop_routes.contains(&hosted.route); + async move { + let route = &hosted.route; + if stop { + tracing::info!( + route = %route, + sandbox_id = %hosted.sandbox_id, + "{REMOVING_WITHDRAWN} {route}" + ); + self.remove(&hosted.route).await; + } else { + tracing::info!( + route = %route, + sandbox_id = %hosted.sandbox_id, + "{RELINQUISHED_TO_OTHER_MEMBER} {route}" + ); + self.host.relinquish(&hosted.route).await; + } + } })) .await; } @@ -322,10 +713,39 @@ impl SharedHost { .any(|hosted| !desired.contains(&hosted.route)) } + /// Keep the recovery deadline a snapshot offers for the instances admitted + /// from now on. Absent offers none; a value the runtime would refuse is + /// ignored, so the runtime default applies. + fn offer_recovery_timeout(&self, offered: Option<&prost_types::Duration>) { + let offer = offered.and_then(|offered| { + let seconds = u32::try_from(offered.seconds) + .ok() + .filter(|_| offered.nanos == 0) + .filter(|seconds| { + openshell_sandbox_backend::boundary_protocol::supervisor_recovery_timeout( + *seconds, + ) + .is_some() + }); + if seconds.is_none() { + tracing::warn!( + seconds = offered.seconds, + "ignoring an out-of-range sandbox recovery deadline from the gateway" + ); + } + seconds + }); + *self + .recovery_timeout_secs + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = offer; + } + /// Launch material of `registration` with its private files in `dir`. fn instance_launch( &self, registration: SandboxInstanceRegistration, + takeover: Option, dir: &Path, ) -> Result<(crate::SandboxInstanceBindings, SandboxInstanceBody), ControlError> { let failure = |error: miette::Report| ControlError::Unavailable(error.to_string()); @@ -388,6 +808,11 @@ impl SharedHost { ca_private_key: Some(ca_private_key), directory: Some(trust_dir), }, + takeover, + supervisor_recovery_timeout_secs: *self + .recovery_timeout_secs + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), }; crate::sandbox_instance_host::sandbox_instance_launch( Box::new(crate::backend_setup::RoutedBackendSetup::new( @@ -401,12 +826,77 @@ impl SharedHost { } } -fn registration_digest( - registration: &SandboxInstanceRegistration, +/// Admission fingerprint of a pool instance: its launch material and its +/// ownership term (the launch term, or a takeover term with its granted +/// supervisor process). Credentials minted for a term differ between gateway +/// replicas and reconnects, so they are not part of it: a host keeps the +/// instance it runs for the same term whichever gateway it hears from. +fn pool_fingerprint( + launch_material: &[u8], + launch_epoch: u64, + takeover: Option<&SandboxInstanceTakeover>, ) -> Result<[u8; 32], ControlError> { - let encoded = serde_json::to_vec(registration) - .map_err(|error| ControlError::Invalid(format!("encode registration: {error}")))?; - openshell_crypto::sha256(&encoded).map_err(|error| ControlError::Unavailable(error.to_string())) + let (auth_epoch, supervisor_instance_id) = takeover.map_or((launch_epoch, ""), |takeover| { + ( + takeover.auth_epoch, + takeover.supervisor_instance_id.as_str(), + ) + }); + let mut term = Vec::with_capacity(launch_material.len() + 96); + term.extend_from_slice(b"openshell-pool-instance-v1\0"); + term.extend_from_slice(&auth_epoch.to_be_bytes()); + term.extend_from_slice(supervisor_instance_id.as_bytes()); + term.push(0); + term.extend_from_slice(launch_material); + openshell_crypto::sha256(&term).map_err(|error| ControlError::Unavailable(error.to_string())) +} + +/// Credentials of one takeover term and the grant that lets a new instance +/// take its generation over. +struct TakeoverTerm { + auth_epoch: openshell_core::jwt::CredentialEpoch, + gateway_token: openshell_core::jwt::SecretJwt, + sandbox_token: openshell_core::jwt::SecretJwt, + boundary: BoundaryTakeover, +} + +/// Validate a takeover term that follows the launch term `launch_epoch`. A +/// term without credentials, minted by another gateway replica, is +/// [`ControlError::Deferred`]: an instance already running for the term is +/// kept, and a new one waits for credentials from the gateway this host +/// hears from. +fn takeover_term( + launch_epoch: openshell_core::jwt::CredentialEpoch, + takeover: &SandboxInstanceTakeover, +) -> Result { + let invalid = |what: &str| ControlError::Invalid(format!("sandbox instance takeover {what}")); + let auth_epoch = openshell_core::jwt::CredentialEpoch::new(takeover.auth_epoch) + .map_err(|_| invalid("names no ownership term"))?; + if auth_epoch <= launch_epoch { + return Err(invalid("does not follow the launch term")); + } + let supervisor_instance_id = takeover + .supervisor_instance_id + .parse::() + .map_err(|_| invalid("names no supervisor process"))?; + if takeover.grant.is_empty() { + return Err(ControlError::Deferred( + "the takeover credentials of this ownership term have not arrived".to_string(), + )); + } + let secret = |token: &str| { + openshell_core::jwt::SecretJwt::parse(token.to_string()) + .map_err(|_| invalid("carries a malformed credential")) + }; + Ok(TakeoverTerm { + auth_epoch, + gateway_token: secret(&takeover.gateway_token)?, + sandbox_token: secret(&takeover.sandbox_token)?, + boundary: BoundaryTakeover { + supervisor_instance_id, + grant: secret(&takeover.grant)?, + }, + }) } #[cfg(unix)] @@ -477,6 +967,7 @@ impl GatewayLink { fn status_message(host: &SharedHost, applied_revision: u64) -> HostMessage { HostMessage { payload: Some(host_message::Payload::Status(HostStatus { + leaving: host.is_leaving(), applied_revision, instances: host .list() @@ -494,6 +985,58 @@ enum StreamEnd { Lost { accepted: bool }, } +/// What a gateway sends a host on its assignment stream. +type GatewayInbound = std::pin::Pin< + Box> + Send>, +>; + +/// Opens a host's assignment stream: hands the gateway the host's outbound +/// messages and returns the gateway's. +trait GatewayConnector: Send + Sync { + /// Connect and attach, each stage bounded by `deadline`. + fn open( + &self, + outbound: tokio::sync::mpsc::Receiver, + deadline: tokio::time::Instant, + ) -> std::pin::Pin> + Send + '_>>; +} + +impl GatewayConnector for GatewayLink { + fn open( + &self, + outbound: tokio::sync::mpsc::Receiver, + deadline: tokio::time::Instant, + ) -> std::pin::Pin> + Send + '_>> { + Box::pin(async move { + if !self.endpoint.starts_with("https://") { + return Err("the supervisor token requires an https gateway endpoint".to_string()); + } + let token = self + .read_token() + .map_err(|error| format!("no usable supervisor token: {error}"))?; + let connect = openshell_core::grpc_client::connect_gateway_channel(&self.endpoint); + let channel = tokio::time::timeout_at(deadline, connect) + .await + .map_err(|_| "connecting to the gateway timed out".to_string())? + .map_err(|error| format!("could not connect to the gateway: {error}"))?; + let mut client = SupervisorHostClient::with_interceptor(channel, token.interceptor()) + .max_decoding_message_size(MAX_GATEWAY_MESSAGE_BYTES) + .max_encoding_message_size(MAX_HOST_MESSAGE_BYTES); + let attach = client.attach(tokio_stream::wrappers::ReceiverStream::new(outbound)); + let response = tokio::time::timeout_at(deadline, attach) + .await + .map_err(|_| { + "gateway did not answer the supervisor host stream in time".to_string() + })? + .map_err(|status| { + format!("gateway refused the supervisor host stream: {status}") + })?; + let inbound: GatewayInbound = Box::pin(response.into_inner()); + Ok(inbound) + }) + } +} + /// Keep one assignment stream open to the gateway until `shutdown` /// completes, reconnecting with backoff. Hosted instances run on across /// disconnects. @@ -501,124 +1044,356 @@ pub async fn attach_to_gateway( host: Arc, link: GatewayLink, shutdown: impl Future, +) { + attach_with(host, Arc::new(link), shutdown).await; +} + +async fn attach_with( + host: Arc, + connector: Arc, + shutdown: impl Future, ) { let instance_id = uuid::Uuid::new_v4().to_string(); let mut connection_epoch = 0_u64; - let mut backoff = RECONNECT_BACKOFF_MIN; + let timings = StreamTimings::pool(); + let reconciler = Reconciler::start(host.clone()); + // Streams lost since the last accepted stream, counting its own loss; + // the backoff grows with it. + let mut failed = 0_u32; + // Attempts since the last accepted stream that were never accepted; the + // attempt bound depends on it, so the attempt after the covered + // unanswered ones is still quick. + let mut unanswered = 0_u32; tokio::pin!(shutdown); loop { connection_epoch += 1; let hello = HostHello { instance_id: instance_id.clone(), connection_epoch, + member: Some(PoolMember { + member_id: host.pool.member_id.clone(), + member_uid: host.pool.member_uid.clone(), + runtime_address: host.pool.runtime_address.to_string(), + }), }; + let attempt_timeout = timings.attempt_timeout(unanswered); let StreamEnd::Lost { accepted } = tokio::select! { - end = run_stream(&host, &link, hello) => end, - () = &mut shutdown => return, + end = run_stream(&host, connector.as_ref(), &reconciler, hello, &timings, attempt_timeout) => end, + () = &mut shutdown => return reconciler.shut_down().await, }; if accepted { - backoff = RECONNECT_BACKOFF_MIN; + failed = 0; + unanswered = 0; + } else { + unanswered = unanswered.saturating_add(1); } + failed = failed.saturating_add(1); + let backoff = timings.backoff(failed, jitter_sample()); tracing::info!( - retry_in_secs = backoff.as_secs(), + retry_in_ms = u64::try_from(backoff.as_millis()).unwrap_or(u64::MAX), "gateway assignment stream lost; hosted sandbox instances keep running" ); tokio::select! { () = tokio::time::sleep(backoff) => {} - () = &mut shutdown => return, + () = &mut shutdown => return reconciler.shut_down().await, } - backoff = (backoff * 2).min(RECONNECT_BACKOFF_MAX); } } -/// One connection and stream to the gateway. +/// One unit of registry work. +enum ReconcileJob { + Snapshot(AssignmentSnapshot), + /// Retry deferred admissions and withdrawals whose cleanup failed. + Retry, +} + +/// A finished registry job of stream `stream`; `revision` is the snapshot +/// revision it applied, if any. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct Applied { + sequence: u64, + stream: u64, + revision: Option, + /// The job ran, rather than refusing a malformed snapshot. + ok: bool, +} + +#[derive(Default)] +struct ReconcilerState { + /// The stream that may queue work; older streams' work is dropped. + stream: u64, + queued: Option<(u64, ReconcileJob)>, + busy: bool, + closed: bool, +} + +/// Applies assignment snapshots to the registry one at a time, across +/// streams. /// -/// Snapshots are reconciled on a separate task so heartbeats and status -/// reports keep flowing while instances stop; at most one reconciliation runs -/// at a time, only the newest unapplied snapshot is kept, and the stream does -/// not end until a running reconciliation has finished. -async fn run_stream(host: &Arc, link: &GatewayLink, hello: HostHello) -> StreamEnd { - let lost = |accepted| StreamEnd::Lost { accepted }; - if !link.endpoint.starts_with("https://") { - tracing::error!("the supervisor token requires an https gateway endpoint"); - return lost(false); +/// A stream hands the worker the newest snapshot it received; one still +/// queued when a newer stream begins is dropped, since its session is gone. +/// A reconciliation the worker began runs to completion even when its stream +/// ends, so registry mutation stays serialized while a new stream connects, +/// reports status, and queues its own snapshots without waiting for it. +struct Reconciler { + state: Mutex, + wake: tokio::sync::Notify, + done: watch::Sender, +} + +impl Reconciler { + /// Start the worker; it stops once the returned handle is shut down or + /// dropped and the job in flight, if any, has finished. + fn start(host: Arc) -> ReconcilerHandle { + let reconciler = Arc::new(Self { + state: Mutex::new(ReconcilerState::default()), + wake: tokio::sync::Notify::new(), + done: watch::Sender::new(Applied::default()), + }); + let worker = tokio::spawn(reconciler.clone().run(host)); + ReconcilerHandle { + reconciler, + worker: Some(worker), + } } - let token = match link.read_token() { - Ok(token) => token, - Err(error) => { - tracing::warn!(%error, "no usable supervisor token; will retry"); - return lost(false); + + /// Accept no more work and drop what is queued; the job in flight, if + /// any, still finishes. + fn close(&self) { + let mut state = self.state(); + state.closed = true; + state.queued = None; + drop(state); + self.wake.notify_one(); + } + + fn state(&self) -> std::sync::MutexGuard<'_, ReconcilerState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Begin a stream: work still queued for earlier streams is dropped. + fn begin_stream(&self) -> u64 { + let mut state = self.state(); + state.stream += 1; + state.queued = None; + state.stream + } + + /// Queue `snapshot` for `stream`, replacing whatever it queued before. + fn submit(&self, stream: u64, snapshot: AssignmentSnapshot) { + let mut state = self.state(); + if state.stream == stream && !state.closed { + state.queued = Some((stream, ReconcileJob::Snapshot(snapshot))); + drop(state); + self.wake.notify_one(); } - }; - let connect = openshell_core::grpc_client::connect_gateway_channel(&link.endpoint); - let channel = match tokio::time::timeout(CONNECT_TIMEOUT, connect).await { - Ok(Ok(channel)) => channel, - Ok(Err(error)) => { - tracing::warn!(%error, "could not connect to the gateway"); - return lost(false); + } + + /// Queue a retry for `stream` when the worker has nothing to do. + fn submit_retry(&self, stream: u64) { + let mut state = self.state(); + if state.stream == stream && !state.closed && !state.busy && state.queued.is_none() { + state.queued = Some((stream, ReconcileJob::Retry)); + drop(state); + self.wake.notify_one(); } - Err(_) => { - tracing::warn!("connecting to the gateway timed out"); - return lost(false); + } + + async fn run(self: Arc, host: Arc) { + loop { + let job = { + let mut state = self.state(); + match state.queued.take() { + Some(job) => { + state.busy = true; + Some(job) + } + None if state.closed => return, + None => None, + } + }; + let Some((stream, job)) = job else { + self.wake.notified().await; + continue; + }; + let (revision, ok) = match job { + ReconcileJob::Snapshot(snapshot) => match host.reconcile(&snapshot).await { + Ok(()) => (Some(snapshot.revision), true), + Err(error) => { + tracing::warn!(%error, "refused a malformed assignment snapshot"); + (None, false) + } + }, + ReconcileJob::Retry => { + host.retry_pending().await; + (None, true) + } + }; + self.state().busy = false; + self.done.send_modify(|applied| { + *applied = Applied { + sequence: applied.sequence + 1, + stream, + revision, + ok, + }; + }); } - }; - let mut client = SupervisorHostClient::with_interceptor(channel, token.interceptor()) - .max_decoding_message_size(MAX_GATEWAY_MESSAGE_BYTES) - .max_encoding_message_size(MAX_HOST_MESSAGE_BYTES); + } +} + +/// Owns the worker of an attach loop. Shutting it down waits for the job in +/// flight, so nothing the worker admits outlives the attachment; dropping it +/// lets the worker finish that job and stop on its own. +struct ReconcilerHandle { + reconciler: Arc, + worker: Option>, +} + +impl ReconcilerHandle { + /// Stop the worker at process shutdown: queued work is dropped and the + /// reconciliation in flight, if any, finishes before this returns, so the + /// host's own shutdown sees every instance the worker admitted. + async fn shut_down(mut self) { + self.reconciler.close(); + if let Some(worker) = self.worker.take() { + let _ = worker.await; + } + } +} + +impl std::ops::Deref for ReconcilerHandle { + type Target = Reconciler; + + fn deref(&self) -> &Reconciler { + &self.reconciler + } +} + +impl Drop for ReconcilerHandle { + fn drop(&mut self) { + self.reconciler.close(); + } +} + +/// The assignments of `snapshot` a host admits, and the routes it leaves +/// unstarted. Hosted instances the snapshot assigns or keeps stay; new ones +/// are admitted in route order while fewer than [`MAX_ASSIGNED_INSTANCES`] +/// would run. +fn within_capacity<'a>( + snapshot: &'a AssignmentSnapshot, + hosted: &std::collections::BTreeSet, +) -> (Vec<&'a SandboxInstanceAssignment>, Vec<&'a str>) { + let assigned = snapshot + .instances + .iter() + .map(|assignment| assignment.route.as_str()) + .collect::>(); + let staying = hosted + .iter() + .filter(|route| assigned.contains(route.as_str()) || snapshot.keep_routes.contains(*route)) + .count(); + let mut room = MAX_ASSIGNED_INSTANCES.saturating_sub(staying); + let mut fresh = snapshot + .instances + .iter() + .filter(|assignment| !hosted.contains(&assignment.route)) + .collect::>(); + fresh.sort_by(|left, right| left.route.cmp(&right.route)); + let mut admitted = snapshot + .instances + .iter() + .filter(|assignment| hosted.contains(&assignment.route)) + .collect::>(); + let mut unstarted = Vec::new(); + for assignment in fresh { + if room > 0 { + room -= 1; + admitted.push(assignment); + } else { + unstarted.push(assignment.route.as_str()); + } + } + (admitted, unstarted) +} + +/// One connection and stream to the gateway. +/// +/// One deadline bounds connecting, attaching, and being accepted. Snapshots +/// go to the attach loop's reconcile worker, so heartbeats and status reports +/// keep flowing while instances stop, and the stream ends without waiting for +/// a reconciliation in flight. A pool member also bounds every send by its +/// silence budget, so a gateway that stops reading is abandoned as promptly +/// as one that stops writing. +async fn run_stream( + host: &Arc, + connector: &dyn GatewayConnector, + reconciler: &Reconciler, + hello: HostHello, + timings: &StreamTimings, + attempt_timeout: Duration, +) -> StreamEnd { + let lost = |accepted| StreamEnd::Lost { accepted }; + // One deadline covers saying hello, connecting, attaching, and being + // accepted, so a slow stage leaves less time for the next; the reconnect + // budget counts on it. + let deadline = tokio::time::Instant::now() + attempt_timeout; + let stream = reconciler.begin_stream(); let (tx, rx) = tokio::sync::mpsc::channel::(OUTBOUND_QUEUE); - let send = |message: HostMessage| { + // A send gives up by the ordinary bound, or once the stream has been + // silent for its whole silence budget. + let send_deadline = |last_heard: tokio::time::Instant| { + (tokio::time::Instant::now() + SEND_TIMEOUT).min(last_heard + timings.silence_timeout) + }; + let send = |message: HostMessage, deadline: tokio::time::Instant| { let tx = tx.clone(); - async move { tx.send_timeout(message, SEND_TIMEOUT).await.is_ok() } + async move { + tokio::time::timeout_at(deadline, tx.send(message)) + .await + .is_ok_and(|sent| sent.is_ok()) + } }; - if !send(HostMessage { - payload: Some(host_message::Payload::Hello(hello)), - }) + if !send( + HostMessage { + payload: Some(host_message::Payload::Hello(hello)), + }, + deadline, + ) .await { return lost(false); } - let attach = client.attach(tokio_stream::wrappers::ReceiverStream::new(rx)); - let mut inbound = match tokio::time::timeout(CONNECT_TIMEOUT, attach).await { - Ok(Ok(response)) => response.into_inner(), - Ok(Err(status)) => { - tracing::warn!(%status, "gateway refused the supervisor host stream"); - return lost(false); - } - Err(_) => { - tracing::warn!("gateway did not answer the supervisor host stream in time"); + let opened = tokio::time::timeout_at(deadline, connector.open(rx, deadline)) + .await + .unwrap_or_else(|_| Err("connecting to the gateway timed out".to_string())); + let mut inbound = match opened { + Ok(inbound) => inbound, + Err(error) => { + tracing::warn!(%error, "could not open the supervisor host stream"); return lost(false); } }; + let accept_deadline = deadline; let mut changes = host.changes(); changes.mark_unchanged(); + let mut leaving = host.leaving.subscribe(); + leaving.mark_unchanged(); + let mut finished = reconciler.done.subscribe(); + finished.mark_unchanged(); let mut applied_revision = 0_u64; // Highest snapshot revision received on this stream, applied or not, so a // snapshot that arrives behind a newer in-flight one is never applied. let mut received_revision = 0_u64; let mut accepted = false; - let accept_deadline = tokio::time::Instant::now() + CONNECT_TIMEOUT; let mut last_heard = tokio::time::Instant::now(); let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL); heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - // The reconciliation in flight and the newest snapshot waiting for it. - // A retry of withdrawals whose cleanup failed runs through the same slot, - // so at most one reconciliation touches the registry at a time. - let mut applying: Option, ControlError>>> = None; - let mut queued: Option = None; - let start = |snapshot: AssignmentSnapshot| { - let host = host.clone(); - tokio::spawn(async move { - host.reconcile(&snapshot).await?; - Ok(Some(snapshot.revision)) - }) - }; - let retry_withdrawn = || { - let host = host.clone(); - tokio::spawn(async move { - host.remove_withdrawn().await; - Ok(None) - }) - }; + // A pool member checks for silence and for acceptance on its own, short + // cadence; a lone host checks on each heartbeat. + let mut silence_check = tokio::time::interval(timings.silence_check); + silence_check.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let end = loop { tokio::select! { message = inbound.next() => { @@ -634,9 +1409,22 @@ async fn run_stream(host: &Arc, link: &GatewayLink, hello: HostHello match payload { Some(host_gateway_message::Payload::Accepted(accepted_message)) => { accepted = true; + host.gateway_connected.send_replace(true); tracing::info!(session = %accepted_message.session_id, "gateway accepted the supervisor host"); + // A pool gateway detects failed members by silence, + // so it asks for more frequent heartbeats. + if let Some(interval) = accepted_message + .heartbeat_interval + .and_then(|interval| Duration::try_from(interval).ok()) + .filter(|interval| { + *interval >= MIN_HEARTBEAT_INTERVAL && *interval < HEARTBEAT_INTERVAL + }) + { + heartbeat = tokio::time::interval(interval); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + } // Report what runs before any snapshot arrives. - if !send(status_message(host, applied_revision)).await { + if !send(status_message(host, applied_revision), send_deadline(last_heard)).await { break lost(accepted); } } @@ -654,71 +1442,995 @@ async fn run_stream(host: &Arc, link: &GatewayLink, hello: HostHello continue; } received_revision = snapshot.revision; - if applying.is_some() { - queued = Some(snapshot); - } else { - applying = Some(start(snapshot)); - } + reconciler.submit(stream, snapshot); } Some(host_gateway_message::Payload::Heartbeat(_)) | None => {} } } - applied = async { applying.as_mut().expect("polled only while applying").await }, if applying.is_some() => { - applying = None; - match applied { - Ok(Ok(revision)) => { - if let Some(revision) = revision { - applied_revision = revision; - } - changes.mark_unchanged(); - if !send(status_message(host, applied_revision)).await { - break lost(accepted); - } - } - Ok(Err(error)) => { - tracing::warn!(%error, "refused a malformed assignment snapshot"); - } - Err(error) => { - tracing::error!(%error, "assignment reconciliation task failed"); - } + worked = finished.changed() => { + if worked.is_err() { + break lost(accepted); } - if let Some(snapshot) = queued.take() { - applying = Some(start(snapshot)); + let applied = *finished.borrow_and_update(); + if applied.stream == stream + && let Some(revision) = applied.revision + { + applied_revision = revision; + } + // A refused snapshot changes nothing to report. + if applied.ok { + changes.mark_unchanged(); + if accepted && !send(status_message(host, applied_revision), send_deadline(last_heard)).await { + break lost(accepted); + } } } notified = changes.changed() => { if notified.is_err() { break lost(accepted); } - if accepted && !send(status_message(host, applied_revision)).await { + if accepted && !send(status_message(host, applied_revision), send_deadline(last_heard)).await { break lost(accepted); } } - _ = heartbeat.tick() => { - if !accepted && tokio::time::Instant::now() > accept_deadline { - tracing::warn!("gateway did not accept the supervisor host in time"); + // The gateway moves a leaving member's instances at once. + _ = leaving.changed() => { + if accepted && !send(status_message(host, applied_revision), send_deadline(last_heard)).await { break lost(accepted); } - if last_heard.elapsed() > GATEWAY_SILENCE_TIMEOUT { + } + // Acceptance is given up exactly at the attempt's deadline. + () = tokio::time::sleep_until(accept_deadline), if !accepted => { + tracing::warn!("gateway did not accept the supervisor host in time"); + break lost(accepted); + } + _ = silence_check.tick() => { + if last_heard.elapsed() > timings.silence_timeout { tracing::warn!("gateway assignment stream went silent; reconnecting"); break lost(accepted); } + } + _ = heartbeat.tick() => { let beat = HostMessage { payload: Some(host_message::Payload::Heartbeat(HostHeartbeat {})) }; - if !send(beat).await { + if !send(beat, send_deadline(last_heard)).await { break lost(accepted); } - // Withdrawn instances whose cleanup failed are retried at the - // heartbeat cadence, never concurrently with a snapshot. - if applying.is_none() && queued.is_none() && host.has_withdrawn_leftovers() { - applying = Some(retry_withdrawn()); + // Deferred admissions and withdrawn instances whose cleanup + // failed are retried at the heartbeat cadence, never + // concurrently with a snapshot. + if host.has_pending() { + reconciler.submit_retry(stream); } } } }; - // Finish the reconciliation this stream started before another stream - // can begin one; the queued snapshot belongs to a session that is gone. - if let Some(applying) = applying { - let _ = applying.await; - } + host.gateway_connected.send_replace(false); end } + +#[cfg(test)] +mod tests { + use super::*; + use crate::sandbox_instance_host::SandboxInstanceAdmission; + use openshell_core::jwt::CredentialEpoch; + use openshell_sandbox_backend::boundary_protocol::SupervisorInstanceId; + + /// A hosted instance whose stop path reports on `stopped`. + fn admit_supervising(host: &SharedHost, route: &str, stopped: tokio::sync::mpsc::Sender<()>) { + let bindings = crate::SandboxInstanceBindings { + gateway_credentials: + openshell_core::grpc_client::GatewayCredentials::from_gateway_token( + "pool-test-token", + ) + .unwrap(), + event_context: Arc::new(crate::sandbox_instance_event_context( + Some(route), + Some(route), + String::new(), + )), + ocsf_json: openshell_ocsf::OcsfJsonSettings::default(), + log_sink: None, + denial_log: + openshell_supervisor_network::policy_local::DenialLogSource::SandboxInstance( + openshell_supervisor_network::policy_local::SandboxInstanceDenialLog::default(), + ), + }; + let body: SandboxInstanceBody = Box::new(move |controls| { + Box::pin(async move { + controls.readiness.send_replace(true); + controls.shutdown.await; + let _ = stopped.send(()).await; + Ok(0) + }) + }); + host.host + .admit( + SandboxInstanceAdmission { + key: route.to_string(), + fingerprint: [7; 32], + sandbox_id: route.to_string(), + }, + |_| Ok::<_, ControlError>((bindings, body)), + ) + .unwrap(); + } + + /// A pool snapshot sets the recovery deadline offered to the runtimes of + /// instances admitted after it; none, or one the runtime would refuse, + /// leaves the runtime default. + #[tokio::test] + async fn snapshots_set_the_recovery_deadline_offered_to_new_instances() { + let root = tempfile::tempdir().unwrap(); + let host = SharedHost::start(SharedHostConfig { + gateway_endpoint: "https://gateway.invalid".to_string(), + instance_root: root.path().join("instances"), + runtime_listen: "127.0.0.1:0".parse().unwrap(), + removal_grace: Duration::from_secs(10), + pool: PoolIdentity { + member_id: "member-1".to_string(), + member_uid: "uid-1".to_string(), + runtime_address: "127.0.0.1:5500".parse().unwrap(), + }, + }) + .unwrap(); + for (revision, (offered, kept)) in [ + (Some(60), Some(60)), + (None, None), + (Some(301), None), + (Some(10), Some(10)), + ] + .into_iter() + .enumerate() + { + host.reconcile(&AssignmentSnapshot { + revision: u64::try_from(revision).unwrap() + 1, + pool: true, + recovery_timeout: offered + .map(|seconds| prost_types::Duration { seconds, nanos: 0 }), + ..AssignmentSnapshot::default() + }) + .await + .unwrap(); + assert_eq!(*host.recovery_timeout_secs.lock().unwrap(), kept); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn pool_members_keep_release_or_stop_unassigned_instances_as_snapshots_say() { + let root = tempfile::tempdir().unwrap(); + let host = SharedHost::start(SharedHostConfig { + gateway_endpoint: "https://gateway.invalid".to_string(), + instance_root: root.path().join("instances"), + runtime_listen: "127.0.0.1:0".parse().unwrap(), + removal_grace: Duration::from_secs(10), + pool: PoolIdentity { + member_id: "member-1".to_string(), + member_uid: "uid-1".to_string(), + runtime_address: "127.0.0.1:5500".parse().unwrap(), + }, + }) + .unwrap(); + let routes = [ + "rt-kept-0123456789", + "rt-released-012345", + "rt-stopped-0123456", + ]; + let mut stops = Vec::new(); + for route in routes { + let (stopped, stops_of) = tokio::sync::mpsc::channel(1); + admit_supervising(&host, route, stopped); + stops.push(stops_of); + } + tokio::time::timeout(Duration::from_secs(10), async { + while host + .list() + .iter() + .filter(|status| status.phase == SandboxInstancePhaseReport::Supervising) + .count() + < routes.len() + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("instances supervise"); + + // A gateway replica whose desired set lacks the kept instance, such + // as one this member just reconnected to, still names it as owned. + host.reconcile(&AssignmentSnapshot { + revision: 1, + pool: true, + stop_routes: vec![routes[2].to_string()], + keep_routes: vec![routes[0].to_string()], + ..AssignmentSnapshot::default() + }) + .await + .unwrap(); + + // The kept instance is untouched; the released one ended without + // stopping its workload; only the withdrawn one was stopped. + assert_eq!( + host.status(routes[0]).map(|status| status.phase), + Some(SandboxInstancePhaseReport::Supervising) + ); + assert!(host.status(routes[1]).is_none()); + assert_eq!( + stops[1].try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected), + "a released instance never runs its stop path" + ); + assert!(host.status(routes[2]).is_none()); + assert_eq!( + stops[2].try_recv(), + Ok(()), + "a withdrawn instance is stopped" + ); + assert!(stops[0].try_recv().is_err()); + + // A snapshot that is not a pool's is refused. + assert!( + host.reconcile(&AssignmentSnapshot { + revision: 2, + ..AssignmentSnapshot::default() + }) + .await + .is_err() + ); + assert_eq!( + host.status(routes[0]).map(|status| status.phase), + Some(SandboxInstancePhaseReport::Supervising) + ); + host.shut_down().await; + assert!(host.list().is_empty()); + assert_eq!( + stops[0].try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected), + "a pool member relinquishes what it still hosts at shutdown" + ); + } + + /// Console output of the shorthand layer the supervisor logs through. + #[derive(Clone, Default)] + struct Captured(Arc>>); + + impl std::io::Write for Captured { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + /// The shorthand console format keeps only an event's message, so the + /// lines that say which member supervises a route carry the route in + /// their text. + #[tokio::test] + async fn assignment_lifecycle_lines_name_their_route_in_shorthand_output() { + use tracing_subscriber::layer::SubscriberExt as _; + let captured = Captured::default(); + let _logging = + tracing::subscriber::set_default(tracing_subscriber::registry().with( + openshell_ocsf::OcsfShorthandLayer::new(captured.clone()).with_non_ocsf(true), + )); + let root = tempfile::tempdir().unwrap(); + let host = SharedHost::start(SharedHostConfig { + gateway_endpoint: "https://gateway.invalid".to_string(), + instance_root: root.path().join("instances"), + runtime_listen: "127.0.0.1:0".parse().unwrap(), + removal_grace: Duration::from_secs(10), + pool: PoolIdentity { + member_id: "member-1".to_string(), + member_uid: "uid-1".to_string(), + runtime_address: "127.0.0.1:5500".parse().unwrap(), + }, + }) + .unwrap(); + let (released, stopped) = ("rt-released-0123456789", "rt-stopped-0123456789"); + let mut stops = Vec::new(); + for route in [released, stopped] { + let (stop, stops_of) = tokio::sync::mpsc::channel(1); + admit_supervising(&host, route, stop); + stops.push(stops_of); + } + host.reconcile(&AssignmentSnapshot { + revision: 1, + pool: true, + stop_routes: vec![stopped.to_string()], + ..AssignmentSnapshot::default() + }) + .await + .unwrap(); + assert_eq!( + stops[1].try_recv(), + Ok(()), + "the withdrawn instance stopped" + ); + host.shut_down().await; + let output = String::from_utf8(captured.0.lock().unwrap().clone()).unwrap(); + for line in [ + format!("{RELINQUISHED_TO_OTHER_MEMBER} {released}"), + format!("{REMOVING_WITHDRAWN} {stopped}"), + ] { + assert!(output.contains(&line), "no {line:?} in:\n{output}"); + } + // The structured fields alone would not show the route. + assert!(!output.contains("route="), "{output}"); + } + + #[test] + fn hosts_reconnect_inside_the_membership_budget() { + use openshell_supervisor_control::pool_timing; + let pool = StreamTimings::pool(); + assert_eq!(pool.silence_timeout, pool_timing::HOST_SILENCE_TIMEOUT); + assert_eq!(pool.silence_check, pool_timing::HOST_SILENCE_CHECK); + // The covered attempts after a lost stream are quick; later ones get + // the ordinary bound, so a slow but responsive gateway is reached. + for failed in 0..=pool_timing::COVERED_UNANSWERED_ATTEMPTS { + assert_eq!( + pool.attempt_timeout(failed), + pool_timing::HOST_ATTEMPT_TIMEOUT + ); + } + assert_eq!( + pool.attempt_timeout(pool_timing::COVERED_UNANSWERED_ATTEMPTS + 1), + CONNECT_TIMEOUT + ); + for failed in 1..=12 { + for sample in [0.0, 0.5, 0.999_999, jitter_sample()] { + let backoff = pool.backoff(failed, sample); + assert!(backoff >= pool_timing::HOST_RECONNECT_BACKOFF_MIN); + assert!(backoff <= pool_timing::host_backoff_bound(failed)); + } + } + assert_eq!( + pool.backoff(1, 0.0), + pool_timing::HOST_RECONNECT_BACKOFF_MIN + ); + assert_eq!( + pool.backoff(9, 0.0), + pool_timing::HOST_RECONNECT_BACKOFF_MAX + ); + let sample = jitter_sample(); + assert!((0.0..1.0).contains(&sample)); + } + + #[tokio::test(flavor = "multi_thread")] + async fn an_oversized_snapshot_still_stops_redirects_and_keeps_within_capacity() { + let root = tempfile::tempdir().unwrap(); + let host = SharedHost::start(SharedHostConfig { + gateway_endpoint: "https://gateway.invalid".to_string(), + instance_root: root.path().join("instances"), + runtime_listen: "127.0.0.1:0".parse().unwrap(), + removal_grace: Duration::from_secs(10), + pool: PoolIdentity { + member_id: "member-1".to_string(), + member_uid: "uid-1".to_string(), + runtime_address: "127.0.0.1:5500".parse().unwrap(), + }, + }) + .unwrap(); + let (hosted_route, stopped_route) = ("rt-hosted-0123456789", "rt-stopped-0123456789"); + let (stopped, mut stops) = tokio::sync::mpsc::channel(1); + let (kept, mut kept_stops) = tokio::sync::mpsc::channel(1); + admit_supervising(&host, hosted_route, kept); + admit_supervising(&host, stopped_route, stopped); + tokio::time::timeout(Duration::from_secs(10), async { + while host + .list() + .iter() + .filter(|status| status.phase == SandboxInstancePhaseReport::Supervising) + .count() + < 2 + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("instances supervise"); + + // The hosted instance plus a full host's worth of new ones: one more + // than the host runs. + let assignment = |route: String| SandboxInstanceAssignment { + sandbox_id: route.clone(), + route, + generation: "generation".to_string(), + launch_material: b"{}".to_vec(), + takeover: None, + }; + let mut instances = vec![assignment(hosted_route.to_string())]; + instances.extend( + (0..MAX_ASSIGNED_INSTANCES) + .map(|index| assignment(format!("rt-new-{index:04}-0123456789"))), + ); + let elsewhere = "rt-elsewhere-0123456789"; + let owner: SocketAddr = "10.0.0.9:5500".parse().unwrap(); + host.reconcile(&AssignmentSnapshot { + revision: 1, + pool: true, + instances, + stop_routes: vec![stopped_route.to_string()], + route_directory: vec![openshell_core::proto::supervisor_host::v1::RouteOwner { + route: elsewhere.to_string(), + runtime_address: owner.to_string(), + }], + ..AssignmentSnapshot::default() + }) + .await + .expect("an oversized snapshot is applied, not refused"); + + // The listed stop ran, the hosted instance stays, and the directory + // redirects runtimes of routes owned elsewhere. + assert_eq!(stops.try_recv(), Ok(())); + assert!(host.status(stopped_route).is_none()); + assert_eq!( + host.status(hosted_route).map(|status| status.phase), + Some(SandboxInstancePhaseReport::Supervising) + ); + assert!(kept_stops.try_recv().is_err()); + let mut stream = tokio::net::TcpStream::connect(host.router.local_addr()) + .await + .unwrap(); + assert_eq!( + openshell_sandbox_backend::runtime_routed::negotiate_route(&mut stream, elsewhere) + .await + .unwrap(), + openshell_sandbox_backend::runtime_routed::RouteReply::Redirect(owner) + ); + // Only up to capacity was attempted: the hosted instance and all but + // the last new one in route order; the last stays unstarted. + let rejected = host.rejected.lock().unwrap().clone(); + assert_eq!(rejected.len(), MAX_ASSIGNED_INSTANCES); + let last = format!("rt-new-{:04}-0123456789", MAX_ASSIGNED_INSTANCES - 1); + assert!(!rejected.contains_key(&last)); + assert!(rejected.contains_key("rt-new-0000-0123456789")); + assert!(host.list().iter().all(|status| status.route != last)); + host.shut_down().await; + } + + /// How one connection attempt to a [`FakeGateway`] goes. + enum Attempt { + /// Never opens. + Hang, + /// Opens after a delay. + After(Duration), + /// Opens at once. + Now, + } + + /// A stream the fake gateway opened, driven by the test. + struct Opened { + /// When the host asked for it. + at: tokio::time::Instant, + outbound: tokio::sync::mpsc::Receiver, + inbound: tokio::sync::mpsc::Sender>, + } + + /// Gateway replicas behind one Service: each attempt goes as planned, and + /// every opened stream is handed to the test. + struct FakeGateway { + plan: Mutex>, + calls: Mutex>, + opened: tokio::sync::mpsc::UnboundedSender, + } + + impl GatewayConnector for FakeGateway { + fn open( + &self, + outbound: tokio::sync::mpsc::Receiver, + _deadline: tokio::time::Instant, + ) -> std::pin::Pin> + Send + '_>> + { + let at = tokio::time::Instant::now(); + self.calls.lock().unwrap().push(at); + let attempt = self + .plan + .lock() + .unwrap() + .pop_front() + .unwrap_or(Attempt::Now); + Box::pin(async move { + match attempt { + Attempt::Hang => std::future::pending::<()>().await, + Attempt::After(delay) => tokio::time::sleep(delay).await, + Attempt::Now => {} + } + let (inbound, stream) = tokio::sync::mpsc::channel(1); + let _ = self.opened.send(Opened { + at, + outbound, + inbound, + }); + let stream: GatewayInbound = + Box::pin(tokio_stream::wrappers::ReceiverStream::new(stream)); + Ok(stream) + }) + } + } + + fn gateway_message(payload: host_gateway_message::Payload) -> HostGatewayMessage { + HostGatewayMessage { + payload: Some(payload), + } + } + + fn accepted() -> HostGatewayMessage { + gateway_message(host_gateway_message::Payload::Accepted( + openshell_core::proto::supervisor_host::v1::HostAccepted { + session_id: "session".to_string(), + heartbeat_interval: Some(prost_types::Duration { + seconds: 1, + nanos: 0, + }), + }, + )) + } + + fn gateway_heartbeat() -> HostGatewayMessage { + gateway_message(host_gateway_message::Payload::Heartbeat( + openshell_core::proto::supervisor_host::v1::HostGatewayHeartbeat {}, + )) + } + + fn pool_snapshot(revision: u64) -> HostGatewayMessage { + gateway_message(host_gateway_message::Payload::Assignments( + AssignmentSnapshot { + revision, + pool: true, + ..AssignmentSnapshot::default() + }, + )) + } + + /// A pool host attached to `plan`, and the streams it opens. + fn attach_pool_host( + plan: Vec, + ) -> ( + Arc, + Arc, + tokio::sync::mpsc::UnboundedReceiver, + tempfile::TempDir, + ) { + let (host, gateway, streams, root, _attachment) = + attach_pool_host_until(plan, std::future::pending::<()>()); + (host, gateway, streams, root) + } + + /// A pool host attached to `plan` until `shutdown`, the streams it opens, + /// and the attachment. + fn attach_pool_host_until( + plan: Vec, + shutdown: impl Future + Send + 'static, + ) -> ( + Arc, + Arc, + tokio::sync::mpsc::UnboundedReceiver, + tempfile::TempDir, + tokio::task::JoinHandle<()>, + ) { + let root = tempfile::tempdir().unwrap(); + let host = SharedHost::start(SharedHostConfig { + gateway_endpoint: "https://gateway.invalid".to_string(), + instance_root: root.path().join("instances"), + runtime_listen: "127.0.0.1:0".parse().unwrap(), + removal_grace: Duration::from_secs(10), + pool: PoolIdentity { + member_id: "member-1".to_string(), + member_uid: "uid-1".to_string(), + runtime_address: "127.0.0.1:5500".parse().unwrap(), + }, + }) + .unwrap(); + let (opened, streams) = tokio::sync::mpsc::unbounded_channel(); + let gateway = Arc::new(FakeGateway { + plan: Mutex::new(plan.into()), + calls: Mutex::new(Vec::new()), + opened, + }); + let attachment = tokio::spawn(attach_with(host.clone(), gateway.clone(), shutdown)); + (host, gateway, streams, root, attachment) + } + + fn is_hello(message: Option) -> bool { + matches!( + message.and_then(|message| message.payload), + Some(host_message::Payload::Hello(_)) + ) + } + + fn status_of(message: Option) -> Option { + match message.and_then(|message| message.payload) { + Some(host_message::Payload::Status(status)) => Some(status), + _ => None, + } + } + + /// The next status the host reports, past its heartbeats. + async fn next_status( + outbound: &mut tokio::sync::mpsc::Receiver, + ) -> Option { + loop { + match outbound.recv().await?.payload { + Some(host_message::Payload::Status(status)) => return Some(status), + Some(host_message::Payload::Heartbeat(_)) => {} + _ => return None, + } + } + } + + /// Accept a stream and keep it answering: heartbeats every second, and + /// every status the host reports forwarded to the test. + fn serve( + mut stream: Opened, + ) -> ( + tokio::sync::mpsc::UnboundedReceiver, + tokio::sync::mpsc::Sender>, + ) { + let (statuses, reported) = tokio::sync::mpsc::unbounded_channel(); + let gateway = stream.inbound.clone(); + let inbound = stream.inbound.clone(); + tokio::spawn(async move { + if inbound.send(Ok(accepted())).await.is_err() { + return; + } + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + if inbound.send(Ok(gateway_heartbeat())).await.is_err() { + return; + } + } + }); + tokio::spawn(async move { + while let Some(message) = stream.outbound.recv().await { + if let Some(status) = status_of(Some(message)) + && statuses.send(status).is_err() + { + return; + } + } + }); + (reported, gateway) + } + + /// The part of the reconnect budget the host itself spends: everything + /// but the gateway's record refresh age and its two refresh rounds. + fn host_budget(unanswered: u32) -> Duration { + use openshell_supervisor_control::pool_timing; + pool_timing::reconnect_budget(unanswered) + .saturating_sub(pool_timing::MEMBER_REFRESH_INTERVAL) + .saturating_sub(pool_timing::RECORD_WRITE_BOUND * 2) + } + + /// A replica that accepted and then went silent is abandoned, and a + /// responsive replica is reached within the budget even when the covered + /// attempts reach the silent replica again; the new stream says hello and + /// reports status at once. + #[tokio::test(start_paused = true)] + async fn a_silent_gateway_is_left_for_a_responsive_one_within_the_budget() { + use openshell_supervisor_control::pool_timing; + let covered = pool_timing::COVERED_UNANSWERED_ATTEMPTS; + let mut plan = vec![Attempt::Now]; + plan.extend((0..covered).map(|_| Attempt::Hang)); + plan.push(Attempt::Now); + let (_host, gateway, mut streams, _root) = attach_pool_host(plan); + let mut silent = streams.recv().await.unwrap(); + assert!(is_hello(silent.outbound.recv().await)); + silent.inbound.send(Ok(accepted())).await.unwrap(); + assert!(next_status(&mut silent.outbound).await.is_some()); + let went_silent = tokio::time::Instant::now(); + + let mut responsive = streams.recv().await.unwrap(); + let reconnected = responsive.at - went_silent; + assert!( + reconnected <= host_budget(covered), + "reached a responsive replica {reconnected:?} after the silence, over {:?}", + host_budget(covered) + ); + assert_eq!(gateway.calls.lock().unwrap().len(), 2 + covered as usize); + assert!(is_hello(responsive.outbound.recv().await)); + responsive.inbound.send(Ok(accepted())).await.unwrap(); + let accepted_at = tokio::time::Instant::now(); + assert!(next_status(&mut responsive.outbound).await.is_some()); + assert_eq!( + tokio::time::Instant::now(), + accepted_at, + "status is reported at once" + ); + } + + /// After an accepted stream is lost, every attempt through the one after + /// the covered unanswered attempts, the one the budget counts on reaching + /// a responsive replica, gets the quick bound; only later ones get the + /// ordinary bound. + #[tokio::test(start_paused = true)] + async fn the_attempt_after_the_covered_unanswered_ones_is_still_quick() { + use openshell_supervisor_control::pool_timing; + let covered = pool_timing::COVERED_UNANSWERED_ATTEMPTS as usize; + let mut plan = vec![Attempt::Now]; + // The covered unanswered attempts, the one after them, and one more. + plan.extend((0..covered + 2).map(|_| Attempt::Hang)); + plan.push(Attempt::Now); + let (_host, gateway, mut streams, _root) = attach_pool_host(plan); + let mut silent = streams.recv().await.unwrap(); + assert!(is_hello(silent.outbound.recv().await)); + silent.inbound.send(Ok(accepted())).await.unwrap(); + assert!(next_status(&mut silent.outbound).await.is_some()); + streams.recv().await.unwrap(); + let calls = gateway.calls.lock().unwrap().clone(); + assert_eq!(calls.len(), covered + 4); + // Reconnect attempt `n` (1-based) runs from calls[n] to calls[n + 1]. + for attempt in 1..=covered + 2 { + let took = calls[attempt + 1] - calls[attempt]; + let backoff = pool_timing::host_backoff_bound(u32::try_from(attempt + 1).unwrap()); + let bound = if attempt <= covered + 1 { + pool_timing::HOST_ATTEMPT_TIMEOUT + } else { + CONNECT_TIMEOUT + }; + assert!( + took >= bound && took <= bound + backoff, + "reconnect attempt {attempt} took {took:?}, expected its {bound:?} bound" + ); + } + } + + /// Connecting, attaching, and being accepted share one deadline: a + /// stream that opens late and is never accepted is abandoned when the + /// attempt's deadline passes, not a full timeout after it opened. + #[tokio::test(start_paused = true)] + async fn one_deadline_bounds_connecting_and_being_accepted() { + use openshell_supervisor_control::pool_timing; + let late = pool_timing::HOST_ATTEMPT_TIMEOUT.mul_f64(0.7); + let (_host, gateway, mut streams, _root) = + attach_pool_host(vec![Attempt::After(late), Attempt::Now]); + let unaccepted = streams.recv().await.unwrap(); + let next = streams.recv().await.unwrap(); + let attempt = next.at - unaccepted.at; + assert!( + attempt <= pool_timing::HOST_ATTEMPT_TIMEOUT + pool_timing::host_backoff_bound(1), + "the unaccepted attempt took {attempt:?}" + ); + assert!(attempt >= pool_timing::HOST_ATTEMPT_TIMEOUT); + assert_eq!(gateway.calls.lock().unwrap().len(), 2); + } + + /// A replica that keeps writing but stops reading fills the host's + /// outbound queue; the blocked send gives up within the silence budget + /// rather than the ordinary send bound. + #[tokio::test(start_paused = true)] + async fn a_gateway_that_stops_reading_is_left_within_the_silence_budget() { + use openshell_supervisor_control::pool_timing; + let (_host, _gateway, mut streams, _root) = + attach_pool_host(vec![Attempt::Now, Attempt::Now]); + let unread = streams.recv().await.unwrap(); + let inbound = unread.inbound.clone(); + tokio::spawn(async move { + if inbound.send(Ok(accepted())).await.is_err() { + return; + } + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + if inbound.send(Ok(gateway_heartbeat())).await.is_err() { + return; + } + } + }); + // The host fills its queue to the replica that never reads it. + while unread.outbound.len() < OUTBOUND_QUEUE { + tokio::time::sleep(Duration::from_millis(10)).await; + } + let full = tokio::time::Instant::now(); + let next = streams.recv().await.unwrap(); + let left = next.at - full; + assert!( + left <= pool_timing::HOST_SILENCE_TIMEOUT + + pool_timing::HEARTBEAT_INTERVAL + + pool_timing::host_backoff_bound(1), + "a blocked send held the host {left:?}" + ); + assert!(left < SEND_TIMEOUT); + } + + /// A reconciliation stalled past the membership TTL keeps neither a new + /// stream from connecting and reporting nor a second reconciliation from + /// waiting its turn: the new stream's snapshot is applied after it, never + /// beside it. + #[tokio::test(start_paused = true)] + async fn a_stalled_reconciliation_neither_delays_reconnecting_nor_runs_twice() { + use openshell_supervisor_control::pool_timing; + let (host, _gateway, mut streams, _root) = + attach_pool_host(vec![Attempt::Now, Attempt::Now]); + let (release, gate) = tokio::sync::oneshot::channel(); + *host.reconcile_probe.gate.lock().unwrap() = Some(gate); + let mut stalled = streams.recv().await.unwrap(); + assert!(is_hello(stalled.outbound.recv().await)); + stalled.inbound.send(Ok(accepted())).await.unwrap(); + assert!(next_status(&mut stalled.outbound).await.is_some()); + stalled.inbound.send(Ok(pool_snapshot(7))).await.unwrap(); + let went_silent = tokio::time::Instant::now(); + // The worker takes the snapshot and stalls; the replica stays silent. + while host + .reconcile_probe + .running + .load(std::sync::atomic::Ordering::SeqCst) + == 0 + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let responsive = tokio::time::timeout(pool_timing::MEMBER_TTL, streams.recv()) + .await + .expect("reconnected while the reconciliation stalls") + .unwrap(); + assert!(responsive.at - went_silent <= host_budget(0)); + let accepted_at = tokio::time::Instant::now(); + let (mut reported, gateway) = serve(responsive); + let first = reported.recv().await.unwrap(); + assert_eq!(tokio::time::Instant::now(), accepted_at); + assert_eq!(first.applied_revision, 0); + + // The new replica's snapshot waits for the stalled one, well past the + // TTL, while the stream stays up. + gateway.send(Ok(pool_snapshot(1))).await.unwrap(); + tokio::time::sleep(pool_timing::MEMBER_TTL + Duration::from_secs(5)).await; + assert_eq!( + host.reconcile_probe + .peak + .load(std::sync::atomic::Ordering::SeqCst), + 1 + ); + release.send(()).unwrap(); + let applied = loop { + let status = reported.recv().await.unwrap(); + if status.applied_revision == 1 { + break status; + } + }; + assert_eq!(applied.applied_revision, 1); + assert_eq!( + host.reconcile_probe + .peak + .load(std::sync::atomic::Ordering::SeqCst), + 1 + ); + } + + /// Process shutdown ends the attachment only once the reconciliation in + /// flight has finished, so the host's own shutdown sees everything it + /// admitted; snapshots still queued behind it are dropped. + #[tokio::test(start_paused = true)] + async fn shutdown_waits_for_the_reconciliation_in_flight_and_drops_queued_work() { + use std::sync::atomic::Ordering; + let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); + let (host, _gateway, mut streams, _root, attachment) = + attach_pool_host_until(vec![Attempt::Now], async { + let _ = stopped.await; + }); + let (release, gate) = tokio::sync::oneshot::channel(); + *host.reconcile_probe.gate.lock().unwrap() = Some(gate); + let stream = streams.recv().await.unwrap(); + let (mut reported, gateway) = serve(stream); + // Accepted: the host reports before any snapshot. + reported.recv().await.unwrap(); + gateway.send(Ok(pool_snapshot(1))).await.unwrap(); + while host.reconcile_probe.running.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + // A newer snapshot waits behind the stalled one. + gateway.send(Ok(pool_snapshot(2))).await.unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + stop.send(()).unwrap(); + tokio::time::sleep(Duration::from_secs(30)).await; + assert!( + !attachment.is_finished(), + "the attachment ended while a reconciliation was still admitting" + ); + release.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), attachment) + .await + .expect("the attachment ends once the reconciliation finishes") + .unwrap(); + assert_eq!(host.reconcile_probe.running.load(Ordering::SeqCst), 0); + assert_eq!( + host.reconcile_probe.entered.load(Ordering::SeqCst), + 1, + "the queued snapshot must not run after shutdown" + ); + } + + fn term(auth_epoch: u64, instance: &str, minted_by: &str) -> SandboxInstanceTakeover { + SandboxInstanceTakeover { + auth_epoch, + supervisor_instance_id: instance.to_string(), + grant: format!("grant.{minted_by}.signature"), + gateway_token: format!("gateway.{minted_by}.signature"), + gateway_expires_at: 10, + sandbox_token: format!("sandbox.{minted_by}.signature"), + sandbox_expires_at: 20, + } + } + + #[test] + fn pool_instances_are_fingerprinted_by_their_term_not_their_credentials() { + let instance = SupervisorInstanceId::new().to_string(); + let fingerprint = |launch: &[u8], takeover: Option<&SandboxInstanceTakeover>| { + pool_fingerprint(launch, 1, takeover).expect("fingerprint") + }; + let minted_here = term(2, &instance, "gateway-a"); + // Another gateway replica mints different bytes for the same term, + // and one that minted nothing yet sends the term alone. + let minted_elsewhere = term(2, &instance, "gateway-b"); + let unminted = SandboxInstanceTakeover { + auth_epoch: 2, + supervisor_instance_id: instance.clone(), + ..SandboxInstanceTakeover::default() + }; + let current = fingerprint(b"launch", Some(&minted_here)); + assert_eq!(current, fingerprint(b"launch", Some(&minted_elsewhere))); + assert_eq!(current, fingerprint(b"launch", Some(&unminted))); + // Another term, process, or launch is another instance. + for other in [ + fingerprint(b"launch", None), + fingerprint(b"launch", Some(&term(3, &instance, "gateway-a"))), + fingerprint( + b"launch", + Some(&term( + 2, + &SupervisorInstanceId::new().to_string(), + "gateway-a", + )), + ), + fingerprint(b"other launch", Some(&minted_here)), + ] { + assert_ne!(current, other); + } + } + + #[test] + fn takeover_terms_follow_the_launch_term_and_wait_for_their_credentials() { + let launch = CredentialEpoch::new(1).expect("launch epoch"); + let instance = SupervisorInstanceId::new(); + let granted = takeover_term(launch, &term(2, &instance.to_string(), "gateway-a")) + .expect("granted term"); + assert_eq!(granted.auth_epoch.get(), 2); + assert_eq!( + granted.boundary.supervisor_instance_id.to_string(), + instance.to_string() + ); + assert_eq!( + granted.gateway_token.expose_secret(), + "gateway.gateway-a.signature" + ); + assert_eq!( + granted.sandbox_token.expose_secret(), + "sandbox.gateway-a.signature" + ); + + let unminted = SandboxInstanceTakeover { + auth_epoch: 2, + supervisor_instance_id: instance.to_string(), + ..SandboxInstanceTakeover::default() + }; + assert!(matches!( + takeover_term(launch, &unminted), + Err(ControlError::Deferred(_)) + )); + for refused in [ + term(1, &instance.to_string(), "gateway-a"), + term(0, &instance.to_string(), "gateway-a"), + term(2, "not-a-process", "gateway-a"), + SandboxInstanceTakeover { + gateway_token: "has whitespace".to_string(), + ..term(2, &instance.to_string(), "gateway-a") + }, + ] { + assert!(matches!( + takeover_term(launch, &refused), + Err(ControlError::Invalid(_)) + )); + } + } +} diff --git a/deploy/helm/openshell/README.md b/deploy/helm/openshell/README.md index fcb7a108e4..4038412654 100644 --- a/deploy/helm/openshell/README.md +++ b/deploy/helm/openshell/README.md @@ -590,6 +590,8 @@ discovery endpoint or its TLS CA. | supervisor.podSecurityContext.runAsNonRoot | bool | `true` | Require the shared supervisor to run as a non-root user. | | supervisor.podSecurityContext.runAsUser | int | `65534` | UID assigned to the shared supervisor pod. | | supervisor.podSecurityContext.seccompProfile.type | string | `"RuntimeDefault"` | Seccomp profile applied to the shared supervisor pod. | +| supervisor.recoveryTimeoutSeconds | int | `60` | Seconds a sandbox workload stays frozen after losing its pool member's supervision before the sandbox runtime terminates it, rendered as [openshell.gateway.supervisor_pool] recovery_timeout_secs. Must be an integer from 10 to 300. Null leaves the gateway's pool default of 60 seconds. | +| supervisor.replicas | int | `1` | Shared supervisor pool members, used only when gatewayConfig selects supervisor_topology = "gateway". The gateway always coordinates the shared supervisor as a high-availability pool: when a member fails or leaves, another member takes over its sandboxes without restarting their workloads. 1 has no failover; use 2 or more to survive losing a member. Only the supervisor Deployment's replica count depends on this value, so changing it neither restarts the gateway nor replaces running members. | | supervisor.sandboxRuntime.boundaryPort | int | `5500` | Workload boundary TLS listener port. | | tolerations | list | `[]` | Tolerations for the gateway pod. | | upstreamProxy | object | `{"authAllowInsecure":false,"authSecret":{"key":"","name":""},"caBundle":{"configMapName":"","key":"ca.crt"},"connectByHostname":false,"noProxy":"","url":""}` | Operator-owned corporate forward proxy for policy-approved TLS egress from Kubernetes sandboxes. The workload cannot select or override it. | diff --git a/deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml b/deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml new file mode 100644 index 0000000000..7a68522f56 --- /dev/null +++ b/deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# E2E overlay: supervise new sandbox generations with a shared supervisor pool +# that starts with one member, so the test can scale it to two members and +# back with `helm upgrade --set supervisor.replicas=N` without changing the +# gateway configuration or the supervisor Pod template. +gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + +supervisor: + replicas: 1 + recoveryTimeoutSeconds: 60 diff --git a/deploy/helm/openshell/ci/values-supervisor-ha.yaml b/deploy/helm/openshell/ci/values-supervisor-ha.yaml new file mode 100644 index 0000000000..4a1dda47ab --- /dev/null +++ b/deploy/helm/openshell/ci/values-supervisor-ha.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# E2E overlay: supervise new sandbox generations with a high-availability pool +# of two shared supervisors, so a failed or leaving member's sandboxes move to +# the other member. Layer it on a workspace-mode overlay to cover managed or +# operator namespaces. +gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + +supervisor: + replicas: 2 + recoveryTimeoutSeconds: 60 diff --git a/deploy/helm/openshell/ci/values-supervisor.yaml b/deploy/helm/openshell/ci/values-supervisor.yaml index 0361016d1b..3e530a9419 100644 --- a/deploy/helm/openshell/ci/values-supervisor.yaml +++ b/deploy/helm/openshell/ci/values-supervisor.yaml @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# E2E overlay: supervise new sandbox generations with one supervisor process -# shared by the gateway. Layer it on a workspace-mode overlay +# E2E overlay: supervise new sandbox generations with a one-member shared +# supervisor pool. Layer it on a workspace-mode overlay # to cover managed or operator namespaces. gatewayConfig: openshell.drivers.kubernetes: diff --git a/deploy/helm/openshell/ci/values-tls-port-forward.yaml b/deploy/helm/openshell/ci/values-tls-port-forward.yaml new file mode 100644 index 0000000000..a1f27d578d --- /dev/null +++ b/deploy/helm/openshell/ci/values-tls-port-forward.yaml @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# E2E overlay: keep gateway TLS on over the Skaffold values, for topologies +# that require it, such as the shared supervisor. The harness reaches the +# gateway through its port-forward with the chart-generated client +# certificate (OPENSHELL_E2E_KUBE_TLS=1); the server certificate names +# 127.0.0.1. +server: + disableTls: false diff --git a/deploy/helm/openshell/skaffold.yaml b/deploy/helm/openshell/skaffold.yaml index 2bbf0acc1b..3ad4d3b8bf 100644 --- a/deploy/helm/openshell/skaffold.yaml +++ b/deploy/helm/openshell/skaffold.yaml @@ -169,6 +169,18 @@ profiles: - op: add path: /deploy/helm/releases/1/valuesFiles/- value: ci/values-gateway.yaml + # Shared supervisor pool: two supervisor replicas that take over each other's + # sandboxes. The topology requires gateway TLS, so this profile turns it back + # on over ci/values-skaffold.yaml. + - name: shared-supervisor-high-availability + patches: + - op: add + path: /deploy/helm/releases/0/valuesFiles/- + value: ci/values-supervisor-ha.yaml + - op: add + path: /deploy/helm/releases/0/setValues + value: + server.disableTls: "false" - name: credential-driver-kubernetes-secrets patches: - op: add diff --git a/deploy/helm/openshell/templates/_helpers.tpl b/deploy/helm/openshell/templates/_helpers.tpl index 44a1a8a8bb..25f885205e 100644 --- a/deploy/helm/openshell/templates/_helpers.tpl +++ b/deploy/helm/openshell/templates/_helpers.tpl @@ -387,6 +387,77 @@ app.kubernetes.io/instance: {{ .Release.Name }} app.kubernetes.io/component: supervisor {{- end }} +{{/* +A chart value as a whole number, or nothing when it is not one. Helm parses +YAML numbers as float64, so whole floats count; booleans, strings, and +fractions do not. +*/}} +{{- define "openshell.wholeNumber" -}} +{{- if has (kindOf .) (list "int" "int8" "int16" "int32" "int64" "uint" "uint8" "uint16" "uint32" "uint64" "float32" "float64") -}} +{{- $whole := int64 . -}} +{{- if eq (float64 .) (float64 $whole) -}}{{- $whole -}}{{- end -}} +{{- end -}} +{{- end }} + +{{/* +Shared supervisor replica count. A missing value (for example after `helm +upgrade --reuse-values` from a release that predates it) means 1. +*/}} +{{- define "openshell.supervisorReplicas" -}} +{{- $supervisor := .Values.supervisor | default dict -}} +{{- $value := get $supervisor "replicas" -}} +{{- if or (not (hasKey $supervisor "replicas")) (kindIs "invalid" $value) -}} +1 +{{- else -}} +{{- $replicas := include "openshell.wholeNumber" $value -}} +{{- if or (eq $replicas "") (lt (int64 $replicas) 1) -}} +{{- fail "supervisor.replicas must be an integer of at least 1" -}} +{{- end -}} +{{- $replicas -}} +{{- end -}} +{{- end }} + +{{/* +Seconds a sandbox runtime keeps its workload frozen after losing a pool +member's supervision, or nothing when unset so the gateway applies its pool +default. +*/}} +{{- define "openshell.supervisorRecoveryTimeoutSeconds" -}} +{{- $supervisor := .Values.supervisor | default dict -}} +{{- $value := get $supervisor "recoveryTimeoutSeconds" -}} +{{- if and (hasKey $supervisor "recoveryTimeoutSeconds") (not (kindIs "invalid" $value)) (ne (toString $value) "") -}} +{{- $timeout := include "openshell.wholeNumber" $value -}} +{{- if or (eq $timeout "") (lt (int64 $timeout) 10) (gt (int64 $timeout) 300) -}} +{{- fail "supervisor.recoveryTimeoutSeconds must be an integer between 10 and 300, or null for the default" -}} +{{- end -}} +{{- $timeout -}} +{{- end -}} +{{- end }} + +{{/* +The gateway's [openshell.gateway.supervisor_pool] tuning table as YAML, or +nothing. The shared supervisor always runs as a high-availability pool that +the gateway coordinates; the chart owns the table only to set the recovery +deadline (supervisor.recoveryTimeoutSeconds; none leaves the gateway default). +It never depends on supervisor.replicas, so scaling never changes the gateway +configuration. Renders nothing without the gateway topology. Validates the +supervisor values whenever the topology is selected, and refuses a +hand-written table in either TOML form. +*/}} +{{- define "openshell.supervisorPoolConfig" -}} +{{- $gatewayConfig := .Values.gatewayConfig | default dict -}} +{{- if or (hasKey $gatewayConfig "openshell.gateway.supervisor_pool") (hasKey (get $gatewayConfig "openshell.gateway" | default dict) "supervisor_pool") -}} +{{- fail "supervisor_pool is rendered by the chart; set supervisor.replicas and supervisor.recoveryTimeoutSeconds" -}} +{{- end -}} +{{- if include "openshell.supervisorEnabled" . -}} +{{- $_ := include "openshell.supervisorReplicas" . -}} +{{- $recoveryTimeout := include "openshell.supervisorRecoveryTimeoutSeconds" . -}} +{{- if $recoveryTimeout -}} +{{- toYaml (dict "recovery_timeout_secs" (int64 $recoveryTimeout)) -}} +{{- end -}} +{{- end -}} +{{- end }} + {{/* Namespace selector for the sandbox namespaces of this gateway's workspace mode: the shared namespace, the managed namespaces it labels, or the @@ -654,6 +725,7 @@ Validate chart values that Helm would otherwise accept silently. {{- if hasKey (get (.Values.gatewayConfig | default dict) "openshell.drivers.kubernetes" | default dict) "shared_supervisor" -}} {{- fail "shared_supervisor is rendered by the chart; set only supervisor_topology" -}} {{- end -}} +{{- $_ := include "openshell.supervisorPoolConfig" . -}} {{- $workloadKind := include "openshell.workloadKind" . -}} {{- $workload := .Values.workload | default dict -}} {{- $maxReplicas := int (include "openshell.maxReplicas" .) -}} diff --git a/deploy/helm/openshell/templates/_toml.tpl b/deploy/helm/openshell/templates/_toml.tpl index 121c93159d..2e54294ed3 100644 --- a/deploy/helm/openshell/templates/_toml.tpl +++ b/deploy/helm/openshell/templates/_toml.tpl @@ -304,6 +304,12 @@ owner: server.*. Override any gatewayConfig copies before serializing TOML. */}} {{- end -}} {{- $_ := set $config "openshell.gateway.tls" $gatewayTls -}} {{- end -}} +{{/* The chart owns the shared supervisor pool; see +openshell.supervisorPoolConfig. */}} +{{- $supervisorPool := include "openshell.supervisorPoolConfig" . | fromYaml -}} +{{- if $supervisorPool -}} +{{- $_ := set $config "openshell.gateway.supervisor_pool" $supervisorPool -}} +{{- end -}} {{/* A nested YAML map and its dotted top-level table would serialize to conflicting TOML definitions. Fail before rendering an invalid ConfigMap. */}} {{- range $tableName := keys $config | sortAlpha -}} diff --git a/deploy/helm/openshell/templates/supervisor.yaml b/deploy/helm/openshell/templates/supervisor.yaml index d6027d2449..8ecac54784 100644 --- a/deploy/helm/openshell/templates/supervisor.yaml +++ b/deploy/helm/openshell/templates/supervisor.yaml @@ -12,12 +12,13 @@ {{- $endpoint := get $kubernetes "grpc_endpoint" }} {{- $endpointUrl := urlParse $endpoint }} {{- $gatewayTls := and (not .Values.server.disableTls) .Values.server.tls.clientTlsSecretName }} -# One supervisor process hosting a sandbox instance per sandbox generation. -# The process dials the gateway over TLS with the supervisor token the gateway -# mints for it and receives its assignments from it; nothing but sandbox -# runtimes and kubelet probes connect to it. The Pod holds no Kubernetes -# identity of its own: each instance acts only with the credentials the -# gateway issued for its generation. +{{- $replicas := int64 (include "openshell.supervisorReplicas" .) }} +# One supervisor process per replica hosts a sandbox instance for each sandbox +# generation placed on it. The process dials the gateway over TLS with the +# supervisor token the gateway mints for it and receives its assignments from +# it; nothing but sandbox runtimes and kubelet probes connect to it. The Pod +# holds no Kubernetes identity of its own: each instance acts only with the +# credentials the gateway issued for its generation. apiVersion: apps/v1 kind: Deployment metadata: @@ -27,10 +28,19 @@ metadata: {{- include "openshell.labels" . | nindent 4 }} app.kubernetes.io/component: supervisor spec: - # Instances live in this process; a second replica could not serve them. - replicas: 1 + # The replicas form a high-availability pool. Each instance lives in the + # process hosting it; the pool rolls one member at a time, starting a new + # member before an old one leaves, and the gateway hands a failed or leaving + # member's instances to the remaining members. One replica has no failover + # but scales like any other count: only the replica count depends on the + # pool's size, so scaling changes nothing else here or in the gateway + # configuration. + replicas: {{ $replicas }} strategy: - type: Recreate + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 selector: matchLabels: {{- include "openshell.supervisorSelectorLabels" . | nindent 6 }} @@ -41,7 +51,9 @@ spec: spec: automountServiceAccountToken: false enableServiceLinks: false - terminationGracePeriodSeconds: 30 + # A stopping member announces that it is leaving, and the gateway moves + # its instances to the remaining members before it exits. + terminationGracePeriodSeconds: 60 {{- with ($imagePullSecrets | uniq) }} imagePullSecrets: {{- range . }} @@ -52,6 +64,17 @@ spec: securityContext: {{- toYaml . | nindent 8 }} {{- end }} + # Prefer separate nodes, so losing one node leaves members to take over + # its instances. + affinity: + podAntiAffinity: + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 100 + podAffinityTerm: + labelSelector: + matchLabels: + {{- include "openshell.supervisorSelectorLabels" . | nindent 20 }} + topologyKey: kubernetes.io/hostname containers: - name: supervisor image: {{ get $kubernetes "supervisor_image" | quote }} @@ -81,6 +104,21 @@ spec: - name: OPENSHELL_GATEWAY_TLS_SERVER_NAME value: {{ get $endpointUrl "hostname" | quote }} {{- end }} + # Pool membership: the Pod name and UID identify this member, and + # other members redirect the runtimes of its instances to its Pod + # IP. + - name: OPENSHELL_SUPERVISOR_POOL_MEMBER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: OPENSHELL_SUPERVISOR_POOL_MEMBER_UID + valueFrom: + fieldRef: + fieldPath: metadata.uid + - name: OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP + valueFrom: + fieldRef: + fieldPath: status.podIP ports: - name: runtime containerPort: {{ get $shared "runtime_port" }} @@ -205,4 +243,20 @@ metadata: {{- include "openshell.labels" . | nindent 4 }} app.kubernetes.io/component: supervisor type: Opaque +--- +# Voluntary disruptions, such as node drains, evict one member at a time, so +# the remaining members can take over its instances. +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: {{ $name }} + namespace: {{ .Release.Namespace }} + labels: + {{- include "openshell.labels" . | nindent 4 }} + app.kubernetes.io/component: supervisor +spec: + maxUnavailable: 1 + selector: + matchLabels: + {{- include "openshell.supervisorSelectorLabels" . | nindent 6 }} {{- end }} diff --git a/deploy/helm/openshell/tests/supervisor_test.yaml b/deploy/helm/openshell/tests/supervisor_test.yaml index 0686a800ec..6ef71da4b4 100644 --- a/deploy/helm/openshell/tests/supervisor_test.yaml +++ b/deploy/helm/openshell/tests/supervisor_test.yaml @@ -13,7 +13,7 @@ tests: asserts: - hasDocuments: count: 0 - - it: deploys one instance host without Kubernetes identity in the gateway namespace + - it: deploys the supervisor pool without Kubernetes identity in the gateway namespace template: templates/supervisor.yaml set: gatewayConfig: @@ -22,7 +22,7 @@ tests: supervisor_image: example.com/supervisor:shared asserts: - hasDocuments: - count: 4 + count: 5 - equal: path: kind value: Deployment @@ -261,6 +261,352 @@ tests: - matchRegex: path: data["gateway.toml"] pattern: 'shared_supervisor = \{ runtime_port = 5500, service_name = "openshell-supervisor", service_namespace = "gateway", token_secret_name = "openshell-supervisor-token" \}' + - it: runs one replica as a high-availability pool by default + template: templates/supervisor.yaml + set: + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - hasDocuments: + count: 5 + - equal: + path: spec.replicas + value: 1 + documentIndex: 0 + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 + documentIndex: 0 + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 60 + documentIndex: 0 + - exists: + path: spec.template.spec.affinity.podAntiAffinity + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_MEMBER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_MEMBER_UID + valueFrom: + fieldRef: + fieldPath: metadata.uid + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + documentIndex: 0 + - equal: + path: spec.maxUnavailable + value: 1 + documentSelector: + path: kind + value: PodDisruptionBudget + - it: runs two or more replicas as a high-availability pool + template: templates/supervisor.yaml + set: + supervisor: + replicas: 3 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - hasDocuments: + count: 5 + - equal: + path: spec.replicas + value: 3 + documentIndex: 0 + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 + documentIndex: 0 + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 60 + documentIndex: 0 + - equal: + path: spec.template.spec.affinity + value: + podAntiAffinity: + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 100 + podAffinityTerm: + labelSelector: + matchLabels: + app.kubernetes.io/name: openshell-supervisor + app.kubernetes.io/instance: openshell + app.kubernetes.io/component: supervisor + topologyKey: kubernetes.io/hostname + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_MEMBER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_MEMBER_UID + valueFrom: + fieldRef: + fieldPath: metadata.uid + documentIndex: 0 + - contains: + path: spec.template.spec.containers[0].env + content: + name: OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + documentIndex: 0 + - equal: + path: spec.selector + value: + app.kubernetes.io/name: openshell-supervisor + app.kubernetes.io/instance: openshell + app.kubernetes.io/component: supervisor + documentIndex: 1 + - equal: + path: spec.ingress[0].ports + value: + - protocol: TCP + port: 5500 + documentIndex: 2 + - equal: + path: apiVersion + value: policy/v1 + documentSelector: + path: kind + value: PodDisruptionBudget + - equal: + path: metadata.name + value: openshell-supervisor + documentSelector: + path: kind + value: PodDisruptionBudget + - equal: + path: spec + value: + maxUnavailable: 1 + selector: + matchLabels: + app.kubernetes.io/name: openshell-supervisor + app.kubernetes.io/instance: openshell + app.kubernetes.io/component: supervisor + documentSelector: + path: kind + value: PodDisruptionBudget + - it: protects the pool from voluntary disruption without a NetworkPolicy + template: templates/supervisor.yaml + set: + networkPolicy: + enabled: false + supervisor: + replicas: 2 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - hasDocuments: + count: 4 + - equal: + path: spec.maxUnavailable + value: 1 + documentSelector: + path: kind + value: PodDisruptionBudget + - it: renders the pool's recovery deadline in the gateway table by default + template: templates/gateway-config.yaml + set: + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: 'shared_supervisor = \{ runtime_port = 5500, service_name = "openshell-supervisor", service_namespace = "gateway", token_secret_name = "openshell-supervisor-token" \}' + - matchRegex: + path: data["gateway.toml"] + pattern: '\[openshell\.gateway\.supervisor_pool\]\nrecovery_timeout_secs = 60\n' + - notMatchRegex: + path: data["gateway.toml"] + pattern: 'supervisor_pool\]\n[^\[]*(enabled|high_availability)' + - it: renders no pool tuning when the recovery timeout is null, leaving the gateway's defaults + template: templates/gateway-config.yaml + set: + supervisor: + recoveryTimeoutSeconds: null + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: 'supervisor_topology = "gateway"' + - notMatchRegex: + path: data["gateway.toml"] + pattern: supervisor_pool|recovery_timeout_secs + - it: renders a configured recovery timeout into the supervisor pool table + template: templates/gateway-config.yaml + set: + supervisor: + replicas: 2 + recoveryTimeoutSeconds: 45 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - matchRegex: + path: data["gateway.toml"] + pattern: '\[openshell\.gateway\.supervisor_pool\]\nrecovery_timeout_secs = 45\n' + - matchRegex: + path: data["gateway.toml"] + pattern: 'shared_supervisor = \{ runtime_port = 5500, service_name = "openshell-supervisor", service_namespace = "gateway", token_secret_name = "openshell-supervisor-token" \}' + - it: renders no supervisor pool without the gateway topology + template: templates/gateway-config.yaml + set: + supervisor: + replicas: 2 + recoveryTimeoutSeconds: 120 + asserts: + - notMatchRegex: + path: data["gateway.toml"] + pattern: supervisor_pool|recovery_timeout_secs|shared_supervisor + - it: rejects a hand-written supervisor pool table + template: templates/supervisor.yaml + set: + supervisor: + replicas: 2 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + openshell.gateway.supervisor_pool: + enabled: true + asserts: + - failedTemplate: + errorMessage: supervisor_pool is rendered by the chart; set supervisor.replicas and supervisor.recoveryTimeoutSeconds + - it: rejects a hand-written supervisor pool inside the gateway table + template: templates/supervisor.yaml + set: + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + openshell.gateway: + supervisor_pool: + enabled: true + asserts: + - failedTemplate: + errorMessage: supervisor_pool is rendered by the chart; set supervisor.replicas and supervisor.recoveryTimeoutSeconds + - it: rejects a hand-written supervisor pool without the gateway topology + template: templates/supervisor.yaml + set: + gatewayConfig: + openshell.gateway.supervisor_pool: + enabled: true + asserts: + - failedTemplate: + errorMessage: supervisor_pool is rendered by the chart; set supervisor.replicas and supervisor.recoveryTimeoutSeconds + - it: rejects zero supervisor replicas + template: templates/supervisor.yaml + set: + supervisor: + replicas: 0 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.replicas must be an integer of at least 1 + - it: rejects fractional supervisor replicas + template: templates/supervisor.yaml + set: + supervisor: + replicas: 1.5 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.replicas must be an integer of at least 1 + - it: rejects supervisor replicas given as a string + template: templates/supervisor.yaml + set: + supervisor: + replicas: "2" + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.replicas must be an integer of at least 1 + - it: rejects a recovery timeout below 10 seconds + template: templates/supervisor.yaml + set: + supervisor: + recoveryTimeoutSeconds: 9 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.recoveryTimeoutSeconds must be an integer between 10 and 300, or null for the default + - it: rejects a recovery timeout above 300 seconds + template: templates/supervisor.yaml + set: + supervisor: + replicas: 2 + recoveryTimeoutSeconds: 301 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.recoveryTimeoutSeconds must be an integer between 10 and 300, or null for the default + - it: rejects a fractional recovery timeout + template: templates/supervisor.yaml + set: + supervisor: + recoveryTimeoutSeconds: 30.5 + gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + asserts: + - failedTemplate: + errorMessage: supervisor.recoveryTimeoutSeconds must be an integer between 10 and 300, or null for the default + - it: ignores supervisor pool values without the gateway topology + template: templates/supervisor.yaml + set: + supervisor: + replicas: 0 + recoveryTimeoutSeconds: 5 + asserts: + - hasDocuments: + count: 0 - it: rejects a hand-written shared supervisor connection template: templates/supervisor.yaml set: diff --git a/deploy/helm/openshell/values.yaml b/deploy/helm/openshell/values.yaml index 04fc1152a4..067f56a76a 100644 --- a/deploy/helm/openshell/values.yaml +++ b/deploy/helm/openshell/values.yaml @@ -69,6 +69,20 @@ supervisor: tag: "" # -- Supervisor image digest. When set, this takes precedence over tag. digest: "" + # -- Shared supervisor pool members, used only when gatewayConfig selects + # supervisor_topology = "gateway". The gateway always coordinates the shared + # supervisor as a high-availability pool: when a member fails or leaves, + # another member takes over its sandboxes without restarting their + # workloads. 1 has no failover; use 2 or more to survive losing a member. + # Only the supervisor Deployment's replica count depends on this value, so + # changing it neither restarts the gateway nor replaces running members. + replicas: 1 + # -- Seconds a sandbox workload stays frozen after losing its pool member's + # supervision before the sandbox runtime terminates it, rendered as + # [openshell.gateway.supervisor_pool] recovery_timeout_secs. Must be an + # integer from 10 to 300. Null leaves the gateway's pool default of 60 + # seconds. + recoveryTimeoutSeconds: 60 # Pod security context for the shared supervisor Deployment, rendered only # when gatewayConfig selects supervisor_topology = "gateway". On OpenShift, # set the user and group IDs to null so the restricted-v2 SCC can assign the diff --git a/deploy/helm/test-supervisor-pool-scaling.sh b/deploy/helm/test-supervisor-pool-scaling.sh new file mode 100755 index 0000000000..148090bcbf --- /dev/null +++ b/deploy/helm/test-supervisor-pool-scaling.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# The shared supervisor always runs as a high-availability pool, which scales +# between one and more replicas without changing anything but the supervisor +# Deployment's replica count: not the gateway configuration (so the gateway +# does not roll), not the supervisor Pod template (so running members are not +# replaced), and not the PodDisruptionBudget or Service. The renders differ +# only in supervisor.replicas. +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd -P)" +chart="${repo_root}/deploy/helm/openshell" +work_dir="$(mktemp -d)" +trap 'rm -rf "${work_dir}"' EXIT + +render() { + local name="$1" + shift + helm template scaling "${chart}" \ + --namespace scaling-namespace \ + --set agentSandbox.preflight.enabled=false \ + --set-string 'gatewayConfig.openshell\.drivers\.kubernetes.supervisor_topology=gateway' \ + "$@" >"${work_dir}/${name}.yaml" +} + +# Every rendered document except Secrets, whose generated values differ +# between renders, with the supervisor Deployment's replica count removed. +without_replica_count() { + yq ea -P \ + 'select(. != null and .kind != "Secret") + | (select(.kind == "Deployment" and .metadata.name == "scaling-openshell-supervisor") | .spec) |= del(.replicas)' \ + "${work_dir}/$1.yaml" >"${work_dir}/$1.normalized.yaml" +} + +supervisor_replicas() { + yq ea -e -r \ + 'select(.kind == "Deployment" and .metadata.name == "scaling-openshell-supervisor") | .spec.replicas' \ + "${work_dir}/$1.yaml" +} + +fail() { + echo "supervisor pool scaling: $*" >&2 + exit 1 +} + +render one --set supervisor.replicas=1 +render two --set supervisor.replicas=2 +render three --set supervisor.replicas=3 + +[[ "$(supervisor_replicas one)" == "1" ]] || fail "one-replica pool renders $(supervisor_replicas one) replicas" +[[ "$(supervisor_replicas two)" == "2" ]] || fail "two-replica pool renders $(supervisor_replicas two) replicas" +[[ "$(supervisor_replicas three)" == "3" ]] || fail "three-replica pool renders $(supervisor_replicas three) replicas" + +for name in one two three; do + without_replica_count "${name}" +done + +for name in two three; do + if ! diff -u "${work_dir}/one.normalized.yaml" "${work_dir}/${name}.normalized.yaml"; then + fail "scaling the pool from 1 to ${name} replicas changed more than the replica count" + fi +done + +# The pool table, with the chart's default recovery deadline, is present at +# every count, and the gateway's config checksum (which rolls the gateway when +# it changes) is the same. +for name in one two three; do + yq ea -e -r 'select(.kind == "ConfigMap" and (.data | has("gateway.toml"))) | .data."gateway.toml"' \ + "${work_dir}/${name}.yaml" >"${work_dir}/${name}.toml" + grep -q '^\[openshell\.gateway\.supervisor_pool\]$' "${work_dir}/${name}.toml" \ + || fail "the ${name}-replica pool renders no supervisor_pool table" + grep -q '^recovery_timeout_secs = 60$' "${work_dir}/${name}.toml" \ + || fail "the ${name}-replica pool renders no recovery deadline" +done + +echo "supervisor pool scaling: only the replica count changes across 1, 2, and 3 replicas" diff --git a/docs/how-it-works/gateways/configuration.mdx b/docs/how-it-works/gateways/configuration.mdx index 7f16929759..e89e19cc5e 100644 --- a/docs/how-it-works/gateways/configuration.mdx +++ b/docs/how-it-works/gateways/configuration.mdx @@ -716,6 +716,33 @@ Multi-replica gateways relay session-bound traffic to the replica that owns a sa The gateway refuses to start when a peer endpoint is configured on a multi-replica backend and peer authentication is unavailable, because the replica could not serve relays from its peers. That happens when the pod identity variables are missing, when the in-cluster Kubernetes client cannot be built, or when the gateway is not running in a cluster. A plaintext `http://` peer endpoint is also rejected when the gateway serves TLS; use it only when the gateway itself serves plaintext. +## Supervisor Pool + +With the Kubernetes driver's `supervisor_topology = "gateway"`, the gateway +always runs the shared supervisor as a high-availability pool, whatever its +replica count. It places each sandbox on one pool member, records its owner +durably, and moves the sandboxes of a member that fails or leaves to a live +member without restarting their workloads. +`[openshell.gateway.supervisor_pool]` is optional tuning of that pool; without +it the defaults apply. The gateway refuses to start with the table when no +in-process compute driver uses the shared supervisor. + +```toml +[openshell.gateway.supervisor_pool] +recovery_timeout_secs = 60 +``` + +| Field | Default | Behavior | +|---|---|---| +| `recovery_timeout_secs` | `60` | Seconds, from 10 to 300, that a sandbox runtime keeps its workload frozen after losing its pool member before it terminates the workload. The pool's supervisors offer this limit to each sandbox runtime when they attach. | + +The Helm chart renders `recovery_timeout_secs` from +`supervisor.recoveryTimeoutSeconds` (60 by default; null renders no table) and +rejects a hand-written table. The table does not depend on the replica count, +so scaling the pool does not change the gateway configuration or restart the +gateway. For failover behavior, refer to +[Shared Supervisor High Availability](/kubernetes/high-availability#shared-supervisor-high-availability). + ## Driver References ### External Resource Admission @@ -899,6 +926,10 @@ sandbox_runtime_image_pull_policy = "if_not_present" # Defaults to the gateway version; override to pin a specific build. # supervisor_image = "ghcr.io/nvidia/openshell/supervisor:" supervisor_image_pull_policy = "if_not_present" +# Supervise new sandboxes with one shared supervisor Deployment instead of a +# supervisor Pod per sandbox. Helm deploys it and renders shared_supervisor. +# supervisor_topology = "gateway" +# shared_supervisor = { service_name = "openshell-supervisor", service_namespace = "openshell", runtime_port = 5500, token_secret_name = "openshell-supervisor-token" } # Optional corporate forward proxy for policy-approved TLS egress. The sandbox # workload cannot select or override these settings. http:// and https:// proxy @@ -1008,6 +1039,12 @@ In managed workspace mode, the Kubernetes driver copies each explicitly named namespace on sandbox creation. Shared and operator modes require the Secret to already exist in the sandbox namespace. +With `supervisor_topology = "gateway"`, `shared_supervisor` tells the gateway +how to reach the shared supervisor that the Helm chart deploys. The chart +renders `shared_supervisor` and rejects a hand-written one. To run the shared +supervisor as a high-availability pool, configure the gateway's +[supervisor pool](#supervisor-pool). + For token-exchange provider profiles, the gateway also needs access to its own SPIFFE Workload API socket. In Helm deployments, set `server.providerTokenGrants.spiffe.enabled=true`; the chart mounts the socket diff --git a/docs/kubernetes/high-availability.mdx b/docs/kubernetes/high-availability.mdx index 8b4d467b4e..3349620ecf 100644 --- a/docs/kubernetes/high-availability.mdx +++ b/docs/kubernetes/high-availability.mdx @@ -4,7 +4,7 @@ title: "High Availability" sidebar-title: "High Availability" description: "Run multiple OpenShell gateway replicas on Kubernetes with shared PostgreSQL, authenticated peer routing, capacity metrics, and optional autoscaling." -keywords: "Generative AI, Cybersecurity, Kubernetes, High Availability, HA, Gateway, PostgreSQL, Replicas, Failover, Autoscaling, HPA, Metrics" +keywords: "Generative AI, Cybersecurity, Kubernetes, High Availability, HA, Gateway, PostgreSQL, Replicas, Failover, Autoscaling, HPA, Metrics, Shared Supervisor" position: 3 --- @@ -504,6 +504,84 @@ 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. +## Shared Supervisor High Availability + +With `supervisor_topology = "gateway"`, one shared supervisor Deployment in the +release namespace supervises new sandboxes instead of a supervisor Pod per +sandbox. The topology requires gateway TLS and does not support upstream proxy +settings or provider SPIFFE identity. The Deployment always runs as a pool of +supervisors that the gateway coordinates. By default it has one member and no +failover. Run two or more replicas when sandboxes must stay supervised through +the loss of a supervisor Pod or node. + +Add these values to your release: + +```yaml +gatewayConfig: + openshell.drivers.kubernetes: + supervisor_topology: gateway + +supervisor: + replicas: 2 + # Optional. From 10 to 300 seconds; defaults to 60. + recoveryTimeoutSeconds: 90 +``` + +The chart enables the gateway's +[supervisor pool](/how-it-works/gateways/configuration#supervisor-pool) and +prefers to schedule the supervisor Pods on different nodes. It replaces them +one at a time during rollouts, and its PodDisruptionBudget lets voluntary +disruptions, such as node drains, evict one supervisor Pod at a time. + +Each sandbox runs on one pool member, and a healthy member keeps its +sandboxes. When a member's process exits, its node fails, or it shuts down, the +gateway moves each of its sandboxes to a live member, first revoking every +credential the previous member held for that sandbox. The sandbox runtime +freezes the workload when it loses its supervisor and resumes it after the +replacement confirms the takeover. + +A takeover does not restart the sandbox. The sandbox keeps its Pod, its running +processes and their PIDs, and its in-memory state. The replacement adopts the +agent and exec processes and keeps the sandbox CA, SSH host key, and +[binary identity](/security/best-practices#binary-identity-binding) pins. +Proxied connections and interactive exec and SSH sessions that are open during +the failure break, and clients must reconnect. + +`supervisor.recoveryTimeoutSeconds` limits how long a workload stays frozen. +The pool's supervisors offer the limit to each sandbox runtime when +they attach. If no replacement confirms the takeover in time, the sandbox +runtime terminates the workload instead of letting it run unsupervised. When +unset, the limit is 60 seconds. The gateway takes up to 22 seconds to detect a +failed node, or 30 seconds right after a gateway replica starts, and the +replacement then needs time to take over, so keep the limit well above that. +A pool fails over only to a live member, so keep at least two members running. + +With one member, a supervisor failure terminates the workloads it supervises +once the limit passes. The gateway configuration and supervisor Pod template do +not depend on the replica count, so changing `supervisor.replicas` neither +restarts the gateway nor replaces running members: + +```shell +helm upgrade openshell -n openshell --reuse-values --set supervisor.replicas=2 +helm upgrade openshell -n openshell --reuse-values --set supervisor.replicas=1 +``` + +Scaling up adds a member and moves no running sandbox; new sandboxes can land +on the new member. Scaling down removes a member, which hands its sandboxes to +the remaining members before it exits, within its 60-second grace period. The +sandboxes keep running; open proxied connections and exec or SSH sessions on +the handed-off sandboxes break and must reconnect. + +Check the supervisor Pods and their disruption budget: + +```shell +kubectl -n openshell get deployment,poddisruptionbudget openshell-supervisor +kubectl -n openshell get pods -l app.kubernetes.io/name=openshell-supervisor -o wide +``` + +If you changed the chart's generated names with `nameOverride` or +`fullnameOverride`, use the rendered `-supervisor` name instead. + ## Next Steps - To expose the gateway through a highly available data path, refer to diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index 1dc6f0b1cf..d903c2e6e2 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -38,6 +38,11 @@ e2e-kubernetes-workspace-managed = ["e2e-kubernetes"] e2e-kubernetes-workspace-operator = ["e2e-kubernetes"] # One supervisor process shared by the gateway's sandboxes. e2e-kubernetes-shared-supervisor = ["e2e-kubernetes", "e2e-host-gateway"] +# A high-availability pool of shared supervisors. +e2e-kubernetes-shared-supervisor-ha = ["e2e-kubernetes", "e2e-host-gateway"] +# Scaling a one-member high-availability pool of shared supervisors to two +# members and back. It upgrades the release, so it runs on its own release. +e2e-kubernetes-shared-supervisor-scaling = ["e2e-kubernetes", "e2e-host-gateway"] e2e-podman = ["e2e", "e2e-host-gateway", "e2e-local-container-driver"] e2e-podman-gpu = ["e2e-podman", "e2e-gpu"] e2e-oidc-pkce = [] diff --git a/e2e/rust/src/harness/sandbox.rs b/e2e/rust/src/harness/sandbox.rs index 0d6ff621dd..c0f0c1aedc 100644 --- a/e2e/rust/src/harness/sandbox.rs +++ b/e2e/rust/src/harness/sandbox.rs @@ -232,10 +232,20 @@ impl SandboxGuard { /// Unlike [`SandboxGuard::create_keep`], this does not open an attachment, /// which lets tests control competing and reconnecting clients directly. pub async fn create_detached_main(command: &[&str]) -> Result { + Self::create_detached_main_with_args(&[], command).await + } + + /// Like [`SandboxGuard::create_detached_main`], but forwards extra flags + /// to `sandbox create` before the `-- ` separator. + pub async fn create_detached_main_with_args( + create_args: &[&str], + command: &[&str], + ) -> Result { let mut cmd = openshell_cmd(); cmd.arg("sandbox").arg("create").arg("--detach"); - add_unique_name_if_missing(&mut cmd, &[]); - add_test_image_if_missing(&mut cmd, &[]); + add_unique_name_if_missing(&mut cmd, create_args); + add_test_image_if_missing(&mut cmd, create_args); + cmd.args(create_args); cmd.arg("--") .args(command) .stdout(Stdio::piped()) @@ -586,10 +596,13 @@ impl SandboxGuard { cmd.arg(arg); } // Never share the test runner's stdin: parallel test processes share - // its open file description, and the command needs no input. + // its open file description, and the command needs no input. A caller + // that bounds the exec with a timeout drops this future; the CLI then + // goes with it instead of lingering. cmd.stdin(Stdio::null()) .stdout(Stdio::piped()) - .stderr(Stdio::piped()); + .stderr(Stdio::piped()) + .kill_on_drop(true); let output = cmd .output() @@ -693,6 +706,45 @@ impl SandboxGuard { } } +impl SandboxGuard { + /// Delete the sandbox and report whether that worked. On failure the + /// guard still deletes it when dropped. + /// + /// # Errors + /// + /// Returns the CLI's output when the delete fails. + pub async fn delete(&mut self) -> Result<(), String> { + if self.cleaned_up { + return Ok(()); + } + if let Some(ref mut child) = self.child { + let _ = child.kill().await; + let _ = child.wait().await; + } + let output = openshell_cmd() + .arg("sandbox") + .arg("delete") + .arg(&self.name) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await + .map_err(|e| format!("failed to spawn openshell sandbox delete: {e}"))?; + if !output.status.success() { + return Err(format!( + "sandbox delete {} failed (exit {:?}):\n{}{}", + self.name, + output.status.code(), + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + )); + } + self.cleaned_up = true; + Ok(()) + } +} + impl Drop for SandboxGuard { fn drop(&mut self) { if self.cleaned_up { diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs new file mode 100644 index 0000000000..896549305b --- /dev/null +++ b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs @@ -0,0 +1,389 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "e2e-kubernetes-shared-supervisor-ha")] + +//! A high-availability pool of shared supervisors (`supervisor_topology = +//! "gateway"` with two or more replicas). +//! +//! Killing the supervisor Pod that owns a sandbox instance, abruptly so it +//! cannot announce that it is leaving, moves the instance to another pool +//! member once its membership expires, and that member adopts the running +//! workload. The workload Pod, runtime generation, main process, and the +//! process's in-memory state survive; exec and policy-mediated networking work +//! again; and a sandbox on another member keeps running and answering +//! throughout, under the same owner. + +mod shared_supervisor; + +use std::process::Stdio; +use std::time::Duration; + +use shared_supervisor::{ + HostEndpoint, SUPERVISOR_CONTAINER, call_host, counts_past, create, gateway_log_lines, + generation, is_takeover, kubectl, kubectl_command, kubectl_json, main_process, owner_of, + policy, release_namespace, supervisor_pods, wait_for_owner, wait_until_attached, +}; +use tokio::io::AsyncBufReadExt as _; + +/// Bound on the takeover: membership expiry, a reconciliation pass, the +/// replacement's startup, and the CLI's own retries. +const RECOVERY_DEADLINE: Duration = Duration::from_mins(3); +/// Sandboxes created, at most, until two land on different members. The ring +/// places each sandbox by its random server-assigned ID, so placements can +/// cluster on one member for many sandboxes in a row; the bound must not +/// assume they alternate (eight in a row on one member was seen live). +const PLACEMENT_ATTEMPTS: usize = 32; +/// A member that announces it is leaving is drained within a reconciliation +/// pass (2 s); an abruptly lost member is taken over only once its record +/// expires (15 s TTL, refreshed every 2 s). A move this late cannot come from +/// a leave. +const EXPIRY_EVIDENCE: Duration = Duration::from_secs(10); +/// What a supervisor logs as it begins to leave the pool. +const LEAVING: &str = "pool member is leaving"; + +/// Run `command` in a cluster node's container; k3d and kind name node +/// containers after their nodes. +async fn node_exec(node: &str, command: &[&str]) -> Result { + let output = tokio::process::Command::new("docker") + .arg("exec") + .arg(node) + .args(command) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await + .map_err(|error| format!("docker exec {node} {command:?}: {error}"))?; + if output.status.success() { + Ok(String::from_utf8_lossy(&output.stdout).into_owned()) + } else { + Err(format!( + "docker exec {node} {command:?}: {}", + String::from_utf8_lossy(&output.stderr).trim() + )) + } +} + +/// SIGKILL the supervisor container of `pod` through its node's container +/// runtime, so it cannot announce that it is leaving, then force-delete the +/// Pod so its member does not come back in place. Kills the task with +/// `ctr`, or failing that the container's process with `kill -KILL`; panics +/// when neither succeeds. Returns how the container was killed. +async fn kill_abruptly(pod: &str) -> String { + let namespace = release_namespace(); + let description = kubectl_json(&["-n", &namespace, "get", "pod", pod, "-o", "json"]).await; + let node = description["spec"]["nodeName"] + .as_str() + .unwrap_or_else(|| panic!("{pod} has no node: {description}")) + .to_string(); + let container = description["status"]["containerStatuses"] + .as_array() + .into_iter() + .flatten() + .find(|status| status["name"] == SUPERVISOR_CONTAINER) + .and_then(|status| status["containerID"].as_str()) + .and_then(|id| id.split_once("://").map(|(_, id)| id.to_string())) + .unwrap_or_else(|| panic!("{pod} has no running {SUPERVISOR_CONTAINER} container")); + let mut failures = Vec::new(); + let mut killed = None; + // k3s ships `ctr` as a multi-call alias; kind ships it on its own. + for ctr in [&["ctr"][..], &["k3s", "ctr"][..]] { + let kill = [ + ctr, + &[ + "-n", "k8s.io", "tasks", "kill", "--signal", "SIGKILL", &container, + ], + ] + .concat(); + match node_exec(&node, &kill).await { + Ok(_) => { + killed = Some(format!("ctr tasks kill --signal SIGKILL on {node}")); + break; + } + Err(error) => failures.push(error), + } + } + if killed.is_none() { + for ctr in [&["ctr"][..], &["k3s", "ctr"][..]] { + let list = [ctr, &["-n", "k8s.io", "tasks", "ls"]].concat(); + let pid = match node_exec(&node, &list).await { + Ok(tasks) => tasks.lines().find_map(|line| { + let mut fields = line.split_whitespace(); + (fields.next() == Some(container.as_str())) + .then(|| fields.next()) + .flatten() + .filter(|pid| pid.parse::().is_ok_and(|pid| pid > 1)) + .map(str::to_string) + }), + Err(error) => { + failures.push(error); + continue; + } + }; + let Some(pid) = pid else { + failures.push(format!("no running task {container} on {node}")); + continue; + }; + match node_exec(&node, &["kill", "-KILL", &pid]).await { + Ok(_) => { + killed = Some(format!("kill -KILL {pid} on {node}")); + break; + } + Err(error) => failures.push(error), + } + } + } + let killed = killed.unwrap_or_else(|| { + panic!( + "could not hard-kill supervisor {pod}; a graceful stop proves nothing: {failures:#?}" + ) + }); + kubectl(&[ + "-n", + &namespace, + "delete", + "pod", + pod, + "--grace-period=0", + "--force", + "--wait=false", + ]) + .await; + killed +} + +/// Lines a supervisor Pod logs from its start until its container ends. +struct FollowedLogs { + lines: std::sync::Arc>>, + ended: tokio::sync::oneshot::Receiver<()>, + _child: tokio::process::Child, +} + +impl FollowedLogs { + /// Follow `pod`'s supervisor log, once it shows the admission of `route`. + async fn start(pod: &str, route: &str) -> Self { + let mut child = kubectl_command() + .args([ + "-n", + &release_namespace(), + "logs", + "--follow", + pod, + "-c", + SUPERVISOR_CONTAINER, + ]) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true) + .spawn() + .expect("spawn kubectl logs --follow"); + let stdout = child.stdout.take().expect("piped logs"); + let lines = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let (end, ended) = tokio::sync::oneshot::channel(); + let collected = lines.clone(); + tokio::spawn(async move { + let mut reader = tokio::io::BufReader::new(stdout).lines(); + while let Ok(Some(line)) = reader.next_line().await { + collected.lock().unwrap().push(line); + } + let _ = end.send(()); + }); + let logs = Self { + lines, + ended, + _child: child, + }; + tokio::time::timeout(Duration::from_secs(30), async { + while !logs.any(|line| line.contains(route)) { + tokio::time::sleep(Duration::from_millis(200)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{pod}'s followed log never showed {route}")); + logs + } + + fn any(&self, matches: impl Fn(&str) -> bool) -> bool { + self.lines.lock().unwrap().iter().any(|line| matches(line)) + } + + /// Everything logged, once the container ended. + async fn ended(mut self, pod: &str) -> Vec { + tokio::time::timeout(Duration::from_secs(30), &mut self.ended) + .await + .unwrap_or_else(|_| panic!("{pod}'s log did not end after the kill")) + .ok(); + self.lines.lock().unwrap().clone() + } +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them() { + let members = supervisor_pods().await; + assert!( + members.len() >= 2, + "the supervisor pool needs at least two running members: {members:?}" + ); + wait_until_attached(&members).await; + let endpoint = HostEndpoint::start().await; + let policy = policy(endpoint.port); + + // Two sandboxes on different members: create until the pool spreads one + // away from the first, and remove the extra ones. + let mut first = create(&policy).await; + let first_generation = generation(&first.name).await; + let first_owner = wait_for_owner(&first_generation.route).await; + let mut extras = Vec::new(); + let mut placed = None; + for _ in 0..PLACEMENT_ATTEMPTS { + let candidate = create(&policy).await; + let candidate_generation = generation(&candidate.name).await; + let candidate_owner = wait_for_owner(&candidate_generation.route).await; + if candidate_owner == first_owner { + extras.push(candidate); + } else { + placed = Some((candidate, candidate_generation, candidate_owner)); + break; + } + } + for mut extra in extras { + extra.cleanup().await; + } + let (mut second, second_generation, second_owner) = placed.unwrap_or_else(|| { + panic!( + "the pool placed {} sandboxes on {first_owner} alone; members: {members:?}", + PLACEMENT_ATTEMPTS + 1 + ) + }); + eprintln!( + "{} on {first_owner}, {} on {second_owner}", + first.name, second.name + ); + assert_eq!(call_host(&first, endpoint.port).await, "HA_OK"); + let before = main_process(&first).await.expect("main process before"); + let second_before = main_process(&second) + .await + .expect("second main process before"); + + // The owner is lost abruptly: its supervisor is killed without a chance + // to announce that it is leaving or to hand anything over. + let victim_logs = FollowedLogs::start(&first_owner, &first_generation.route).await; + let killed_at = tokio::time::Instant::now(); + let killed = kill_abruptly(&first_owner).await; + eprintln!("supervisor {first_owner} killed: {killed}"); + + // The sandbox on the surviving member keeps answering, and counting, + // while the other is recovered. + let during = main_process(&second) + .await + .expect("a sandbox of the surviving member answers during the failover"); + tokio::time::sleep(Duration::from_secs(3)).await; + let later = main_process(&second).await.expect("still answering"); + assert_eq!(during.pid, later.pid); + assert!( + later.counter > during.counter, + "the surviving member's workload must keep running: {during:?} then {later:?}" + ); + + // The route moves only once the lost member's record expires, which a + // leaving member would not wait for. + let moved_after = tokio::time::timeout(RECOVERY_DEADLINE, async { + loop { + if owner_of(&first_generation.route) + .await + .is_some_and(|owner| owner != first_owner) + { + return killed_at.elapsed(); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + }) + .await + .unwrap_or_else(|_| { + panic!( + "no other member admitted {} within {RECOVERY_DEADLINE:?}", + first.name + ) + }); + // The cause is the member's failure, not a leave: the lost supervisor + // logged nothing about leaving before it died, a gateway took the route + // over from it, and the move waited for its record to expire. + let victim_lines = victim_logs.ended(&first_owner).await; + let left = victim_lines + .iter() + .filter(|line| line.contains(LEAVING)) + .collect::>(); + assert!( + left.is_empty(), + "{first_owner} announced that it was leaving, so the kill ({killed}) was not abrupt: {left:?}" + ); + // The gateway names the cause it saw: the member had no live record, + // not a leaving flag or a restarted process. + let takeover = gateway_log_lines(&first_generation.route).await; + assert!( + takeover.iter().any(|line| is_takeover( + line, + &first_generation.route, + &first_owner, + "expired" + )), + "no gateway took {} over from {first_owner} because its record expired: {takeover:?}", + first_generation.route + ); + assert!( + moved_after >= EXPIRY_EVIDENCE, + "{} moved {moved_after:?} after the kill, sooner than its member's record could expire", + first.name + ); + + // The lost owner's sandbox comes back through its new owner. + let after = tokio::time::timeout(RECOVERY_DEADLINE, async { + loop { + if let Ok(after) = main_process(&first).await { + return after; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await + .unwrap_or_else(|_| { + panic!( + "{} was not recovered within {RECOVERY_DEADLINE:?}", + first.name + ) + }); + + // Same Pod, same generation, same process: the workload was adopted, not + // restarted, and keeps counting on from the state that lives only in its + // memory. + assert_eq!(generation(&first.name).await, first_generation); + assert_eq!(after.pid, before.pid, "the main process was restarted"); + assert_eq!( + after.start_ticks, before.start_ticks, + "a new process reused the main process's PID" + ); + let after = counts_past(&first, &before).await; + let new_owner = wait_for_owner(&first_generation.route).await; + assert_ne!(new_owner, first_owner, "the lost Pod cannot own the route"); + + // Mediated networking works again through the replacement's policy, and + // the workload is running again. + assert_eq!(call_host(&first, endpoint.port).await, "HA_OK"); + counts_past(&first, &after).await; + + // The other sandbox is still its original process, under its original + // owner. + assert_eq!(generation(&second.name).await, second_generation); + assert_eq!( + wait_for_owner(&second_generation.route).await, + second_owner, + "the surviving member's sandbox must not move" + ); + assert_eq!(call_host(&second, endpoint.port).await, "HA_OK"); + let second_after = counts_past(&second, &second_before).await; + assert_eq!(second_after.start_ticks, second_before.start_ticks); + + first.cleanup().await; + second.cleanup().await; +} diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_scaling.rs b/e2e/rust/tests/kubernetes_shared_supervisor_scaling.rs new file mode 100644 index 0000000000..71223b5125 --- /dev/null +++ b/e2e/rust/tests/kubernetes_shared_supervisor_scaling.rs @@ -0,0 +1,640 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "e2e-kubernetes-shared-supervisor-scaling")] + +//! Scaling the shared supervisor pool (`supervisor_topology = "gateway"`, which +//! always runs as a high-availability pool) from one member to two and back to +//! one, with `helm upgrade --set supervisor.replicas=N` of the installed +//! release. +//! +//! The release starts as a one-member pool (ci/values-supervisor-ha-scaling.yaml). +//! Scaling up adds a member and moves nothing: a live owner keeps its +//! sandboxes. Scaling down removes a member that owns a sandbox; it leaves +//! the pool, and the remaining member adopts its sandboxes. Throughout, every +//! sandbox keeps its workload Pod, runtime generation, main process, and the +//! counter that lives only in that process's memory; exec and +//! policy-mediated networking answer, with failures bounded to the hand-off; +//! and neither the gateway nor the remaining supervisor Pod is replaced, +//! since scaling the pool changes only the supervisor Deployment's replica +//! count. +//! +//! The test upgrades the release it runs against, so it needs a release of +//! its own (`mise run e2e:kubernetes:shared-supervisor-scaling`). + +mod shared_supervisor; + +use std::collections::BTreeMap; +use std::process::Stdio; +use std::time::Duration; + +use futures_util::future::join_all; +use openshell_e2e::harness::sandbox::SandboxGuard; +use shared_supervisor::{ + GATEWAY_SELECTOR, HostEndpoint, MainProcess, SUPERVISOR_COMPONENT, TAKEOVER, call_host, + counts_past, create, gateway_log_lines, generation, helm_command, is_takeover, kubectl_json, + main_process, owner_of, policy, release_namespace, supervisor_pod_items, supervisor_pods, + try_call_host, wait_for_owner, wait_until_attached, +}; +use tokio::time::Instant; + +/// Sandboxes created, at most, until one lands on the new member. The ring +/// places each sandbox by its random server-assigned ID, so placements can +/// cluster on one member for many sandboxes in a row; the bound must not +/// assume they alternate (eight in a row on one member was seen live). +const PLACEMENT_ATTEMPTS: usize = 32; +/// Bound on one `helm upgrade --wait`. +const HELM_TIMEOUT: Duration = Duration::from_mins(6); +/// Bound on the pool reaching its new size once Helm returns: a removed +/// member leaves within its 60-second grace period. +const CONVERGENCE_TIMEOUT: Duration = Duration::from_mins(3); +/// Interval between probes of each sandbox. +const PROBE_INTERVAL: Duration = Duration::from_secs(2); +/// Probing continues this long after the pool reached its new size. +const SETTLE: Duration = Duration::from_secs(15); +/// Longest a sandbox whose supervisor stays may go without a successful +/// probe: brief reconnects only. +const STAYING_OUTAGE: Duration = Duration::from_secs(30); +/// Longest a sandbox handed to another member may go without a successful +/// probe; within the pool's 60-second recovery timeout. +const MOVING_OUTAGE: Duration = Duration::from_mins(1); +const POOL_MEMBER_ENV: &str = "OPENSHELL_SUPERVISOR_POOL_MEMBER_ID"; + +fn required_env(name: &str) -> String { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| { + panic!( + "{name} must be set; e2e/with-kube-gateway.sh sets it for the release it installs" + ) + }) +} + +/// The release's shared supervisor Deployment. +async fn supervisor_deployment() -> serde_json::Value { + let list = kubectl_json(&[ + "-n", + &release_namespace(), + "get", + "deployments", + "-l", + SUPERVISOR_COMPONENT, + "-o", + "json", + ]) + .await; + let [deployment] = list["items"] + .as_array() + .expect("Deployment items") + .as_slice() + else { + panic!("expected one shared supervisor Deployment: {list}"); + }; + deployment.clone() +} + +/// What scaling the pool must leave alone: the gateway workloads (their +/// generation and config checksum), the gateway Pods (UID and restarts), and +/// the supervisor Pod template. +#[derive(Debug, PartialEq, Eq)] +struct Untouched { + gateway_workloads: Vec<(String, i64, String)>, + gateway_pods: Vec<(String, String, i64)>, + supervisor_template: serde_json::Value, +} + +async fn untouched() -> Untouched { + let namespace = release_namespace(); + let selector = format!("{GATEWAY_SELECTOR},app.kubernetes.io/component!=supervisor"); + let workloads = kubectl_json(&[ + "-n", + &namespace, + "get", + "statefulsets,deployments", + "-l", + &selector, + "-o", + "json", + ]) + .await; + let mut gateway_workloads = workloads["items"] + .as_array() + .expect("workload items") + .iter() + .map(|workload| { + ( + format!( + "{}/{}", + workload["kind"].as_str().unwrap_or_default(), + workload["metadata"]["name"].as_str().unwrap_or_default() + ), + workload["metadata"]["generation"].as_i64().unwrap_or(-1), + workload["spec"]["template"]["metadata"]["annotations"]["checksum/gateway-config"] + .as_str() + .unwrap_or_default() + .to_string(), + ) + }) + .collect::>(); + gateway_workloads.sort(); + assert!( + !gateway_workloads.is_empty(), + "no gateway workload matches {selector}" + ); + let pods = kubectl_json(&[ + "-n", + &namespace, + "get", + "pods", + "-l", + GATEWAY_SELECTOR, + "-o", + "json", + ]) + .await; + let mut gateway_pods = pods["items"] + .as_array() + .expect("Pod items") + .iter() + .map(|pod| { + ( + pod["metadata"]["name"] + .as_str() + .unwrap_or_default() + .to_string(), + pod["metadata"]["uid"] + .as_str() + .unwrap_or_default() + .to_string(), + pod["status"]["containerStatuses"] + .as_array() + .into_iter() + .flatten() + .filter_map(|status| status["restartCount"].as_i64()) + .sum(), + ) + }) + .collect::>(); + gateway_pods.sort(); + assert!( + !gateway_pods.is_empty(), + "no gateway Pods match {GATEWAY_SELECTOR}" + ); + Untouched { + gateway_workloads, + gateway_pods, + supervisor_template: supervisor_deployment().await["spec"]["template"].clone(), + } +} + +/// UIDs of the running supervisor Pods, by name. +async fn running_members() -> BTreeMap { + let running = supervisor_pods().await; + supervisor_pod_items() + .await + .iter() + .filter_map(|pod| { + let name = pod["metadata"]["name"].as_str()?.to_string(); + running.contains(&name).then(|| { + ( + name, + pod["metadata"]["uid"] + .as_str() + .unwrap_or_default() + .to_string(), + ) + }) + }) + .collect() +} + +/// Upgrade the release to `replicas` pool members with its own values and +/// the exact chart it was installed from, then wait until exactly that many +/// members run and no removed member is still leaving. +async fn scale_pool(replicas: usize) -> BTreeMap { + let release = required_env("OPENSHELL_E2E_KUBE_RELEASE"); + let chart = required_env("OPENSHELL_E2E_KUBE_CHART"); + let namespace = release_namespace(); + let replicas_value = format!("supervisor.replicas={replicas}"); + let args = [ + "upgrade", + release.as_str(), + chart.as_str(), + "--namespace", + namespace.as_str(), + "--reuse-values", + "--set", + replicas_value.as_str(), + "--wait", + "--timeout", + "5m", + ]; + let output = tokio::time::timeout( + HELM_TIMEOUT, + helm_command() + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output(), + ) + .await + .unwrap_or_else(|_| panic!("helm {args:?} did not finish within {HELM_TIMEOUT:?}")) + .unwrap_or_else(|error| panic!("spawn helm {args:?}: {error}")); + assert!( + output.status.success(), + "helm {args:?} failed: {}{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + tokio::time::timeout(CONVERGENCE_TIMEOUT, async { + loop { + let pods = supervisor_pod_items().await; + let members = running_members().await; + if members.len() == replicas && pods.len() == replicas { + return members; + } + tokio::time::sleep(Duration::from_secs(1)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("the supervisor pool did not settle at {replicas} members")) +} + +/// What the probes of one sandbox saw during a transition. +#[derive(Debug, Default)] +struct Samples { + exec_ok: usize, + exec_failures: Vec, + network_ok: usize, + network_failures: Vec, + /// Longest time without a successful exec probe. + exec_outage: Duration, + /// Longest time without a successful network probe. + network_outage: Duration, + last: Option, +} + +/// Probe `sandbox` every [`PROBE_INTERVAL`] until `stop`: its main process +/// must stay `baseline`'s and its counter must never go back; record how +/// long exec and the policy-mediated call to the host endpoint fail. +async fn sample( + sandbox: &SandboxGuard, + baseline: &MainProcess, + port: u16, + mut stop: tokio::sync::watch::Receiver, +) -> Samples { + let mut samples = Samples::default(); + let mut last_exec = Instant::now(); + let mut last_network = Instant::now(); + let mut counter = baseline.counter; + loop { + match main_process(sandbox).await { + Ok(current) => { + assert_eq!( + (current.pid, current.start_ticks), + (baseline.pid, baseline.start_ticks), + "{}'s main process was replaced: {baseline:?} then {current:?}", + sandbox.name + ); + assert!( + current.counter >= counter, + "{}'s in-memory counter went back from {counter} to {}", + sandbox.name, + current.counter + ); + counter = current.counter; + samples.exec_outage = samples.exec_outage.max(last_exec.elapsed()); + last_exec = Instant::now(); + samples.exec_ok += 1; + samples.last = Some(current); + } + Err(error) => samples.exec_failures.push(error), + } + match try_call_host(sandbox, port).await { + Ok(result) if result == "HA_OK" => { + samples.network_outage = samples.network_outage.max(last_network.elapsed()); + last_network = Instant::now(); + samples.network_ok += 1; + } + Ok(result) => samples.network_failures.push(result), + Err(error) => samples.network_failures.push(error), + } + if *stop.borrow() { + break; + } + tokio::select! { + () = tokio::time::sleep(PROBE_INTERVAL) => {} + _ = stop.changed() => {} + } + } + samples.exec_outage = samples.exec_outage.max(last_exec.elapsed()); + samples.network_outage = samples.network_outage.max(last_network.elapsed()); + samples +} + +/// Run `transition` while probing every sandbox, and keep probing for +/// [`SETTLE`] after it finished. +async fn while_probing( + sandboxes: &[(&SandboxGuard, MainProcess)], + port: u16, + transition: impl Future, +) -> (T, Vec) { + let (stop, stopped) = tokio::sync::watch::channel(false); + let samplers = join_all( + sandboxes + .iter() + .map(|(sandbox, baseline)| sample(sandbox, baseline, port, stopped.clone())), + ); + let driver = async { + let outcome = transition.await; + tokio::time::sleep(SETTLE).await; + stop.send_replace(true); + outcome + }; + tokio::join!(driver, samplers) +} + +/// Check what the probes of `sandbox` saw against `bound`. +fn check_samples(phase: &str, sandbox: &SandboxGuard, samples: &Samples, bound: Duration) { + eprintln!( + "{phase}: {} exec {} ok / {} failed (longest outage {:?}), network {} ok / {} failed (longest outage {:?})", + sandbox.name, + samples.exec_ok, + samples.exec_failures.len(), + samples.exec_outage, + samples.network_ok, + samples.network_failures.len(), + samples.network_outage + ); + assert!( + samples.exec_ok > 0 && samples.network_ok > 0, + "{phase}: {} never answered: {samples:?}", + sandbox.name + ); + assert!( + samples.exec_outage <= bound, + "{phase}: exec on {} failed for {:?}, longer than {bound:?}: {:?}", + sandbox.name, + samples.exec_outage, + samples.exec_failures + ); + assert!( + samples.network_outage <= bound, + "{phase}: networking from {} failed for {:?}, longer than {bound:?}: {:?}", + sandbox.name, + samples.network_outage, + samples.network_failures + ); +} + +/// A sandbox under test and what it must keep. +struct Tracked { + sandbox: SandboxGuard, + generation: shared_supervisor::Generation, + baseline: MainProcess, +} + +impl Tracked { + async fn new(sandbox: SandboxGuard) -> Self { + let generation = generation(&sandbox.name).await; + let baseline = tokio::time::timeout(Duration::from_mins(1), async { + loop { + if let Ok(baseline) = main_process(&sandbox).await { + return baseline; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{}'s main process did not answer", sandbox.name)); + Self { + sandbox, + generation, + baseline, + } + } + + /// Same workload Pod, generation, and main process, which counted past + /// its baseline; networking answers. Moves the baseline forward. + async fn check_continuity(&mut self, phase: &str, port: u16) { + assert_eq!( + generation(&self.sandbox.name).await, + self.generation, + "{phase}: {}'s workload Pod or runtime generation changed", + self.sandbox.name + ); + self.baseline = counts_past(&self.sandbox, &self.baseline).await; + assert_eq!( + call_host(&self.sandbox, port).await, + "HA_OK", + "{phase}: {} cannot reach the host endpoint through its policy", + self.sandbox.name + ); + } +} + +/// Delete every sandbox in `tracked`, then fail with every delete error. +async fn delete_all(tracked: Vec) { + let mut failures = Vec::new(); + for mut sandbox in tracked.into_iter().map(|tracked| tracked.sandbox) { + if let Err(error) = sandbox.delete().await { + failures.push(error); + } + } + assert!( + failures.is_empty(), + "could not delete sandboxes: {failures:#?}" + ); +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running() { + // A one-member pool: pool mode is chosen explicitly, not by the count. + let initial = running_members().await; + assert_eq!( + initial.len(), + 1, + "the release must start with one supervisor pool member: {initial:?}" + ); + let first_member = initial.keys().next().expect("one member").clone(); + wait_until_attached(std::slice::from_ref(&first_member)).await; + let before = untouched().await; + assert!( + before.supervisor_template["spec"]["containers"] + .as_array() + .into_iter() + .flatten() + .flat_map(|container| container["env"].as_array().into_iter().flatten()) + .any(|env| env["name"] == POOL_MEMBER_ENV), + "the shared supervisor Pod template carries no pool member identity; install the chart from this branch" + ); + let endpoint = HostEndpoint::start().await; + let policy = policy(endpoint.port); + + // Phase 1: one member supervises a sandbox. + let mut a = Tracked::new(create(&policy).await).await; + assert_eq!(wait_for_owner(&a.generation.route).await, first_member); + a.check_continuity("one member", endpoint.port).await; + + // Phase 2: scale to two members. The new member joins; nothing moves. + let (two_members, samples) = while_probing( + &[(&a.sandbox, a.baseline.clone())], + endpoint.port, + scale_pool(2), + ) + .await; + check_samples("scale up", &a.sandbox, &samples[0], STAYING_OUTAGE); + assert_eq!( + two_members.get(&first_member), + initial.get(&first_member), + "scaling up replaced the first member: {two_members:?}" + ); + let second_member = two_members + .keys() + .find(|member| **member != first_member) + .expect("a second member") + .clone(); + wait_until_attached(&[first_member.clone(), second_member.clone()]).await; + assert_eq!( + untouched().await, + before, + "scaling up changed the gateway or the supervisor template" + ); + assert_eq!( + wait_for_owner(&a.generation.route).await, + first_member, + "a live owner keeps its sandbox when the pool grows" + ); + assert!( + gateway_log_lines(&a.generation.route) + .await + .iter() + .all(|line| !line.contains(TAKEOVER)), + "scaling up moved {}", + a.sandbox.name + ); + a.check_continuity("two members", endpoint.port).await; + + // A new sandbox lands on the new member; others the ring places on the + // first member are removed again. + let mut b = None; + for _ in 0..PLACEMENT_ATTEMPTS { + let candidate = Tracked::new(create(&policy).await).await; + if wait_for_owner(&candidate.generation.route).await == second_member { + b = Some(candidate); + break; + } + delete_all(vec![candidate]).await; + } + let mut b = b.unwrap_or_else(|| { + panic!( + "the pool placed {PLACEMENT_ATTEMPTS} new sandboxes on {first_member}, none on {second_member}" + ) + }); + b.check_continuity("two members", endpoint.port).await; + eprintln!( + "{} on {first_member}, {} on {second_member}", + a.sandbox.name, b.sandbox.name + ); + + // Phase 3: scale back to one member. Each member owns a sandbox, so the + // member that leaves hands at least one live sandbox to the other. + let owners_before = BTreeMap::from([ + (a.generation.route.clone(), first_member.clone()), + (b.generation.route.clone(), second_member.clone()), + ]); + let (members, samples) = while_probing( + &[ + (&a.sandbox, a.baseline.clone()), + (&b.sandbox, b.baseline.clone()), + ], + endpoint.port, + scale_pool(1), + ) + .await; + let (survivor, survivor_uid) = members.into_iter().next().expect("one member"); + let leaver = if survivor == first_member { + second_member.clone() + } else { + first_member.clone() + }; + assert_eq!( + two_members.get(&survivor), + Some(&survivor_uid), + "scaling down replaced the remaining member" + ); + eprintln!("{leaver} left the pool; {survivor} remains"); + wait_until_attached(std::slice::from_ref(&survivor)).await; + assert_eq!( + untouched().await, + before, + "scaling down changed the gateway or the supervisor template" + ); + let handed_off = owners_before + .iter() + .filter(|(_, owner)| **owner == leaver) + .map(|(route, _)| route.clone()) + .collect::>(); + assert!( + !handed_off.is_empty(), + "{leaver} owned no sandbox to hand off" + ); + for (tracked, samples) in [&a, &b].into_iter().zip(&samples) { + let moved = handed_off.contains(&tracked.generation.route); + check_samples( + if moved { + "scale down (handed off)" + } else { + "scale down (stayed)" + }, + &tracked.sandbox, + samples, + if moved { MOVING_OUTAGE } else { STAYING_OUTAGE }, + ); + } + for route in &handed_off { + assert_eq!( + wait_for_owner(route).await, + survivor, + "the remaining member must adopt {route}" + ); + let takeover = gateway_log_lines(route).await; + assert!( + takeover + .iter() + .any(|line| is_takeover(line, route, &leaver, "leaving")), + "no gateway handed {route} over from {leaver} as a leaving member: {takeover:?}" + ); + } + for (route, owner) in &owners_before { + if *owner == survivor { + assert_eq!( + owner_of(route).await.as_deref(), + Some(survivor.as_str()), + "the remaining member's own sandbox must stay with it" + ); + } + } + a.check_continuity("back to one member", endpoint.port) + .await; + b.check_continuity("back to one member", endpoint.port) + .await; + + // A sandbox created after scaling down runs on the remaining member. + let mut c = Tracked::new(create(&policy).await).await; + assert_eq!(wait_for_owner(&c.generation.route).await, survivor); + c.check_continuity("back to one member", endpoint.port) + .await; + for tracked in [&mut a, &mut b] { + tracked.check_continuity("final", endpoint.port).await; + } + assert_eq!( + untouched().await, + before, + "the gateway or the supervisor template changed" + ); + + delete_all(vec![a, b, c]).await; +} diff --git a/e2e/rust/tests/shared_supervisor/mod.rs b/e2e/rust/tests/shared_supervisor/mod.rs new file mode 100644 index 0000000000..55ee3c715e --- /dev/null +++ b/e2e/rust/tests/shared_supervisor/mod.rs @@ -0,0 +1,808 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Helpers shared by the shared supervisor pool tests: the release's +//! Kubernetes objects, which pool member supervises a sandbox, and the +//! sandbox workload whose main process counts in memory. +//! +//! Each test binary uses a subset. +#![allow(dead_code)] + +use std::process::Stdio; +use std::time::Duration; + +use openshell_e2e::harness::output::strip_ansi; +use openshell_e2e::harness::sandbox::SandboxGuard; +use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _}; +use tokio::net::{TcpListener, TcpStream}; + +pub const TEST_HOST: &str = "host.openshell.internal"; +pub const SUPERVISOR_COMPONENT: &str = "app.kubernetes.io/component=supervisor"; +const INSTANCE_ROUTE_ANNOTATION: &str = "openshell.ai/sandbox-runtime-supervisor-session"; +const GENERATION_ANNOTATION: &str = "openshell.ai/sandbox-runtime-generation"; +/// The workload's main process: it records its PID once and a counter that +/// lives only in its memory every second. A restarted process would start +/// counting from one again. +pub const COUNTER_MAIN: &str = "echo $$ > /tmp/ha-main.pid; i=0; while true; do i=$((i+1)); echo $i > /tmp/ha-counter; sleep 1; done"; +/// Gateway Pods of the release: the chart's selector labels name the gateway +/// `openshell`, while the supervisor and certgen Pods carry their own names. +pub const GATEWAY_SELECTOR: &str = "app.kubernetes.io/name=openshell"; +pub const GATEWAY_CONTAINER: &str = "openshell-gateway"; +pub const SUPERVISOR_CONTAINER: &str = "supervisor"; +/// What the gateway logs when it moves an instance to another member. +pub const TAKEOVER: &str = "taking over a shared sandbox instance from a failed supervisor"; +/// Supervisor lines that change which member supervises a route; each names +/// the route in its text after the phrase. They must match the messages in +/// `crates/openshell-supervisor/src/shared_host.rs`. +const ADMITTED: &str = "admitted assigned sandbox instance"; +const RELEASES: [&str; 3] = [ + "removing withdrawn sandbox instance", + "relinquishing a sandbox instance owned by another pool member", + "relinquishing a sandbox instance for its later ownership term", +]; +/// What a supervisor logs when the gateway accepts its stream, and when any +/// stream ends; the later of the two says whether it is attached now. +const ATTACHED: &str = "gateway accepted the supervisor host"; +const DETACHED: &str = "gateway assignment stream lost"; +/// Bound on a supervisor attaching to the gateway. A Pod is TCP-ready before +/// it can authenticate: on a fresh install the gateway writes the supervisor +/// token into the chart's initially empty Secret, and the kubelet projects +/// the update into the Pod only after its sync period (about 90 s observed). +pub const ATTACH_TIMEOUT: Duration = Duration::from_mins(3); +/// After a member attaches, the gateway writes its membership record within +/// two bounded refresh rounds (2 x 500 ms); placement reads that record. +const MEMBERSHIP_RECORD_BOUND: Duration = Duration::from_secs(2); +/// Bound on one exec probe of a sandbox, so a hung CLI cannot stall a test. +pub const EXEC_PROBE_TIMEOUT: Duration = Duration::from_secs(20); +/// Bound on one policy-mediated network probe from a sandbox. +pub const NETWORK_PROBE_TIMEOUT: Duration = Duration::from_secs(40); + +pub fn release_namespace() -> String { + std::env::var("OPENSHELL_E2E_KUBE_NAMESPACE").unwrap_or_else(|_| "openshell".to_string()) +} + +/// The kubeconfig and context the harness uses, when it names them. +fn kube_target() -> (Option, Option) { + let set = |name: &str| { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) + }; + ( + set("OPENSHELL_E2E_KUBECONFIG"), + set("OPENSHELL_E2E_KUBE_CONTEXT_ACTIVE"), + ) +} + +pub fn kubectl_command() -> tokio::process::Command { + let mut command = tokio::process::Command::new("kubectl"); + let (kubeconfig, context) = kube_target(); + if let Some(kubeconfig) = kubeconfig { + command.arg("--kubeconfig").arg(kubeconfig); + } + if let Some(context) = context { + command.arg("--context").arg(context); + } + command.stdin(Stdio::null()).kill_on_drop(true); + command +} + +pub fn helm_command() -> tokio::process::Command { + let mut command = tokio::process::Command::new("helm"); + let (kubeconfig, context) = kube_target(); + if let Some(kubeconfig) = kubeconfig { + command.arg("--kubeconfig").arg(kubeconfig); + } + if let Some(context) = context { + command.arg("--kube-context").arg(context); + } + command.stdin(Stdio::null()).kill_on_drop(true); + command +} + +pub async fn kubectl(args: &[&str]) -> String { + let output = tokio::time::timeout( + Duration::from_mins(2), + kubectl_command() + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output(), + ) + .await + .unwrap_or_else(|_| panic!("kubectl {args:?} did not finish within two minutes")) + .unwrap_or_else(|error| panic!("spawn kubectl {args:?}: {error}")); + assert!( + output.status.success(), + "kubectl {args:?} failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + String::from_utf8_lossy(&output.stdout).into_owned() +} + +pub async fn kubectl_json(args: &[&str]) -> serde_json::Value { + serde_json::from_str(&kubectl(args).await).expect("parse kubectl JSON") +} + +/// The Sandbox resource of an `OpenShell` sandbox. +async fn sandbox_resource(name: &str) -> serde_json::Value { + let list = kubectl_json(&[ + "get", + "sandboxes.agents.x-k8s.io", + "--all-namespaces", + "-l", + &format!("openshell.ai/sandbox-name={name}"), + "-o", + "json", + ]) + .await; + let [resource] = list["items"].as_array().expect("Sandbox items").as_slice() else { + panic!("expected one Sandbox resource for {name}: {list}"); + }; + resource.clone() +} + +/// Identity of a sandbox's running generation as Kubernetes records it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Generation { + pub route: String, + pub runtime: String, + pub workload_pod_uid: String, +} + +pub async fn generation(name: &str) -> Generation { + let resource = sandbox_resource(name).await; + let annotations = &resource["metadata"]["annotations"]; + let namespace = resource["metadata"]["namespace"] + .as_str() + .expect("Sandbox namespace"); + let resource_name = resource["metadata"]["name"].as_str().expect("Sandbox name"); + let resource_uid = resource["metadata"]["uid"].as_str().expect("Sandbox UID"); + let pods = kubectl_json(&[ + "-n", + namespace, + "get", + "pods", + "-l", + "openshell.ai/boundary-role=workload", + "-o", + "json", + ]) + .await; + let workload = pods["items"] + .as_array() + .expect("Pod items") + .iter() + .find(|pod| { + pod["metadata"]["name"] == resource_name + || pod["metadata"]["ownerReferences"] + .as_array() + .into_iter() + .flatten() + .any(|owner| owner["uid"] == resource_uid) + }) + .unwrap_or_else(|| panic!("no workload Pod for {name}: {pods}")); + Generation { + route: annotations[INSTANCE_ROUTE_ANNOTATION] + .as_str() + .unwrap_or_else(|| panic!("{name} is not on the shared supervisor: {annotations}")) + .to_string(), + runtime: annotations[GENERATION_ANNOTATION] + .as_str() + .unwrap_or_else(|| panic!("{name} has no runtime generation: {annotations}")) + .to_string(), + workload_pod_uid: workload["metadata"]["uid"] + .as_str() + .expect("workload Pod UID") + .to_string(), + } +} + +/// The release's shared supervisor Pods, as JSON. +pub async fn supervisor_pod_items() -> Vec { + let list = kubectl_json(&[ + "-n", + &release_namespace(), + "get", + "pods", + "-l", + SUPERVISOR_COMPONENT, + "-o", + "json", + ]) + .await; + list["items"].as_array().expect("Pod items").clone() +} + +/// Names of the running shared supervisor Pods. +pub async fn supervisor_pods() -> Vec { + supervisor_pod_items() + .await + .iter() + .filter(|pod| { + pod["metadata"]["deletionTimestamp"].is_null() && pod["status"]["phase"] == "Running" + }) + .map(|pod| { + pod["metadata"]["name"] + .as_str() + .expect("Pod name") + .to_string() + }) + .collect() +} + +/// Whether `text` holds `phrase` followed by exactly `value` as a whole +/// token: a route or member name that merely starts with `value` does not +/// count. +fn names(text: &str, phrase: &str, value: &str) -> bool { + let needle = format!("{phrase}{value}"); + text.match_indices(&needle).any(|(start, _)| { + let before = text[..start].chars().next_back(); + let after = text[start + needle.len()..].chars().next(); + let token = |c: char| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.'; + (phrase.ends_with(' ') || before.is_none_or(|c| !token(c))) + && after.is_none_or(|c| !token(c)) + }) +} + +/// What one supervisor log line says about who supervises `route`: +/// `Some(true)` for an admission, `Some(false)` for a removal or release, +/// and `None` for anything else, including lines that only mention the route +/// (such as an SSH relay's socket path). +pub fn ownership_event(line: &str, route: &str) -> Option { + let line = strip_ansi(line); + if names(&line, &format!("{ADMITTED} "), route) { + Some(true) + } else if RELEASES + .iter() + .any(|phrase| names(&line, &format!("{phrase} "), route)) + { + Some(false) + } else { + None + } +} + +/// Whether a supervisor's log says it supervises `route` now: its latest +/// ownership event for the route is an admission. +pub fn supervises(logs: &str, route: &str) -> bool { + logs.lines() + .rev() + .find_map(|line| ownership_event(line, route)) + .unwrap_or(false) +} + +/// Whether a gateway log line records moving `route` from member `from` +/// for `cause`, with the gateway's styling removed. +pub fn is_takeover(line: &str, route: &str, from: &str, cause: &str) -> bool { + let line = strip_ansi(line); + line.contains(TAKEOVER) + && names(&line, "route=", route) + && names(&line, "from=", from) + && names(&line, "cause=", cause) +} + +/// The supervisor Pod that currently supervises `route`, judged from what +/// each running member logged about it last. +pub async fn owner_of(route: &str) -> Option { + for pod in supervisor_pods().await { + let Ok(Ok(logs)) = tokio::time::timeout( + Duration::from_secs(30), + kubectl_command() + .args([ + "-n", + &release_namespace(), + "logs", + &pod, + "-c", + SUPERVISOR_CONTAINER, + ]) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .output(), + ) + .await + else { + continue; + }; + if supervises(&String::from_utf8_lossy(&logs.stdout), route) { + return Some(pod); + } + } + None +} + +/// The latest attach or detach line `pod`'s supervisor logged, if any. +async fn attachment(pod: &str) -> Option { + let output = tokio::time::timeout( + Duration::from_secs(30), + kubectl_command() + .args([ + "-n", + &release_namespace(), + "logs", + pod, + "-c", + SUPERVISOR_CONTAINER, + ]) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .output(), + ) + .await + .ok()? + .ok()?; + String::from_utf8_lossy(&output.stdout) + .lines() + .rev() + .find(|line| line.contains(ATTACHED) || line.contains(DETACHED)) + .map(str::to_string) +} + +/// Wait until every one of `members` is attached to the gateway now (its +/// latest stream event is an acceptance, not a loss), and its membership +/// record has had time to be written. Pod readiness only means the process +/// listens; it says nothing about authenticating to the gateway. +pub async fn wait_until_attached(members: &[String]) { + let mut last = std::collections::BTreeMap::new(); + let attached = tokio::time::timeout(ATTACH_TIMEOUT, async { + loop { + let mut all = true; + for member in members { + let event = attachment(member).await; + all &= event.as_deref().is_some_and(|line| line.contains(ATTACHED)); + last.insert(member.clone(), event); + } + if all { + return; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await; + assert!( + attached.is_ok(), + "supervisor pool members did not attach to the gateway within {ATTACH_TIMEOUT:?}; latest stream event per member: {last:#?}" + ); + tokio::time::sleep(MEMBERSHIP_RECORD_BOUND).await; +} + +pub async fn wait_for_owner(route: &str) -> String { + tokio::time::timeout(Duration::from_mins(1), async { + loop { + if let Some(owner) = owner_of(route).await { + return owner; + } + tokio::time::sleep(Duration::from_secs(1)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("no shared supervisor admitted route {route}")) +} + +/// PID, start time, and counter of a sandbox's main process. +#[derive(Debug, Clone)] +pub struct MainProcess { + pub pid: u32, + pub start_ticks: u64, + pub counter: u64, +} + +/// One bounded exec probe of the main process. +pub async fn main_process(sandbox: &SandboxGuard) -> Result { + let output = tokio::time::timeout( + EXEC_PROBE_TIMEOUT, + sandbox.exec(&[ + "sh", + "-c", + "pid=$(cat /tmp/ha-main.pid) && echo pid=$pid && \ + echo start=$(cut -d')' -f2 /proc/$pid/stat | awk '{print $20}') && \ + echo counter=$(cat /tmp/ha-counter)", + ]), + ) + .await + .map_err(|_| format!("exec did not finish within {EXEC_PROBE_TIMEOUT:?}"))??; + let field = |key: &str| { + output + .lines() + .find_map(|line| line.trim().strip_prefix(key)) + .and_then(|value| value.trim().parse::().ok()) + .ok_or_else(|| format!("no {key} in main process report:\n{output}")) + }; + Ok(MainProcess { + pid: u32::try_from(field("pid=")?).map_err(|error| error.to_string())?, + start_ticks: field("start=")?, + counter: field("counter=")?, + }) +} + +/// Wait until the same main process as `previous` has counted past it. A +/// successful read alone does not prove a tick happened since the last one. +pub async fn counts_past(sandbox: &SandboxGuard, previous: &MainProcess) -> MainProcess { + tokio::time::timeout(Duration::from_mins(1), async { + loop { + if let Ok(current) = main_process(sandbox).await { + assert_eq!(current.pid, previous.pid, "the main process was restarted"); + assert_eq!( + current.start_ticks, previous.start_ticks, + "a new process reused the main process's PID" + ); + if current.counter > previous.counter { + return current; + } + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{} stopped counting after {previous:?}", sandbox.name)) +} + +/// Host endpoint that answers every request with `HA_OK`. +pub struct HostEndpoint { + pub port: u16, + task: tokio::task::JoinHandle<()>, +} + +impl HostEndpoint { + pub async fn start() -> Self { + let listener = TcpListener::bind(("0.0.0.0", 0)) + .await + .expect("bind host endpoint"); + let port = listener.local_addr().expect("endpoint address").port(); + let task = tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + tokio::spawn(async move { + let _ = answer(stream).await; + }); + } + }); + Self { port, task } + } +} + +impl Drop for HostEndpoint { + fn drop(&mut self) { + self.task.abort(); + } +} + +/// The host endpoint's only response. +const RESPONSE: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nConnection: close\r\n\r\nHA_OK"; + +/// Read one request head and answer it; returns the head as received. +async fn answer(stream: TcpStream) -> std::io::Result { + let mut reader = tokio::io::BufReader::new(stream); + let mut head = String::new(); + loop { + let mut line = String::new(); + if reader.read_line(&mut line).await? == 0 { + break; + } + head.push_str(&line); + if line == "\r\n" { + break; + } + } + reader.get_mut().write_all(RESPONSE).await?; + Ok(head) +} + +pub fn policy(port: u16) -> tempfile::NamedTempFile { + let mut file = tempfile::Builder::new() + .suffix(".yaml") + .tempfile() + .expect("create policy file"); + std::io::Write::write_all( + &mut file, + format!( + r"version: 1 +filesystem_policy: + include_workdir: true + read_only: [/usr, /lib, /proc, /dev/urandom, /app, /etc, /var/log] + read_write: [/sandbox, /tmp, /dev/null] +landlock: + compatibility: best_effort +process: + run_as_user: sandbox + run_as_group: sandbox +network_policies: + ha: + name: ha + endpoints: + - host: {TEST_HOST} + port: {port} + protocol: rest + access: full + binaries: + - path: /usr/bin/python* + - path: /usr/local/bin/python* + - path: /sandbox/.uv/python/*/bin/python* +" + ) + .as_bytes(), + ) + .expect("write policy file"); + file +} + +/// Python that sends one HTTP request to `host:port` and prints `HA_OK`, +/// `UNEXPECTED`, or `DENIED`. The policy enforces the request's authority, so +/// the Host header names the port as well as the host. The Rust raw string +/// leaves `\r\n` for Python's bytes literal to turn into CRLF. +fn host_call_script(host: &str, port: u16) -> String { + format!( + r#" +import socket +try: + with socket.create_connection(({host:?}, {port}), timeout=15) as sock: + sock.sendall(b"GET / HTTP/1.1\r\nHost: {host}:{port}\r\nConnection: close\r\n\r\n") + sock.settimeout(15) + response = b"" + while chunk := sock.recv(4096): + response += chunk + print("HA_OK" if b"HA_OK" in response else "UNEXPECTED") +except OSError: + print("DENIED") +"# + ) +} + +/// One bounded call from the sandbox to the host endpoint through its policy: +/// `HA_OK`, `UNEXPECTED`, `DENIED`, or `NO_RESULT`. +pub async fn try_call_host(sandbox: &SandboxGuard, port: u16) -> Result { + let script = host_call_script(TEST_HOST, port); + let output = tokio::time::timeout( + NETWORK_PROBE_TIMEOUT, + sandbox.exec(&["python3", "-c", &script]), + ) + .await + .map_err(|_| format!("network probe did not finish within {NETWORK_PROBE_TIMEOUT:?}"))??; + Ok(output + .lines() + .map(str::trim) + .find(|line| matches!(*line, "HA_OK" | "UNEXPECTED" | "DENIED")) + .unwrap_or("NO_RESULT") + .to_string()) +} + +/// Call the host endpoint from the sandbox through its policy. +pub async fn call_host(sandbox: &SandboxGuard, port: u16) -> String { + try_call_host(sandbox, port) + .await + .unwrap_or_else(|error| panic!("call host from {}: {error}", sandbox.name)) +} + +/// Gateway log lines that mention `route`, from every gateway Pod, without +/// the styling the gateway's console format puts between field names, `=`, +/// and values. +pub async fn gateway_log_lines(route: &str) -> Vec { + let namespace = release_namespace(); + let pods = kubectl_json(&[ + "-n", + &namespace, + "get", + "pods", + "-l", + GATEWAY_SELECTOR, + "-o", + "json", + ]) + .await; + let pods = pods["items"] + .as_array() + .expect("Pod items") + .iter() + .map(|pod| { + pod["metadata"]["name"] + .as_str() + .expect("Pod name") + .to_string() + }) + .collect::>(); + assert!(!pods.is_empty(), "no gateway Pods match {GATEWAY_SELECTOR}"); + let mut lines = Vec::new(); + for pod in pods { + let logs = kubectl(&["-n", &namespace, "logs", &pod, "-c", GATEWAY_CONTAINER]).await; + lines.extend( + logs.lines() + .map(strip_ansi) + .filter(|line| line.contains(route)), + ); + } + lines +} + +/// Create a sandbox whose main process counts in memory, and wait until it +/// counts. +pub async fn create(policy: &tempfile::NamedTempFile) -> SandboxGuard { + let policy = policy.path().to_str().expect("UTF-8 policy path"); + let sandbox = SandboxGuard::create_detached_main_with_args( + &["--policy", policy], + &["sh", "-c", COUNTER_MAIN], + ) + .await + .expect("create sandbox"); + tokio::time::timeout(Duration::from_mins(2), async { + while main_process(&sandbox).await.is_err() { + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{} never started counting", sandbox.name)); + sandbox +} + +#[cfg(test)] +mod tests { + use super::*; + + const ROUTE: &str = "rt-17999333d7754f3e957b9cf2fca0ee29"; + + fn supervisor_line(message: &str) -> String { + format!("2026-10-23T06:07:20.104Z INFO openshell_supervisor::shared_host: {message}") + } + + #[test] + fn an_admission_in_shorthand_output_names_its_owner() { + let logs = [ + supervisor_line("gateway accepted the supervisor host"), + supervisor_line(&format!("{ADMITTED} {ROUTE}")), + ] + .join("\n"); + assert!(supervises(&logs, ROUTE)); + } + + #[test] + fn lines_that_only_mention_the_route_change_nothing() { + // SSH relay activity after the admission names the route only in a + // socket path. + let logs = [ + supervisor_line(&format!("{ADMITTED} {ROUTE}")), + supervisor_line(&format!( + "SSH relay opened socket=/run/openshell/sessions/{ROUTE}/ssh.sock" + )), + supervisor_line(&format!( + "SSH relay closed /run/openshell/sessions/{ROUTE}/ssh.sock" + )), + ] + .join("\n"); + assert!(supervises(&logs, ROUTE)); + assert_eq!(ownership_event(logs.lines().last().unwrap(), ROUTE), None); + } + + #[test] + fn a_removal_or_release_ends_supervision() { + for release in RELEASES.iter().take(2) { + let logs = [ + supervisor_line(&format!("{ADMITTED} {ROUTE}")), + supervisor_line(&format!("{release} {ROUTE}")), + supervisor_line(&format!( + "SSH relay closed /run/openshell/sessions/{ROUTE}/ssh.sock" + )), + ] + .join("\n"); + assert!(!supervises(&logs, ROUTE), "{release}"); + } + } + + #[test] + fn adopting_a_later_term_keeps_supervising() { + let logs = [ + supervisor_line(&format!("{ADMITTED} {ROUTE}")), + supervisor_line(&format!("{} {ROUTE}", RELEASES[2])), + supervisor_line(&format!("{ADMITTED} {ROUTE}")), + ] + .join("\n"); + assert!(supervises(&logs, ROUTE)); + } + + #[test] + fn another_route_sharing_a_prefix_is_not_this_route() { + let logs = supervisor_line(&format!("{ADMITTED} {ROUTE}0")); + assert!(!supervises(&logs, ROUTE)); + assert!(!supervises("", ROUTE)); + } + + /// A takeover line as the gateway's styled console prints it: italic + /// field names, a dimmed `=`, plain values. + fn styled_takeover(route: &str, from: &str, to: &str, cause: &str) -> String { + let field = |key: &str, value: &str| format!("\x1b[3m{key}\x1b[0m\x1b[2m=\x1b[0m{value}"); + format!( + "\x1b[2m2026-10-23T06:09:01.512Z\x1b[0m \x1b[32m INFO\x1b[0m \x1b[2mopenshell_server::supervisor_pool\x1b[0m\x1b[2m:\x1b[0m {TAKEOVER} {} {} {} {} {} {}", + field("sandbox_id", "sandbox-1"), + field("route", route), + field("from", from), + field("to", to), + field("term", "2"), + field("cause", cause), + ) + } + + /// The sandbox's request reaches a listener as a well-formed HTTP request + /// whose authority names the port, and the endpoint's answer reads as + /// `HA_OK`: the exact bytes both ways, without a cluster. + #[tokio::test] + async fn the_host_call_sends_a_request_with_its_authority_and_reads_the_answer() { + let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let served = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + answer(stream).await.unwrap() + }); + let output = match tokio::process::Command::new("python3") + .args(["-c", &host_call_script("127.0.0.1", port)]) + .stdin(Stdio::null()) + .output() + .await + { + Ok(output) => output, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + eprintln!("python3 is not installed; skipping"); + return; + } + Err(error) => panic!("run python3: {error}"), + }; + let head = tokio::time::timeout(Duration::from_secs(20), served) + .await + .expect("the request arrives") + .unwrap(); + assert_eq!( + head, + format!("GET / HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n") + ); + assert_eq!( + String::from_utf8_lossy(&output.stdout).trim(), + "HA_OK", + "{}", + String::from_utf8_lossy(&output.stderr) + ); + // The response's Content-Length is its body's. + let (headers, body) = std::str::from_utf8(RESPONSE) + .unwrap() + .split_once("\r\n\r\n") + .unwrap(); + assert!(headers.contains(&format!("Content-Length: {}", body.len()))); + } + + #[test] + fn a_styled_takeover_line_matches_its_route_member_and_cause() { + let line = styled_takeover( + ROUTE, + "openshell-supervisor-abc-1", + "openshell-supervisor-abc-2", + "leaving", + ); + assert!(!line.contains("cause=leaving"), "the fixture is styled"); + assert!(is_takeover( + &line, + ROUTE, + "openshell-supervisor-abc-1", + "leaving" + )); + assert!(!is_takeover( + &line, + ROUTE, + "openshell-supervisor-abc-1", + "expired" + )); + assert!(!is_takeover( + &line, + ROUTE, + "openshell-supervisor-abc-2", + "leaving" + )); + assert!(!is_takeover( + &line, + ROUTE, + "openshell-supervisor-abc", + "leaving" + )); + assert!(!is_takeover( + &line, + "rt-other", + "openshell-supervisor-abc-1", + "leaving" + )); + } +} diff --git a/e2e/with-kube-gateway.sh b/e2e/with-kube-gateway.sh index b13dbc380b..13e58b10bf 100755 --- a/e2e/with-kube-gateway.sh +++ b/e2e/with-kube-gateway.sh @@ -32,6 +32,10 @@ # files, relative to the repository root or absolute, to layer additional chart # configuration on top of ci/values-skaffold.yaml. # +# Set OPENSHELL_E2E_KUBE_TLS=1 with an overlay that keeps gateway TLS on (for +# example ci/values-tls-port-forward.yaml) to reach the gateway over mTLS +# through the port-forward with the chart-generated client certificate. +# # Image source: # - Ephemeral k3d mode builds local # `openshell/{gateway,sandbox,supervisor}:${IMAGE_TAG}` @@ -753,6 +757,11 @@ run_scenario() { export OPENSHELL_E2E_KUBE_CONTEXT="${KUBE_CONTEXT}" export OPENSHELL_E2E_KUBE_NAMESPACE="${NAMESPACE}" export OPENSHELL_E2E_KUBE_RELEASE="${RELEASE_NAME}" + # Tests that upgrade the release (for example to scale the shared + # supervisor pool) reuse its values against this exact chart, through the + # kubeconfig this harness uses. + export OPENSHELL_E2E_KUBE_CHART="${ROOT}/deploy/helm/openshell" + export OPENSHELL_E2E_KUBECONFIG="${KUBECONFIG:-}" export OPENSHELL_PROVISION_TIMEOUT="${OPENSHELL_PROVISION_TIMEOUT:-300}" e2e_import_example_provider_profiles \ @@ -1465,6 +1474,27 @@ else # OpenShift: reach the gateway over the passthrough Route with mTLS so the # SSH-relay `sandbox connect` suites work (port-forward stalls them). openshift_register_route_gateway || exit 1 + elif [ "${OPENSHELL_E2E_KUBE_TLS:-0}" = "1" ]; then + # Topologies that require gateway TLS: reach the gateway over mTLS through + # the port-forward with the chart-generated client certificate. The + # server certificate names 127.0.0.1. + start_gateway_portforward || exit 1 + kube_pki_dir="${WORKDIR}/kube-client-pki" + mkdir -p "${kube_pki_dir}/client" + kctl -n "${NAMESPACE}" get secret openshell-client-tls \ + -o jsonpath='{.data.ca\.crt}' | base64 -d >"${kube_pki_dir}/ca.crt" + kctl -n "${NAMESPACE}" get secret openshell-client-tls \ + -o jsonpath='{.data.tls\.crt}' | base64 -d >"${kube_pki_dir}/client/tls.crt" + kctl -n "${NAMESPACE}" get secret openshell-client-tls \ + -o jsonpath='{.data.tls\.key}' | base64 -d >"${kube_pki_dir}/client/tls.key" + GATEWAY_NAME="openshell-e2e-kube-${LOCAL_PORT}" + GATEWAY_ENDPOINT="https://127.0.0.1:${LOCAL_PORT}" + e2e_register_mtls_gateway \ + "${XDG_CONFIG_HOME}" \ + "${GATEWAY_NAME}" \ + "${GATEWAY_ENDPOINT}" \ + "${LOCAL_PORT}" \ + "${kube_pki_dir}" else # Vanilla Kubernetes: reach the gateway in plaintext over port-forward. start_gateway_portforward || exit 1 @@ -1491,6 +1521,11 @@ else export OPENSHELL_E2E_KUBE_CONTEXT="${KUBE_CONTEXT}" export OPENSHELL_E2E_KUBE_NAMESPACE="${NAMESPACE}" export OPENSHELL_E2E_KUBE_RELEASE="${RELEASE_NAME}" + # Tests that upgrade the release (for example to scale the shared + # supervisor pool) reuse its values against this exact chart, through the + # kubeconfig this harness uses. + export OPENSHELL_E2E_KUBE_CHART="${ROOT}/deploy/helm/openshell" + export OPENSHELL_E2E_KUBECONFIG="${KUBECONFIG:-}" export OPENSHELL_PROVISION_TIMEOUT="${OPENSHELL_PROVISION_TIMEOUT:-300}" e2e_import_example_provider_profiles \ diff --git a/examples/isolation-backend-v8/src/lifecycle.rs b/examples/isolation-backend-v8/src/lifecycle.rs index 9391e237fd..f80a48575e 100644 --- a/examples/isolation-backend-v8/src/lifecycle.rs +++ b/examples/isolation-backend-v8/src/lifecycle.rs @@ -87,6 +87,7 @@ impl Runtime { let registry = Arc::new(SandboxConnectionRegistry::new( launch.boundary.session_id, launch.rotation, + launch.epoch, )); Ok(Arc::new(Self { launch, @@ -234,8 +235,10 @@ impl Runtime { supervisor_instance_id, .. } => { + // This backend does not offer supervisor takeover, so a + // replacement supervisor is refused rather than staged. self.registry - .attach(principal, *supervisor_instance_id) + .attach(principal, *supervisor_instance_id, None) .map_err(|e| Status::failed_precondition(e.to_string()))?; } Q::Confirm => self @@ -298,6 +301,7 @@ impl Runtime { snapshot: wire::SessionSnapshotWire { generation: self.launch.boundary.generation.clone(), processes: vec![], + bound_policy: state.policy.clone().map(Box::new), }, } } diff --git a/proto/sandbox_protocol.proto b/proto/sandbox_protocol.proto index 40a35cf096..d0befff52b 100644 --- a/proto/sandbox_protocol.proto +++ b/proto/sandbox_protocol.proto @@ -113,6 +113,17 @@ message DiscoverPolicy {} message Attach { string supervisor_instance_id = 1; SandboxPolicy policy = 2; + // Gateway-signed supervisor takeover grant. Present only when this + // supervisor process replaces the current owner of the generation at a + // higher credential epoch; send it only to a runtime advertising + // openshell.sandbox.supervisor-takeover. The runtime stages the replacement + // and promotes it at Confirm, fencing every older connection and epoch. + optional string supervisor_takeover = 3 [(openshell.options.v1.secret) = true]; + // How long the runtime should keep the workload frozen after losing this + // connection's supervision, offered by the supervisor in whole seconds. The + // runtime refuses an offer outside 10 to 300 seconds and applies its + // default (30 seconds) without one. + google.protobuf.Duration supervisor_recovery_timeout = 4; } // Success asserts that the backend can enforce the attached policy and trusted // workload identity. Unsupported requirements fail before StartAgent, which @@ -329,6 +340,10 @@ message BoundaryConfirmation { message SessionSnapshot { string generation = 1; repeated ProcessSnapshot processes = 2; + // Launch controls the runtime bound at its first attach. A replacement + // supervisor reattaches with exactly these; they never change for a + // running generation. + SandboxPolicy bound_policy = 3; } message ProcessSnapshot { string process_id = 1; diff --git a/proto/supervisor_host.proto b/proto/supervisor_host.proto index 2747dd5749..3b5a892742 100644 --- a/proto/supervisor_host.proto +++ b/proto/supervisor_host.proto @@ -18,6 +18,12 @@ import "options.proto"; // instances and sends it as complete snapshots; the host reconciles its // registry to each snapshot and reports what it observes. Disconnecting // never stops hosted instances; only a later snapshot that omits them does. +// +// A host that names a pool member in its hello belongs to a supervisor pool +// (high availability). The gateway places each instance on one member, sends +// each member only the instances it owns, and moves an instance only when its +// owner fails or leaves. A pool host releases an instance a snapshot omits +// without stopping its workload; only `stop_routes` stops workloads. service SupervisorHost { // Supervisor-initiated assignment stream. Only the gateway-minted // supervisor token may call it. @@ -55,6 +61,19 @@ message HostHello { // Counter the host increments for every connection of one instance_id, so // the gateway can order reconnects. uint64 connection_epoch = 2; + // Present when the host is a member of a supervisor pool. + PoolMember member = 3; +} + +// One member of a supervisor pool, as its deployment identifies it. +message PoolMember { + // Stable member name (the Pod name). One process serves a member at a time. + string member_id = 1; + // Identity of the member's current incarnation (the Pod UID). + string member_uid = 2; + // Address (IP:port) of the member's routed runtime listener. Other members + // redirect sandbox runtimes here for the instances this member owns. + string runtime_address = 3; } // The gateway accepted the host and will send assignment snapshots. @@ -70,11 +89,38 @@ message HostRejected { } // Complete desired set of sandbox instances. Each snapshot replaces the -// previous one; an instance absent from a snapshot must stop. +// previous one; an instance absent from a snapshot must stop, except on a +// pool host, which releases it without stopping its workload. message AssignmentSnapshot { // Monotonic within the session, starting at 1. uint64 revision = 1; repeated SandboxInstanceAssignment instances = 2; + // The host is a pool member: release omitted instances instead of + // stopping them, and stop only the routes listed in `stop_routes`. + bool pool = 3; + // Pool hosts only: routes whose sandbox is stopping or deleted; the host + // stops their workloads if it still hosts them. + repeated string stop_routes = 4; + // Pool hosts only: where each instance placed in the pool is supervised, + // so a sandbox runtime that reaches this host is redirected to its owner. + repeated RouteOwner route_directory = 5; + // Pool hosts only: routes this host's process owns whose launch material + // this gateway replica has not loaded yet. The host keeps supervising them + // unchanged; it neither releases nor stops them, and starts nothing for + // them. + repeated string keep_routes = 6; + // Pool hosts only: how long, in whole seconds, a sandbox runtime keeps its + // workload frozen after losing this host's supervision before terminating + // it. The host offers it to each instance's runtime when it attaches; + // absent keeps the runtime default. + google.protobuf.Duration recovery_timeout = 7; +} + +// The pool member that supervises one route. +message RouteOwner { + string route = 1; + // Runtime listener address (IP:port) of the owning member. + string runtime_address = 2; } // One running generation of a stable Sandbox resource, assigned to the host. @@ -88,6 +134,30 @@ message SandboxInstanceAssignment { // instance registration (runtime descriptor, gateway-issued credentials, // proxy CA, main process). Never logged or forwarded. bytes launch_material = 4 [(openshell.options.v1.secret) = true]; + // Present when the host supervises the instance under a later ownership + // term than its launch: it replaces the previous owner with these + // credentials and adopts the running workload. + SandboxInstanceTakeover takeover = 5; +} + +// Credentials of one ownership term after the launch term. The gateway +// advanced the sandbox's credential epoch to `auth_epoch` before granting it, +// which revoked every credential of earlier terms. A term is identified by +// `auth_epoch` and `supervisor_instance_id`: each gateway replica mints its +// own credentials for it, and may mint them again, so a host that already +// supervises the term keeps doing so whatever credentials arrive. A term sent +// without credentials (empty `grant`) was minted by another gateway replica; +// a host that does not supervise it yet waits for credentials. +message SandboxInstanceTakeover { + uint64 auth_epoch = 1; + // Supervisor process identity the runtime pins for this term. + string supervisor_instance_id = 2; + // Gateway-signed takeover grant presented when attaching to the runtime. + string grant = 3 [(openshell.options.v1.secret) = true]; + string gateway_token = 4 [(openshell.options.v1.secret) = true]; + int64 gateway_expires_at = 5; + string sandbox_token = 6 [(openshell.options.v1.secret) = true]; + int64 sandbox_expires_at = 7; } // Everything the host observes, sent after reconciling each snapshot and @@ -99,6 +169,9 @@ message HostStatus { // Every hosted instance, including ones that ended but are still being // removed. repeated SandboxInstanceStatus instances = 2; + // A pool host is shutting down: the gateway moves its instances to other + // members at once. Its workloads keep running. + bool leaving = 3; } // Lifecycle phase of a hosted sandbox instance. diff --git a/rfc/0012-isolation-backend/README.md b/rfc/0012-isolation-backend/README.md index 9febf8d763..74a16029d8 100644 --- a/rfc/0012-isolation-backend/README.md +++ b/rfc/0012-isolation-backend/README.md @@ -109,6 +109,18 @@ may reconnect that process and re-confirm its binding, but a replacement supervisor process cannot resume an existing runtime generation with launch credentials. Planned upgrades stop the sandbox and start a fresh generation. +> **Implementation note ([#4860](https://github.com/NVIDIA/OpenShell/issues/4860)).** +> The built-in OpenShell runtime adds an explicit, capability-gated exception +> for highly available shared supervisors +> (`openshell.sandbox.supervisor-takeover`, protocol 1.1). A replacement +> process may resume a generation only with a gateway-signed takeover grant +> for a later credential epoch, which the gateway issues after it has advanced +> the sandbox's persisted epoch and revoked every credential of the previous +> owner. The runtime stages the replacement, promotes it at confirmation, and +> fences the previous owner; launch credentials alone still never resume a +> generation, and backends without the capability keep this contract. See +> the `openshell-sandbox` and `openshell-sandbox-backend` READMEs. + ### Provisioning Provisioning runs on the control plane, and three rules hold for every sandbox: diff --git a/skills/debug-openshell-cluster/SKILL.md b/skills/debug-openshell-cluster/SKILL.md index 40351c998e..e1d9349c9b 100644 --- a/skills/debug-openshell-cluster/SKILL.md +++ b/skills/debug-openshell-cluster/SKILL.md @@ -1143,6 +1143,119 @@ done Any `no` result can strand recovery before workload Pod creation. Upgrade the OpenShell chart rather than treating missing workload Pods as Pod failures. +#### Shared supervisor high availability + +With `supervisor_topology = "gateway"`, the gateway always coordinates the +shared supervisor as a pool, and the `-supervisor` Deployment runs +`supervisor.replicas` members (default 1, which has no failover). The chart +renders only `[openshell.gateway.supervisor_pool] recovery_timeout_secs` (the +driver's `shared_supervisor` table carries no pool settings), rolls one member +at a time, prefers separate nodes, gives each member a 60-second grace period to leave, +and adds a `-supervisor` PodDisruptionBudget with `maxUnavailable: 1`. +Each member takes its identity from `OPENSHELL_SUPERVISOR_POOL_MEMBER_ID`, +`OPENSHELL_SUPERVISOR_POOL_MEMBER_UID`, and `OPENSHELL_SUPERVISOR_POOL_ADVERTISE_IP` +(Pod name, UID, and IP). Only the Deployment's replica count depends on +`supervisor.replicas`: scaling it must not change the +gateway ConfigMap or the supervisor Pod template, so neither the gateway nor +the remaining members restart. A member removed by scaling down announces +that it is leaving (`pool member is leaving` in its log), and the gateway logs +`cause=leaving` as it moves each of its sandboxes. A scale step that restarts +the gateway or every supervisor Pod means something besides +`supervisor.replicas` changed in the release values. + +```bash +kubectl -n openshell get deployment,poddisruptionbudget,pod -l app.kubernetes.io/component=supervisor -o wide +kubectl -n openshell get configmap openshell-config -o jsonpath='{.data.gateway\.toml}' | grep -E -A2 'supervisor_topology|shared_supervisor|supervisor_pool' +kubectl -n openshell logs -l app.kubernetes.io/name=openshell,app.kubernetes.io/instance=openshell \ + -c openshell-gateway --prefix --tail=500 | grep -E 'supervisor pool|taking over a shared sandbox instance' +kubectl -n openshell logs -l app.kubernetes.io/name=openshell-supervisor --prefix --tail=500 \ + | grep -Ei 'takeover|sandbox instance' +kubectl -n logs -l openshell.ai/sandbox-id=,openshell.ai/boundary-role=workload \ + --tail=200 | grep -E 'sandbox-supervisor-|Supervisor takeover' +``` + +Each sandbox instance runs on one member, and a live owner keeps it. A +member's console log names an instance by its route (the Sandbox's +`openshell.ai/sandbox-runtime-supervisor-session` annotation) at the end of +`admitted assigned sandbox instance `, `removing withdrawn sandbox +instance `, and `relinquishing a sandbox instance ... `; the +member whose latest such line for a route is an admission supervises it. Other +lines that mention the route, such as SSH relay socket paths, say nothing +about ownership. Gateway console fields are styled, so strip ANSI escapes +(for example `sed 's/\x1b\[[0-9;]*m//g'`) before matching `cause=leaving`. The +gateway logs `shared supervisor pool enabled; reconciling pool membership and +ownership` at startup. When an owner's process exits, its node is lost, or it +leaves, the gateway logs `taking over a shared sandbox instance from a failed +supervisor` with `from`, `to`, `term`, and `cause` (`expired`: the member's +record expired; `leaving`: it announced leaving; `restarted`: its process +changed, and a member whose process restarted re-adopts its own instances +first), revokes the previous owner's +credentials for that sandbox, and grants the instance to a live member. +Runtimes dial the supervisor Service, and a member that does not own the +instance redirects the runtime to the owner's Pod IP on the runtime port, so +the workload must reach every supervisor Pod IP. The replacement and the +sandbox runtime both log `Supervisor takeover confirmed; previous owner +fenced`. The sandbox keeps its Pod, generation, PIDs, and memory; open proxied +connections and exec or SSH streams break and must reconnect. + +The sandbox runtime writes these OCSF findings to the workload Pod log: + +- `sandbox-supervisor-connection-lost`: the runtime lost its supervisor and + froze the workload. +- `sandbox-supervisor-recovery-expired`: no replacement confirmed within the + recovery timeout, so the runtime terminated the workload. +- `sandbox-supervisor-takeover-denied`: the runtime refused a supervisor + without a valid grant for the current term, such as a fenced owner returning. + +The recovery timeout is the gateway's +`[openshell.gateway.supervisor_pool] recovery_timeout_secs` (Helm +`supervisor.recoveryTimeoutSeconds`, 10 to 300 seconds, default 60). Each pool +supervisor offers it to a sandbox runtime when it attaches. A graceful leave hands instances over at +once. +A lost node takes up to about 22 seconds to detect, before the replacement +starts: the gateway closes the member's silent stream after 5 seconds, the +member's record expires 15 seconds after its last refresh, and a reconcile pass +runs every 2 seconds. A short timeout +terminates workloads that would have recovered. A gateway replica that just +started waits 30 seconds before it takes over instances whose owner looks +failed, so members cut off by a gateway restart reconnect first. With one +replica, or when no live member takes over before the timeout, the workload +ends in `sandbox-supervisor-recovery-expired`. + +`supervisor pool reconciliation failed; retrying` means a gateway replica could +not read or write the pool's store records: `supervisor_pool_member` (one per +member, refreshed every 2 seconds by the gateway replica serving its newest +stream, independently of reconciliation; older than 15 seconds means dead), +`supervisor_instance_assignment` (one per instance the pool must run, written by +the replica that launched it; it holds the handle of the launch material, which +the gateway's credential store keeps), `supervisor_instance_owner` (one per +instance, written with compare-and-swap, naming the owner and its term), +`supervisor_pool_capacity` (one per member, listing the routes whose owner +record names it; changed in the same transaction as each owner record, so a +member is never charged more than 256 instances across replicas), and +`supervisor_instance_report` (one per instance and ownership term, the owner's +latest report of that term, recorded by the replica serving it). `supervisor +pool member refresh failed` or `did not finish in time` means the replica could +not keep a member's record fresh; if it persists for 15 seconds the member's +instances move. Every gateway replica +serves the pool from these records, so a sandbox launched through one replica +reaches a member connected to another. `could not load a supervisor pool +instance's launch material` points at the credential store instead. Check the +database. To read the records on PostgreSQL without changing them: + +```sql +SELECT object_type, id, + (extract(epoch FROM now()) * 1000)::bigint - updated_at_ms AS age_ms, + convert_from(payload, 'UTF8') AS record +FROM objects +WHERE object_type IN ('supervisor_pool_member', 'supervisor_instance_assignment', + 'supervisor_instance_owner', 'supervisor_pool_capacity', + 'supervisor_instance_report') +ORDER BY object_type, id; +``` + +Do not edit or delete these records; the gateway fences takeovers with them. + #### Corporate upstream proxy When the deployment routes sandbox egress through a corporate HTTP forward @@ -1282,7 +1395,8 @@ credential failures. | Gateway or `-certgen` hook pod rejected with `violates PodSecurity "restricted:latest"` | The rendered container lacks `runAsNonRoot` or `seccompProfile`: an older chart version does not set them, or an override in `securityContext` / `pkiInitJob.securityContext` unset them | `kubectl -n openshell get events --field-selector reason=FailedCreate`; upgrade the chart, then remove overrides that unset those fields | | 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 `-supervisor` Deployment never creates pods after setting `supervisor_topology = "gateway"` (SCC `runAsUser`/`fsGroup` not in the namespace range) | `supervisor.podSecurityContext` defaults `runAsUser`, `runAsGroup`, and `fsGroup` to 65534, outside the namespace's assigned UID/GID ranges | `kubectl -n openshell get events --field-selector reason=FailedCreate`; compare with the `openshift.io/sa.scc.*` annotations from `oc describe ns openshell`; set `supervisor.podSecurityContext.runAsUser`, `.runAsGroup`, and `.fsGroup` to `null`, as `deploy/helm/openshell/ci/values-openshift-scc.yaml` does, and keep `runAsNonRoot` and `seccompProfile` | -| Shared sandboxes stay `Starting` or report `unavailable` readiness with `supervisor_topology = "gateway"`, gateway logs say no shared supervisor host is connected | The `-supervisor` Pod cannot reach or authenticate to the gateway: wrong `OPENSHELL_TLS_CA`, the `-supervisor-token` Secret is missing or still empty (gateway logs `wrote the supervisor token Secret` once it has delivered a token; `no usable supervisor token` in the supervisor means it has not), or the gateway was deployed without TLS or without the sandbox-JWT signing key | `kubectl -n openshell logs deployment/-supervisor` for `gateway accepted the supervisor host` or a refusal; gateway logs for `supervisor host attached`, `wrote the supervisor token Secret`, or `failed to write the supervisor token Secret`; `kubectl -n openshell get secret -supervisor-token -o jsonpath='{.metadata.annotations}'` shows the token expiry; the gateway refuses to start the topology without TLS or `gateway_jwt` | +| Shared sandboxes stay `Starting` or report `unavailable` readiness with `supervisor_topology = "gateway"`, gateway logs say no shared supervisor host is connected | The `-supervisor` Pod cannot reach or authenticate to the gateway: wrong `OPENSHELL_TLS_CA`, the `-supervisor-token` Secret is missing or still empty (gateway logs `wrote the supervisor token Secret` once it has delivered a token; `no usable supervisor token` in the supervisor means it has not; right after a fresh install the kubelet can take about 90 seconds to project the written token into the already-ready Pod, so launches can be refused for no live pool member until the supervisor logs `gateway accepted the supervisor host`), or the gateway was deployed without TLS or without the sandbox-JWT signing key | `kubectl -n openshell logs deployment/-supervisor` for `gateway accepted the supervisor host` or a refusal (a refusal or startup failure naming `OPENSHELL_SUPERVISOR_POOL_MEMBER_ID`, `_MEMBER_UID`, or `_ADVERTISE_IP` means the Pod lacks its member identity; the chart sets them from the Pod name, UID, and IP); gateway logs for `supervisor host attached`, `wrote the supervisor token Secret`, or `failed to write the supervisor token Secret`; `kubectl -n openshell get secret -supervisor-token -o jsonpath='{.metadata.annotations}'` shows the token expiry; the gateway refuses to start the topology without TLS or `gateway_jwt` | +| Shared sandbox workloads terminate after a supervisor Pod or node failure (`sandbox-supervisor-recovery-expired`) | One supervisor replica (no member to fail over to), no live pool member took over, or the recovery timeout is shorter than detection plus takeover | `supervisor.replicas` and the `-supervisor` PodDisruptionBudget; gateway `taking over a shared sandbox instance` logs; workload Pod `sandbox-supervisor-*` findings; `supervisor.recoveryTimeoutSeconds`. See "Shared supervisor high availability" in Step 6 | | OpenShift sandbox pod fails to start (`unable to validate against any security context constraint`) | An explicit sandbox UID/GID override or custom security context conflicts with the namespace's SCC range | Inspect pod events and the namespace's `openshift.io/sa.scc.uid-range` and `openshift.io/sa.scc.supplemental-groups` annotations; remove incompatible `sandbox_uid` / `sandbox_gid` overrides so the driver uses the namespace-assigned identity. Current sandbox pods use `restricted-v2` without a privileged SCC grant; see the [OpenShift guide](https://docs.nvidia.com/openshell/latest/kubernetes/openshift) and [RFC 0012](https://github.com/NVIDIA/OpenShell/blob/main/rfc/0012-isolation-backend/README.md#motivation). | | 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 | | Vault credential driver returns HTTP 403 / `Vault Kubernetes auth denied the configured role` on provider create | Vault's `auth/kubernetes` method or the gateway login role is not provisioned, or the role is not bound to the gateway service account and namespace | In Vault: `bao auth enable kubernetes` and `bao write auth/kubernetes/config kubernetes_host=... kubernetes_ca_cert=@...`; ensure the login role's `bound_service_account_names`/`bound_service_account_namespaces` match the gateway SA and namespace and its policy grants the credential paths | diff --git a/tasks/helm.toml b/tasks/helm.toml index feaf624e90..047cdad602 100644 --- a/tasks/helm.toml +++ b/tasks/helm.toml @@ -67,6 +67,7 @@ run = """ deploy/helm/test-split-ownership.sh deploy/helm/test-gateway-config-parser.sh deploy/helm/test-gateway-resource-coherence.sh + deploy/helm/test-supervisor-pool-scaling.sh """ run_windows = "echo Skipping helm:test: Helm validation is not part of the native Windows lane." diff --git a/tasks/test.toml b/tasks/test.toml index a5a2a03928..30e8982599 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -268,10 +268,25 @@ env = { OPENSHELL_E2E_CREDENTIAL_DRIVERS = "1", OPENSHELL_E2E_CREDENTIAL_DRIVER run = "e2e/rust/e2e-kubernetes.sh" ["e2e:kubernetes:shared-supervisor"] -description = "Run Kubernetes e2e with one supervisor shared by the gateway's sandboxes (gateway supervisor topology)" -env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor" } -# Prepend the supervisor overlay to any caller overlays, such as a workspace mode. -run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' +description = "Run Kubernetes e2e with a one-member shared supervisor pool for the gateway's sandboxes (gateway supervisor topology)" +env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor" } +# Prepend the supervisor overlay to any caller overlays, such as a workspace +# mode. The shared supervisor requires gateway TLS, which the Skaffold values +# turn off. +run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' + +["e2e:kubernetes:shared-supervisor-ha"] +description = "Run Kubernetes e2e that kills the owning member of a two-replica shared supervisor pool and checks its sandboxes are adopted, not restarted" +env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_ha", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-ha" } +# The shared supervisor requires gateway TLS, which the Skaffold values turn off. +run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor-ha.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' + +["e2e:kubernetes:shared-supervisor-scaling"] +description = "Run Kubernetes e2e that scales a one-member shared supervisor pool to two members and back to one with Helm, checking sandboxes are handed off, not restarted" +env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_scaling", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-scaling" } +# Its own release: the test upgrades it, so it never runs beside the abrupt-loss +# test. The shared supervisor requires gateway TLS, which the Skaffold values turn off. +run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' ["e2e:kubernetes:workspace-operator"] description = "Run Kubernetes e2e with operator workspace mode (pre-provisioned per-workspace namespaces)"