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(()) +}