sbx-mgr: suspend writes only over the state it expects

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

@@ -225,8 +225,8 @@ edge_table! {
 
     // Pause and resident resume.
 
-    /// `pause` and the suspend auto-pause, after `transition(Pausing)` on the
-    /// read.
+    /// `pause` after `transition(Pausing)` on the read, and the suspend
+    /// auto-pause after reading `started`.
     PAUSE_BEGIN: Request, of(&[Started]) => Pausing;
     /// `pause` and the suspend auto-pause, after node-manager confirms the
     /// pause.

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

@@ -224,6 +224,13 @@ pub(super) async fn service_cleanup_via_node_with_options(
 
 pub(super) use crate::db::CleanupCompletion;
 
+#[cfg_attr(
+    not(test),
+    expect(
+        dead_code,
+        reason = "the update_state parity tests call it; production writes use update_state_from"
+    )
+)]
 pub(super) fn require_state_write(
     result: rusqlite::Result<crate::db::SandboxStateWriteOutcome>,
     context: &'static str,

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

@@ -10,9 +10,7 @@ use super::activation_support::{
 use super::cleanup::{
     AsyncCleanupRequest, CleanupAdmission, CleanupDispatchDisposition, SynchronousCleanupRequest,
 };
-use super::cleanup_support::{
-    CleanupCompletion, CleanupLockTryResult, OwnedCleanupFenceGuard, require_state_write,
-};
+use super::cleanup_support::{CleanupCompletion, CleanupLockTryResult, OwnedCleanupFenceGuard};
 use super::create_support::SUSPEND_CHECKPOINT;
 use super::network_runtime::*;
 use super::recovery::SuspendSnapshotPublicationRecoveryOrigin;
@@ -54,9 +52,9 @@ pub(super) use self::artifact_backup::{
     reupload_precondition_for_cluster_api_env, reupload_resolution_to_config,
 };
 use self::backup_submission::suspend_snapshot_metadata;
-use self::pause_resume::pause_with_reconciliation;
 #[cfg(test)]
 use self::pause_resume::{PauseReconciliationOutcome, pause_reconciliation_outcome};
+use self::pause_resume::{pause_with_reconciliation, prepare_guest_for_pause};
 #[cfg(test)]
 pub(super) use self::suspend::classify_suspend_after_wake;
 use self::suspend::emit_suspend_failed_event;

crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/pause_resume.rs

@@ -37,7 +37,7 @@ fn roll_back_pause_best_effort(
 }
 
 #[cfg(test)]
-type GuestPreparation = Arc<
+pub(super) type GuestPreparation = Arc<
     dyn Fn(
             String,
         ) -> futures_util::future::BoxFuture<
@@ -49,13 +49,13 @@ type GuestPreparation = Arc<
 
 #[cfg(test)]
 tokio::task_local! {
-    /// Replaces pause's guest snapshot preparation, which dials the host's
-    /// per-sandbox spaced-host socket, and lets a test act between pause's
-    /// state writes.
-    static GUEST_PREPARATION: GuestPreparation;
+    /// Replaces the guest snapshot preparation of pause and the suspend
+    /// auto-pause, which dials the host's per-sandbox spaced-host socket, and
+    /// lets a test act between their state writes.
+    pub(super) static GUEST_PREPARATION: GuestPreparation;
 }
 
-async fn prepare_guest_for_pause(
+pub(super) async fn prepare_guest_for_pause(
     sandbox_id: &str,
     budget: GuestSnapshotPreparationBudget,
 ) -> Result<(), crate::guest_filesystem_flush::GuestFlushError> {

crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend.rs

@@ -2,7 +2,7 @@
 //! and suspend preparation once the locks are held.
 
 use super::*;
-use crate::service::cleanup_support::StateWriteRejected;
+use crate::service::cleanup_support::{StateWriteRejected, require_state_write_from};
 
 fn completed_suspend_retry_response(
     sandbox: &crate::db::SandboxRow,
@@ -354,22 +354,26 @@ impl SandboxManagerService {
                 Status::internal(error)
             })?;
             self.ensure_teardown_suspend_owner_for_durability(&sandbox.id, context.durability)?;
+            // The locks held here do not exclude a writer that takes none of
+            // them, such as warm-snapshot Start, so each auto-pause write names
+            // the state it expects instead of overwriting whatever it finds.
             {
                 let db = self.db_conn()?;
-                require_state_write(
-                    crate::db::update_state(&db, &sandbox.id, SandboxState::Pausing),
+                require_state_write_from(
+                    crate::db::update_state_from(&db, &sandbox.id, &crate::db::edges::PAUSE_BEGIN),
                     "begin checkpoint pause",
                 )?;
             }
             self.watchdog.record_lifecycle_edge(&sandbox.id).await;
             let guest_preparation_budget = context.guest_preparation_budget();
-            if let Err(e) =
-                prepare_guest_for_vm_snapshot_via_spaced_host(&sandbox.id, guest_preparation_budget)
-                    .await
-            {
+            if let Err(e) = prepare_guest_for_pause(&sandbox.id, guest_preparation_budget).await {
                 let db = self.db_conn()?;
-                if let Err(error) = require_state_write(
-                    crate::db::update_state(&db, &sandbox.id, SandboxState::Started),
+                if let Err(error) = require_state_write_from(
+                    crate::db::update_state_from(
+                        &db,
+                        &sandbox.id,
+                        &crate::db::edges::PAUSE_ROLLBACK,
+                    ),
                     "roll back checkpoint pause preparation",
                 ) {
                     tracing::warn!(sandbox_id = %sandbox.id, %error, "failed to roll back checkpoint pause state");
@@ -442,8 +446,12 @@ impl SandboxManagerService {
             }
             {
                 let db = self.db_conn()?;
-                require_state_write(
-                    crate::db::update_state(&db, &sandbox.id, SandboxState::Paused),
+                require_state_write_from(
+                    crate::db::update_state_from(
+                        &db,
+                        &sandbox.id,
+                        &crate::db::edges::PAUSE_COMPLETE,
+                    ),
                     "complete checkpoint pause",
                 )?;
             }
@@ -528,3 +536,6 @@ impl SandboxManagerService {
         Ok((blocking_local_artifact_id, publication))
     }
 }
+
+#[cfg(test)]
+mod tests;