greentic_flow/cache/

disk.rs

use std::fs;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result, bail};

use crate::cache::PruneReport;
use crate::cache::engine_profile::EngineProfile;
use crate::cache::keys::ArtifactKey;
use crate::cache::metadata::ArtifactMetadata;

#[derive(Clone, Debug)]
pub struct DiskCache {
    root: PathBuf,
    profile: EngineProfile,
    disk_max_bytes: Option<u64>,
}

impl DiskCache {
    pub fn new(root: PathBuf, profile: EngineProfile, disk_max_bytes: Option<u64>) -> Self {
        Self {
            root,
            profile,
            disk_max_bytes,
        }
    }

    pub fn root(&self) -> &Path {
        &self.root
    }

    pub fn try_read(&self, key: &ArtifactKey) -> Result<Option<Vec<u8>>> {
        let paths = self.paths_for(key)?;
        if !paths.meta_path.exists() {
            if paths.artifact_path.exists() {
                let _ = fs::remove_file(&paths.artifact_path);
            }
            return Ok(None);
        }
        let meta = match fs::read_to_string(&paths.meta_path) {
            Ok(raw) => match serde_json::from_str::<ArtifactMetadata>(&raw) {
                Ok(meta) => meta,
                Err(_) => {
                    self.delete_entry(&paths)?;
                    return Ok(None);
                }
            },
            Err(_) => {
                self.delete_entry(&paths)?;
                return Ok(None);
            }
        };
        if meta.validate_for_profile(&self.profile).is_err() {
            self.delete_entry(&paths)?;
            return Ok(None);
        }
        if meta.wasm_digest != key.wasm_digest {
            self.delete_entry(&paths)?;
            return Ok(None);
        }
        if !paths.artifact_path.exists() {
            self.delete_entry(&paths)?;
            return Ok(None);
        }
        let artifact_bytes = fs::read(&paths.artifact_path).ok();
        let Some(artifact_bytes) = artifact_bytes else {
            self.delete_entry(&paths)?;
            return Ok(None);
        };
        if meta.artifact_bytes != artifact_bytes.len() as u64 {
            self.delete_entry(&paths)?;
            return Ok(None);
        }
        self.update_access(&paths, meta).ok();
        Ok(Some(artifact_bytes))
    }

    pub fn write_atomic(
        &self,
        key: &ArtifactKey,
        bytes: &[u8],
        meta: &ArtifactMetadata,
    ) -> Result<()> {
        meta.validate_for_profile(&self.profile)
            .context("cache metadata does not match engine profile")?;
        if meta.wasm_digest != key.wasm_digest {
            bail!("cache metadata digest does not match artifact key");
        }
        let paths = self.paths_for(key)?;
        fs::create_dir_all(&paths.artifacts_dir).with_context(|| {
            format!(
                "failed to create cache dir {}",
                paths.artifacts_dir.display()
            )
        })?;
        fs::create_dir_all(&paths.tmp_dir).with_context(|| {
            format!("failed to create cache tmp dir {}", paths.tmp_dir.display())
        })?;
        let tmp_artifact = paths.tmp_path("artifact");
        let tmp_meta = paths.tmp_path("meta");
        fs::write(&tmp_artifact, bytes)
            .with_context(|| format!("failed to write {}", tmp_artifact.display()))?;
        let meta_json =
            serde_json::to_vec_pretty(meta).context("failed to serialize cache metadata")?;
        fs::write(&tmp_meta, meta_json)
            .with_context(|| format!("failed to write {}", tmp_meta.display()))?;
        fs::rename(&tmp_artifact, &paths.artifact_path)
            .with_context(|| format!("failed to rename {}", paths.artifact_path.display()))?;
        fs::rename(&tmp_meta, &paths.meta_path)
            .with_context(|| format!("failed to rename {}", paths.meta_path.display()))?;
        Ok(())
    }

    pub fn approx_size_bytes(&self) -> Result<u64> {
        let artifacts_dir = self.root.join("artifacts");
        if !artifacts_dir.exists() {
            return Ok(0);
        }
        let mut total = 0u64;
        for entry in fs::read_dir(&artifacts_dir)
            .with_context(|| format!("failed to read {}", artifacts_dir.display()))?
        {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|ext| ext.to_str()) != Some("cwasm") {
                continue;
            }
            let size = entry.metadata().map(|m| m.len()).unwrap_or(0);
            total = total.saturating_add(size);
        }
        Ok(total)
    }

    pub fn prune_to_limit(&self, dry_run: bool) -> Result<PruneReport> {
        let Some(limit) = self.disk_max_bytes else {
            return Ok(PruneReport {
                removed_entries: 0,
                removed_bytes: 0,
            });
        };
        let artifacts_dir = self.root.join("artifacts");
        if !artifacts_dir.exists() {
            return Ok(PruneReport {
                removed_entries: 0,
                removed_bytes: 0,
            });
        }
        let canonical_artifacts_dir = artifacts_dir
            .canonicalize()
            .with_context(|| format!("failed to canonicalize {}", artifacts_dir.display()))?;
        let mut entries = Vec::new();
        let mut total_bytes = 0u64;
        for entry in fs::read_dir(&artifacts_dir)
            .with_context(|| format!("failed to read {}", artifacts_dir.display()))?
        {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
                continue;
            }
            let raw = match fs::read_to_string(&path) {
                Ok(raw) => raw,
                Err(_) => continue,
            };
            let meta: ArtifactMetadata = match serde_json::from_str(&raw) {
                Ok(meta) => meta,
                Err(_) => continue,
            };
            let access = meta.last_access_time();
            let artifact_path = path.with_extension("cwasm");
            let canonical_artifact = match artifact_path.canonicalize() {
                Ok(path) if path.starts_with(&canonical_artifacts_dir) => path,
                _ => continue,
            };
            let size = fs::metadata(&canonical_artifact)
                .map(|m| m.len())
                .unwrap_or(0);
            total_bytes = total_bytes.saturating_add(size);
            entries.push((access, meta, canonical_artifact, path, size));
        }
        entries.sort_by_key(|(access, _, _, _, _)| {
            access.map(|ts| ts.timestamp()).unwrap_or(i64::MIN)
        });
        let mut removed_entries = 0u64;
        let mut removed_bytes = 0u64;
        let mut remaining = total_bytes;
        for (_access, _meta, artifact_path, meta_path, size) in entries {
            if remaining <= limit {
                break;
            }
            if !dry_run {
                let _ = fs::remove_file(&artifact_path);
                let _ = fs::remove_file(&meta_path);
            }
            removed_entries = removed_entries.saturating_add(1);
            removed_bytes = removed_bytes.saturating_add(size);
            remaining = remaining.saturating_sub(size);
        }
        Ok(PruneReport {
            removed_entries,
            removed_bytes,
        })
    }

    pub fn artifact_count(&self) -> Result<u64> {
        let artifacts_dir = self.root.join("artifacts");
        if !artifacts_dir.exists() {
            return Ok(0);
        }
        let mut count = 0u64;
        for entry in fs::read_dir(&artifacts_dir)
            .with_context(|| format!("failed to read {}", artifacts_dir.display()))?
        {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|ext| ext.to_str()) == Some("cwasm") {
                count = count.saturating_add(1);
            }
        }
        Ok(count)
    }

    pub fn delete(&self, key: &ArtifactKey) -> Result<()> {
        let paths = self.paths_for(key)?;
        self.delete_entry(&paths)
    }

    fn update_access(&self, paths: &DiskPaths, mut meta: ArtifactMetadata) -> Result<()> {
        meta.touch();
        let tmp = paths.tmp_path("meta");
        let json = serde_json::to_vec_pretty(&meta)?;
        let _ = fs::create_dir_all(&paths.tmp_dir);
        fs::write(&tmp, json).ok();
        let _ = fs::rename(&tmp, &paths.meta_path);
        Ok(())
    }

    fn delete_entry(&self, paths: &DiskPaths) -> Result<()> {
        let _ = fs::remove_file(&paths.artifact_path);
        let _ = fs::remove_file(&paths.meta_path);
        Ok(())
    }

    fn paths_for(&self, key: &ArtifactKey) -> Result<DiskPaths> {
        if key.engine_profile_id != self.profile.engine_profile_id {
            bail!("artifact key engine_profile_id mismatch");
        }
        let artifacts_dir = self.root.join("artifacts");
        let tmp_dir = self.root.join("tmp");
        let name = digest_to_filename(&key.wasm_digest);
        let artifact_path = artifacts_dir.join(format!("{}.cwasm", name));
        let meta_path = artifacts_dir.join(format!("{}.json", name));
        Ok(DiskPaths {
            artifacts_dir,
            tmp_dir,
            artifact_path,
            meta_path,
        })
    }
}

struct DiskPaths {
    artifacts_dir: PathBuf,
    tmp_dir: PathBuf,
    artifact_path: PathBuf,
    meta_path: PathBuf,
}

impl DiskPaths {
    fn tmp_path(&self, suffix: &str) -> PathBuf {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        let pid = std::process::id();
        self.tmp_dir.join(format!("tmp_{}_{}_{}", pid, now, suffix))
    }
}

fn digest_to_filename(digest: &str) -> String {
    let sanitized: String = digest
        .chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_') {
                ch
            } else {
                '_'
            }
        })
        .collect();
    let trimmed = sanitized.trim_matches('_');
    if trimmed.is_empty() {
        "artifact".to_string()
    } else {
        trimmed.to_string()
    }
}