From b70c03b5fe994d51c60a57f6e827f202b53dab15 Mon Sep 17 00:00:00 2001 From: Abhinav Chordia Date: Mon, 2 Nov 2026 21:23:07 +0000 Subject: [PATCH] sandbox-store: move the restore-lease and pod-tracking sandbox row writes into the store The Kubernetes substrate's store wrote the control-plane `sandboxes` row from inline SQL in `restore_lease.rs` and `pod_tracking.rs`. All twelve statements move into two new `pplx-space-sandbox-store` modules, one function per transition: - `restore_lease`: the user and evacuation restore claims, the markerless `restoring` lease claim, the owned `restoring -> started` completion, the lease release, and the three reverts of an abandoned claim. - `pod_tracking`: both forms of the unplace after pod teardown (placement clock bumped or preserved), the markerless `restoring` recovery, and the `started` recovery onto a current suspend snapshot (execution-retired and failed-recovery proofs, with its artifact-marker SET fragment). Statement text is unchanged byte for byte, and each keeps its construction, including the portability and routability fragments it interpolates. Binds keep their order and Postgres types; the claims keep their unused parameters, so each statement is prepared with the same parameter list as before. The claims, the completion and the release take any Postgres executor and still run on the pool; the reverts and the pod-tracking writes take the caller's transaction, since they rely on the row lock the caller takes. The revert's lock and evidence read, the ownership snapshot, the transactions and the execution writes stay in the Kubernetes store, whose methods now delegate. Each function's doc names the fact it writes, the guard it relies on, and what its caller must have done first. The moved statements' `sql!` locators now name `pplx_space_sandbox_store::restore_lease` and `pplx_space_sandbox_store::pod_tracking`. The one test that finds a statement by its locator, the claim-herd test that counts lock-blocked user restore claims in `pg_stat_activity`, now matches the new module path. The `sandbox_writes` ratchet drops both entries (`pod_tracking.rs` 5, `restore_lease.rs` 8). --- architecture-baseline.toml | 2 - .../controllers/substrate/kubernetes/store.rs | 15 +- .../kubernetes/store/pod_tracking.rs | 171 +-------- .../kubernetes/store/restore_lease.rs | 231 ++--------- .../controllers/substrate/kubernetes/tests.rs | 2 +- crates/libs/sandbox/store/src/lib.rs | 2 + crates/libs/sandbox/store/src/pod_tracking.rs | 250 ++++++++++++ .../libs/sandbox/store/src/restore_lease.rs | 359 ++++++++++++++++++ 8 files changed, 671 insertions(+), 361 deletions(-) create mode 100644 crates/libs/sandbox/store/src/pod_tracking.rs create mode 100644 crates/libs/sandbox/store/src/restore_lease.rs diff --git a/architecture-baseline.toml b/architecture-baseline.toml index 45ad6f17c9..b28dd9c2af 100644 --- a/architecture-baseline.toml +++ b/architecture-baseline.toml @@ -150,8 +150,6 @@ "crates/libs/api/src/controllers/sandbox/responses.rs" = 3 "crates/libs/api/src/controllers/sandbox/updates.rs" = 2 "crates/libs/api/src/controllers/substrate/direct.rs" = 6 -"crates/libs/api/src/controllers/substrate/kubernetes/store/pod_tracking.rs" = 5 -"crates/libs/api/src/controllers/substrate/kubernetes/store/restore_lease.rs" = 8 "crates/libs/api/src/controllers/substrate/mod.rs" = 1 "crates/libs/cluster/controller/server/src/drained_node_terminator/store.rs" = 1 "crates/libs/cluster/controller/server/src/node_drain/store.rs" = 15 diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/store.rs b/crates/libs/api/src/controllers/substrate/kubernetes/store.rs index 14e3e4c122..52beea88ff 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/store.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/store.rs @@ -26,10 +26,9 @@ use pplx_space_events::{ use pplx_space_metrics::sandbox::api::RouteIdentityDecisionReason; use pplx_space_net_common::NetworkPolicyV1; use pplx_space_sandbox_portability::{ - clear_current_local_artifact_marker_assignments, failed_repair_revert_eligibility, fresh_ready_snapshot_exists, missing_suspending_pod_revert_eligibility, no_newer_full_snapshot, restore_retry_source_snapshot_is_current, restore_retry_suspended_revert_snapshot_exists, - snapshot_is_current_local_artifact, started_suspend_revert_eligibility, + snapshot_is_current_local_artifact, }; use sqlx::{PgConnection, PgPool, Postgres, Row, Transaction, postgres::PgRow}; use uuid::Uuid; @@ -175,12 +174,6 @@ pub(super) enum PlacementClock { Preserve, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum StartedSuspendRecoveryProof { - ExecutionRetired, - FailedRecovery { generation: ExecutionGeneration }, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum StartedSuspendRecoveryOperation { PodTrackingCleanup, @@ -697,12 +690,6 @@ impl RestoreClaimRevert { } } -#[derive(sqlx::FromRow)] -struct RestoreClaimRow { - lifecycle: String, - restore_lease_token: Option, -} - fn recovery_lease_ttl_secs() -> i64 { i64::try_from(crate::controllers::substrate::NODE_RPC_TIMEOUT_SECS) .expect("node RPC timeout seconds fit i64") diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/store/pod_tracking.rs b/crates/libs/api/src/controllers/substrate/kubernetes/store/pod_tracking.rs index 807105495c..b836467306 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/store/pod_tracking.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/store/pod_tracking.rs @@ -2,6 +2,7 @@ //! started-suspend recovery that runs after execution cleanup. use super::*; +use pplx_space_sandbox_store::pod_tracking; impl K8sStore { pub(in crate::controllers::substrate::kubernetes) async fn forget_pod_tracking( @@ -533,29 +534,14 @@ impl K8sStore { // retire keeps `placement_changed_at` so the legacy snapshot-freshness // gate doesn't strand a markerless row's ready snapshot. See // [`PlacementClock`]. - let query = match placement_clock { - PlacementClock::Bump => sqlx::query(sql!( - "UPDATE sandboxes - SET placement = 'unplaced', - active_node_id = NULL, - target_node_id = NULL, - placement_changed_at = NOW() - WHERE id = $1" - )), - PlacementClock::Preserve => sqlx::query(sql!( - "UPDATE sandboxes - SET placement = 'unplaced', - active_node_id = NULL, - target_node_id = NULL - WHERE id = $1" - )), - }; - query - .bind(sandbox_id) - .execute(&mut **tx) - .await - .db_context("failed to mark sandbox unplaced")?; - Ok(()) + match placement_clock { + PlacementClock::Bump => pod_tracking::unplace_after_pod_teardown(tx, sandbox_id).await, + PlacementClock::Preserve => { + pod_tracking::unplace_after_pod_teardown_preserving_placement_clock(tx, sandbox_id) + .await + } + } + .db_context("failed to mark sandbox unplaced") } pub(super) async fn recover_markerless_restoring_after_execution_cleanup_in( @@ -566,37 +552,9 @@ impl K8sStore { // Callers hold the sandbox row lock and retire the authoritative // execution earlier in this transaction. Execution creation takes the // same lock, so a successor cannot appear between this proof and commit. - let safe_revert_snapshot = restore_retry_suspended_revert_snapshot_exists("s"); - let query = format!( - "UPDATE sandboxes s - SET lifecycle = 'suspended', - placement = 'unplaced', - active_node_id = NULL, - target_node_id = NULL, - placement_changed_at = NULL, - current_local_artifact_id = NULL, - current_local_artifact_kind = NULL, - current_local_artifact_origin = NULL, - command_ready_pending = FALSE, - restore_lease_expires_at = NULL, - restore_lease_token = NULL - WHERE s.id = $1 - AND s.lifecycle = 'restoring' - AND s.evacuation_restore_target IS NULL - AND NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = s.id - AND e.active - ) - AND {safe_revert_snapshot}" - ); - sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) - .bind(sandbox_id) - .execute(&mut **tx) + pod_tracking::recover_markerless_restoring_after_execution_cleanup(tx, sandbox_id) .await - .db_context("failed to recover markerless restoring sandbox after execution cleanup")?; - Ok(()) + .db_context("failed to recover markerless restoring sandbox after execution cleanup") } pub(super) async fn recover_started_after_suspend_cleanup_in( @@ -604,12 +562,9 @@ impl K8sStore { tx: &mut Transaction<'_, Postgres>, sandbox_id: SandboxId, ) -> Result { - self.recover_started_with_current_suspend_snapshot_in( - tx, - sandbox_id, - StartedSuspendRecoveryProof::ExecutionRetired, - ) - .await + pod_tracking::recover_started_after_execution_retired(tx, sandbox_id) + .await + .db_context("failed to recover started sandbox with current suspend snapshot") } pub(super) async fn recover_started_after_failed_recovery_in( @@ -618,102 +573,8 @@ impl K8sStore { sandbox_id: SandboxId, generation: ExecutionGeneration, ) -> Result { - self.recover_started_with_current_suspend_snapshot_in( - tx, - sandbox_id, - StartedSuspendRecoveryProof::FailedRecovery { generation }, - ) - .await - } - - async fn recover_started_with_current_suspend_snapshot_in( - &self, - tx: &mut Transaction<'_, Postgres>, - sandbox_id: SandboxId, - proof: StartedSuspendRecoveryProof, - ) -> Result { - let (generation, execution_proof) = match proof { - StartedSuspendRecoveryProof::ExecutionRetired => ( - None, - "NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = s.id - AND e.active - ) - AND NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = s.id - AND e.pod_ip IS NOT NULL - )", - ), - StartedSuspendRecoveryProof::FailedRecovery { generation } => ( - Some(generation.get()), - "EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = s.id - AND e.generation = args.generation - AND e.active - AND e.state = 'recovery_failed' - AND e.pod_uid IS NULL - AND e.pod_ip IS NULL - AND e.k8s_node_name IS NULL - ) - AND NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions newer - WHERE newer.sandbox_id = s.id - AND newer.active - AND newer.generation > args.generation - )", - ), - }; - let eligibility = match proof { - StartedSuspendRecoveryProof::ExecutionRetired => { - started_suspend_revert_eligibility("s") - } - StartedSuspendRecoveryProof::FailedRecovery { .. } => { - failed_repair_revert_eligibility("s") - } - }; - let clear_marker = eligibility.clear_marker; - let eligible = eligibility.eligible; - let artifact_marker_assignments = - clear_current_local_artifact_marker_assignments("s", "pre.clear_marker"); - let query = format!( - "WITH args AS ( - SELECT $2::BIGINT AS generation - ), - pre AS ( - SELECT id, - {clear_marker} AS clear_marker - FROM sandboxes s - WHERE s.id = $1 - ) - UPDATE sandboxes s - SET lifecycle = 'suspended', - placement = 'unplaced', - active_node_id = NULL, - target_node_id = NULL, - placement_changed_at = NULL, - command_ready_pending = FALSE, - {artifact_marker_assignments} - FROM pre, args - WHERE s.id = pre.id - AND s.lifecycle = 'started' - AND s.placement = 'unplaced' - AND s.active_node_id IS NULL - AND {execution_proof} - AND {eligible}" - ); - let result = sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) - .bind(sandbox_id) - .bind(generation) - .execute(&mut **tx) + pod_tracking::recover_started_after_failed_recovery(tx, sandbox_id, generation) .await - .db_context("failed to recover started sandbox with current suspend snapshot")?; - Ok(result.rows_affected() > 0) + .db_context("failed to recover started sandbox with current suspend snapshot") } } diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/store/restore_lease.rs b/crates/libs/api/src/controllers/substrate/kubernetes/store/restore_lease.rs index 176d74e897..a48e0bb17f 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/store/restore_lease.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/store/restore_lease.rs @@ -2,6 +2,7 @@ //! rollback reverts, and preparing a sandbox for wake. use super::*; +use pplx_space_sandbox_store::restore_lease; impl K8sStore { /// Post-CAS-loss snapshot of the restore gate. Fresh execution work keeps @@ -354,24 +355,16 @@ impl K8sStore { // `target_node_id` and the artifact marker name the wake // source and stay; the lifecycle triggers keep the marker on // `suspended`. `placement_changed_at` is cleared, as the - // markerless restoring recovery in `pod_tracking.rs` does for - // its rows: a + // markerless restoring recovery + // (`pplx_space_sandbox_store::pod_tracking`) does for its + // rows: a // restore whose pod bound moved it past the proven // pre-restore source, and both the suspended selector and // `durably_restorable` take a markerless snapshot only when // it is newer than the last placement. - sqlx::query(sql!( - "UPDATE sandboxes - SET lifecycle = 'suspended', - placement = 'unplaced', - active_node_id = NULL, - placement_changed_at = NULL - WHERE id = $1" - )) - .bind(sandbox_id) - .execute(&mut *tx) - .await - .db_context("failed to revert sandbox lifecycle to suspended")?; + restore_lease::revert_restore_claim_retiring_placement(&mut tx, sandbox_id) + .await + .db_context("failed to revert sandbox lifecycle to suspended")?; RestoreClaimRevert::ExecutionRetired { pod_name } } (None, true) if !any_active => { @@ -379,26 +372,15 @@ impl K8sStore { // clear as above. An execution the claim reused rather than // created keeps the clock, since nothing proves its pod // holds no state of its own. - sqlx::query(sql!( - "UPDATE sandboxes - SET lifecycle = 'suspended', - placement_changed_at = NULL - WHERE id = $1" - )) - .bind(sandbox_id) - .execute(&mut *tx) - .await - .db_context("failed to revert sandbox lifecycle to suspended")?; + restore_lease::revert_restore_claim_clearing_placement_clock(&mut tx, sandbox_id) + .await + .db_context("failed to revert sandbox lifecycle to suspended")?; RestoreClaimRevert::LifecycleOnly } (claimed, _) => { - sqlx::query(sql!( - "UPDATE sandboxes SET lifecycle = 'suspended' WHERE id = $1" - )) - .bind(sandbox_id) - .execute(&mut *tx) - .await - .db_context("failed to revert sandbox lifecycle to suspended")?; + restore_lease::revert_restore_claim_lifecycle(&mut tx, sandbox_id) + .await + .db_context("failed to revert sandbox lifecycle to suspended")?; if claimed.is_some() || any_active { RestoreClaimRevert::ExecutionLeftActive } else { @@ -471,97 +453,26 @@ impl K8sStore { // `stopped` sandbox is user-owned, and one still on its node would get // two homes. It claims to `restoring` like the suspend path; see // `RestoreClaim::FreshFromStopped`. - let restore_lease_token = (!require_evacuation_marker).then(Uuid::now_v7); let restore_claim = if require_evacuation_marker { - sqlx::query_as::<_, RestoreClaimRow>(sql!( - "WITH candidate AS ( - SELECT id, lifecycle - FROM sandboxes - WHERE id = $1 - AND ( - LOWER(TRIM(lifecycle)) IN ('suspended', 'restoring') - OR ( - LOWER(TRIM(lifecycle)) = 'stopped' - AND evacuation_restore_target = 'stopped' - AND active_node_id IS NULL - AND $5::text = 'stopped' - ) - ) - AND evacuation_restore_target IS NOT NULL - FOR UPDATE - ), - updated AS ( - UPDATE sandboxes s - SET lifecycle = 'restoring' - FROM candidate c - WHERE s.id = c.id - RETURNING c.lifecycle - ) - SELECT u.lifecycle, args.restore_lease_token - FROM updated u - CROSS JOIN ( - SELECT $2::bigint AS ttl_secs, $3::uuid AS restore_lease_token - ) args - WHERE args.ttl_secs >= 0" - )) + restore_lease::claim_evacuation_restore( + &self.pool, + sandbox_id, + recovery_lease_ttl_secs(), + super::super::RESTORE_EXECUTION_OWNERSHIP_TTL_SECS, + evacuation_target, + ) + .await } else { - // Preserve `lifecycle_changed_at` across `restoring -> restoring` - // re-drives: restore selection uses it as the immutable boundary - // between the original source window and artifacts published by a - // failed restore branch. The lifecycle trigger still sets it for a - // fresh `suspended -> restoring` claim. The restore lease protects - // an active re-drive from the transient-lifecycle sweeper. - sqlx::query_as::<_, RestoreClaimRow>(sql!( - "WITH candidate AS ( - SELECT id, lifecycle - FROM sandboxes - WHERE id = $1 - AND ( - LOWER(TRIM(lifecycle)) = 'suspended' - OR ( - LOWER(TRIM(lifecycle)) = 'restoring' - AND evacuation_restore_target IS NULL - AND ( - restore_lease_expires_at IS NULL - OR restore_lease_expires_at <= NOW() - ) - ) - ) - AND ( - restore_lease_expires_at IS NULL - OR restore_lease_expires_at <= NOW() - ) - AND NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = sandboxes.id - AND e.active - AND e.state IN ('creating', 'created', 'recovering', 'running') - AND e.updated_at > - NOW() - make_interval(secs => $4) - ) - FOR UPDATE - ), - updated AS ( - UPDATE sandboxes s - SET lifecycle = 'restoring', - evacuation_restore_target = NULL, - restore_lease_expires_at = NOW() + ($2::bigint * INTERVAL '1 second'), - restore_lease_token = $3 - FROM candidate c - WHERE s.id = c.id - RETURNING c.lifecycle - ) - SELECT lifecycle, $3::uuid AS restore_lease_token FROM updated" - )) + restore_lease::claim_user_restore( + &self.pool, + sandbox_id, + recovery_lease_ttl_secs(), + Uuid::now_v7(), + super::super::RESTORE_EXECUTION_OWNERSHIP_TTL_SECS, + evacuation_target, + ) + .await } - .bind(sandbox_id.as_uuid()) - .bind(recovery_lease_ttl_secs()) - .bind(restore_lease_token) - .bind(super::super::RESTORE_EXECUTION_OWNERSHIP_TTL_SECS) - .bind(evacuation_target.map(SandboxLifecycle::as_str)) - .fetch_optional(&self.pool) - .await .db_context("failed to begin sandbox restore")?; let Some(restore_claim) = restore_claim else { @@ -643,43 +554,13 @@ impl K8sStore { &self, sandbox_id: SandboxId, ) -> Result, SubstrateError> { - let restore_lease_token = Uuid::now_v7(); - let restore_claim = sqlx::query_as::<_, RestoreClaimRow>(sql!( - "WITH candidate AS ( - SELECT id, lifecycle - FROM sandboxes - WHERE id = $1 - AND LOWER(TRIM(lifecycle)) = 'restoring' - AND evacuation_restore_target IS NULL - AND ( - restore_lease_expires_at IS NULL - OR restore_lease_expires_at <= NOW() - ) - AND NOT EXISTS ( - SELECT 1 - FROM sandbox_k8s_executions e - WHERE e.sandbox_id = sandboxes.id - AND e.active - AND e.state IN ('creating', 'created', 'recovering', 'running') - AND e.updated_at > NOW() - make_interval(secs => $4) - ) - FOR UPDATE - ), - updated AS ( - UPDATE sandboxes s - SET restore_lease_expires_at = NOW() + ($2::bigint * INTERVAL '1 second'), - restore_lease_token = $3 - FROM candidate c - WHERE s.id = c.id - RETURNING c.lifecycle - ) - SELECT lifecycle, $3::uuid AS restore_lease_token FROM updated" - )) - .bind(sandbox_id.as_uuid()) - .bind(recovery_lease_ttl_secs()) - .bind(restore_lease_token) - .bind(super::super::RESTORE_EXECUTION_OWNERSHIP_TTL_SECS) - .fetch_optional(&self.pool) + let restore_claim = restore_lease::claim_markerless_restoring_lease( + &self.pool, + sandbox_id, + recovery_lease_ttl_secs(), + Uuid::now_v7(), + super::super::RESTORE_EXECUTION_OWNERSHIP_TTL_SECS, + ) .await .db_context("failed to claim restoring sandbox restore lease")?; @@ -743,28 +624,9 @@ impl K8sStore { sandbox_id: SandboxId, restore_lease_token: Uuid, ) -> Result { - let routable = - pplx_space_common_svc::sql::sandbox_routability::routable_execution_or_non_k8s( - "sandboxes", - "$3", - ); - let query = format!( - "UPDATE sandboxes - SET lifecycle = 'started', placement = 'on_node' - WHERE id = $1 - AND lifecycle = 'restoring' - AND restore_lease_token = $2 - AND {routable}" - ); - let completed = sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) - .bind(sandbox_id.as_uuid()) - .bind(restore_lease_token) - .bind(pplx_space_common_core::types::SubstrateKind::Kubernetes) - .execute(&self.pool) + restore_lease::complete_owned_user_restore(&self.pool, sandbox_id, restore_lease_token) .await - .db_context("failed to complete owned sandbox restore")? - .rows_affected(); - Ok(completed == 1) + .db_context("failed to complete owned sandbox restore") } pub(in crate::controllers::substrate::kubernetes) async fn release_user_restore_lease( @@ -776,18 +638,9 @@ impl K8sStore { // but retain the paired token and attempt clock at NOW(). The transient // sweeper requires a full quiet window after that clock before changing // `restoring`, while the next retry atomically replaces the stale pair. - sqlx::query(sql!( - "UPDATE sandboxes - SET restore_lease_expires_at = NOW() - WHERE id = $1 - AND restore_lease_token = $2" - )) - .bind(sandbox_id.as_uuid()) - .bind(restore_lease_token) - .execute(&self.pool) - .await - .db_context("failed to release sandbox restore lease")?; - Ok(()) + restore_lease::release_user_restore_lease(&self.pool, sandbox_id, restore_lease_token) + .await + .db_context("failed to release sandbox restore lease") } pub(in crate::controllers::substrate::kubernetes) async fn prepare_wake_sandbox( diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/tests.rs b/crates/libs/api/src/controllers/substrate/kubernetes/tests.rs index fdb88f2e57..a798ebfa35 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/tests.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/tests.rs @@ -6131,7 +6131,7 @@ mod kubernetes_tests { FROM pg_stat_activity WHERE datname = current_database() AND wait_event_type = 'Lock' - AND query LIKE '%kubernetes::store%' + AND query LIKE '%pplx_space_sandbox_store::restore_lease:%' AND query LIKE '%WITH candidate AS (%'", ) .fetch_one(&pool) diff --git a/crates/libs/sandbox/store/src/lib.rs b/crates/libs/sandbox/store/src/lib.rs index d2ace8905a..b2ad7ea33e 100644 --- a/crates/libs/sandbox/store/src/lib.rs +++ b/crates/libs/sandbox/store/src/lib.rs @@ -105,7 +105,9 @@ pub mod delete; pub mod fork_source_lifecycle; pub mod ingest; pub mod lost; +pub mod pod_tracking; pub mod repair; +pub mod restore_lease; pub mod suspend; pub mod sweeper; pub mod wake; diff --git a/crates/libs/sandbox/store/src/pod_tracking.rs b/crates/libs/sandbox/store/src/pod_tracking.rs new file mode 100644 index 0000000000..b55ab710b7 --- /dev/null +++ b/crates/libs/sandbox/store/src/pod_tracking.rs @@ -0,0 +1,250 @@ +//! Guarded writes the Kubernetes substrate makes to `sandboxes` when it +//! forgets a sandbox's tracked pod: unplacing the row, and recovering a +//! `restoring` or `started` row to `suspended` once its execution's cleanup +//! leaves nothing running. +//! +//! The Kubernetes store keeps everything around these statements: the +//! sandbox row lock (`SELECT ... FOR UPDATE`), the transaction, the +//! execution-row writes, the warm-pool slot delete, and the event a recovery +//! emits. Every function here takes the caller's transaction because it +//! relies on the row lock the caller took in it. What else the caller must +//! have done first differs per function; see each function's doc. + +use pplx_space_common_core::types::{ExecutionGeneration, SandboxId}; + +use pplx_space_common_core::sql; +use pplx_space_sandbox_portability::{ + clear_current_local_artifact_marker_assignments, failed_repair_revert_eligibility, + restore_retry_suspended_revert_snapshot_exists, started_suspend_revert_eligibility, +}; +use sqlx::{Postgres, Transaction}; + +/// Unplace a sandbox whose pod tracking was torn down, clearing its active +/// and target nodes and bumping `placement_changed_at` to `NOW()`. Writes an +/// **operation step** (ends the row's placement after pod teardown). +/// +/// For delete-shaped teardowns, where nothing is left to keep recoverable. +/// Matches on the id alone: no lifecycle predicate of its own. It relies on +/// the caller's row lock, and on the caller having already retired the +/// execution it tore down, or confirmed that none is active, earlier in this +/// transaction. +pub async fn unplace_after_pod_teardown( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes + SET placement = 'unplaced', + active_node_id = NULL, + target_node_id = NULL, + placement_changed_at = NOW() + WHERE id = $1" + )) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// [`unplace_after_pod_teardown`] for a recoverability-preserving retire: +/// `placement_changed_at` is kept, so a markerless row's ready snapshot stays +/// fresh against it. Writes an **operation step**, as +/// [`unplace_after_pod_teardown`] does. +/// +/// Matches on the id alone, as [`unplace_after_pod_teardown`] does. +pub async fn unplace_after_pod_teardown_preserving_placement_clock( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes + SET placement = 'unplaced', + active_node_id = NULL, + target_node_id = NULL + WHERE id = $1" + )) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// Revert a markerless `restoring` row whose execution cleanup left no +/// active execution to `suspended`, unplaced, with its placement clock, +/// artifact marker, `command_ready_pending` and restore lease cleared. +/// Writes an **inference** (the control plane concluding the restore is +/// abandoned). +/// +/// Guarded on `restoring`, no evacuation marker, no active execution, and +/// [`restore_retry_suspended_revert_snapshot_exists`]. The caller holds the +/// sandbox row lock, which execution creation also takes, so no successor +/// execution can appear between this proof and commit. Callers retire the +/// execution, or confirm that none is active, earlier in this transaction. +pub async fn recover_markerless_restoring_after_execution_cleanup( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + let safe_revert_snapshot = restore_retry_suspended_revert_snapshot_exists("s"); + let query = format!( + "UPDATE sandboxes s + SET lifecycle = 'suspended', + placement = 'unplaced', + active_node_id = NULL, + target_node_id = NULL, + placement_changed_at = NULL, + current_local_artifact_id = NULL, + current_local_artifact_kind = NULL, + current_local_artifact_origin = NULL, + command_ready_pending = FALSE, + restore_lease_expires_at = NULL, + restore_lease_token = NULL + WHERE s.id = $1 + AND s.lifecycle = 'restoring' + AND s.evacuation_restore_target IS NULL + AND NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = s.id + AND e.active + ) + AND {safe_revert_snapshot}" + ); + sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// Revert an unplaced `started` row to `suspended` when no execution is +/// active and no execution row still has a `pod_ip`, and its current suspend +/// snapshot is still eligible ([`started_suspend_revert_eligibility`]). +/// Execution retirement clears `pod_ip`, so this does not prove that no pod +/// ever had an IP. Writes an **inference**. Returns whether the row moved. +/// +/// The caller holds the sandbox row lock and has already retired the +/// execution it tore down in this transaction. +pub async fn recover_started_after_execution_retired( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result { + recover_started_with_current_suspend_snapshot( + tx, + sandbox_id, + StartedSuspendRecoveryProof::ExecutionRetired, + ) + .await +} + +/// Revert an unplaced `started` row to `suspended` when its active execution +/// at `generation` is a `recovery_failed` one that tracks no pod and no newer +/// execution is active, and its snapshot passes +/// [`failed_repair_revert_eligibility`]. Writes an **inference**. Returns +/// whether the row moved. +/// +/// The caller holds the sandbox row lock. The `recovery_failed` execution +/// must still be active when this runs; the caller retires it afterwards, in +/// the same transaction, only when this returns `true`. +pub async fn recover_started_after_failed_recovery( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, + generation: ExecutionGeneration, +) -> Result { + recover_started_with_current_suspend_snapshot( + tx, + sandbox_id, + StartedSuspendRecoveryProof::FailedRecovery { generation }, + ) + .await +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StartedSuspendRecoveryProof { + ExecutionRetired, + FailedRecovery { generation: ExecutionGeneration }, +} + +async fn recover_started_with_current_suspend_snapshot( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, + proof: StartedSuspendRecoveryProof, +) -> Result { + let (generation, execution_proof) = match proof { + StartedSuspendRecoveryProof::ExecutionRetired => ( + None, + "NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = s.id + AND e.active + ) + AND NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = s.id + AND e.pod_ip IS NOT NULL + )", + ), + StartedSuspendRecoveryProof::FailedRecovery { generation } => ( + Some(generation.get()), + "EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = s.id + AND e.generation = args.generation + AND e.active + AND e.state = 'recovery_failed' + AND e.pod_uid IS NULL + AND e.pod_ip IS NULL + AND e.k8s_node_name IS NULL + ) + AND NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions newer + WHERE newer.sandbox_id = s.id + AND newer.active + AND newer.generation > args.generation + )", + ), + }; + let eligibility = match proof { + StartedSuspendRecoveryProof::ExecutionRetired => started_suspend_revert_eligibility("s"), + StartedSuspendRecoveryProof::FailedRecovery { .. } => failed_repair_revert_eligibility("s"), + }; + let clear_marker = eligibility.clear_marker; + let eligible = eligibility.eligible; + let artifact_marker_assignments = + clear_current_local_artifact_marker_assignments("s", "pre.clear_marker"); + let query = format!( + "WITH args AS ( + SELECT $2::BIGINT AS generation + ), + pre AS ( + SELECT id, + {clear_marker} AS clear_marker + FROM sandboxes s + WHERE s.id = $1 + ) + UPDATE sandboxes s + SET lifecycle = 'suspended', + placement = 'unplaced', + active_node_id = NULL, + target_node_id = NULL, + placement_changed_at = NULL, + command_ready_pending = FALSE, + {artifact_marker_assignments} + FROM pre, args + WHERE s.id = pre.id + AND s.lifecycle = 'started' + AND s.placement = 'unplaced' + AND s.active_node_id IS NULL + AND {execution_proof} + AND {eligible}" + ); + let result = sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) + .bind(sandbox_id) + .bind(generation) + .execute(&mut **tx) + .await?; + Ok(result.rows_affected() > 0) +} diff --git a/crates/libs/sandbox/store/src/restore_lease.rs b/crates/libs/sandbox/store/src/restore_lease.rs new file mode 100644 index 0000000000..64c4231475 --- /dev/null +++ b/crates/libs/sandbox/store/src/restore_lease.rs @@ -0,0 +1,359 @@ +//! Guarded writes of the Kubernetes substrate's restore lease on `sandboxes`: +//! the user and evacuation restore claims, the markerless re-drive's lease +//! claim, the owned completion and the lease release, and the reverts of an +//! abandoned claim. +//! +//! The restore lease is the pair `restore_lease_expires_at` / +//! `restore_lease_token`. A user claim writes both; the evacuation claim is +//! tokenless and is fenced by the `evacuation_restore_target` marker +//! instead. Leaving `restoring` clears the pair through the +//! `sandboxes_clear_restore_lease` trigger. +//! +//! The Kubernetes store keeps everything around these statements: the +//! revert's row lock (`SELECT ... FOR UPDATE`) and its evidence read, the +//! transactions, the execution retirement beside the revert, and the +//! ownership snapshot it reads when a claim matches nothing. A claim or +//! completion that matches no row is an outcome, not an error. + +use pplx_space_common_core::types::{SandboxId, SandboxLifecycle, SubstrateKind}; + +use pplx_space_common_core::sql; +use pplx_space_sandbox_routability::routable_execution_or_non_k8s; +use sqlx::{PgExecutor, Postgres, Transaction}; +use uuid::Uuid; + +/// What a restore claim returns: the lifecycle the row had before the claim, +/// and the lease token the claim wrote (`None` for the tokenless evacuation +/// claim). +#[derive(Debug, Clone, sqlx::FromRow)] +pub struct RestoreClaimRow { + pub lifecycle: String, + pub restore_lease_token: Option, +} + +/// The evacuation restore claim: flip a row that carries the node-drain +/// `evacuation_restore_target` marker to `restoring`, atomically on that +/// marker. Writes an **operation step** (the claim). Returns the row's +/// lifecycle before the claim, or `None` when the claim lost. +/// +/// Takes a `suspended` or `restoring` row with the marker still set (a +/// `restoring` one is crash recovery after the controller died between this +/// claim and finalize), or a released stopped evacuation: lifecycle +/// `stopped`, marker `'stopped'`, no `active_node_id`, and only when +/// `evacuation_target` is [`SandboxLifecycle::Stopped`]. It writes no lease: +/// the marker fences it, and the claim's `restore_lease_token` is always +/// `NULL`. `lease_ttl_secs` is bound as `$2` but not written: it only +/// filters the returned row. The `UPDATE` in the statement's `WITH` runs +/// whatever its value, so a negative value still moves the row to +/// `restoring` and returns `None`. Pass a non-negative value. +/// +/// `execution_ownership_ttl_secs` is bound as `$4`, which this statement +/// doesn't read; it keeps the parameter list the statement has always been +/// prepared with. +pub async fn claim_evacuation_restore( + executor: impl PgExecutor<'_>, + sandbox_id: SandboxId, + lease_ttl_secs: i64, + execution_ownership_ttl_secs: i64, + evacuation_target: Option, +) -> Result, sqlx::Error> { + sqlx::query_as::<_, RestoreClaimRow>(sql!( + "WITH candidate AS ( + SELECT id, lifecycle + FROM sandboxes + WHERE id = $1 + AND ( + LOWER(TRIM(lifecycle)) IN ('suspended', 'restoring') + OR ( + LOWER(TRIM(lifecycle)) = 'stopped' + AND evacuation_restore_target = 'stopped' + AND active_node_id IS NULL + AND $5::text = 'stopped' + ) + ) + AND evacuation_restore_target IS NOT NULL + FOR UPDATE + ), + updated AS ( + UPDATE sandboxes s + SET lifecycle = 'restoring' + FROM candidate c + WHERE s.id = c.id + RETURNING c.lifecycle + ) + SELECT u.lifecycle, args.restore_lease_token + FROM updated u + CROSS JOIN ( + SELECT $2::bigint AS ttl_secs, $3::uuid AS restore_lease_token + ) args + WHERE args.ttl_secs >= 0" + )) + .bind(sandbox_id) + .bind(lease_ttl_secs) + .bind(None::) + .bind(execution_ownership_ttl_secs) + .bind(evacuation_target.map(SandboxLifecycle::as_str)) + .fetch_optional(executor) + .await +} + +/// The user restore claim: flip a `suspended` row, or a markerless +/// `restoring` row, to `restoring` under a new restore lease. Writes an +/// **operation step** (the claim). Returns the row's lifecycle before the +/// claim and the lease token it wrote, or `None` when the claim lost. +/// +/// In both cases the row must hold no live restore lease, and no active +/// execution in `creating`/`created`/`recovering`/`running` may have been +/// updated within `execution_ownership_ttl_secs`. A `restoring` row that +/// still carries the `evacuation_restore_target` marker is a drain restore in +/// flight and matches nothing. The claim clears that marker, so on a marked +/// `suspended` row it is a user takeover: the drain's marker-required claim +/// then loses. +/// +/// The claim preserves `lifecycle_changed_at` across +/// `restoring -> restoring` re-drives: restore selection uses it as the +/// immutable boundary between the original source window and artifacts +/// published by a failed restore branch. The lifecycle trigger still sets it +/// for a fresh `suspended -> restoring` claim. The restore lease protects an +/// active re-drive from the transient-lifecycle sweeper, with one exception: +/// when the sweeper's restoring-episode deadline is configured, an overdue +/// episode (`sweeper::restoring_episode_overdue`) passes the lease and owner +/// guards, so the sweeper may revert the row under a live lease. +/// +/// `evacuation_target` is bound as `$5`, which this statement doesn't read; +/// it keeps the parameter list the statement has always been prepared with. +pub async fn claim_user_restore( + executor: impl PgExecutor<'_>, + sandbox_id: SandboxId, + lease_ttl_secs: i64, + restore_lease_token: Uuid, + execution_ownership_ttl_secs: i64, + evacuation_target: Option, +) -> Result, sqlx::Error> { + sqlx::query_as::<_, RestoreClaimRow>(sql!( + "WITH candidate AS ( + SELECT id, lifecycle + FROM sandboxes + WHERE id = $1 + AND ( + LOWER(TRIM(lifecycle)) = 'suspended' + OR ( + LOWER(TRIM(lifecycle)) = 'restoring' + AND evacuation_restore_target IS NULL + AND ( + restore_lease_expires_at IS NULL + OR restore_lease_expires_at <= NOW() + ) + ) + ) + AND ( + restore_lease_expires_at IS NULL + OR restore_lease_expires_at <= NOW() + ) + AND NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = sandboxes.id + AND e.active + AND e.state IN ('creating', 'created', 'recovering', 'running') + AND e.updated_at > + NOW() - make_interval(secs => $4) + ) + FOR UPDATE + ), + updated AS ( + UPDATE sandboxes s + SET lifecycle = 'restoring', + evacuation_restore_target = NULL, + restore_lease_expires_at = NOW() + ($2::bigint * INTERVAL '1 second'), + restore_lease_token = $3 + FROM candidate c + WHERE s.id = c.id + RETURNING c.lifecycle + ) + SELECT lifecycle, $3::uuid AS restore_lease_token FROM updated" + )) + .bind(sandbox_id) + .bind(lease_ttl_secs) + .bind(Some(restore_lease_token)) + .bind(execution_ownership_ttl_secs) + .bind(evacuation_target.map(SandboxLifecycle::as_str)) + .fetch_optional(executor) + .await +} + +/// Take the restore lease of a markerless `restoring` row without moving its +/// lifecycle, so a retry may re-drive it. Writes an **operation step** (the +/// lease claim). Returns the claimed row, or `None` when the row isn't a +/// markerless `restoring` row, its lease is still live, or an active +/// execution did work within `execution_ownership_ttl_secs`. +pub async fn claim_markerless_restoring_lease( + executor: impl PgExecutor<'_>, + sandbox_id: SandboxId, + lease_ttl_secs: i64, + restore_lease_token: Uuid, + execution_ownership_ttl_secs: i64, +) -> Result, sqlx::Error> { + sqlx::query_as::<_, RestoreClaimRow>(sql!( + "WITH candidate AS ( + SELECT id, lifecycle + FROM sandboxes + WHERE id = $1 + AND LOWER(TRIM(lifecycle)) = 'restoring' + AND evacuation_restore_target IS NULL + AND ( + restore_lease_expires_at IS NULL + OR restore_lease_expires_at <= NOW() + ) + AND NOT EXISTS ( + SELECT 1 + FROM sandbox_k8s_executions e + WHERE e.sandbox_id = sandboxes.id + AND e.active + AND e.state IN ('creating', 'created', 'recovering', 'running') + AND e.updated_at > NOW() - make_interval(secs => $4) + ) + FOR UPDATE + ), + updated AS ( + UPDATE sandboxes s + SET restore_lease_expires_at = NOW() + ($2::bigint * INTERVAL '1 second'), + restore_lease_token = $3 + FROM candidate c + WHERE s.id = c.id + RETURNING c.lifecycle + ) + SELECT lifecycle, $3::uuid AS restore_lease_token FROM updated" + )) + .bind(sandbox_id) + .bind(lease_ttl_secs) + .bind(restore_lease_token) + .bind(execution_ownership_ttl_secs) + .fetch_optional(executor) + .await +} + +/// Commit `restoring -> started` (and `on_node`) for the restore whose lease +/// `restore_lease_token` holds. Writes an **operation step** (the claim's +/// completion). Returns whether the row moved. +/// +/// The row must still be `restoring` under that token, with a routable +/// active execution behind it, or a `managed_by` other than Kubernetes +/// ([`routable_execution_or_non_k8s`]). +/// A row a rollback, the sweeper, a suspend or a delete has moved on matches +/// nothing. +pub async fn complete_owned_user_restore( + executor: impl PgExecutor<'_>, + sandbox_id: SandboxId, + restore_lease_token: Uuid, +) -> Result { + let routable = routable_execution_or_non_k8s("sandboxes", "$3"); + let query = format!( + "UPDATE sandboxes + SET lifecycle = 'started', placement = 'on_node' + WHERE id = $1 + AND lifecycle = 'restoring' + AND restore_lease_token = $2 + AND {routable}" + ); + let completed = sqlx::query(sqlx::AssertSqlSafe(sql!(query.as_str()))) + .bind(sandbox_id) + .bind(restore_lease_token) + .bind(SubstrateKind::Kubernetes) + .execute(executor) + .await? + .rows_affected(); + Ok(completed == 1) +} + +/// Expire the restore lease `restore_lease_token` holds, keeping the token +/// and setting the expiry to the database's `NOW()`. Writes an **operation +/// step** (the lease's release). +/// +/// Guarded by the token alone, with no lifecycle predicate: a lease a newer +/// claim replaced, or one the lease-clear trigger dropped, matches nothing. +pub async fn release_user_restore_lease( + executor: impl PgExecutor<'_>, + sandbox_id: SandboxId, + restore_lease_token: Uuid, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes + SET restore_lease_expires_at = NOW() + WHERE id = $1 + AND restore_lease_token = $2" + )) + .bind(sandbox_id) + .bind(restore_lease_token) + .execute(executor) + .await?; + Ok(()) +} + +/// Revert an abandoned restore claim to `suspended` whose execution the +/// caller retired in this transaction: unplace the row and clear +/// `placement_changed_at`. Writes an **operation step** (the claim's +/// compensation). +/// +/// No lifecycle or lease predicate of its own; it relies on the caller's row +/// lock, taken on `restoring` (and, when the claim wrote a lease token, on +/// that token), and on the caller's proof that the pre-restore source is still restorable. +pub async fn revert_restore_claim_retiring_placement( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes + SET lifecycle = 'suspended', + placement = 'unplaced', + active_node_id = NULL, + placement_changed_at = NULL + WHERE id = $1" + )) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// Revert an abandoned restore claim with no execution at all to +/// `suspended`, clearing `placement_changed_at`. Writes an **operation step** +/// (the claim's compensation). +/// +/// No lifecycle or lease predicate of its own; it relies on the caller's row +/// lock and on the caller's proof that the pre-restore source is still +/// restorable. +pub async fn revert_restore_claim_clearing_placement_clock( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes + SET lifecycle = 'suspended', + placement_changed_at = NULL + WHERE id = $1" + )) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// Revert an abandoned restore claim's lifecycle to `suspended`, leaving +/// its execution, placement and placement clock as the restore left them. +/// Writes an **operation step** (the claim's compensation). +/// +/// No lifecycle or lease predicate of its own; it relies on the caller's row +/// lock. +pub async fn revert_restore_claim_lifecycle( + tx: &mut Transaction<'_, Postgres>, + sandbox_id: SandboxId, +) -> Result<(), sqlx::Error> { + sqlx::query(sql!( + "UPDATE sandboxes SET lifecycle = 'suspended' WHERE id = $1" + )) + .bind(sandbox_id) + .execute(&mut **tx) + .await?; + Ok(()) +}