use std::{sync::Arc, time::SystemTime};
use opendal::{raw::Timestamp, Operator};
use rattler_networking::retry_policies::default_retry_policy;
use retry_policies::{RetryDecision, RetryPolicy};
use tokio::sync::RwLock;
use crate::{IndexedPackageRecord, RepodataFileMetadata};
#[derive(Debug, Clone)]
struct CachedPackage {
record: IndexedPackageRecord,
etag: Option<String>,
last_modified: Option<Timestamp>,
}
#[derive(Debug)]
pub(crate) enum CacheResult {
Hit(Box<IndexedPackageRecord>),
Miss {
etag: Option<String>,
last_modified: Option<Timestamp>,
},
}
#[derive(Debug, Clone, Default)]
pub struct PackageRecordCache {
inner: Arc<RwLock<ahash::HashMap<String, CachedPackage>>>,
}
impl PackageRecordCache {
pub fn new() -> Self {
Self::default()
}
pub(crate) async fn get_or_stat(
&self,
op: &Operator,
path: &str,
) -> opendal::Result<CacheResult> {
let metadata = match op.stat(path).await {
Ok(m) => m,
Err(e) => {
tracing::warn!(
"Failed to stat file during cache lookup for {}: {}",
path,
e
);
return Err(e);
}
};
let current_etag = metadata.etag().map(str::to_owned);
let current_last_modified = metadata.last_modified();
let cached = {
let guard = self.inner.read().await;
guard.get(path).cloned()
};
if let Some(cached) = cached {
if let (Some(cached_etag), Some(ref current_etag)) = (&cached.etag, ¤t_etag) {
if cached_etag == current_etag {
tracing::debug!("Cache hit for {} (etag validated)", path);
return Ok(CacheResult::Hit(Box::new(cached.record)));
} else {
tracing::debug!(
"Cache entry for {} has mismatched etag, treating as miss",
path
);
return Ok(CacheResult::Miss {
etag: Some(current_etag.clone()),
last_modified: current_last_modified,
});
}
}
if let (Some(cached_modified), Some(current_modified)) =
(cached.last_modified, current_last_modified)
{
if cached_modified == current_modified {
tracing::debug!("Cache hit for {} (last_modified validated)", path);
return Ok(CacheResult::Hit(Box::new(cached.record)));
} else {
tracing::debug!(
"Cache entry for {} has mismatched last_modified, treating as miss",
path
);
return Ok(CacheResult::Miss {
etag: current_etag,
last_modified: current_last_modified,
});
}
}
tracing::debug!(
"Cache entry for {} cannot be validated (no metadata), treating as miss",
path
);
} else {
tracing::debug!("Cache miss for {} (not in cache)", path);
}
Ok(CacheResult::Miss {
etag: current_etag,
last_modified: current_last_modified,
})
}
pub(crate) async fn insert(
&self,
path: &str,
record: IndexedPackageRecord,
etag: Option<String>,
last_modified: Option<Timestamp>,
) {
let cached = CachedPackage {
record,
etag,
last_modified,
};
let mut guard = self.inner.write().await;
guard.insert(path.to_string(), cached);
}
}
pub async fn read_package_with_retry(
op: &Operator,
path: &str,
initial_metadata: RepodataFileMetadata,
) -> opendal::Result<(opendal::Buffer, RepodataFileMetadata)> {
let retry_policy = default_retry_policy();
let mut current_try = 0;
let mut metadata = initial_metadata;
loop {
let request_start_time = SystemTime::now();
match crate::utils::read_with_metadata_check(op, path, &metadata).await {
Ok(buffer) => return Ok((buffer, metadata)),
Err(e) if e.kind() == opendal::ErrorKind::Unsupported => {
tracing::debug!(
"Conditional reads not supported for {}, using simple read",
path
);
let buffer = op.read(path).await?;
return Ok((buffer, metadata));
}
Err(e) if e.kind() == opendal::ErrorKind::ConditionNotMatch => {
match retry_policy.should_retry(request_start_time, current_try) {
RetryDecision::Retry { execute_after } => {
let duration = execute_after
.duration_since(SystemTime::now())
.unwrap_or_default();
tracing::debug!(
"File {} changed between stat and read (attempt {}), retrying in {:?}",
path,
current_try + 1,
duration
);
tokio::time::sleep(duration).await;
current_try += 1;
let fresh_metadata = op.stat(path).await?;
metadata = RepodataFileMetadata {
etag: fresh_metadata.etag().map(str::to_owned),
last_modified: fresh_metadata.last_modified(),
file_existed: true,
precondition_checks: crate::PreconditionChecks::Enabled,
};
}
RetryDecision::DoNotRetry => {
tracing::warn!(
"Max retries exceeded for reading {} due to concurrent modifications",
path
);
return Err(e);
}
}
}
Err(e) => {
return Err(e);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_creation() {
let cache = PackageRecordCache::new();
assert!(cache.inner.try_read().is_ok());
}
#[test]
fn test_cache_default() {
let cache = PackageRecordCache::default();
assert!(cache.inner.try_read().is_ok());
}
}