sbx-mgr: snapshot completion and rollback write only over the state they expect

architecture-baseline.toml

@@ -88,8 +88,8 @@
 "crates/libs/sbx/mgr/server/src/service/create.rs" = 1195
 "crates/libs/sbx/mgr/server/src/service/grpc_core.rs" = 1087
 "crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend_snapshot.rs" = 1016
-"crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs" = 1640
-"crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs" = 1835
+"crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs" = 1636
+"crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs" = 1831
 "crates/libs/sbx/mgr/server/src/service/mod.rs" = 1757
 "crates/libs/sbx/mgr/server/src/service/network_runtime.rs" = 1112
 "crates/libs/sbx/mgr/server/src/service/recovery.rs" = 2017

crates/libs/sbx/mgr/server/src/db/state_write.rs

@@ -251,18 +251,22 @@ edge_table! {
 
     // Checkpoints and template base snapshots.
 
-    /// `do_checkpoint` and `create_template_base_snapshot` (from `paused`):
-    /// `update_state(Checkpointing)` checked and written under one guard.
+    /// `begin_checkpoint_state` from `do_checkpoint`, checked and written under
+    /// one guard.
     CHECKPOINT_BEGIN: Request, of(&[Started, Paused]) => Checkpointing;
+    /// `begin_template_base_snapshot_state` from
+    /// `create_template_base_snapshot`, checked and written under one guard.
+    TEMPLATE_SNAPSHOT_BEGIN: Request, of(&[Paused]) => Checkpointing;
     /// `restore_checkpoint_state` from `finish_checkpoint_success`.
     CHECKPOINT_COMPLETE_STARTED: Completion, of(&[Checkpointing]) => Started;
     /// `restore_checkpoint_state` from `finish_checkpoint_success`, and
-    /// `create_template_base_snapshot`'s `update_state(Paused)`.
+    /// `complete_template_base_snapshot_state`.
     CHECKPOINT_COMPLETE_PAUSED: Completion, of(&[Checkpointing]) => Paused;
     /// `restore_checkpoint_state_best_effort` after a failed checkpoint.
     CHECKPOINT_ROLLBACK_STARTED: Compensation, of(&[Checkpointing]) => Started;
     /// `restore_checkpoint_state_best_effort`, and
-    /// `create_template_base_snapshot`'s `update_state_best_effort(Paused)`.
+    /// `rollback_template_base_snapshot_state` after a failed template base
+    /// snapshot.
     CHECKPOINT_ROLLBACK_PAUSED: Compensation, of(&[Checkpointing]) => Paused;
 
     // Suspend.
@@ -274,9 +278,9 @@ edge_table! {
     SUSPEND_COMPLETE: Completion, of(&[Suspending]) => Suspended;
     /// `rollback_suspending_to_paused` from the suspend snapshot path.
     SUSPEND_ROLLBACK: Compensation, of(&[Suspending]) => Paused;
-    /// The full process-snapshot rollback in `grpc_spaced_process_snapshots`:
-    /// `update_state(Suspended)` after its own stop, so `start_from_checkpoint`
-    /// can restore the snapshot.
+    /// `restage_rolled_back_full_snapshot`: the full process-snapshot rollback
+    /// restages its stopped row, so `start_from_checkpoint` can restore the
+    /// snapshot. A row that is already `suspended` is left as it is.
     PROCESS_SNAPSHOT_ROLLBACK_RESTAGE: Completion, of(&[Stopped]) => Suspended;
 
     // Stop and delete.
@@ -429,14 +433,6 @@ edge_table! {
     PAUSE_ROLLBACK_AFTER_STOP: KnownRace, of(&[Stopped]) => Started;
     /// Pause completion (`update_state(Paused)`) after a concurrent Stop.
     PAUSE_COMPLETE_AFTER_STOP: KnownRace, of(&[Stopped]) => Paused;
-    /// The full process-snapshot rollback's `update_state(Suspended)` after a
-    /// Stop admitted while it restaged the snapshot. Its own stop leaves no
-    /// runtime, so the second Stop is reached only in the window between
-    /// Stop's runtime gate and the cleanup fence, or with a leaked runtime.
-    PROCESS_SNAPSHOT_ROLLBACK_AFTER_STOP: KnownRace, of(&[Stopping]) => Suspended;
-    /// `create_template_base_snapshot`'s `update_state(Paused)` after a Stop
-    /// completed during capture.
-    TEMPLATE_SNAPSHOT_COMPLETE_AFTER_STOP: KnownRace, of(&[Stopped]) => Paused;
 }
 
 pub use edges::EDGES;

crates/libs/sbx/mgr/server/src/db/tests/state_write_fixtures.rs

@@ -658,7 +658,7 @@ pub(super) fn insert_suspended_with_artifact(conn: &Connection) -> (String, Loca
     (id, artifact)
 }
 
-pub(super) fn decorate(conn: &Connection, id: &str, receipt: SandboxState) {
+pub(crate) fn decorate(conn: &Connection, id: &str, receipt: SandboxState) {
     set_failure_reason(conn, id, "seeded failure");
     seed_receipt(conn, id, receipt);
     backdate(conn, id);

crates/libs/sbx/mgr/server/src/db/tests/state_write_table.rs

@@ -61,13 +61,12 @@ const ALLOWLIST: &[(SandboxState, SandboxState, Disagreement, &str)] = &[
     (Started,         Suspending,      AllowedUnwritten,    "suspend pauses a started row first"),
     (Started,         Quarantined,     AllowedNotRequested, RECOVERY_QUARANTINE),
     (Stopping,        Stopped,         AllowedNotRequested, COMPLETION_ONLY),
-    (Stopping,        Suspended,       WriteForbidden,      "known race: process-snapshot rollback in the gate-to-fence window"),
     (Stopping,        CleanupRequired, WriteForbidden,      "an interrupted stop is handed to cleanup"),
     (Stopping,        CleanupFailed,   WriteForbidden,      "stop cleanup failed"),
     (Stopping,        Quarantined,     AllowedNotRequested, RECOVERY_QUARANTINE),
     (Stopped,         Started,         WriteForbidden,      "known race: pause rollback after Stop"),
     (Stopped,         Stopping,        RequestForbidden,    STOP_SETTLED_SOURCE),
-    (Stopped,         Paused,          WriteForbidden,      "known races: pause and template-snapshot completion after Stop"),
+    (Stopped,         Paused,          WriteForbidden,      "known race: pause completion after Stop"),
     (Stopped,         Suspended,       RequestForbidden,    "warm start and process-snapshot rollback restage a stopped row"),
     (Stopped,         Quarantined,     AllowedUnwritten,    NO_QUARANTINE_WRITER),
     (Deleting,        CleanupRequired, WriteForbidden,      "an interrupted delete is handed to cleanup"),

crates/libs/sbx/mgr/server/src/service/cleanup_support.rs

@@ -331,10 +331,6 @@ impl StateWriteRejected {
 /// or an error is logged and leaves the row as it is, and only an applied
 /// write marks the sandbox dirty. `None` is a poisoned connection or a
 /// database error.
-#[cfg_attr(
-    not(test),
-    expect(dead_code, reason = "the update_state_from conversions adopt it")
-)]
 pub(super) fn update_state_from_best_effort(
     db: &Arc<Mutex<rusqlite::Connection>>,
     sandbox_id: &str,

crates/libs/sbx/mgr/server/src/service/cleanup_support/tests.rs

@@ -355,14 +355,14 @@ fn best_effort_refusals_change_nothing() {
 
 #[test]
 fn best_effort_database_error_is_none_and_changes_nothing() {
-    let edge: &'static Edge = &edges::TEMPLATE_SNAPSHOT_COMPLETE_AFTER_STOP;
+    let edge: &'static Edge = &edges::ORPHAN_SUSPENDED_INSERT;
     let db = Arc::new(Mutex::new(crate::db::test_conn()));
     let id = SandboxId::new().to_string();
     insert_test_sandbox(&db.lock().unwrap(), &id, SandboxState::Stopped);
     let before = snapshot(&db.lock().unwrap());
 
     assert_eq!(update_state_from_best_effort(&db, &id, edge), None);
-    assert_unchanged("known-race edge", &before, &db.lock().unwrap());
+    assert_unchanged("insert edge", &before, &db.lock().unwrap());
 }
 
 /// A poisoned connection is `None`, changes nothing and signals the failure:

crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs

@@ -6,7 +6,7 @@ use super::cleanup::{
     CleanupCompletionDisposition, CleanupDispatchDisposition, ObservedCleanupClaim,
 };
 use super::cleanup_support::{
-    CleanupCompletion, cleanup_sandbox_via_node, read_node_execution_fence, require_state_write,
+    CleanupCompletion, cleanup_sandbox_via_node, read_node_execution_fence,
 };
 use super::create::StartFromCheckpointLaunch;
 use super::create_support::*;
@@ -20,6 +20,7 @@ use pplx_space_events::{
     emit_sandbox_event,
 };
 use pplx_space_sbx_mgr_client::operations as spaced_ops;
+use rollback_state::restage_rolled_back_full_snapshot;
 use std::str::FromStr;
 use tracing::Instrument;
 
@@ -1687,14 +1688,7 @@ impl SandboxManagerService {
                             "snapshot.json",
                         )?;
 
-                        // Set state to Suspended so start_from_checkpoint can proceed
-                        {
-                            let db = service.db_conn()?;
-                            require_state_write(
-                                crate::db::update_state(&db, &copy_id, SandboxState::Suspended),
-                                "complete process snapshot suspension",
-                            )?;
-                        }
+                        restage_rolled_back_full_snapshot(&*service.db_conn()?, &copy_id)?;
                         Ok(artifact_guard)
                     })
                     .await?;
@@ -1781,6 +1775,8 @@ fn start_process_options(
     })
 }
 
+mod rollback_state;
+
 #[cfg(test)]
 mod start_process_options_tests;
 

crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots/rollback_state.rs

@@ -0,0 +1,27 @@
+//! The full process-snapshot rollback's write to the node-local SQLite
+//! `sandboxes.state`.
+
+use super::super::cleanup_support::{StateWriteRejected, state_write_result};
+use crate::db::edges;
+use rusqlite::Connection;
+use tonic::Status;
+
+/// Restage a full process-snapshot rollback's `stopped` row as `suspended`, so
+/// `start_from_checkpoint` can restore the snapshot. A row that is already
+/// `suspended` needs no write. Any other row was moved after the rollback
+/// checked it, for example by a Stop admitted after the rollback's own stop,
+/// and the rollback reports that it was superseded instead of overwriting it.
+pub(super) fn restage_rolled_back_full_snapshot(
+    db: &Connection,
+    sandbox_id: &str,
+) -> Result<(), Status> {
+    const CONTEXT: &str = "complete process snapshot suspension";
+    let edge = &edges::PROCESS_SNAPSHOT_ROLLBACK_RESTAGE;
+    let outcome = crate::db::update_state_from(db, sandbox_id, edge)
+        .map_err(|error| Status::internal(format!("{CONTEXT}: {error}")))?;
+    match state_write_result(outcome) {
+        Ok(_) => Ok(()),
+        Err(StateWriteRejected::FromState { found }) if found == edge.to() => Ok(()),
+        Err(rejected) => Err(rejected.into_status(CONTEXT)),
+    }
+}

crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots/tests.rs

@@ -343,3 +343,136 @@ fn rollback_replaces_optional_sidecars_or_reports_unremovable_stale_evidence() {
         }
     }
 }
+
+fn sandbox_state(conn: &rusqlite::Connection, id: &str) -> SandboxState {
+    crate::db::get_sandbox(conn, id).unwrap().unwrap().state
+}
+
+/// Stop admission as the cleanup fence runs it, with a non-blocking Stop's
+/// persisted intent.
+fn admit_stop(conn: &rusqlite::Connection, id: &str) {
+    let tx = crate::db::immediate_transaction(conn).unwrap();
+    crate::db::update_state_superseding_activation_in_transaction(&tx, id, SandboxState::Stopping)
+        .unwrap();
+    assert!(crate::db::set_cleanup_intent(&tx, id, true, CleanupCompletion::MarkStopped).unwrap());
+    tx.commit().unwrap();
+}
+
+fn complete_stop(conn: &rusqlite::Connection, id: &str) {
+    crate::db::complete_row_intent_cleanup_for_job(
+        conn,
+        id,
+        CleanupCompletion::MarkStopped,
+        None,
+        None,
+        None,
+    )
+    .unwrap();
+}
+
+/// The full rollback holds the artifact lock, which Stop does not take, while
+/// it rolls the disk back and copies the snapshot. A Stop admitted in that
+/// window must not be overwritten by the restage to `suspended`: it keeps its
+/// `stopping` row and its completion still settles it.
+#[test]
+fn full_rollback_restage_leaves_a_stop_admitted_after_its_own_stop() {
+    let conn = crate::db::test_conn();
+    let id = SandboxId::new().to_string();
+    crate::service::tests::service_tests::insert_test_sandbox_with_runtime(
+        &conn,
+        &id,
+        SandboxState::Started,
+    );
+    // The rollback's own stop.
+    admit_stop(&conn, &id);
+    complete_stop(&conn, &id);
+    assert_eq!(sandbox_state(&conn, &id), SandboxState::Stopped);
+    admit_stop(&conn, &id);
+    let before = crate::db::tests::database_rows(&conn);
+
+    let error = restage_rolled_back_full_snapshot(&conn, &id)
+        .expect_err("the restage must not overwrite a concurrent Stop");
+
+    assert_eq!(error.code(), tonic::Code::Aborted);
+    assert_eq!(
+        error.message(),
+        "complete process snapshot suspension: sandbox state write was superseded: sandbox is stopping"
+    );
+    assert_eq!(crate::db::tests::database_rows(&conn), before);
+    complete_stop(&conn, &id);
+    let row = crate::db::get_sandbox(&conn, &id).unwrap().unwrap();
+    assert_eq!(row.state, SandboxState::Stopped);
+    assert_eq!(row.cleanup_completion, None);
+}
+
+/// The restage writes only over a `stopped` row and accepts one that is
+/// already `suspended` without a write. Every other row keeps every table
+/// unchanged and the rollback reports why. Rows carry a failure reason, a
+/// receipt for another target and backdated timestamps, which an in-place
+/// rewrite of a `suspended` row would clear or restamp.
+#[test]
+fn full_rollback_restage_writes_only_over_a_stopped_row() {
+    enum Expected {
+        Restaged,
+        AlreadySuspended,
+        Refused(String),
+    }
+    for state in SandboxState::all() {
+        let expected = match state {
+            SandboxState::Stopped => Expected::Restaged,
+            SandboxState::Suspended => Expected::AlreadySuspended,
+            // The generic guard: a deleting row owns its `delete_row` cleanup.
+            SandboxState::Deleting => Expected::Refused(
+                "complete process snapshot suspension: sandbox state write was not applied: \
+                 superseded_by_cleanup_episode"
+                    .to_string(),
+            ),
+            SandboxState::Quarantined => Expected::Refused(
+                "complete process snapshot suspension: sandbox state write was not applied: \
+                 blocked_by_quarantine"
+                    .to_string(),
+            ),
+            SandboxState::Creating
+            | SandboxState::Created
+            | SandboxState::Starting
+            | SandboxState::Started
+            | SandboxState::Stopping
+            | SandboxState::Pausing
+            | SandboxState::Paused
+            | SandboxState::Resuming
+            | SandboxState::Checkpointing
+            | SandboxState::Suspending
+            | SandboxState::Restoring
+            | SandboxState::CleanupRequired
+            | SandboxState::CleanupFailed
+            | SandboxState::Failed => Expected::Refused(format!(
+                "complete process snapshot suspension: sandbox state write was superseded: \
+                 sandbox is {state}"
+            )),
+        };
+        let conn = crate::db::test_conn();
+        let id = SandboxId::new().to_string();
+        crate::service::tests::service_tests::insert_test_sandbox_with_runtime(&conn, &id, state);
+        crate::db::tests::state_write_fixtures::decorate(&conn, &id, SandboxState::Started);
+        let before = crate::db::tests::database_rows(&conn);
+
+        let result = restage_rolled_back_full_snapshot(&conn, &id);
+
+        match expected {
+            Expected::Restaged => {
+                result.unwrap();
+                assert_eq!(sandbox_state(&conn, &id), SandboxState::Suspended);
+            }
+            Expected::AlreadySuspended => {
+                result.unwrap();
+                assert_eq!(crate::db::tests::database_rows(&conn), before, "{state}");
+            }
+            Expected::Refused(message) => {
+                let error = result.expect_err("a refused restage must fail");
+                assert_eq!(error.code(), tonic::Code::Aborted, "{state}");
+                assert_eq!(error.message(), message, "{state}");
+                assert_eq!(crate::db::tests::database_rows(&conn), before, "{state}");
+            }
+        }
+    }
+}

crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs

@@ -1,12 +1,12 @@
 use pplx_space_common_core::types::LocalArtifactId;
 
-use super::activation_support::update_state_best_effort;
 use super::cache_cleanup::is_sandbox_running;
-use super::cleanup_support::require_state_write;
 use super::create_support::*;
 use super::tail_support::*;
 use super::template_snapshot_support::{
+    begin_template_base_snapshot_state, complete_template_base_snapshot_state,
     compute_template_base_snapshot_hash, publish_template_base_snapshot_dir,
+    rollback_template_base_snapshot_state,
 };
 use super::validation::{canonical_batched_sandbox_ids, snapshot_type_from_proto};
 use super::*;
@@ -1785,10 +1785,7 @@ impl SandboxManagerService {
                 .state
                 .transition(SandboxState::Checkpointing)
                 .map_err(Status::from)?;
-            require_state_write(
-                crate::db::update_state(&db, &sandbox.id, SandboxState::Checkpointing),
-                "begin template checkpoint",
-            )?;
+            begin_template_base_snapshot_state(&db, &sandbox.id)?;
             (sandbox, content_hash)
         };
 
@@ -1900,13 +1897,7 @@ impl SandboxManagerService {
 
         match snapshot_result {
             Ok((snapshot_hash, snapshot_memory_mib)) => {
-                {
-                    let db = self.db_conn()?;
-                    require_state_write(
-                        crate::db::update_state(&db, &sandbox.id, SandboxState::Paused),
-                        "complete template checkpoint",
-                    )?;
-                }
+                complete_template_base_snapshot_state(&*self.db_conn()?, &sandbox.id)?;
                 tracing::debug!(
                     sandbox_id = %sandbox.id,
                     template = %sandbox.template,
@@ -1936,7 +1927,7 @@ impl SandboxManagerService {
                     "template base snapshot checkpoint failed"
                 );
                 let _ = std::fs::remove_dir_all(&staging_dir);
-                update_state_best_effort(&self.db, &sandbox.id, SandboxState::Paused);
+                rollback_template_base_snapshot_state(&self.db, &sandbox.id);
                 Err(status)
             }
         }