use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CacheEntry {
pub status: u16,
pub headers: HashMap<String, String>,
#[serde(with = "barbacane_plugin_sdk::types::base64_body")]
pub body: Option<Vec<u8>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<CacheMetadata>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CacheMetadata {
pub created_at: u64,
pub expires_at: u64,
pub ttl_remaining: u64,
}
struct InternalEntry {
entry: CacheEntry,
expires_at: Instant,
created_at: Instant,
_ttl_secs: u64,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CacheResult {
pub hit: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub entry: Option<CacheEntry>,
}
#[derive(Clone)]
pub struct ResponseCache {
entries: Arc<RwLock<HashMap<String, InternalEntry>>>,
cleanup_interval: Duration,
last_cleanup: Arc<RwLock<Instant>>,
}
impl Default for ResponseCache {
fn default() -> Self {
Self::new()
}
}
impl ResponseCache {
pub fn new() -> Self {
Self {
entries: Arc::new(RwLock::new(HashMap::new())),
cleanup_interval: Duration::from_secs(60),
last_cleanup: Arc::new(RwLock::new(Instant::now())),
}
}
pub fn get(&self, key: &str) -> CacheResult {
self.maybe_cleanup();
let entries = self.entries.read();
let now = Instant::now();
if let Some(internal) = entries.get(key) {
if internal.expires_at > now {
let ttl_remaining = internal.expires_at.saturating_duration_since(now).as_secs();
let now_unix = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let mut entry = internal.entry.clone();
entry.metadata = Some(CacheMetadata {
created_at: now_unix - internal.created_at.elapsed().as_secs(),
expires_at: now_unix + ttl_remaining,
ttl_remaining,
});
return CacheResult {
hit: true,
entry: Some(entry),
};
}
}
CacheResult {
hit: false,
entry: None,
}
}
pub fn set(&self, key: &str, entry: CacheEntry, ttl_secs: u64) {
let now = Instant::now();
let expires_at = now + Duration::from_secs(ttl_secs);
let internal = InternalEntry {
entry,
expires_at,
created_at: now,
_ttl_secs: ttl_secs,
};
let mut entries = self.entries.write();
entries.insert(key.to_string(), internal);
}
pub fn invalidate(&self, key: &str) {
let mut entries = self.entries.write();
entries.remove(key);
}
pub fn clear(&self) {
let mut entries = self.entries.write();
entries.clear();
}
fn maybe_cleanup(&self) {
let now = Instant::now();
{
let last = self.last_cleanup.read();
if now.duration_since(*last) < self.cleanup_interval {
return;
}
}
if let Some(mut last) = self.last_cleanup.try_write() {
if now.duration_since(*last) >= self.cleanup_interval {
*last = now;
if let Some(mut entries) = self.entries.try_write() {
entries.retain(|_, v| v.expires_at > now);
}
}
}
}
pub fn stats(&self) -> CacheStats {
let entries = self.entries.read();
let now = Instant::now();
let valid_count = entries.values().filter(|e| e.expires_at > now).count();
CacheStats {
total_entries: entries.len(),
valid_entries: valid_count,
}
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub total_entries: usize,
pub valid_entries: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_miss() {
let cache = ResponseCache::new();
let result = cache.get("test-key");
assert!(!result.hit);
assert!(result.entry.is_none());
}
#[test]
fn test_cache_hit() {
let cache = ResponseCache::new();
let entry = CacheEntry {
status: 200,
headers: HashMap::new(),
body: Some(b"test body".to_vec()),
metadata: None,
};
cache.set("test-key", entry, 60);
let result = cache.get("test-key");
assert!(result.hit);
assert!(result.entry.is_some());
let cached = result.entry.unwrap();
assert_eq!(cached.status, 200);
assert_eq!(cached.body, Some(b"test body".to_vec()));
}
#[test]
fn test_cache_invalidate() {
let cache = ResponseCache::new();
let entry = CacheEntry {
status: 200,
headers: HashMap::new(),
body: None,
metadata: None,
};
cache.set("test-key", entry, 60);
assert!(cache.get("test-key").hit);
cache.invalidate("test-key");
assert!(!cache.get("test-key").hit);
}
#[test]
fn test_cache_stats() {
let cache = ResponseCache::new();
let entry = CacheEntry {
status: 200,
headers: HashMap::new(),
body: None,
metadata: None,
};
cache.set("key1", entry.clone(), 60);
cache.set("key2", entry.clone(), 60);
cache.set("key3", entry, 60);
let stats = cache.stats();
assert_eq!(stats.total_entries, 3);
assert_eq!(stats.valid_entries, 3);
}
#[test]
fn test_cache_binary_body_roundtrip() {
let cache = ResponseCache::new();
let binary_body = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; let entry = CacheEntry {
status: 200,
headers: {
let mut h = HashMap::new();
h.insert("content-type".to_string(), "image/png".to_string());
h
},
body: Some(binary_body.clone()),
metadata: None,
};
cache.set("binary-key", entry, 60);
let cached = cache.get("binary-key").entry.unwrap();
assert_eq!(cached.body, Some(binary_body));
}
#[test]
fn test_cache_entry_json_roundtrip() {
let entry = CacheEntry {
status: 200,
headers: HashMap::new(),
body: Some(vec![0x00, 0xFF, 0x80]),
metadata: None,
};
let json = serde_json::to_string(&entry).unwrap();
let decoded: CacheEntry = serde_json::from_str(&json).unwrap();
assert_eq!(decoded.body, entry.body);
}
#[test]
fn test_cache_entry_json_none_body() {
let entry = CacheEntry {
status: 204,
headers: HashMap::new(),
body: None,
metadata: None,
};
let json = serde_json::to_string(&entry).unwrap();
let decoded: CacheEntry = serde_json::from_str(&json).unwrap();
assert!(decoded.body.is_none());
}
}