Link received payload objects into the core cache instead of copying

heliopause/src/export/transfer.rs

@@ -741,9 +741,11 @@ fn verify_cached_payload(
 
 /// Adopt a received cache payload object under its signed payload identity.
 ///
-/// The staged file is a complete companion container; its payload bytes are
-/// verified against the declared identity by the cache insertion path, so a
-/// corrupt or mislabelled object never enters the cache.
+/// The staged file is a complete, synced companion container; its payload bytes
+/// are verified against the declared identity by the cache insertion path, so a
+/// corrupt or mislabelled object never enters the cache. A canonical wrapper is
+/// linked into the cache as is; otherwise, or when the filesystem refuses the
+/// link, the slice is copied.
 fn install_payload_object(
     cache: &Path,
     id: helios_core::core::PayloadId,
@@ -758,7 +760,10 @@ fn install_payload_object(
             "received payload object length differs from its declared group",
         ));
     }
-    let object = root.store_slice(Some(id), &slice)?;
+    let object = match root.adopt_companion(id, &slice, &partial.display().to_string())? {
+        Some(object) => object,
+        None => root.store_slice(Some(id), &slice)?,
+    };
     if object.length() != length {
         return Err(invalid("received payload object has the wrong length"));
     }
@@ -1038,6 +1043,9 @@ fn receive_rank(
                         fs::Permissions::from_mode(0o444)
                     })?;
                 }
+                // The synced file may become a cache object itself: close the
+                // writable descriptor before it is published.
+                drop(file);
                 if info.payload.is_some() {
                     let id = helios_core::core::PayloadId::parse(
                         info.payload.as_deref().unwrap_or_default(),
@@ -1444,6 +1452,59 @@ mod tests {
         (output, id)
     }
 
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn received_canonical_wrapper_is_adopted_and_others_are_copied() {
+        use std::os::unix::fs::MetadataExt;
+        let directory = Directory::new();
+        let payload: Vec<u8> = (0..4096u32).map(|i| (i % 97) as u8).collect();
+        let (_, id) = cache_backed_output(&directory, &payload);
+        let sent = directory.0.join("sender-cache").join("blobs").join(id.file_name());
+        let cache = directory.0.join("cache");
+        helios_core::core::cache::Root::open(&cache).unwrap();
+
+        // The receiver writes the sender's object file under the cache's blobs directory.
+        let partial = cache.join("blobs").join(".rank-blob-canonical");
+        fs::copy(&sent, &partial).unwrap();
+        let object = install_payload_object(&cache, id, payload.len() as u64, &partial).unwrap();
+        assert_eq!(
+            fs::metadata(object.path()).unwrap().ino(),
+            fs::metadata(&partial).unwrap().ino(),
+            "a canonical received wrapper was copied instead of adopted"
+        );
+        drop(object);
+
+        // Another framing of the same payload is stored through its slice in a new file.
+        let other = directory.0.join("other-cache");
+        helios_core::core::cache::Root::open(&other).unwrap();
+        let framed = other.join("blobs").join(".rank-blob-framed");
+        let slice = helios_core::core::container::Container::open(&sent)
+            .unwrap()
+            .tensor(helios_core::core::container::COMPANION_TENSOR)
+            .unwrap();
+        let mut writer = helios_core::core::container::Writer::new(
+            fs::File::create(&framed).unwrap(),
+            helios_core::core::container::PAYLOAD_FORMAT,
+        )
+        .unwrap();
+        writer
+            .add(
+                helios_core::core::container::COMPANION_TENSOR,
+                slice.len(),
+                helios_core::core::container::Source::Slice(slice),
+            )
+            .unwrap();
+        writer
+            .add("extra.tail", 3, helios_core::core::container::Source::Bytes(vec![1, 2, 3]))
+            .unwrap();
+        writer.finish().unwrap();
+        let copied = install_payload_object(&other, id, payload.len() as u64, &framed).unwrap();
+        assert_ne!(
+            fs::metadata(copied.path()).unwrap().ino(),
+            fs::metadata(&framed).unwrap().ino()
+        );
+    }
+
     #[test]
     fn repeated_identical_payload_groups_share_one_object_and_round_trip() {
         let directory = Directory::new();

helios/src/core/cache.rs

@@ -45,8 +45,8 @@ thread_local! {
 }
 use crate::Result;
 use crate::core::container::{
-    COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, open_relative,
-    validate_relative_path, write_companion,
+    COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, link_companion,
+    open_relative, validate_relative_path, write_companion,
 };
 
 const MAX_NAME: usize = 128;
@@ -1503,11 +1503,7 @@ impl Root {
             )
             .into());
         }
-        if let Some(object) = self.payload(id)? {
-            if object.length() != slice.len() {
-                return Err(format!("cached object {id} has a different length").into());
-            }
-            verify_slice(object.slice(), id, None)?;
+        if let Some(object) = self.cached_slice(id, slice.len())? {
             return Ok(object);
         }
         let transaction = Transaction::begin(self)?;
@@ -1529,6 +1525,48 @@ impl Root {
         result
     }
 
+    /// Adopt the complete, durable companion wrapper file that holds `slice` as
+    /// a payload object without copying it.
+    ///
+    /// The opened file is linked into staging only when it is byte-for-byte
+    /// the wrapper `store_slice` would write; the linked file must then hash to
+    /// `declared`, and publication verifies the installed object again. `None`
+    /// means the file is not that canonical wrapper or the filesystem refuses
+    /// the link, and the caller stores its slice instead. A mismatched identity
+    /// is an error. An already cached identity returns the cached object. The
+    /// caller must have synced the file under `slice` and must not modify it
+    /// afterwards: it becomes the object. `label` names the payload in errors.
+    pub fn adopt_companion(
+        &self,
+        declared: PayloadId,
+        slice: &Slice,
+        label: &str,
+    ) -> Result<Option<PayloadObject>> {
+        let transaction = Transaction::begin(self)?;
+        let staged = transaction.staging_file("wrapped")?;
+        if !link_companion(slice, declared, &staged, label)? {
+            return Ok(None);
+        }
+        let cached = self.cached_slice(declared, slice.len());
+        if !matches!(cached, Ok(None)) {
+            let _ = fs::remove_file(&staged);
+            return cached;
+        }
+        self.publish_payload(&transaction, &staged, declared, slice.len()).map(Some)
+    }
+
+    /// A cached object for `id`, verified, when one is already installed.
+    fn cached_slice(&self, id: PayloadId, length: u64) -> Result<Option<PayloadObject>> {
+        let Some(object) = self.payload(id)? else {
+            return Ok(None);
+        };
+        if object.length() != length {
+            return Err(format!("cached object {id} has a different length").into());
+        }
+        verify_slice(object.slice(), id, None)?;
+        Ok(Some(object))
+    }
+
     /// Insert an optional retention source while distinguishing a corrupt
     /// source from publication failure. A missing/corrupt source yields
     /// `Ok(None)`; object-store or staging I/O errors are returned.
@@ -4855,6 +4893,74 @@ mod tests {
         assert_eq!(root.collect().unwrap().objects, 1);
     }
 
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn adopt_companion_links_canonical_wrappers_and_declines_others() {
+        use std::os::unix::fs::MetadataExt;
+        let temporary = Temporary::new();
+        let root = temporary.root();
+        let bytes: Vec<u8> = (0..2 * 1024 * 1024 + 77).map(|i| (i % 241) as u8).collect();
+        let raw = temporary.0.join("slice.bin");
+        fs::write(&raw, &bytes).unwrap();
+        let slice = Slice::new(Arc::new(File::open(&raw).unwrap()), 0, bytes.len() as u64).unwrap();
+        let id = PayloadId::of(&bytes);
+        let adopt = |declared: PayloadId, path: &Path| {
+            let slice = Container::open(path).unwrap().tensor(COMPANION_TENSOR).unwrap();
+            root.adopt_companion(declared, &slice, "test")
+        };
+        let staging_is_empty = || {
+            fs::read_dir(root.root.join("staging").join("wrapped"))
+                .unwrap()
+                .count()
+                == 0
+        };
+
+        // A framing that is not the canonical wrapper is left to store_slice.
+        let framed = temporary.0.join("framed.safetensors");
+        let mut writer = Writer::new(File::create(&framed).unwrap(), PAYLOAD_FORMAT).unwrap();
+        writer.add(COMPANION_TENSOR, bytes.len() as u64, Source::Slice(slice.clone())).unwrap();
+        writer.add("extra.tail", 3, Source::Bytes(vec![1, 2, 3])).unwrap();
+        writer.finish().unwrap();
+        assert!(adopt(id, &framed).unwrap().is_none());
+
+        // Same length as the canonical wrapper, but its header padding differs.
+        let padded = temporary.0.join("padded.safetensors");
+        write_companion(&mut File::create(&padded).unwrap(), &slice).unwrap();
+        let mut wrapper = fs::read(&padded).unwrap();
+        let header_end = 8 + u64::from_le_bytes(wrapper[..8].try_into().unwrap()) as usize;
+        assert_eq!(wrapper[header_end - 1], b' ');
+        wrapper[header_end - 1] = b'\n';
+        fs::write(&padded, &wrapper).unwrap();
+        assert!(adopt(id, &padded).unwrap().is_none());
+        assert!(staging_is_empty());
+
+        let received = temporary.0.join("received.safetensors");
+        write_companion(&mut File::create(&received).unwrap(), &slice).unwrap();
+        assert!(adopt(PayloadId::of(b"other"), &received).is_err());
+        assert!(root.payload(id).unwrap().is_none());
+        assert!(staging_is_empty());
+
+        // The canonical wrapper becomes the object itself: no copy is written.
+        let object = adopt(id, &received).unwrap().unwrap();
+        assert_eq!(object.id(), id);
+        assert_eq!(object.length(), bytes.len() as u64);
+        assert_eq!(
+            fs::metadata(object.path()).unwrap().ino(),
+            fs::metadata(&received).unwrap().ino()
+        );
+        fs::remove_file(&received).unwrap();
+        let mut copy = vec![0; bytes.len()];
+        object.slice().read_at(0, &mut copy).unwrap();
+        assert_eq!(copy, bytes);
+
+        // A second arrival of a cached identity returns the cached object.
+        let again_path = temporary.0.join("again.safetensors");
+        write_companion(&mut File::create(&again_path).unwrap(), &slice).unwrap();
+        let again = adopt(id, &again_path).unwrap().unwrap();
+        assert_eq!(again.path(), object.path());
+        assert!(staging_is_empty());
+    }
+
     #[test]
     #[cfg(unix)]
     fn legacy_import_streams_large_companions_and_publishes_once() {

helios/src/core/container.rs

@@ -491,6 +491,105 @@ pub(crate) fn companion_prefix(nbytes: u64) -> Result<Vec<u8>> {
     Ok(prefix)
 }
 
+/// Hard-link the file under `payload` to `to` when that file is exactly the
+/// deterministic companion wrapper of the payload, then verify the linked file
+/// against `id`. Returns false, creating no file, when the file is not that
+/// wrapper or the filesystem refuses the link; the caller then copies. An
+/// existing `to` is an error. If the linked file fails verification, the link
+/// is removed and the error returned; `label` names the payload in errors.
+#[cfg(target_os = "linux")]
+pub(crate) fn link_companion(payload: &Slice, id: PayloadId, to: &Path, label: &str) -> Result<bool> {
+    use std::os::unix::ffi::OsStrExt;
+    use std::os::unix::fs::MetadataExt;
+    use std::os::unix::io::AsRawFd;
+
+    // I/O failures are state errors; a linked file that is not the payload's wrapper
+    // or identity, and an unrepresentable path, are manifest errors.
+    let failed = |context: &str, error: io::Error| -> crate::Error {
+        format!("{context}: {error}").into()
+    };
+
+    let nbytes = payload.len();
+    let prefix = companion_prefix(nbytes)?;
+    let source = payload.file();
+    let metadata = source
+        .metadata()
+        .map_err(|e| failed("inspect payload source", e))?;
+    let wrapper_len = prefix.len() as u64 + nbytes;
+    if payload.offset() != prefix.len() as u64
+        || !metadata.is_file()
+        || metadata.len() != wrapper_len
+    {
+        return Ok(false);
+    }
+    let mut header = vec![0u8; prefix.len()];
+    // std `FileExt` keeps the short-read message that directory publication reports.
+    std::os::unix::fs::FileExt::read_exact_at(&**source, &mut header, 0)
+        .map_err(|e| failed("read payload wrapper", e))?;
+    if header != prefix {
+        return Ok(false);
+    }
+    if let Some(parent) = to.parent() {
+        std::fs::create_dir_all(parent).map_err(|e| failed("create payload directory", e))?;
+    }
+    // The descriptor link names the pinned file itself, not a path that could
+    // have been replaced since the source was opened.
+    let from = std::ffi::CString::new(format!("/proc/self/fd/{}", source.as_raw_fd()))
+        .map_err(|_| invalid("payload source descriptor path"))?;
+    let destination = std::ffi::CString::new(to.as_os_str().as_bytes())
+        .map_err(|_| invalid(format!("payload path {label:?} contains NUL")))?;
+    // SAFETY: Both arguments are NUL-terminated paths that outlive the call.
+    let status = unsafe {
+        libc::linkat(
+            libc::AT_FDCWD,
+            from.as_ptr(),
+            libc::AT_FDCWD,
+            destination.as_ptr(),
+            libc::AT_SYMLINK_FOLLOW,
+        )
+    };
+    if status != 0 {
+        let error = io::Error::last_os_error();
+        // An existing destination is a publication conflict, as for a copy.
+        if error.kind() == io::ErrorKind::AlreadyExists {
+            return Err(failed("link companion", error));
+        }
+        return Ok(false);
+    }
+    let result = (|| -> Result<()> {
+        let linked = File::open(to).map_err(|e| failed("open linked companion", e))?;
+        let linked_metadata = linked
+            .metadata()
+            .map_err(|e| failed("inspect linked companion", e))?;
+        if (linked_metadata.dev(), linked_metadata.ino()) != (metadata.dev(), metadata.ino()) {
+            return Err(invalid("linked companion is not the payload source file"));
+        }
+        // Verify the file published under `to`, not the pre-link source reads.
+        let mut header = vec![0u8; prefix.len()];
+        std::os::unix::fs::FileExt::read_exact_at(&linked, &mut header, 0)
+            .map_err(|e| failed("read linked companion", e))?;
+        if header != prefix || linked_metadata.len() != wrapper_len {
+            return Err(invalid("linked companion is not the payload wrapper"));
+        }
+        let slice = Slice::new(Arc::new(linked), prefix.len() as u64, nbytes)?;
+        if slice.digest()? != id {
+            return Err(invalid(format!(
+                "payload {label} bytes do not match the assigned identity"
+            )));
+        }
+        Ok(())
+    })();
+    if result.is_err() {
+        let _ = std::fs::remove_file(to);
+    }
+    result.map(|()| true)
+}
+
+#[cfg(not(target_os = "linux"))]
+pub(crate) fn link_companion(_payload: &Slice, _id: PayloadId, _to: &Path, _label: &str) -> Result<bool> {
+    Ok(false)
+}
+
 fn write_source<W: Write>(
     output: &mut W,
     name: &str,