powercut 0.1.0

//! A dependency-free scratch directory.
//!
//! `powercut` needs somewhere to materialise crash states, and pulling in a
//! temp-file crate for `create_dir_all` plus a `Drop` impl is not worth a
//! dependency in a crate whose pitch includes "zero dependencies".

use std::io;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

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

/// A uniquely named directory under the system temp dir, removed on drop.
///
/// Set the environment variable `POWERCUT_KEEP=1` to leave every scratch
/// directory behind for inspection; the path is printed to stderr when it is
/// kept, because a directory you cannot find is not a debugging aid.
///
/// ```
/// use powercut::ScratchDir;
///
/// let path;
/// {
///     let dir = ScratchDir::new("doc-scratch").unwrap();
///     path = dir.path().to_path_buf();
///     std::fs::write(path.join("x"), b"hi").unwrap();
///     assert!(path.exists());
/// }
/// assert!(!path.exists(), "dropped scratch dirs clean up after themselves");
/// ```
#[derive(Debug)]
pub struct ScratchDir {
    path: PathBuf,
    keep: bool,
}

impl ScratchDir {
    /// Create a fresh scratch directory whose name carries `tag`.
    ///
    /// # Errors
    ///
    /// Returns the underlying [`io::Error`] if the directory cannot be created
    /// after a few attempts at a unique name.
    ///
    /// ```
    /// # use powercut::ScratchDir;
    /// let dir = ScratchDir::new("my-case").unwrap();
    /// assert!(dir.path().is_dir());
    /// assert!(dir.path().to_string_lossy().contains("my-case"));
    /// ```
    pub fn new(tag: &str) -> io::Result<Self> {
        let base = std::env::temp_dir();
        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or_default();
        let mut last_err = None;
        for attempt in 0..8u64 {
            let n = COUNTER.fetch_add(1, Ordering::SeqCst);
            let path = base.join(format!(
                "powercut-{}-{}-{nanos:x}-{n}-{attempt}",
                sanitise(tag),
                std::process::id(),
            ));
            match std::fs::create_dir(&path) {
                Ok(()) => {
                    return Ok(Self {
                        path,
                        keep: std::env::var_os("POWERCUT_KEEP").is_some(),
                    })
                }
                Err(e) => last_err = Some(e),
            }
        }
        Err(last_err.unwrap_or_else(|| io::Error::other("could not create a scratch directory")))
    }

    /// The directory's path.
    ///
    /// ```
    /// # use powercut::ScratchDir;
    /// let dir = ScratchDir::new("path-demo").unwrap();
    /// assert!(dir.path().is_absolute());
    /// ```
    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Give up ownership: the directory stays on disk and is returned.
    ///
    /// ```
    /// # use powercut::ScratchDir;
    /// let kept = ScratchDir::new("kept").unwrap().keep();
    /// assert!(kept.exists());
    /// std::fs::remove_dir_all(&kept).unwrap();
    /// ```
    #[must_use]
    pub fn keep(mut self) -> PathBuf {
        self.keep = true;
        self.path.clone()
    }
}

impl Drop for ScratchDir {
    fn drop(&mut self) {
        if self.keep {
            eprintln!("powercut: keeping {}", self.path.display());
        } else {
            let _ = std::fs::remove_dir_all(&self.path);
        }
    }
}

fn sanitise(tag: &str) -> String {
    tag.chars()
        .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
        .collect()
}

/// Copy a regular file, creating parent directories as needed, and return the
/// number of bytes copied. `limit` truncates the copy to at most that length.
pub(crate) fn copy_file(src: &Path, dst: &Path, limit: Option<u64>) -> io::Result<u64> {
    if let Some(parent) = dst.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let copied = std::fs::copy(src, dst)?;
    match limit {
        Some(limit) if limit < copied => {
            let file = std::fs::OpenOptions::new().write(true).open(dst)?;
            file.set_len(limit)?;
            Ok(limit)
        }
        _ => Ok(copied),
    }
}

/// Every regular file under `root`, as paths relative to `root`, sorted.
///
/// Symlinks are skipped and reported by the caller: resolving them would copy
/// bytes from outside the store, and following them blindly is how a harness
/// starts measuring the wrong file.
pub(crate) fn walk_files(root: &Path) -> io::Result<(Vec<PathBuf>, Vec<PathBuf>)> {
    let mut files = Vec::new();
    let mut skipped = Vec::new();
    walk_into(root, root, &mut files, &mut skipped)?;
    files.sort();
    skipped.sort();
    Ok((files, skipped))
}

fn walk_into(
    root: &Path,
    dir: &Path,
    files: &mut Vec<PathBuf>,
    skipped: &mut Vec<PathBuf>,
) -> io::Result<()> {
    for entry in std::fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        let meta = entry.metadata()?;
        let relative = path.strip_prefix(root).unwrap_or(&path).to_path_buf();
        if meta.is_symlink() {
            skipped.push(relative);
        } else if meta.is_dir() {
            walk_into(root, &path, files, skipped)?;
        } else if meta.is_file() {
            files.push(relative);
        } else {
            skipped.push(relative);
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn walk_finds_nested_files_and_skips_symlinks() {
        let dir = ScratchDir::new("walk").unwrap();
        std::fs::create_dir_all(dir.path().join("sub")).unwrap();
        std::fs::write(dir.path().join("a"), b"a").unwrap();
        std::fs::write(dir.path().join("sub/b"), b"b").unwrap();
        #[cfg(unix)]
        std::os::unix::fs::symlink(dir.path().join("a"), dir.path().join("link")).unwrap();

        let (files, skipped) = walk_files(dir.path()).unwrap();
        assert_eq!(files, vec![PathBuf::from("a"), PathBuf::from("sub/b")]);
        #[cfg(unix)]
        assert_eq!(skipped, vec![PathBuf::from("link")]);
        #[cfg(not(unix))]
        assert!(skipped.is_empty());
    }

    #[test]
    fn copy_truncates_to_the_durable_length() {
        let dir = ScratchDir::new("truncate").unwrap();
        let src = dir.path().join("src");
        let dst = dir.path().join("dst");
        std::fs::write(&src, b"0123456789").unwrap();
        assert_eq!(copy_file(&src, &dst, Some(4)).unwrap(), 4);
        assert_eq!(std::fs::read(&dst).unwrap(), b"0123");
    }
}