Preserve live prepared-lower materializations during recovery

crates/lazaret-prepared-lower/src/store.rs

@@ -47,6 +47,7 @@ pub struct PrepareRequest {
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum PrepareRecovery {
     ClearedUnmaterializedIntent,
+    PreservedLiveOwnedStage,
     RemovedOwnedStage,
     PublishedCompletedStage,
 }
@@ -267,6 +268,8 @@ impl PreparedLowerStore {
             )));
         }
         let identities = Identities::new(request).map_err(PreparedLowerError::layer_validation)?;
+        let (stage_relative, stage, owned, device, inode, _stage_lock) = {
+            let _physical_lock = self.physical_lock()?;
             if let Some(existing) = self
                 .available_artifact(identities.artifact_id)
                 .map_err(PreparedLowerError::catalog)?
@@ -293,6 +296,22 @@ impl PreparedLowerStore {
             };
             let (device, inode) =
                 file_identity(&stage).map_err(PreparedLowerError::operating_system)?;
+            let stage_lock = match try_exclusive_flock(
+                stage
+                    .try_clone()
+                    .map_err(PreparedLowerError::operating_system)?,
+            ) {
+                Ok(lock) => lock,
+                Err(error) => {
+                    if self
+                        .remove_exact_path(&stage_relative, device, inode)
+                        .is_ok()
+                    {
+                        let _ = self.catalog.remove_prepare(&intent);
+                    }
+                    return Err(PreparedLowerError::locking(error));
+                }
+            };
             let owned = PrepareOperation {
                 stage_state: PrepareStageState::Owned,
                 device: Some(device),
@@ -308,6 +327,8 @@ impl PreparedLowerStore {
                 }
                 return Err(PreparedLowerError::catalog(error));
             }
+            (stage_relative, stage, owned, device, inode, stage_lock)
+        };
         let materialization_map = if materialize_logical_owners {
             logical_id_map(&request.id_map)
         } else {
@@ -320,7 +341,8 @@ impl PreparedLowerStore {
                 })
                 .and_then(|()| sync_tree(&stage))
         {
-            if self
+            if let Ok(_physical_lock) = self.physical_lock()
+                && self
                     .remove_exact_path(&stage_relative, device, inode)
                     .is_ok()
             {
@@ -343,11 +365,16 @@ impl PreparedLowerStore {
         &self,
         operation: &PrepareOperation,
     ) -> PreparedLowerResult<PrepareRecovery> {
-        self.recover_prepare_inner(operation)
+        let physical_lock = self.physical_lock()?;
+        self.recover_prepare_inner(operation, &physical_lock)
             .map_err(PreparedLowerError::recovery)
     }
 
-    fn recover_prepare_inner(&self, operation: &PrepareOperation) -> Result<PrepareRecovery> {
+    fn recover_prepare_inner(
+        &self,
+        operation: &PrepareOperation,
+        _physical_lock: &ExclusiveFileLock,
+    ) -> Result<PrepareRecovery> {
         let expected_stage = stage_path(operation.operation_id);
         if operation.stage_path.as_bytes() != path_bytes(&expected_stage) {
             bail!("prepare operation stage path is not canonical");
@@ -363,7 +390,15 @@ impl PreparedLowerStore {
             PrepareStageState::Owned => {
                 let device = operation.device.context("owned prepare lacks device")?;
                 let inode = operation.inode.context("owned prepare lacks inode")?;
-                if self.open_relative(&expected_stage)?.is_some() {
+                if let Some(stage) = self.open_relative(&expected_stage)? {
+                    require_identity(&stage, device, inode)?;
+                    let _stage_lock = match try_exclusive_flock(stage) {
+                        Ok(lock) => lock,
+                        Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
+                            return Ok(PrepareRecovery::PreservedLiveOwnedStage);
+                        }
+                        Err(error) => return Err(error.into()),
+                    };
                     self.remove_exact_path(&expected_stage, device, inode)?;
                     self.catalog.remove_prepare(operation)?;
                     return Ok(PrepareRecovery::RemovedOwnedStage);
@@ -391,16 +426,20 @@ impl PreparedLowerStore {
     pub fn recover_preparations(
         &self,
     ) -> PreparedLowerResult<Vec<(PrepareOperationId, PrepareRecovery)>> {
-        self.recover_preparations_inner()
+        let physical_lock = self.physical_lock()?;
+        self.recover_preparations_inner(&physical_lock)
             .map_err(PreparedLowerError::recovery)
     }
 
-    fn recover_preparations_inner(&self) -> Result<Vec<(PrepareOperationId, PrepareRecovery)>> {
+    fn recover_preparations_inner(
+        &self,
+        physical_lock: &ExclusiveFileLock,
+    ) -> Result<Vec<(PrepareOperationId, PrepareRecovery)>> {
         let operations = self.catalog.snapshot()?.prepare_operations;
         operations
             .iter()
             .map(|operation| {
-                self.recover_prepare_inner(operation)
+                self.recover_prepare_inner(operation, physical_lock)
                     .map(|outcome| (operation.operation_id, outcome))
             })
             .collect()

crates/lazaret-prepared-lower/tests/prepared_lower.rs

@@ -5,6 +5,7 @@ use std::os::unix::fs::{MetadataExt, PermissionsExt};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Barrier};
 use std::thread;
+use std::time::{Duration, Instant};
 
 use flate2::Compression;
 use flate2::write::GzEncoder;
@@ -75,6 +76,15 @@ impl Fixture {
             artifact.path.as_bytes().to_vec(),
         )))
     }
+
+    fn independent_store(&self) -> (Catalog, PreparedLowerStore) {
+        let catalog = Catalog::new(
+            self.root.join("catalog.redb"),
+            self.root.join("storage-catalog.lock"),
+        );
+        let store = PreparedLowerStore::initialize(&self.root, catalog.clone()).unwrap();
+        (catalog, store)
+    }
 }
 
 #[test]

crates/lazaret-prepared-lower/tests/prepared_lower/publication.rs

@@ -75,6 +75,115 @@ fn concurrent_publication_reuses_one_immutable_artifact() {
     assert!(snapshot.prepare_operations.is_empty());
 }
 
+#[test]
+fn startup_recovery_preserves_a_stage_owned_by_a_live_materializer() {
+    let fixture = Fixture::new();
+    let first_layer = layer(&[Entry {
+        path: "large",
+        kind: EntryType::Regular,
+        data: &vec![0u8; 16 * 1024 * 1024],
+        target: None,
+        mode: 0o644,
+    }]);
+    let second_layer = layer(&[Entry {
+        path: "marker",
+        kind: EntryType::Regular,
+        data: b"complete",
+        target: None,
+        mode: 0o644,
+    }]);
+    let first_request = fixture.request(
+        42,
+        &[(first_layer, OCI_TAR), (second_layer.clone(), OCI_TAR)],
+    );
+    let second_request = PrepareRequest {
+        operation_id: PrepareOperationId::new([43; 32]),
+        ..first_request.clone()
+    };
+    let second_blob = first_request
+        .layout_path
+        .join("blobs/sha256")
+        .join(hex(&Sha256::digest(&second_layer)));
+    let (_, first_store) = fixture.independent_store();
+    let (second_catalog, second_store) = fixture.independent_store();
+    let first_handle = thread::spawn(move || first_store.prepare(&first_request));
+
+    let deadline = Instant::now() + Duration::from_secs(5);
+    let operation = loop {
+        if let Some(operation) = second_catalog
+            .snapshot()
+            .unwrap()
+            .prepare_operations
+            .into_iter()
+            .find(|operation| operation.stage_state == PrepareStageState::Owned)
+        {
+            break operation;
+        }
+        assert!(
+            !first_handle.is_finished(),
+            "materialization finished before synchronization"
+        );
+        assert!(
+            Instant::now() < deadline,
+            "materialization did not publish stage ownership"
+        );
+        thread::yield_now();
+    };
+
+    std::fs::remove_file(&second_blob).unwrap();
+    rustix::fs::mknodat(
+        rustix::fs::CWD,
+        &second_blob,
+        rustix::fs::FileType::Fifo,
+        rustix::fs::Mode::from_raw_mode(0o600),
+        0,
+    )
+    .unwrap();
+    thread::sleep(Duration::from_millis(25));
+    assert!(
+        !first_handle.is_finished(),
+        "materializer passed the synchronization FIFO"
+    );
+
+    assert_eq!(
+        second_store.recover_preparations().unwrap(),
+        vec![(
+            operation.operation_id,
+            PrepareRecovery::PreservedLiveOwnedStage,
+        )]
+    );
+
+    let writer_path = second_blob.clone();
+    let writer_bytes = second_layer.clone();
+    let writer = thread::spawn(move || {
+        std::fs::OpenOptions::new()
+            .write(true)
+            .open(writer_path)
+            .unwrap()
+            .write_all(&writer_bytes)
+            .unwrap();
+    });
+    let first = first_handle.join().unwrap().unwrap();
+    writer.join().unwrap();
+    std::fs::remove_file(&second_blob).unwrap();
+    std::fs::write(&second_blob, &second_layer).unwrap();
+
+    let second = second_store.prepare(&second_request).unwrap();
+    assert_eq!(first, second);
+    let snapshot = second_catalog.snapshot().unwrap();
+    assert_eq!(snapshot.prepared_artifacts, vec![first]);
+    assert!(snapshot.prepare_operations.is_empty());
+    assert!(
+        fixture
+            .root
+            .join("stages")
+            .read_dir()
+            .unwrap()
+            .next()
+            .is_none()
+    );
+}
+
 #[test]
 fn publication_during_reclaim_preserves_replacement_and_deletes_old_inode() {
     let fixture = Fixture::new();