use anyhow::{Context, Result, anyhow};
use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use std::time::SystemTime;
use std::hash::Hasher;
#[derive(Debug, Clone)]
struct CacheEntry {
wasm_bytes: Vec<u8>,
cached_at: Instant,
last_access: Instant,
access_count: usize,
file_path: Option<PathBuf>,
size: usize,
}
impl CacheEntry {
fn new(wasm_bytes: Vec<u8>) -> Self {
let now: _ = Instant::now();
CacheEntry {
wasm_bytes,
cached_at: now,
last_access: now,
access_count: 0,
file_path: None,
size: 0,
}
}
fn update_access(&mut self) {
self.last_access = Instant::now();
self.access_count += 1;
}
fn age(&self) -> Duration {
self.cached_at.elapsed()
}
fn usage_score(&self) -> f64 {
let age_secs: _ = self.cached_at.elapsed().as_secs_f64();
if age_secs == 0.0 {
self.access_count as f64
} else {
self.access_count as f64 / age_secs
}
}
}
type L1Cache = HashMap<String, Arc<RwLock<CacheEntry>>>;
type L2Cache = HashMap<String, PathBuf>;
#[derive(Debug, Clone, Default)]
pub struct CacheStats {
pub l1_entries: usize,
pub l2_entries: usize,
pub hits: usize,
pub misses: usize,
pub total_modules: usize,
pub total_size: usize,
pub hit_ratio: f64,
pub avg_load_time: Duration,
}
#[derive(Debug, Clone)]
struct CacheConfig {
max_l1_entries: usize,
max_l1_size: usize,
l2_cache_dir: PathBuf,
eviction_threshold: f64,
expiration_time: Duration,
enable_l2: bool,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
max_l1_entries: 1000,
max_l1_size: 100 * 1024 * 1024, l2_cache_dir: PathBuf::from("./wasm_cache"),
eviction_threshold: 0.1,
expiration_time: Duration::from_secs(3600), enable_l2: true,
}
}
}
pub struct WasmModuleCache {
l1_cache: Arc<Mutex<L1Cache>>,
l2_cache: Arc<Mutex<L2Cache>>,
config: CacheConfig,
stats: Arc<Mutex<CacheStats>>,
load_times: Arc<Mutex<Vec<Duration>>>,
}
impl WasmModuleCache {
pub fn new() -> Result<Self> {
Self::new_with_config(CacheConfig::default())
}
pub fn new_with_config(config: CacheConfig) -> Result<Self> {
if config.enable_l2 {
std::fs::create_dir_all(&config.l2_cache_dir)
.context("Failed to create L2 cache directory")?;
}
Ok(WasmModuleCache {
l1_cache: Arc::new(Mutex::new(HashMap::new())),
l2_cache: Arc::new(Mutex::new(HashMap::new())),
config,
stats: Arc::new(Mutex::new(CacheStats::default())),
load_times: Arc::new(Mutex::new(Vec::new())),
})
}
pub fn store_module(&self, module_hash: String, wasm_bytes: Vec<u8>) -> Result<()> {
let entry: _ = Arc::new(Mutex::new(CacheEntry::new(wasm_bytes)));
{
let mut l1 = self.l1_cache.lock().unwrap();
if l1.len() >= self.config.max_l1_entries {
self.evict_l1_cache()?;
}
l1.insert(module_hash.clone(), Arc::clone(entry));
}
if self.config.enable_l2 {
self.store_to_l2(&module_hash, &entry)?;
}
self.update_stats_after_store(&entry.read().unwrap());
Ok(())
}
pub fn load_module(&self, module_hash: &str) -> Result<Vec<u8>> {
let start: _ = Instant::now();
{
let l1: _ = self.l1_cache.lock().unwrap();
if let Some(entry) = l1.get(module_hash) {
let mut entry = entry.write().unwrap();
entry.update_access();
let load_time: _ = start.elapsed();
self.record_load_time(load_time);
self.update_stats_after_hit(true);
return Ok(entry.wasm_bytes.clone());
}
}
if self.config.enable_l2 {
let wasm_bytes: _ = self.load_from_l2(module_hash)?;
if !wasm_bytes.is_empty() {
self.store_module(module_hash.to_string(), wasm_bytes.clone())?;
let load_time: _ = start.elapsed();
self.record_load_time(load_time);
self.update_stats_after_hit(true);
return Ok(wasm_bytes);
}
}
self.update_stats_after_hit(false);
Err(anyhow!("Module not found in cache: {}", module_hash))
}
pub fn contains(&self, module_hash: &str) -> bool {
let l1: _ = self.l1_cache.lock().unwrap();
if l1.contains_key(module_hash) {
return true;
}
if self.config.enable_l2 {
let l2: _ = self.l2_cache.lock().unwrap();
return l2.contains_key(module_hash);
}
false
}
pub fn warmup_cache(&self, modules: Vec<(String, Vec<u8>)>) -> Result<()> {
for (hash, wasm_bytes) in modules {
self.store_module(hash, wasm_bytes)?;
}
Ok(())
}
pub fn clear_cache(&self) -> Result<()> {
{
let mut l1 = self.l1_cache.lock().unwrap();
l1.clear();
}
if self.config.enable_l2 {
let l2: _ = self.l2_cache.lock().unwrap();
for (_, file_path) in l2.iter() {
let _: _ = std::fs::remove_file(file_path);
}
}
{
let mut stats = self.stats.lock().unwrap();
*stats = CacheStats::default();
}
Ok(())
}
pub fn get_stats(&self) -> CacheStats {
let mut stats = self.stats.lock().unwrap();
{
let l1: _ = self.l1_cache.lock().unwrap();
stats.l1_entries = l1.len();
}
if self.config.enable_l2 {
let l2: _ = self.l2_cache.lock().unwrap();
stats.l2_entries = l2.len();
}
let mut total_modules = 0;
let mut total_size = 0;
{
let l1: _ = self.l1_cache.lock().unwrap();
for entry in l1.values() {
let entry = entry.read().unwrap();
total_modules += 1;
total_size += entry.size;
}
}
stats.total_modules = total_modules;
stats.total_size = total_size;
let total_accesses: _ = stats.hits + stats.misses;
stats.hit_ratio = if total_accesses > 0 {
stats.hits as f64 / total_accesses as f64
} else {
0.0
};
{
let load_times: _ = self.load_times.lock().unwrap();
if !load_times.is_empty() {
let total: Duration = load_times.iter().sum();
stats.avg_load_time = Duration::from_nanos(total.as_nanos() as u64 / load_times.len() as u64);
}
}
stats.clone()
}
pub fn calculate_hash(&self, wasm_bytes: &[u8]) -> String {
let mut hasher = Hasher::new();
hasher.update(wasm_bytes);
hasher.finalize().to_hex().to_string()
}
fn evict_l1_cache(&self) -> Result<()> {
let hashes_to_evict: Vec<String> = {
let l1: _ = self.l1_cache.lock().unwrap();
let mut entries: Vec<(String, f64)> = l1.iter()
.map(|(hash, entry)| {
let score: _ = entry.read().unwrap().usage_score();
(hash.clone(), score)
})
.collect();
entries.sort_by(|a, b| {
a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
});
let to_evict: _ = entries.len() / 2;
entries.into_iter().take(to_evict).map(|(hash, _)| hash).collect()
};
{
let mut l1 = self.l1_cache.lock().unwrap();
for hash in hashes_to_evict {
l1.remove(hash.as_str());
}
}
Ok(())
}
fn store_to_l2(&self, module_hash: &str, entry: &Arc<RwLock<CacheEntry>>) -> Result<()> {
let file_path: _ = self.config.l2_cache_dir.join(format!("{}.wasm", module_hash));
let wasm_bytes: _ = {
let entry = entry.read().unwrap();
entry.wasm_bytes.clone()
};
std::fs::write(&file_path, wasm_bytes)
.context("Failed to write to L2 cache")?;
{
let mut l2 = self.l2_cache.lock().unwrap();
l2.insert(module_hash.to_string(), file_path.clone());
}
{
let mut entry = entry.write().unwrap();
entry.file_path = Some(file_path);
}
Ok(())
}
fn load_from_l2(&self, module_hash: &str) -> Result<Vec<u8>> {
let l2: _ = self.l2_cache.lock().unwrap();
if let Some(file_path) = l2.get(module_hash) {
let wasm_bytes: _ = std::fs::read(file_path)
.context("Failed to read from L2 cache")?;
let _: _ = filetime::set_file_atime(
file_path,
filetime::FileTime::now()
);
Ok(wasm_bytes)
} else {
Ok(Vec::new())
}
}
fn record_load_time(&self, load_time: Duration) {
let mut times = self.load_times.lock().unwrap();
times.push(load_time);
if times.len() > 1000 {
times.remove(0);
}
}
fn update_stats_after_hit(&self, hit: bool) {
let mut stats = self.stats.lock().unwrap();
if hit {
stats.hits += 1;
} else {
stats.misses += 1;
}
}
fn update_stats_after_store(&self, entry: &CacheEntry) {
let mut stats = self.stats.lock().unwrap();
stats.total_modules += 1;
stats.total_size += entry.size;
}
pub fn cleanup_expired(&self) -> Result<usize> {
let mut cleaned = 0;
let now: _ = Instant::now();
{
let mut l1 = self.l1_cache.lock().unwrap();
let expired_keys: Vec<String> = l1.iter()
.filter_map(|(hash, entry)| {
let entry = entry.read().unwrap();
if entry.age() > self.config.expiration_time {
Some(hash.clone())
} else {
None
}
})
.collect();
for key in expired_keys {
l1.remove(&key);
cleaned += 1;
}
}
if self.config.enable_l2 {
let mut l2 = self.l2_cache.lock().unwrap();
let sys_now: _ = std::time::SystemTime::now();
let expired_keys: Vec<String> = l2.iter()
.filter_map(|(hash, path)| {
if let Ok(metadata) = std::fs::metadata(path) {
if let Ok(modified) = metadata.modified() {
if let Ok(elapsed) = sys_now.duration_since(modified) {
if elapsed > self.config.expiration_time {
return Some(hash.clone());
}
}
}
}
None
})
.collect();
for key in expired_keys {
if let Some(path) = l2.remove(&key) {
let _: _ = std::fs::remove_file(path);
cleaned += 1;
}
}
}
Ok(cleaned)
}
}
impl Drop for WasmModuleCache {
fn drop(&mut self) {
if self.config.enable_l2 {
let l2: _ = self.l2_cache.lock().unwrap();
if l2.is_empty() {
let _: _ = std::fs::remove_dir_all(&self.config.l2_cache_dir);
}
}
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_cache_creation() {
let cache: _ = WasmModuleCache::new();
assert!(cache.is_ok());
}
#[test]
fn test_cache_store_and_load() {
let cache: _ = WasmModuleCache::new().unwrap();
let wasm_bytes: _ = vec![0, 1, 2, 3, 4];
let hash: _ = cache.calculate_hash(&wasm_bytes);
let result: _ = cache.store_module(hash.clone(), wasm_bytes.clone());
assert!(result.is_ok());
let loaded: _ = cache.load_module(&hash);
assert!(loaded.is_ok());
assert_eq!(loaded.unwrap(), wasm_bytes);
}
#[test]
fn test_cache_contains() {
let cache: _ = WasmModuleCache::new().unwrap();
let wasm_bytes: _ = vec![0, 1, 2, 3, 4];
let hash: _ = cache.calculate_hash(&wasm_bytes);
assert!(!cache.contains(&hash));
cache.store_module(hash.clone(), wasm_bytes).unwrap();
assert!(cache.contains(&hash));
}
#[test]
fn test_cache_miss() {
let cache: _ = WasmModuleCache::new().unwrap();
let result: _ = cache.load_module("nonexistent");
assert!(result.is_err());
}
#[test]
fn test_cache_stats() {
let cache: _ = WasmModuleCache::new().unwrap();
let wasm_bytes: _ = vec![0, 1, 2, 3, 4];
let hash: _ = cache.calculate_hash(&wasm_bytes);
cache.store_module(hash.clone(), wasm_bytes).unwrap();
cache.load_module(&hash).unwrap();
let stats: _ = cache.get_stats();
assert_eq!(stats.l1_entries, 1);
assert!(stats.hits > 0);
assert_eq!(stats.misses, 0);
assert!(stats.hit_ratio > 0.0);
}
#[test]
fn test_cache_clear() {
let cache: _ = WasmModuleCache::new().unwrap();
let wasm_bytes: _ = vec![0, 1, 2, 3, 4];
let hash: _ = cache.calculate_hash(&wasm_bytes);
cache.store_module(hash.clone(), wasm_bytes).unwrap();
cache.clear_cache().unwrap();
assert!(!cache.contains(&hash));
let stats: _ = cache.get_stats();
assert_eq!(stats.l1_entries, 0);
assert_eq!(stats.total_modules, 0);
}
#[test]
fn test_hash_calculation() {
let cache: _ = WasmModuleCache::new().unwrap();
let wasm_bytes: _ = vec![0, 1, 2, 3, 4];
let hash1: _ = cache.calculate_hash(&wasm_bytes);
let wasm_bytes2: _ = vec![0, 1, 2, 3, 5];
let hash2: _ = cache.calculate_hash(&wasm_bytes2);
assert_ne!(hash1, hash2);
}
#[test]
fn test_cache_warmup() {
let cache: _ = WasmModuleCache::new().unwrap();
let modules: _ = vec![
("hash1".to_string(), vec![0, 1, 2]),
("hash2".to_string(), vec![3, 4, 5]),
];
let result: _ = cache.warmup_cache(modules);
assert!(result.is_ok());
let stats: _ = cache.get_stats();
assert_eq!(stats.l1_entries, 2);
}
#[test]
fn test_batch_operations() {
let cache: _ = WasmModuleCache::new().unwrap();
for i in 0..10 {
let wasm_bytes: _ = vec![i as u8; 100];
let hash: _ = cache.calculate_hash(&wasm_bytes);
cache.store_module(hash, wasm_bytes).unwrap();
}
for i in 0..10 {
let wasm_bytes: _ = vec![i as u8; 100];
let hash: _ = cache.calculate_hash(&wasm_bytes);
let result: _ = cache.load_module(&hash);
assert!(result.is_ok());
}
let stats: _ = cache.get_stats();
assert!(stats.hit_ratio > 0.8);
}
}