crank/crates/crank-artifacts/tests/artifact_store_io.rs

use std::{
    fs,
    io::{self, Read},
    os::unix::fs::PermissionsExt,
    path::PathBuf,
    process::{Child, Command, ExitStatus},
    sync::atomic::{AtomicU64, Ordering},
    thread,
    time::Duration,
};

#[cfg(debug_assertions)]
use crank_artifacts::test_support::{FaultAction, set_checkpoint};
use crank_artifacts::{ArtifactError, ArtifactRef, ArtifactStore, StaleTemp};
use sha2::{Digest, Sha256};
#[cfg(debug_assertions)]
use std::sync::{Mutex, OnceLock};

static NEXT_ROOT: AtomicU64 = AtomicU64::new(0);

#[cfg(debug_assertions)]
fn fault_guard() -> std::sync::MutexGuard<'static, ()> {
    static GUARD: OnceLock<Mutex<()>> = OnceLock::new();
    GUARD.get_or_init(|| Mutex::new(())).lock().unwrap()
}

struct TestRoot(PathBuf);

impl TestRoot {
    fn new(name: &str) -> Self {
        let path = std::env::temp_dir().join(format!(
            "crank-artifacts-{name}-{}-{}",
            std::process::id(),
            NEXT_ROOT.fetch_add(1, Ordering::Relaxed)
        ));
        fs::create_dir(&path).unwrap();
        fs::set_permissions(&path, fs::Permissions::from_mode(0o700)).unwrap();
        Self(path)
    }
}

impl Drop for TestRoot {
    fn drop(&mut self) {
        let _ = fs::set_permissions(&self.0, fs::Permissions::from_mode(0o700));
        let _ = fs::remove_dir_all(&self.0);
    }
}

fn wait_for_child(mut child: Child, timeout: Duration) -> ExitStatus {
    let deadline = std::time::Instant::now() + timeout;
    loop {
        if let Some(status) = child.try_wait().unwrap() {
            return status;
        }
        if std::time::Instant::now() >= deadline {
            let _ = child.kill();
            let _ = child.wait();
            panic!("child process did not finish within {timeout:?}");
        }
        thread::sleep(Duration::from_millis(10));
    }
}

#[test]
fn fifo_entries_fail_without_blocking_reads() {
    let root = TestRoot::new("fifo");
    let expected =
        ArtifactRef::from_digest_hex(&format!("{:x}", Sha256::digest(b"fifo bytes"))).unwrap();
    let sha = root.0.join("sha256");
    let shard = sha.join(&expected.digest_hex()[..2]);
    fs::create_dir(&sha).unwrap();
    fs::create_dir(&shard).unwrap();
    fs::set_permissions(&sha, fs::Permissions::from_mode(0o700)).unwrap();
    fs::set_permissions(&shard, fs::Permissions::from_mode(0o700)).unwrap();
    let blob_fifo = shard.join(expected.digest_hex());
    let legacy = root.0.join("legacy");
    fs::create_dir(&legacy).unwrap();
    fs::set_permissions(&legacy, fs::Permissions::from_mode(0o700)).unwrap();
    let legacy_fifo = legacy.join("source");
    for path in [&blob_fifo, &legacy_fifo] {
        let path = std::ffi::CString::new(path.as_os_str().as_encoded_bytes()).unwrap();
        assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o400) }, 0);
    }

    let executable = std::env::current_exe().unwrap();
    for kind in ["blob", "legacy"] {
        let child = Command::new(&executable)
            .args(["--exact", "fifo_read_child", "--nocapture"])
            .env("CRANK_ARTIFACTS_FIFO_ROOT", &root.0)
            .env("CRANK_ARTIFACTS_FIFO_KIND", kind)
            .env("CRANK_ARTIFACTS_FIFO_REF", expected.as_str())
            .spawn()
            .unwrap();
        assert!(wait_for_child(child, Duration::from_secs(2)).success());
    }
}

#[test]
fn fifo_read_child() {
    let Some(root) = std::env::var_os("CRANK_ARTIFACTS_FIFO_ROOT") else {
        return;
    };
    let root = PathBuf::from(root);
    let kind = std::env::var("CRANK_ARTIFACTS_FIFO_KIND").unwrap();
    let expected = ArtifactRef::parse(&std::env::var("CRANK_ARTIFACTS_FIFO_REF").unwrap()).unwrap();
    let store = ArtifactStore::open(&root).unwrap();
    let result = if kind == "blob" {
        store.read(&expected)
    } else {
        store
            .with_legacy_root(root.join("legacy"))
            .read_legacy_file_url(
                &format!("file://{}/legacy/source", root.display()),
                &expected,
            )
    };
    assert_eq!(result.unwrap_err(), ArtifactError::Integrity);
}

struct ShortReader {
    bytes: Vec<u8>,
    offset: usize,
    interrupted: bool,
}

struct FailingReader;

struct HousekeepingReader {
    store: ArtifactStore,
    candidate: Option<StaleTemp>,
    bytes: io::Cursor<Vec<u8>>,
}

impl Read for FailingReader {
    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
        Err(io::Error::other("private reader detail"))
    }
}

impl Read for HousekeepingReader {
    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
        if let Some(candidate) = self.candidate.take() {
            self.store.delete_stale_temp(candidate).unwrap();
        }
        self.bytes.read(buffer)
    }
}

impl Read for ShortReader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if !self.interrupted {
            self.interrupted = true;
            return Err(io::Error::from(io::ErrorKind::Interrupted));
        }
        if self.offset == self.bytes.len() {
            return Ok(0);
        }
        let amount = 1.min(buf.len()).min(self.bytes.len() - self.offset);
        buf[..amount].copy_from_slice(&self.bytes[self.offset..self.offset + amount]);
        self.offset += amount;
        Ok(amount)
    }
}

#[test]
fn streaming_input_handles_short_reads_and_interruption() {
    let root = TestRoot::new("stream");
    let store = ArtifactStore::open(&root.0).unwrap();
    let mut reader = ShortReader {
        bytes: b"short chunk source".to_vec(),
        offset: 0,
        interrupted: false,
    };
    let stored = store.put_reader(&mut reader).unwrap();
    assert_eq!(
        store.read(&stored.artifact_ref).unwrap(),
        b"short chunk source"
    );

    assert_eq!(
        store.put_reader(&mut FailingReader).unwrap_err(),
        ArtifactError::Storage
    );
    for error in [
        ArtifactError::InvalidReference,
        ArtifactError::Integrity,
        ArtifactError::Storage,
        ArtifactError::UnsafeRoot,
    ] {
        let diagnostic = format!("{error} {error:?}");
        assert!(!diagnostic.contains("private reader detail"));
        assert!(!diagnostic.contains(std::env::temp_dir().to_string_lossy().as_ref()));
    }
}

#[test]
fn user_reader_can_run_housekeeping_without_self_deadlock() {
    use std::sync::mpsc;

    let root = TestRoot::new("reader-housekeeping");
    let store = ArtifactStore::open(&root.0).unwrap();
    let stored = store.put(b"reader housekeeping shard").unwrap();
    let temp = root
        .0
        .join("sha256")
        .join(&stored.artifact_ref.digest_hex()[..2])
        .join(".crank-artifact-tmp-v1-00000000000000000000000000000000-1-10");
    fs::write(&temp, b"partial").unwrap();
    let (_, candidates) = store.scan_stale_temps(Duration::ZERO, 128, 1).unwrap();
    let mut reader = HousekeepingReader {
        store: store.clone(),
        candidate: candidates.into_iter().next(),
        bytes: io::Cursor::new(b"reader bytes".to_vec()),
    };
    let writer_store = store.clone();
    let (sender, receiver) = mpsc::channel();
    thread::spawn(move || {
        sender.send(writer_store.put_reader(&mut reader)).unwrap();
    });

    let stored = receiver
        .recv_timeout(Duration::from_secs(2))
        .expect("put_reader deadlocked while its reader ran housekeeping")
        .unwrap();
    assert_eq!(store.read(&stored.artifact_ref).unwrap(), b"reader bytes");
    assert!(!temp.exists());
}

#[test]
#[cfg(debug_assertions)]
fn forked_crash_checkpoints_publish_only_complete_or_absent_blobs() {
    let _guard = fault_guard();
    if let Some(root) = std::env::var_os("CRANK_ARTIFACTS_CHILD_ROOT") {
        set_checkpoint(
            std::env::var("CRANK_ARTIFACTS_TEST_STAGE").unwrap(),
            FaultAction::Exit,
        );
        let store = ArtifactStore::open(PathBuf::from(root)).unwrap();
        let _ = store.put(b"crash-consistent bytes");
        panic!("checkpoint did not terminate the child");
    }

    for stage in ["write", "file_fsync", "publish", "directory_fsync"] {
        let root = TestRoot::new(stage);
        let status = Command::new(std::env::current_exe().unwrap())
            .args([
                "--exact",
                "forked_crash_checkpoints_publish_only_complete_or_absent_blobs",
                "--nocapture",
            ])
            .env("CRANK_ARTIFACTS_CHILD_ROOT", &root.0)
            .env("CRANK_ARTIFACTS_TEST_STAGE", stage)
            .status()
            .unwrap();
        assert_eq!(status.code(), Some(86), "stage={stage}");

        let store = ArtifactStore::open(&root.0).unwrap();
        let expected = ArtifactRef::from_digest_hex(&format!(
            "{:x}",
            Sha256::digest(b"crash-consistent bytes")
        ))
        .unwrap();
        match store.read(&expected) {
            Ok(bytes) => assert_eq!(bytes, b"crash-consistent bytes", "stage={stage}"),
            Err(ArtifactError::NotFound) => {}
            Err(error) => panic!("stage={stage} exposed an invalid final blob: {error}"),
        }
        let stored = store.put(b"crash-consistent bytes").unwrap();
        assert_eq!(stored.artifact_ref, expected, "stage={stage}");
        assert_eq!(
            store.read(&expected).unwrap(),
            b"crash-consistent bytes",
            "stage={stage}"
        );

        let (_, candidates) = store.scan_stale_temps(Duration::ZERO, 128, 16).unwrap();
        for candidate in candidates {
            store.delete_stale_temp(candidate).unwrap();
        }
    }
}

#[test]
#[cfg(debug_assertions)]
fn disk_full_fault_seams_leave_a_retryable_store() {
    let _guard = fault_guard();
    if let Some(root) = std::env::var_os("CRANK_ARTIFACTS_CHILD_ROOT") {
        set_checkpoint(
            std::env::var("CRANK_ARTIFACTS_TEST_STAGE").unwrap(),
            FaultAction::Fail,
        );
        let store = ArtifactStore::open(PathBuf::from(root)).unwrap();
        assert_eq!(
            store.put(b"disk-full seam bytes").unwrap_err(),
            ArtifactError::Storage
        );
        return;
    }

    for stage in ["write", "file_fsync", "publish", "directory_fsync"] {
        let root = TestRoot::new(&format!("full-{stage}"));
        let status = Command::new(std::env::current_exe().unwrap())
            .args([
                "--exact",
                "disk_full_fault_seams_leave_a_retryable_store",
                "--nocapture",
            ])
            .env("CRANK_ARTIFACTS_CHILD_ROOT", &root.0)
            .env("CRANK_ARTIFACTS_TEST_STAGE", stage)
            .status()
            .unwrap();
        assert!(status.success(), "stage={stage}");

        let store = ArtifactStore::open(&root.0).unwrap();
        let expected =
            ArtifactRef::from_digest_hex(&format!("{:x}", Sha256::digest(b"disk-full seam bytes")))
                .unwrap();
        match store.read(&expected) {
            Ok(bytes) => assert_eq!(bytes, b"disk-full seam bytes", "stage={stage}"),
            Err(ArtifactError::NotFound) => {}
            Err(error) => panic!("stage={stage} exposed an invalid final blob: {error}"),
        }
        assert_eq!(
            store.put(b"disk-full seam bytes").unwrap().artifact_ref,
            expected,
            "stage={stage}"
        );
        assert_eq!(
            store.read(&expected).unwrap(),
            b"disk-full seam bytes",
            "stage={stage}"
        );
    }
}

#[test]
fn real_process_write_limit_leaves_no_partial_final_blob() {
    if let Some(root) = std::env::var_os("CRANK_ARTIFACTS_CHILD_ROOT") {
        // Keep the test's host filesystem untouched while making the actual
        // write syscall fail as it would for a quota/disk-full condition.
        unsafe {
            assert_ne!(libc::signal(libc::SIGXFSZ, libc::SIG_IGN), libc::SIG_ERR);
            let limit = libc::rlimit {
                rlim_cur: 1,
                rlim_max: 1,
            };
            assert_eq!(libc::setrlimit(libc::RLIMIT_FSIZE, &limit), 0);
        }
        let store = ArtifactStore::open(PathBuf::from(root)).unwrap();
        assert_eq!(
            store.put(b"real write limit bytes").unwrap_err(),
            ArtifactError::Storage
        );
        return;
    }

    let root = TestRoot::new("real-write-limit");
    let status = Command::new(std::env::current_exe().unwrap())
        .args([
            "--exact",
            "real_process_write_limit_leaves_no_partial_final_blob",
            "--nocapture",
        ])
        .env("CRANK_ARTIFACTS_CHILD_ROOT", &root.0)
        .status()
        .unwrap();
    assert!(status.success());

    let store = ArtifactStore::open(&root.0).unwrap();
    let expected =
        ArtifactRef::from_digest_hex(&format!("{:x}", Sha256::digest(b"real write limit bytes")))
            .unwrap();
    assert_eq!(store.read(&expected).unwrap_err(), ArtifactError::NotFound);
    assert_eq!(
        store.put(b"real write limit bytes").unwrap().artifact_ref,
        expected
    );
    assert_eq!(store.read(&expected).unwrap(), b"real write limit bytes");
}