test: verify real orphan pin repair across owned restore process death

crates/worker-runtime/src/restore_checkpoints.rs

@@ -25,6 +25,7 @@ pub const MAX_RESTORE_CHECKPOINT_ACK_BYTES: usize = 4096;
 pub enum RestoreCheckpointStage {
     PreparedNotCommitted,
     CommittedPreterminal,
+    ActivatedPreterminal,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]

crates/worker-runtime/src/runtime.rs

@@ -2871,6 +2871,12 @@ impl Runtime {
                                 sanitize_worker_create_failure_message(&error),
                             ));
                         }
+                        #[cfg(feature = "restore-checkpoints")]
+                        self.test_restore_checkpoint(
+                            worker_ref,
+                            operation_id,
+                            crate::restore_checkpoints::RestoreCheckpointStage::ActivatedPreterminal,
+                        );
                         if let Err(error) =
                             self.activate_restore_candidate_context(worker_ref, &candidate_context)
                         {

crates/worker-runtime/src/runtime/restore_checkpoint.rs

@@ -75,7 +75,8 @@ impl Runtime {
         };
         let expected = match config.stage {
             RestoreCheckpointStage::PreparedNotCommitted => RestoreCommitFact::NotCommitted,
-            RestoreCheckpointStage::CommittedPreterminal => RestoreCommitFact::Committed,
+            RestoreCheckpointStage::CommittedPreterminal
+            | RestoreCheckpointStage::ActivatedPreterminal => RestoreCommitFact::Committed,
         };
         if config.worker_id != worker_ref.worker_id
             || pending.operation_id != operation_id

crates/worker-runtime/src/runtime/restore_checkpoint_tests.rs

@@ -14,6 +14,7 @@ fn checkpoint_stages_ack_canonical_decision_and_eof_resumes_same_operation_once(
         (RestoreCheckpointStage::PreparedNotCommitted, "release"),
         (RestoreCheckpointStage::CommittedPreterminal, "eof"),
         (RestoreCheckpointStage::CommittedPreterminal, "timeout"),
+        (RestoreCheckpointStage::ActivatedPreterminal, "release"),
     ] {
         let dir = tempfile::tempdir().unwrap();
         let backend = Arc::new(TestExecutionBackend::default());
@@ -79,14 +80,19 @@ fn checkpoint_stages_ack_canonical_decision_and_eof_resumes_same_operation_once(
             match stage {
                 RestoreCheckpointStage::PreparedNotCommitted =>
                     RestoreCheckpointCommit::NotCommitted,
-                RestoreCheckpointStage::CommittedPreterminal => RestoreCheckpointCommit::Committed,
+                RestoreCheckpointStage::CommittedPreterminal
+                | RestoreCheckpointStage::ActivatedPreterminal => RestoreCheckpointCommit::Committed,
             }
         );
         assert_eq!(
             backend.restore_operation_ids.lock().unwrap().as_slice(),
             &[ack.operation_id]
         );
-        assert!(backend.activation_operation_ids.lock().unwrap().is_empty());
+        if stage == RestoreCheckpointStage::ActivatedPreterminal {
+            assert_eq!(backend.activation_operation_ids.lock().unwrap().as_slice(), &[ack.operation_id]);
+        } else {
+            assert!(backend.activation_operation_ids.lock().unwrap().is_empty());
+        }
         assert!(received.try_recv().is_err());
         let canonical = runtime_store(&runtime)
             .read_durable_worker_record(&worker.worker_ref)
@@ -255,6 +261,8 @@ fn checkpoint_verification_requires_canonical_worker_operation_commit_and_absent
             .restore_checkpoint_ack(&config, &worker.worker_ref, operation)
             .is_err()
     );
+    config.stage = RestoreCheckpointStage::ActivatedPreterminal;
+    assert!(runtime.restore_checkpoint_ack(&config, &worker.worker_ref, operation).is_err());
     config.stage = RestoreCheckpointStage::PreparedNotCommitted;
     runtime_store(&runtime).fail_next_worker_reconcile_sync();
     assert!(

crates/worker-runtime/src/worker_backend_restore_candidate_tests.rs

@@ -557,6 +557,40 @@ mod restore_candidate {
         }
     }
 
+    #[test]
+    #[serial_test::serial(worker_allocation)]
+    fn restore_run_reclamation_preserves_live_allocation_owner() {
+        let fixture = Fixture::no_model();
+        let factory = fixture.backend.factory.clone();
+        let request = fixture.request.clone();
+        let run = factory
+            .worker_restore_run_dir(&request.worker_ref, request.operation_id)
+            .unwrap();
+        fs::create_dir_all(&run).unwrap();
+        fs::write(run.join("retained-owner"), b"live allocation").unwrap();
+        let name = ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&request.worker_ref);
+        let allocation = worker::runtime::worker_allocation::install_top_level(
+            name.clone(), std::process::id(), run.join("worker.sock"),
+            Vec::new(), fixture.segment_id,
+        ).unwrap();
+        let failure = fixture.backend.run_on_adapter_runtime(async move {
+            match factory.prepare_restore(request).await {
+                Err(failure) => Ok(failure),
+                Ok(_) => Err("live allocation unexpectedly reclaimed".into()),
+            }
+        }).unwrap();
+        assert_eq!(failure.stage, RuntimeRestoreStage::Controller);
+        assert_eq!(failure.cause, "restore run directory still has an allocation owner");
+        let table = worker::runtime::worker_allocation::LockFileGuard::open(
+            &worker::runtime::worker_allocation::default_allocation_path().unwrap(),
+        ).unwrap();
+        assert_eq!(table.data().find(&name).unwrap().pid, std::process::id());
+        assert_eq!(fs::read(run.join("retained-owner")).unwrap(), b"live allocation");
+        drop(table);
+        drop(allocation);
+        fixture.assert_quiet();
+    }
+
     #[test]
     #[serial_test::serial(worker_allocation)]
     fn session_availability_requires_attachable_current_controller() {