use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tracing::{debug, trace};
use crate::{Error, Result, ensure_dir, get_cache_dir};
pub struct RibbitCache {
base_dir: PathBuf,
default_ttl: Duration,
}
impl RibbitCache {
pub async fn new() -> Result<Self> {
Self::with_ttl(Duration::from_secs(300)).await
}
pub async fn with_ttl(ttl: Duration) -> Result<Self> {
let base_dir = get_cache_dir()?.join("ribbit");
ensure_dir(&base_dir).await?;
debug!(
"Initialized Ribbit cache at: {:?} with TTL: {:?}",
base_dir, ttl
);
Ok(Self {
base_dir,
default_ttl: ttl,
})
}
pub fn cache_path(&self, region: &str, product: &str, endpoint: &str) -> PathBuf {
self.base_dir.join(region).join(product).join(endpoint)
}
pub fn metadata_path(&self, region: &str, product: &str, endpoint: &str) -> PathBuf {
let mut path = self.cache_path(region, product, endpoint);
path.set_extension("meta");
path
}
pub async fn is_valid(&self, region: &str, product: &str, endpoint: &str) -> bool {
let meta_path = self.metadata_path(region, product, endpoint);
if let Ok(metadata) = tokio::fs::read_to_string(&meta_path).await {
if let Ok(timestamp) = metadata.trim().parse::<u64>() {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
return (now - timestamp) < self.default_ttl.as_secs();
}
}
false
}
pub async fn write(
&self,
region: &str,
product: &str,
endpoint: &str,
data: &[u8],
) -> Result<()> {
let path = self.cache_path(region, product, endpoint);
let meta_path = self.metadata_path(region, product, endpoint);
if let Some(parent) = path.parent() {
ensure_dir(parent).await?;
}
let temp_path = path.with_file_name(format!(
"{}.tmp",
path.file_name().unwrap().to_string_lossy()
));
let temp_meta_path = meta_path.with_file_name(format!(
"{}.tmp",
meta_path.file_name().unwrap().to_string_lossy()
));
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
trace!(
"Writing {} bytes to Ribbit cache: {}/{}/{}",
data.len(),
region,
product,
endpoint
);
let write_result = async {
tokio::fs::write(&temp_path, data).await?;
tokio::fs::write(&temp_meta_path, timestamp.to_string()).await?;
tokio::fs::rename(&temp_path, &path).await?;
tokio::fs::rename(&temp_meta_path, &meta_path).await?;
Ok::<(), std::io::Error>(())
}
.await;
if write_result.is_err() {
let _ = tokio::fs::remove_file(&temp_path).await;
let _ = tokio::fs::remove_file(&temp_meta_path).await;
}
write_result?;
Ok(())
}
pub async fn read(&self, region: &str, product: &str, endpoint: &str) -> Result<Vec<u8>> {
if !self.is_valid(region, product, endpoint).await {
return Err(Error::CacheEntryNotFound(format!(
"{region}/{product}/{endpoint}"
)));
}
let path = self.cache_path(region, product, endpoint);
trace!(
"Reading from Ribbit cache: {}/{}/{}",
region, product, endpoint
);
Ok(tokio::fs::read(&path).await?)
}
pub async fn clear_expired(&self) -> Result<()> {
debug!("Clearing expired entries from Ribbit cache");
self.clear_expired_in_dir(&self.base_dir).await
}
fn clear_expired_in_dir<'a>(
&'a self,
dir: &'a Path,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>> {
Box::pin(async move {
let mut entries = tokio::fs::read_dir(dir).await?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if path.is_dir() {
self.clear_expired_in_dir(&path).await?;
} else if path.extension().and_then(|s| s.to_str()) == Some("meta") {
if let Ok(metadata) = tokio::fs::read_to_string(&path).await {
if let Ok(timestamp) = metadata.trim().parse::<u64>() {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
if (now - timestamp) >= self.default_ttl.as_secs() {
let data_path = path.with_extension("");
let _ = tokio::fs::remove_file(&data_path).await;
let _ = tokio::fs::remove_file(&path).await;
trace!("Removed expired cache entry: {:?}", data_path);
}
}
}
}
}
Ok(())
})
}
pub fn base_dir(&self) -> &PathBuf {
&self.base_dir
}
pub fn ttl(&self) -> Duration {
self.default_ttl
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_ribbit_cache_operations() {
let cache = RibbitCache::with_ttl(Duration::from_secs(60))
.await
.unwrap();
let region = "us";
let product = "wow";
let endpoint = "versions";
let data = b"test ribbit response";
cache.write(region, product, endpoint, data).await.unwrap();
assert!(cache.is_valid(region, product, endpoint).await);
let read_data = cache.read(region, product, endpoint).await.unwrap();
assert_eq!(read_data, data);
let _ = tokio::fs::remove_file(cache.cache_path(region, product, endpoint)).await;
let _ = tokio::fs::remove_file(cache.metadata_path(region, product, endpoint)).await;
}
#[tokio::test]
async fn test_ribbit_cache_expiry() {
let cache = RibbitCache::with_ttl(Duration::from_secs(0)).await.unwrap();
let region = "eu";
let product = "wow";
let endpoint = "cdns";
let data = b"test data";
cache.write(region, product, endpoint, data).await.unwrap();
tokio::time::sleep(Duration::from_millis(10)).await;
assert!(!cache.is_valid(region, product, endpoint).await);
assert!(cache.read(region, product, endpoint).await.is_err());
let _ = tokio::fs::remove_file(cache.cache_path(region, product, endpoint)).await;
let _ = tokio::fs::remove_file(cache.metadata_path(region, product, endpoint)).await;
}
}