olympipe-rs 1.1.0

//! On-disk JSON cache used by [`crate::Pipeline::cached_task`].

use crate::pipeline::PipelineError;
use serde::Serialize;
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};

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

/// Lowercase hex SHA-256 of `sha256(step)` concatenated with the `serde_json`
/// bytes of `item`.
///
/// This is the file name (without `.json`) used by [`Pipeline::cached_task`]
/// and [`cache_file`]. `step` is hashed as UTF-8. Changing `step` changes the
/// key; a version suffix such as `"resize-v2"` is the way to drop old entries
/// when the computation changes.
///
/// The bytes that are hashed are exactly what `serde_json` emits. A
/// [`std::collections::HashMap`] may serialize keys in a different order on
/// equal values, so those keys will not match. Prefer
/// [`std::collections::BTreeMap`] or another type with a stable encoding.
///
/// # Errors
///
/// Returns [`PipelineError::TaskFailed`] `{ stage: "cached_task", .. }` when
/// `item` cannot be serialized.
///
/// [`Pipeline::cached_task`]: crate::Pipeline::cached_task
///
/// ```
/// use olympipe_rs::cache_key;
///
/// let key = cache_key("resize-v2", &1).unwrap();
/// assert_eq!(
///     key,
///     "4ea06298ee3793703ef2e38fc4c7fb9a1d28b73b55dd12424ee0321acc166499"
/// );
/// assert_ne!(cache_key("resize-v3", &1).unwrap(), key);
/// ```
#[cfg_attr(docsrs, doc(cfg(feature = "cache")))]
pub fn cache_key<T: Serialize>(step: &str, item: &T) -> Result<String, PipelineError> {
    let json = serde_json::to_vec(item).map_err(task_failed)?;
    Ok(hex_encode(&content_hash(step, &json)))
}

/// Path of the cache entry for `item` at `step`.
///
/// `{cache_dir}/{first 16 hex characters of sha256(step)}/{cache_key}.json`.
/// See [`cache_key`] for what goes into the hash. The directory component is
/// only the step, so one computation's files stay together and a new `step`
/// string uses a different directory.
///
/// # Errors
///
/// Returns [`PipelineError::TaskFailed`] `{ stage: "cached_task", .. }` when
/// `item` cannot be serialized, same as [`cache_key`].
///
/// ```
/// use olympipe_rs::cache_file;
/// use std::path::Path;
///
/// let path = cache_file("cache", "resize-v2", &1).unwrap();
/// assert_eq!(
///     path,
///     Path::new("cache").join("c0b86eb1f7f19838").join(
///         "4ea06298ee3793703ef2e38fc4c7fb9a1d28b73b55dd12424ee0321acc166499.json"
///     )
/// );
/// ```
#[cfg_attr(docsrs, doc(cfg(feature = "cache")))]
pub fn cache_file<T: Serialize>(
    cache_dir: impl AsRef<Path>,
    step: &str,
    item: &T,
) -> Result<PathBuf, PipelineError> {
    let key = cache_key(step, item)?;
    Ok(entry_path(cache_dir.as_ref(), step, &key))
}

/// Hit: return the deserialized value and do not call `compute`.
/// Miss, or a file that cannot be read or deserialized: call `compute`, then
/// try to store the result. A store error is logged at most once per `logged`
/// flag and still returns the value. Serializing `item` fails the caller.
pub(crate) fn resolve<T, U, F>(
    step: &str,
    cache_dir: &Path,
    item: T,
    compute: &F,
    logged: &mut bool,
) -> Result<U, PipelineError>
where
    T: Serialize,
    U: Serialize + DeserializeOwned,
    F: Fn(T) -> U,
{
    let json = serde_json::to_vec(&item).map_err(task_failed)?;
    let path = entry_path(cache_dir, step, &hex_encode(&content_hash(step, &json)));
    if let Some(hit) = read_value(&path) {
        return Ok(hit);
    }
    let value = compute(item);
    if let Err(err) = persist(&path, &value) {
        log_once(logged, &err);
    }
    Ok(value)
}

fn task_failed(err: impl std::fmt::Display) -> PipelineError {
    PipelineError::TaskFailed {
        stage: "cached_task",
        message: err.to_string(),
    }
}

fn entry_path(cache_dir: &Path, step: &str, key: &str) -> PathBuf {
    cache_dir
        .join(step_prefix(step))
        .join(format!("{key}.json"))
}

fn step_prefix(step: &str) -> String {
    let hex = hex_encode(&sha256(step.as_bytes()));
    hex[..16].to_owned()
}

fn content_hash(step: &str, json: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(sha256(step.as_bytes()));
    hasher.update(json);
    finalize(hasher)
}

fn sha256(bytes: &[u8]) -> [u8; 32] {
    finalize(Sha256::new().chain_update(bytes))
}

fn finalize(hasher: Sha256) -> [u8; 32] {
    let digest = hasher.finalize();
    let mut out = [0u8; 32];
    out.copy_from_slice(&digest);
    out
}

fn hex_encode(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        out.push(HEX[usize::from(byte >> 4)] as char);
        out.push(HEX[usize::from(byte & 0x0f)] as char);
    }
    out
}

fn read_value<U: DeserializeOwned>(path: &Path) -> Option<U> {
    let bytes = std::fs::read(path).ok()?;
    serde_json::from_slice(&bytes).ok()
}

fn persist<U: Serialize>(path: &Path, value: &U) -> Result<(), String> {
    let bytes = serde_json::to_vec(value).map_err(|err| err.to_string())?;
    write_atomic(path, &bytes).map_err(|err| err.to_string())
}

fn log_once(logged: &mut bool, err: &dyn std::fmt::Display) {
    if *logged {
        return;
    }
    *logged = true;
    eprintln!("olympipe-rs: [cached_task] cache write failed: {err}");
}

/// Write `bytes` to a unique temporary file in the destination directory, then
/// `rename` it over `path`. Readers only open `path`, so they never see a
/// partial write. `rename` on the same directory replaces an existing file.
fn write_atomic(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
    let parent = path
        .parent()
        .filter(|dir| !dir.as_os_str().is_empty())
        .ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "cache entry path has no directory",
            )
        })?;
    std::fs::create_dir_all(parent)?;
    let file_name = path.file_name().ok_or_else(|| {
        std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "cache entry path has no file name",
        )
    })?;
    let seq = TEMP_SEQ.fetch_add(1, Ordering::Relaxed);
    let mut tmp_name = std::ffi::OsString::from(".");
    tmp_name.push(file_name);
    tmp_name.push(format!(".{}.{seq}.tmp", std::process::id()));
    let tmp = parent.join(tmp_name);

    let written = (|| -> std::io::Result<()> {
        let mut file = std::fs::File::create(&tmp)?;
        file.write_all(bytes)?;
        file.sync_all()?;
        std::fs::rename(&tmp, path)?;
        Ok(())
    })();
    if written.is_err() {
        let _ = std::fs::remove_file(&tmp);
    }
    written
}