use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
#[derive(Clone, Debug)]
struct CachedResponse {
body: Arc<[u8]>,
etag: String,
expires_at: Instant,
}
#[derive(Clone, Debug, Default)]
pub struct CacheStats {
pub total_requests: usize,
pub cache_hits: usize,
pub cache_misses: usize,
pub etag_hits: usize,
}
impl CacheStats {
pub fn hit_ratio(&self) -> f64 {
if self.total_requests == 0 {
0.0
} else {
(self.cache_hits + self.etag_hits) as f64 / self.total_requests as f64
}
}
}
#[derive(Debug, Default)]
struct AtomicCacheStats {
total_requests: AtomicUsize,
cache_hits: AtomicUsize,
cache_misses: AtomicUsize,
etag_hits: AtomicUsize,
}
#[derive(Debug)]
pub struct ApiResponseCache {
cache: Arc<RwLock<HashMap<String, CachedResponse>>>,
stats: Arc<AtomicCacheStats>,
}
impl ApiResponseCache {
pub fn new() -> Self {
Self {
cache: Arc::new(RwLock::new(HashMap::new())),
stats: Arc::new(AtomicCacheStats::default()),
}
}
pub fn get_or_compute<F>(
&self,
key: &str,
ttl: Duration,
compute_fn: F,
) -> Result<(Arc<[u8]>, String), String>
where
F: FnOnce() -> Result<Vec<u8>, String>,
{
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
{
let cache = self.cache.read().unwrap();
if let Some(cached) = cache.get(key) {
if cached.expires_at > Instant::now() {
self.stats.cache_hits.fetch_add(1, Ordering::Relaxed);
return Ok((Arc::clone(&cached.body), cached.etag.clone()));
}
}
}
self.stats.cache_misses.fetch_add(1, Ordering::Relaxed);
let body = compute_fn()?;
let etag = compute_etag(&body);
let body_arc: Arc<[u8]> = Arc::from(body);
let mut cache = self.cache.write().unwrap();
cache.insert(
key.to_string(),
CachedResponse {
body: Arc::clone(&body_arc),
etag: etag.clone(),
expires_at: Instant::now() + ttl,
},
);
Ok((body_arc, etag))
}
pub fn etag_matches(&self, key: &str, client_etag: &str) -> bool {
let cache = self.cache.read().unwrap();
if let Some(cached) = cache.get(key) {
if cached.expires_at > Instant::now() {
return cached.etag == client_etag;
}
}
false
}
pub fn record_etag_hit(&self) {
self.stats.etag_hits.fetch_add(1, Ordering::Relaxed);
}
pub fn stats(&self) -> CacheStats {
CacheStats {
total_requests: self.stats.total_requests.load(Ordering::Relaxed),
cache_hits: self.stats.cache_hits.load(Ordering::Relaxed),
cache_misses: self.stats.cache_misses.load(Ordering::Relaxed),
etag_hits: self.stats.etag_hits.load(Ordering::Relaxed),
}
}
pub fn clear(&self) {
self.cache.write().unwrap().clear();
self.stats.total_requests.store(0, Ordering::Relaxed);
self.stats.cache_hits.store(0, Ordering::Relaxed);
self.stats.cache_misses.store(0, Ordering::Relaxed);
self.stats.etag_hits.store(0, Ordering::Relaxed);
}
pub fn pre_warm<F>(&self, endpoints: Vec<(&str, Duration)>, compute_fn: F) -> Result<(), String>
where
F: Fn(&str) -> Result<Vec<u8>, String>,
{
for (key, ttl) in endpoints {
self.get_or_compute(key, ttl, || compute_fn(key))?;
}
Ok(())
}
pub fn size(&self) -> usize {
self.cache.read().unwrap().len()
}
}
impl Default for ApiResponseCache {
fn default() -> Self {
Self::new()
}
}
impl Clone for ApiResponseCache {
fn clone(&self) -> Self {
Self {
cache: Arc::clone(&self.cache),
stats: Arc::clone(&self.stats),
}
}
}
fn compute_etag(body: &[u8]) -> String {
let crc = crate::protocol::crc32(body);
format!("W/\"{:x}-{:x}\"", crc, body.len())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_miss_computes_value() {
let cache = ApiResponseCache::new();
let body = b"response data".to_vec();
let body_clone = body.clone();
let (result, etag) = cache
.get_or_compute("test", Duration::from_secs(60), || Ok(body_clone))
.unwrap();
assert_eq!(result.as_ref(), body.as_slice());
assert!(!etag.is_empty());
}
#[test]
fn cache_hit_reuses() {
let cache = ApiResponseCache::new();
let body1 = b"response1".to_vec();
let body1_clone = body1.clone();
let (result1, etag1) = cache
.get_or_compute("test", Duration::from_secs(60), || Ok(body1_clone))
.unwrap();
let (result2, etag2) = cache
.get_or_compute(
"test",
Duration::from_secs(60),
|| Ok(b"response2".to_vec()),
)
.unwrap();
assert_eq!(result1, result2);
assert_eq!(etag1, etag2);
assert_eq!(cache.stats().cache_hits, 1);
}
#[test]
fn etag_match_detection() {
let cache = ApiResponseCache::new();
let body = b"data".to_vec();
let body_clone = body.clone();
let (_result, etag) = cache
.get_or_compute("test", Duration::from_secs(60), || Ok(body_clone))
.unwrap();
assert!(cache.etag_matches("test", &etag));
assert!(!cache.etag_matches("test", "wrong-etag"));
}
#[test]
fn expiration_and_recompute() {
let cache = ApiResponseCache::new();
cache
.get_or_compute("test", Duration::from_millis(5), || Ok(b"data1".to_vec()))
.unwrap();
std::thread::sleep(Duration::from_millis(50));
let (result2, _etag2) = cache
.get_or_compute("test", Duration::from_secs(60), || Ok(b"data2".to_vec()))
.unwrap();
assert_eq!(result2.as_ref(), b"data2");
}
#[test]
fn statistics_tracking() {
let cache = ApiResponseCache::new();
for i in 0..5 {
let body = format!("response{}", i).into_bytes();
cache
.get_or_compute(&format!("endpoint{}", i), Duration::from_secs(60), || {
Ok(body)
})
.unwrap();
}
let stats = cache.stats();
assert_eq!(stats.total_requests, 5);
assert_eq!(stats.cache_misses, 5);
assert_eq!(stats.cache_hits, 0);
}
#[test]
fn hit_ratio_calculation() {
let cache = ApiResponseCache::new();
for _ in 0..3 {
cache
.get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
.unwrap();
cache
.get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
.unwrap();
}
let stats = cache.stats();
let ratio = stats.hit_ratio();
assert!(ratio > 0.5);
}
#[test]
fn pre_warm_populates() {
let cache = ApiResponseCache::new();
cache
.pre_warm(
vec![
("models", Duration::from_secs(60)),
("runtime", Duration::from_secs(300)),
],
|_key| Ok(b"warm data".to_vec()),
)
.unwrap();
assert_eq!(cache.size(), 2);
}
#[test]
fn clone_shares_state() {
let cache1 = ApiResponseCache::new();
cache1
.get_or_compute("test", Duration::from_secs(60), || Ok(b"data".to_vec()))
.unwrap();
let cache2 = cache1.clone();
assert_eq!(cache1.size(), cache2.size());
}
}