server: record an exit heard while reconnecting as reported, and fix a test's early terminate (P1R-C04 follow-up)

crates/server/src/mcp.rs

@@ -4575,7 +4575,7 @@ mod independent_tests {
     use codespace_runner::{
         host_boot_id, host_worker, host_worker_with_retention, InProcessRunner, IndependentRunner,
         LinkEvent, RetentionPolicy, RunnerEvents, RunnerIdentity, RunnerOp, RunnerOpResult,
-        StopHandle, UdsRunner,
+        RunnerProcessStatus, StopHandle, UdsRunner,
     };
     use codespace_store::{
         NewProcessRecord, ProcessEnd, ProcessLoss, ProcessOutcome, ProcessPhase, ProcessRecord,
@@ -5903,6 +5903,80 @@ mod independent_tests {
         released(&gateway, "gov").await;
     }
 
+    /// An exit that the gateway hears of while it claims its Runner again, as when the Runner
+    /// sends it right after the Claim's answer, is recorded as the Runner reports it once the
+    /// gateway has claimed the Runner (P1R-C04), not as an end whose kind is unknown.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn an_exit_heard_while_claiming_again_is_recorded_as_the_runner_reports_it() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        let dir = tempfile::tempdir().unwrap();
+        let database = dir.path().join("operations.db");
+        let identity = fake_identity();
+        left_running(
+            &database,
+            &runner_ref(identity.clone()),
+            "proc-ended",
+            Some(1),
+            "demo",
+        );
+        let ended = ProcessId("proc-ended".into());
+        let (client, server) = UnixStream::pair().unwrap();
+        let asked = Arc::new(AtomicUsize::new(0));
+        let acknowledged = Arc::new(std::sync::Mutex::new(Vec::new()));
+        tokio::spawn(fake_runner_answering(server, identity.clone(), {
+            let (asked, acknowledged) = (asked.clone(), acknowledged.clone());
+            move |op| match op {
+                RunnerOp::ProcessStatus { process_id } => {
+                    // The first request is refused as the client refuses it while it claims
+                    // the Runner again.
+                    if asked.fetch_add(1, Ordering::SeqCst) == 0 {
+                        return Some(Err(codespace_runner::runner_reconnecting()));
+                    }
+                    Some(Ok(RunnerOpResult::ProcessStatus(RunnerProcessStatus {
+                        process_id,
+                        state: ProcessState::Exited,
+                        exit_code: Some(4),
+                        termination: Some(ProcessTermination::Exited),
+                        signal: None,
+                        output_total: 0,
+                        output_retained_from: 0,
+                        eof: true,
+                    })))
+                }
+                RunnerOp::Acknowledge { process_ids } => {
+                    acknowledged.lock().unwrap().extend(process_ids);
+                    Some(Ok(RunnerOpResult::Acknowledged))
+                }
+                other => panic!("a request the fake Runner does not answer: {other:?}"),
+            }
+        }));
+        let store = Arc::new(Store::open(&database).unwrap());
+        let (release, exits) = ProcessRecords::release_hook(Leases::new(store.clone()));
+        let uds = UdsRunner::from_stream(client, Arc::new(|_: &str| {}));
+        uds.handshake().await.unwrap();
+        uds.claim(&store.store_id().unwrap()).await.unwrap();
+        let _records = ProcessRecords::start(store, &identity, RuntimeBackend::Uds(uds), exits);
+
+        // The Runner reports that the process exited.
+        release(&ended.0);
+        let record = record_once(&database, &ended, terminal).await;
+        assert_eq!(
+            (record.outcome, record.exit_code),
+            (
+                Some(ProcessOutcome::Ended(ProcessTermination::Exited)),
+                Some(4)
+            ),
+            "{record:?}"
+        );
+        assert_eq!(asked.load(Ordering::SeqCst), 2);
+        let deadline = Instant::now() + Duration::from_secs(10);
+        while acknowledged.lock().unwrap().is_empty() {
+            assert!(Instant::now() < deadline, "the end is not acknowledged");
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+        assert_eq!(*acknowledged.lock().unwrap(), [ended]);
+    }
+
     /// An execution, a patch or an approval resume that the Runner refuses as sent under an
     /// earlier Claim never ran (P1R-C03): the approval is given back granted and can be resumed
     /// again, the patch's operation_key begins afresh, the workspace is free again, and the
@@ -6721,7 +6795,7 @@ mod independent_tests {
         // A record that says this Runner started an execution it never received.
         let this = runner_ref(RunnerIdentity::this_process());
         left_running(&database, &this, "proc-phantom", None, "gov");
-        // A process that a controller started without recording it.
+        // A process that a controller started without recording it. It says when it runs.
         let unrecorded = ProcessId("proc-unrecorded".into());
         {
             let (client, server) = UnixStream::pair().unwrap();
@@ -6736,7 +6810,11 @@ mod independent_tests {
                 &raw,
                 ws,
                 RunnerExecRequest::for_host(
-                    vec!["/bin/sleep".into(), "30".into()],
+                    vec![
+                        "/bin/sh".into(),
+                        "-c".into(),
+                        "echo started; exec /bin/sleep 30".into(),
+                    ],
                     unrecorded.clone(),
                     ws.profile,
                 ),
@@ -6787,6 +6865,24 @@ mod independent_tests {
             .0;
         assert_eq!(status.state, ProcessState::Running);
 
+        // Terminate it once it runs: a sandboxed command whose launcher is killed while its
+        // sandbox is still being set up runs on outside it, and keeps its output open.
+        let deadline = Instant::now() + Duration::from_secs(10);
+        loop {
+            let read = gateway
+                .read_process(Parameters(ReadProcessParams {
+                    process_id: unrecorded.clone(),
+                    cursor: 0,
+                }))
+                .await
+                .unwrap()
+                .0;
+            if read.chunk.contains("started") {
+                break;
+            }
+            assert!(Instant::now() < deadline, "{read:?}");
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
         gateway
             .terminate_process(Parameters(TerminateProcessParams {
                 process_id: unrecorded.clone(),

crates/server/src/records.rs

@@ -13,6 +13,7 @@
 use std::collections::HashSet;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex, PoisonError};
+use std::time::Duration;
 
 use codespace_domain::{
     ApprovalId, ErrorBody, ErrorCode, ExecCommandParams, LostReason, ProcessId, ProcessState,
@@ -28,9 +29,38 @@ use codespace_store::{
 };
 use tokio::sync::{mpsc, watch};
 
+/// How long the records first wait before they ask the Runner again about an exit while the
+/// gateway claims it again, and the longest they wait (P1R-C04).
+const RECLAIMING_RETRY_START: Duration = Duration::from_millis(10);
+const RECLAIMING_RETRY_MAX: Duration = Duration::from_millis(500);
+
 /// The exits that a gateway's release hook hands to its process records.
 pub struct Exits(mpsc::UnboundedReceiver<ProcessId>);
 
+/// The Runner's status of `process_id`, which it reported exited (P1R-C04). The report can reach
+/// the gateway while it claims the Runner again, as when the Runner sends it right after the
+/// Claim's answer, and the request is then refused before it is sent: it is asked again once
+/// the gateway has claimed the Runner, which keeps the ended process until its end is recorded.
+async fn status_after_exit(
+    runner: &RuntimeBackend,
+    process_id: &ProcessId,
+) -> Result<RunnerProcessStatus, RunnerError> {
+    let mut retry = RECLAIMING_RETRY_START;
+    loop {
+        match runner.process_status(process_id).await {
+            Err(err)
+                if err
+                    .as_execution()
+                    .is_some_and(|body| body.code == ErrorCode::RunnerReconnecting) =>
+            {
+                tokio::time::sleep(retry).await;
+                retry = (retry * 2).min(RECLAIMING_RETRY_MAX);
+            }
+            status => return status,
+        }
+    }
+}
+
 /// The workspace leases of a gateway connected to an independent Runner (P1R-C04): each exit
 /// releases its process's lease, and a reconciliation holds a workspace again for a process
 /// that it found running. While a reconciliation runs, the exits are noted, so a process that
@@ -134,7 +164,7 @@ impl ProcessRecords {
         let Exits(mut exited) = exits;
         tokio::spawn(async move {
             while let Some(process_id) = exited.recv().await {
-                let recorded = match runner.process_status(&process_id).await {
+                let recorded = match status_after_exit(&runner, &process_id).await {
                     Ok(status) if status.state == ProcessState::Exited => following.ended(&status),
                     Ok(_) => false,
                     // The Runner reported the exit; how it ended is not known.