galaxy/crates/asset_cache/src/lib.rs

use std::hash::{DefaultHasher, Hash, Hasher};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::Result;
use async_fs::{OpenOptions, create_dir_all};
use bytes::Bytes;
use futures::AsyncWriteExt;
use galaxyui_core::assets::asset_cache::{
    Asset, AssetCache, AssetSource, AssetState, AsyncAssetId, AsyncAssetType,
};
use reqwest::Url;

/// Namespace marker for URL-based async asset sources without persistence.
pub struct UrlAssetWithoutPersistence;
impl AsyncAssetType for UrlAssetWithoutPersistence {}

/// Namespace marker for URL-based async asset sources with persistence.
///
/// This is intentionally separate from `UrlAssetWithoutPersistence` to allow
/// ensure we persist the asset even if we fetched it once already without
/// persistence.
pub struct UrlAssetWithPersistence;
impl AsyncAssetType for UrlAssetWithPersistence {}

/// Creates an [`AssetSource::Async`] that fetches bytes from the given URL
/// without persisting them to the local filesystem.
pub fn url_source(url: impl Into<String>) -> AssetSource {
    let url = url.into();
    let url_for_fetch = url.clone();
    AssetSource::Async {
        id: AsyncAssetId::new::<UrlAssetWithoutPersistence>(url),
        fetch: Arc::new(move || {
            let url = url_for_fetch.clone();
            Box::pin(async move {
                let parsed = Url::parse(&url)?;
                fetch_file_to_memory(parsed).await
            })
        }),
    }
}

/// Creates an [`AssetSource::Async`] that fetches bytes from the given URL,
/// persisting them to a file under `cache_dir` for future reads.
pub fn url_source_with_persistence(url: impl Into<String>, cache_dir: &Path) -> AssetSource {
    let url = url.into();
    let url_for_fetch = url.clone();
    let cache_dir_owned = cache_dir.to_path_buf();
    AssetSource::Async {
        id: AsyncAssetId::new::<UrlAssetWithPersistence>(url),
        fetch: Arc::new(move || {
            let url = url_for_fetch.clone();
            let cache_dir = cache_dir_owned.clone();
            Box::pin(async move {
                let parsed = Url::parse(&url)?;
                let file = get_file_path_for_asset(&parsed, &cache_dir);
                fetch_asset_from_url(parsed, Some(file)).await
            })
        }),
    }
}

/// Extension trait that adds URL-based asset loading to [`AssetCache`].
pub trait AssetCacheExt {
    /// Loads an asset from a URL, optionally persisting the fetched bytes to
    /// a file under `cache_dir` for future cache hits.
    fn load_asset_from_url<T: Asset>(&self, url: &str, cache_dir: Option<&Path>) -> AssetState<T>;
}

impl AssetCacheExt for AssetCache {
    fn load_asset_from_url<T: Asset>(&self, url: &str, cache_dir: Option<&Path>) -> AssetState<T> {
        let source = match cache_dir {
            Some(dir) => url_source_with_persistence(url, dir),
            None => url_source(url),
        };
        self.load_asset(source)
    }
}

/// Fetches a file from the given `url` to memory.
async fn fetch_file_to_memory(url: Url) -> Result<Bytes, anyhow::Error> {
    cfg_if::cfg_if! {
        if #[cfg(target_family = "wasm")] {
            let response = reqwest::get(url).await?;
        } else {
            // On non-web platforms, reqwest expects that it is operating within
            // a Tokio-compatible runtime, so use async-compat to wrap the call
            // so reqwest's expectations are met.
            let response = async_compat::Compat::new(async move { reqwest::get(url).await }).await?;
        }
    }
    let content = response.bytes().await?;
    Ok(content)
}

/// Given a url and a directory where cached artifacts are stored, returns a unique
/// file path for an asset.
fn get_file_path_for_asset(url: &Url, cache_dir: &Path) -> PathBuf {
    // Hash the URL so that we can derive a "safe" file name for it. We need something
    // unique and not too long (most filesystems have a maximum length limit for file
    // names. On MacOS it's 255).
    let mut hasher = DefaultHasher::new();
    url.hash(&mut hasher);
    let digest = hasher.finish();
    // Stringify the bytes in hexadecimal. Be careful not to use base64-digests in file
    // names b/c base64 uses a mix of upper and lowercase chars, which is problematic on
    // case-insensitive filesystems such as MacOS
    let filename = format!("{digest:x}");
    cache_dir.join(filename)
}

async fn persist_bytes(bytes: &Bytes, file: &Path) {
    let Some(parent_folder) = file.parent() else {
        log::error!("attempted to write cache file in filesystem root");
        return;
    };

    if let Err(e) = create_dir_all(parent_folder).await {
        log::error!("Error creating directory for cache files: {e:#}");
    }

    let mut file = match OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(file)
        .await
    {
        Ok(file) => file,
        Err(e) => {
            log::error!("Error opening file: {e:#}");
            return;
        }
    };

    if let Err(e) = file.write_all(bytes).await {
        log::error!("Error writing to file: {e:#}");
    }

    if let Err(e) = file.flush().await {
        log::error!("Error flushing file: {e:#}");
    };
}

async fn fetch_file_and_persist_bytes(url: Url, file: Option<PathBuf>) -> Result<Bytes> {
    let result = fetch_file_to_memory(url).await;

    // If the bytes should be written to a file, do so now.
    if let Ok(bytes) = result.as_ref()
        && let Some(filename) = file
    {
        persist_bytes(bytes, &filename).await;
    }

    result
}

async fn fetch_asset_from_url(url: Url, file: Option<PathBuf>) -> Result<Bytes> {
    match file {
        // If a file path is specified and that file path currently exists in the
        // user's filesystem, read the bytes out of the file.
        Some(filename) if filename.exists() => {
            log::debug!("Reading bytes from cached file: {filename:?}");
            let buffer = async_fs::read(filename.clone()).await?;

            // If buffer is empty, try to fetch from url instead
            if buffer.is_empty() {
                return fetch_file_and_persist_bytes(url, Some(filename)).await;
            }

            Ok(buffer.into())
        }
        // Otherwise, fetch the bytes from the url.
        _ => fetch_file_and_persist_bytes(url, file).await,
    }
}