use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use async_trait::async_trait;
use crate::protocol::CachedArtifact;
use crate::DaemonError;
#[async_trait]
pub trait LocalCacheStore: Send + Sync {
async fn get(&self, content_hash: &str) -> Result<Option<CachedArtifact>, DaemonError>;
async fn put(&self, content_hash: &str, artifact: &CachedArtifact) -> Result<(), DaemonError>;
}
pub struct RealLocalCacheStore {
root: PathBuf,
}
impl RealLocalCacheStore {
#[must_use]
pub fn new(root: PathBuf) -> Self {
Self { root }
}
fn entry_path(&self, content_hash: &str) -> PathBuf {
let shard = if content_hash.len() >= 2 { &content_hash[..2] } else { "00" };
self.root.join(shard).join(format!("{content_hash}.json"))
}
}
#[async_trait]
impl LocalCacheStore for RealLocalCacheStore {
async fn get(&self, content_hash: &str) -> Result<Option<CachedArtifact>, DaemonError> {
let path = self.entry_path(content_hash);
match tokio::fs::read(&path).await {
Ok(bytes) => {
let artifact: CachedArtifact =
serde_json::from_slice(&bytes).map_err(|e| DaemonError::Protocol(e.to_string()))?;
Ok(Some(artifact))
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(DaemonError::Io(e)),
}
}
async fn put(&self, content_hash: &str, artifact: &CachedArtifact) -> Result<(), DaemonError> {
let path = self.entry_path(content_hash);
if let Some(parent) = path.parent() {
tokio::fs::create_dir_all(parent).await?;
}
let body = serde_json::to_vec(artifact).map_err(|e| DaemonError::Protocol(e.to_string()))?;
let tmp_path = path.with_extension("json.tmp");
tokio::fs::write(&tmp_path, &body).await?;
tokio::fs::rename(&tmp_path, &path).await?;
Ok(())
}
}
#[derive(Default)]
pub struct MockLocalCacheStore {
entries: Mutex<BTreeMap<String, CachedArtifact>>,
pub get_calls: Mutex<usize>,
}
impl MockLocalCacheStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_entry(self, content_hash: &str, artifact: CachedArtifact) -> Self {
self.entries.lock().expect("mock mutex poisoned").insert(content_hash.to_string(), artifact);
self
}
#[must_use]
pub fn contains(&self, content_hash: &str) -> bool {
self.entries.lock().expect("mock mutex poisoned").contains_key(content_hash)
}
}
#[async_trait]
impl LocalCacheStore for MockLocalCacheStore {
async fn get(&self, content_hash: &str) -> Result<Option<CachedArtifact>, DaemonError> {
*self.get_calls.lock().expect("mock mutex poisoned") += 1;
Ok(self.entries.lock().expect("mock mutex poisoned").get(content_hash).cloned())
}
async fn put(&self, content_hash: &str, artifact: &CachedArtifact) -> Result<(), DaemonError> {
self.entries
.lock()
.expect("mock mutex poisoned")
.insert(content_hash.to_string(), artifact.clone());
Ok(())
}
}
pub async fn ensure_root_exists(root: &Path) -> Result<(), DaemonError> {
tokio::fs::create_dir_all(root).await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn real_store_put_then_get_round_trips() {
let dir = tempfile::tempdir().unwrap();
let store = RealLocalCacheStore::new(dir.path().to_path_buf());
let artifact = CachedArtifact { image_ref: "example/image:v1".to_string() };
store.put("deadbeef", &artifact).await.unwrap();
let got = store.get("deadbeef").await.unwrap();
assert_eq!(got, Some(artifact));
}
#[tokio::test]
async fn real_store_miss_returns_none_not_an_error() {
let dir = tempfile::tempdir().unwrap();
let store = RealLocalCacheStore::new(dir.path().to_path_buf());
let got = store.get("nope").await.unwrap();
assert_eq!(got, None);
}
#[tokio::test]
async fn mock_store_get_calls_are_counted() {
let store = MockLocalCacheStore::new();
store.get("x").await.unwrap();
store.get("x").await.unwrap();
assert_eq!(*store.get_calls.lock().unwrap(), 2);
}
}