fix(compute): revalidate late exit events after restart (#4235)

Signed-off-by: Simon Scatton <sscatton@nvidia.com>

Author: SDAChess

crates/openshell-server/src/compute/mod.rs

@@ -4434,18 +4434,23 @@ impl ComputeRuntime {
                 SandboxPhase::try_from(sandbox.phase()).unwrap_or(SandboxPhase::Unknown)
             });
 
-        if existing_phase != SandboxPhase::Starting {
+        let ready_container_exit = existing_phase == SandboxPhase::Ready
+            && (driver_snapshot_confirms_stopped(&incoming)
+                || driver_snapshot_reports_runtime_restart(&incoming));
+        if existing_phase != SandboxPhase::Starting && !ready_container_exit {
             return self.apply_sandbox_update_locked(incoming, existing).await;
         }
 
         // Any snapshot can already be queued when StartSandbox moves the
         // durable phase to Starting. In particular, an old-generation Ready
         // event followed by its terminal event can otherwise promote and then
         // stop the new generation before the replacement supervisor connects.
+        // The replacement may already be Ready when the old exit arrives, so
+        // terminal container snapshots need the same check after readiness.
         // Release the global watch lock, wait for that lifecycle operation,
         // and then reread both the driver and store before applying an
         // authoritative observation. Taking the per-sandbox gate only for
-        // this ambiguous phase avoids delaying unrelated watch events behind
+        // these ambiguous snapshots avoids delaying unrelated watch events behind
         // slow lifecycle operations.
         let existing_name = existing_sandbox.as_ref().map_or_else(
             || incoming.name.clone(),
@@ -4475,7 +4480,7 @@ impl ComputeRuntime {
             {
                 warn!(
                     sandbox_id = %incoming.id,
-                    "Could not validate driver snapshot during sandbox start; retaining current sandbox state"
+                    "Could not validate driver snapshot against current runtime; retaining current sandbox state"
                 );
                 return Ok(());
             }
@@ -11604,6 +11609,65 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn container_exit_after_replacement_ready_is_revalidated() {
+        for (reason, live_exited) in [
+            ("ContainerExited", false),
+            ("ContainerExited", true),
+            ("ContainerStopped", false),
+            ("ContainerStopped", true),
+            ("ContainerRuntimeRestart", false),
+            ("ContainerRuntimeRestart", true),
+        ] {
+            let driver = ControlledDriver::new();
+            let mut sandbox =
+                sandbox_record("sb-restarted", "sandbox-restarted", SandboxPhase::Starting);
+            sandbox.status.as_mut().unwrap().restart_count = 1;
+            sandbox.status.as_mut().unwrap().exit_code = Some(17);
+            sandbox.status.as_mut().unwrap().main_process_instance_id = "previous".to_string();
+            let mut exited = ready_driver_sandbox(sandbox.object_id(), sandbox.object_name());
+            exited.status = Some(make_driver_status(make_driver_condition(
+                reason,
+                "container exited",
+            )));
+            let live = if live_exited {
+                exited.clone()
+            } else {
+                ready_driver_sandbox(sandbox.object_id(), sandbox.object_name())
+            };
+            driver.set_get_outcome(ControlledGetOutcome::Sandbox(Box::new(live)));
+            let runtime = test_runtime(driver).await;
+            runtime.store.put_message(&sandbox).await.unwrap();
+            register_test_supervisor_session(&runtime, sandbox.object_id());
+            runtime
+                .supervisor_session_connected(sandbox.object_id(), "replacement")
+                .await
+                .unwrap();
+
+            runtime.apply_sandbox_update(exited).await.unwrap();
+
+            let stored = runtime
+                .store
+                .get_message::<Sandbox>(sandbox.object_id())
+                .await
+                .unwrap()
+                .unwrap();
+            assert_eq!(
+                stored.phase(),
+                if live_exited {
+                    SandboxPhase::Error
+                } else {
+                    SandboxPhase::Ready
+                } as i32,
+                "only a live {reason} may fail the replacement runtime"
+            );
+            let status = stored.status.unwrap();
+            assert_eq!(status.main_process_instance_id, "replacement");
+            assert_eq!(status.restart_count, 1);
+            assert_eq!(status.exit_code, None);
+        }
+    }
+
     #[tokio::test]
     async fn stale_ready_snapshot_queued_before_start_is_revalidated() {
         let driver = ControlledDriver::new();
@@ -13375,14 +13439,18 @@ mod tests {
 
     #[tokio::test]
     async fn unexpected_term_runtime_restart_transitions_to_error() {
-        let runtime = test_runtime(Arc::new(TestDriver::default())).await;
+        let driver = ControlledDriver::new();
+        let runtime = test_runtime(driver.clone()).await;
         let sandbox = sandbox_record("sb-term-exit", "sandbox-term-exit", SandboxPhase::Ready);
         runtime.store.put_message(&sandbox).await.unwrap();
         let mut runtime_restart = ready_driver_sandbox(sandbox.object_id(), sandbox.object_name());
         runtime_restart.status = Some(make_driver_status(make_driver_condition(
             "ContainerRuntimeRestart",
             "container exited with status 143",
         )));
+        driver.set_get_outcome(ControlledGetOutcome::Sandbox(Box::new(
+            runtime_restart.clone(),
+        )));
 
         runtime.apply_sandbox_update(runtime_restart).await.unwrap();
 

e2e/rust/tests/sandbox_lifecycle.rs

@@ -9,9 +9,11 @@ use std::process::Stdio;
 use std::time::Duration;
 
 use openshell_e2e::harness::binary::{openshell_cmd, openshell_tty_cmd};
-use openshell_e2e::harness::cli::{run_cli, wait_for_sandbox_phase};
+use openshell_e2e::harness::cli::{
+    run_cli, wait_for_sandbox_exec_contains, wait_for_sandbox_phase,
+};
 use openshell_e2e::harness::output::{extract_field, strip_ansi};
-use openshell_e2e::harness::sandbox::{SandboxGuard, unique_sandbox_name};
+use openshell_e2e::harness::sandbox::{E2E_WORKLOAD_IMAGE, SandboxGuard, unique_sandbox_name};
 use serial_test::serial;
 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
 use tokio::time::{Instant, sleep};
@@ -1300,45 +1302,81 @@ async fn canonical_main_exit_255_is_not_retried_as_transport_failure() {
 #[tokio::test]
 #[serial(sandbox_lifecycle)]
 async fn on_failure_policy_replaces_runtime_and_preserves_workspace() {
-    const FIRST_MARKER: &str = "initial-main-ready";
     const SCRIPT: &str = r#"
 marker=/sandbox/.openshell-restart-e2e
 if [ -e "$marker" ]; then
   printf 'replacement-%s\n' "$(cat /proc/sys/kernel/random/uuid)" > /sandbox/replacement-run
-  printf 'replacement-main-ready\n'
   sleep 300
 else
   touch "$marker"
   printf 'initial-%s\n' "$(cat /proc/sys/kernel/random/uuid)" > /sandbox/initial-run
-  printf 'initial-main-ready\n'
-  sleep 2
+  while [ ! -e /sandbox/.openshell-restart-release ]; do sleep 0.05; done
   exit 17
 fi
 "#;
 
-    let mut sandbox = SandboxGuard::create_keep_with_args(
-        &["--restart-policy", "on-failure"],
-        &["sh", "-lc", SCRIPT],
-        FIRST_MARKER,
+    let mut sandbox = SandboxGuard::manage_existing(unique_sandbox_name());
+    let (output, exit_code) = run_cli(&[
+        "sandbox",
+        "create",
+        "--name",
+        &sandbox.name,
+        "--from",
+        E2E_WORKLOAD_IMAGE,
+        "--detach",
+        "--restart-policy",
+        "on-failure",
+        "--",
+        "sh",
+        "-lc",
+        SCRIPT,
+    ])
+    .await;
+    assert_eq!(
+        exit_code, 0,
+        "create sandbox with OnFailure policy: {output}"
+    );
+    wait_for_sandbox_exec_contains(
+        &sandbox.name,
+        &["cat", "/sandbox/initial-run"],
+        "initial-",
+        SANDBOX_RESTART_TIMEOUT,
     )
     .await
-    .expect("create sandbox with OnFailure restart policy");
+    .expect("initial main process should write its workspace marker");
+
+    let initial_details = sandbox_details(&sandbox.name).await;
+    let initial_instance = extract_field(&initial_details, "Main process instance")
+        .expect("initial main process instance");
+    sandbox
+        .exec(&["touch", "/sandbox/.openshell-restart-release"])
+        .await
+        .expect("release initial main process");
 
     let deadline = Instant::now() + SANDBOX_RESTART_TIMEOUT;
-    let mut observed_replacement_starting = false;
+    let mut last_runs = Err("replacement is not ready".to_string());
     let final_details = loop {
         let details = sandbox_details(&sandbox.name).await;
-        observed_replacement_starting |= details.contains("Phase: Starting");
-        if observed_replacement_starting
-            && details.contains("Phase: Ready")
-            && details.contains("Restart count: 1")
+        let replacement_instance = extract_field(&details, "Main process instance")
+            .is_some_and(|instance| instance != "-" && instance != initial_instance);
+        if details.contains("Phase: Ready")
+            && extract_field(&details, "Restart count").as_deref() == Some("1")
+            && replacement_instance
         {
-            break details;
+            last_runs = sandbox
+                .exec(&["cat", "/sandbox/initial-run", "/sandbox/replacement-run"])
+                .await;
+            if let Ok(runs) = &last_runs
+                && runs.contains("initial-")
+                && runs.contains("replacement-")
+            {
+                break details;
+            }
         }
         assert!(
             Instant::now() < deadline,
             "sandbox did not complete its policy-driven restart within \
-             {SANDBOX_RESTART_TIMEOUT:?}; last details:\n{details}"
+             {SANDBOX_RESTART_TIMEOUT:?}; last details:\n{details}\nworkspace: {last_runs:?}"
         );
         sleep(Duration::from_millis(250)).await;
     };
@@ -1347,19 +1385,6 @@ fi
         final_details.contains("Restart policy: on-failure"),
         "restart policy should remain visible after replacement:\n{final_details}"
     );
-    let runs = sandbox
-        .exec(&["cat", "/sandbox/initial-run", "/sandbox/replacement-run"])
-        .await
-        .expect("replacement sandbox should retain the first run's workspace");
-    assert!(
-        runs.contains("initial-"),
-        "missing initial run marker:\n{runs}"
-    );
-    assert!(
-        runs.contains("replacement-"),
-        "missing replacement run marker:\n{runs}"
-    );
-
     sandbox.cleanup().await;
 }