use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::fs::File;
use tracing::{debug, info, warn};
use wasmtime::{Config, Engine, Module};
use anyhow::{Context, Result};
use memmap2::{Mmap, MmapOptions};
use std::time::Instant;
use std::path::Path;
use std::path::PathBuf;
use std::num::NonZeroUsize;
#[derive(Debug, Clone)]
pub struct ZeroCopyLoadResult {
pub module_name: String,
pub load_time_ms: f64,
pub file_size_bytes: u64,
pub cache_hit: bool,
pub memory_mapped: bool,
pub precompiled: bool,
}
#[derive(Debug, Clone)]
pub struct LoadStats {
pub total_loads: u64,
pub cache_hits: u64,
pub cache_misses: u64,
pub avg_load_time_ms: f64,
pub cache_hit_rate: f64,
}
pub struct WasmZeroCopyLoader {
engine: Arc<Engine>,
module_cache: Arc<RwLock<LruCache<String, Arc<Module>>>>,
memory_maps: Arc<RwLock<HashMap<String, Arc<Mmap>>>>,
load_stats: Arc<RwLock<LoadStats>>,
cache_size: usize,
prewarm_enabled: bool,
}
impl WasmZeroCopyLoader {
pub fn new(cache_size: usize, prewarm_enabled: bool) -> Result<Self> {
info!("🚀 初始化 WASM 零拷贝加载器 (缓存大小: {}, 预热: {})", cache_size, prewarm_enabled);
let mut config = Config::new();
config
.debug_info(false)
.wasm_threads(true)
.wasm_simd(true)
.parallel_compilation(true);
let engine: _ = Arc::new(Mutex::new(Engine::new(&config)));
let loader: _ = Self {
engine,
module_cache: Arc::new(Mutex::new(LruCache::new(NonZero::new(cache_size).unwrap_or(NonZero::new(100).unwrap())))),
memory_maps: Arc::new(Mutex::new(HashMap::new())),
load_stats: Arc::new(Mutex::new(LoadStats {
total_loads: 0,
cache_hits: 0,
cache_misses: 0,
avg_load_time_ms: 0.0,
cache_hit_rate: 0.0,
})),
cache_size,
prewarm_enabled,
};
if prewarm_enabled {
info!("🔥 预热功能已启用");
}
Ok(loader)
}
pub async fn load_zero_copy(&self, name: &str, file_path: &Path) -> Result<ZeroCopyLoadResult> {
let start_time: _ = std::time::Instant::now();
let cache_hit: _ = self.check_cache(name).await?;
if cache_hit {
let load_time: _ = start_time.elapsed().as_secs_f64() * 1000.0;
self.update_stats(true, load_time).await;
let result: _ = ZeroCopyLoadResult {
module_name: name.to_string(),
load_time_ms: load_time,
file_size_bytes: 0,
cache_hit: true,
memory_mapped: false,
precompiled: true,
};
debug!("✅ 缓存命中: {} (耗时: {:.2}ms)", name, load_time);
return Ok(result);
}
let memory_mapped: _ = self.memory_map_file(name, file_path).await?;
let module: _ = self.create_module_from_memory(name, file_path).await?;
self.cache_module(name, module).await?;
let load_time: _ = start_time.elapsed().as_secs_f64() * 1000.0;
let file_size: _ = tokio::fs::metadata(file_path).await?.len();
self.update_stats(false, load_time).await;
let result: _ = ZeroCopyLoadResult {
module_name: name.to_string(),
load_time_ms: load_time,
file_size_bytes: file_size,
cache_hit: false,
memory_mapped,
precompiled: true,
};
info!("✅ 零拷贝加载完成: {} (耗时: {:.2}ms, 文件大小: {} bytes)",
name, load_time, file_size);
Ok(result)
}
pub async fn prewarm_modules(&self, modules: Vec<(String, PathBuf)>) -> Result<()> {
if !self.prewarm_enabled {
warn!("⚠️ 预热功能未启用");
return Ok(());
}
info!("🔥 开始预热 {} 个模块", modules.len());
let start_time: _ = std::time::Instant::now();
let module_count: _ = modules.len();
let mut success_count = 0;
for (name, path) in modules.into_iter() {
match self.load_zero_copy(&name, &path).await {
Ok(_) => success_count += 1,
Err(_) => debug!("预热失败: {}", name),
}
}
let prewarm_time: _ = start_time.elapsed().as_secs_f64() * 1000.0;
info!("✅ 预热完成: {}/{} 成功 (耗时: {:.2}ms)", success_count, module_count, prewarm_time);
Ok(())
}
pub async fn load_batch(&self, modules: Vec<(String, PathBuf)>) -> Result<Vec<ZeroCopyLoadResult>> {
let module_count: _ = modules.len();
info!("📦 批量加载 {} 个模块", module_count);
let start_time: _ = std::time::Instant::now();
let mut successful_results = Vec::new();
for (name, path) in modules.into_iter() {
match self.load_zero_copy(&name, &path).await {
Ok(load_result) => successful_results.push(load_result),
Err(_) => debug!("加载失败: {}", name),
}
}
let batch_time: _ = start_time.elapsed().as_secs_f64() * 1000.0;
let avg_time: _ = if successful_results.len() > 0 {
batch_time / successful_results.len() as f64
} else {
0.0
};
info!("✅ 批量加载完成: {}/{} 成功 (总耗时: {:.2}ms, 平均: {:.2}ms)",
successful_results.len(), module_count, batch_time, avg_time);
Ok(successful_results)
}
async fn check_cache(&self, name: &str) -> Result<bool> {
let cache: _ = self.module_cache.read().await;
Ok(cache.contains(name))
}
async fn memory_map_file(&self, name: &str, file_path: &Path) -> Result<bool> {
let file: _ = tokio::fs::File::open(file_path).await
.with_context(|| format!("打开文件失败: {:?}", file_path))?;
let metadata: _ = file.metadata().await
.with_context(|| format!("获取文件元数据失败: {:?}", file_path))?;
if metadata.len() < 1024 * 1024 {
return Ok(false);
}
let mmap: _ = unsafe {
MmapOptions::new()
.map(&file)
.with_context(|| format!("内存映射失败: {:?}", file_path))?
};
let mut memory_maps = self.memory_maps.write().await;
memory_maps.insert(name.to_string(), Arc::new(Mutex::new(mmap)),
Ok(true))
}
async fn create_module_from_memory(&self, name: &str, file_path: &Path) -> Result<Arc<Module>> {
let wasm_bytes: _ = tokio::fs::read(file_path).await
.with_context(|| format!("读取文件失败: {:?}", file_path))?;
let module: _ = Module::from_binary(self.engine.as_ref(), &wasm_bytes)
.with_context(|| format!("创建模块失败: {}", name))?;
Ok(Arc::new(Mutex::new(module)),)
}
async fn cache_module(&self, name: &str, module: Arc<Module>) -> Result<()> {
let mut cache = self.module_cache.write().await;
cache.put(name.to_string(), module);
Ok(())
}
async fn update_stats(&self, cache_hit: bool, load_time: f64) {
let mut stats = self.load_stats.write().await;
stats.total_loads += 1;
if cache_hit {
stats.cache_hits += 1;
} else {
stats.cache_misses += 1;
}
let total: _ = stats.total_loads as f64;
stats.avg_load_time_ms = (stats.avg_load_time_ms * (total - 1.0) + load_time) / total;
if stats.total_loads > 0 {
stats.cache_hit_rate = stats.cache_hits as f64 / stats.total_loads as f64 * 100.0;
}
}
pub async fn get_stats(&self) -> LoadStats {
let stats: _ = self.load_stats.read().await;
stats.clone()
}
pub async fn clear_cache(&self) -> Result<()> {
info!("🧹 清理 WASM 模块缓存");
let mut cache = self.module_cache.write().await;
cache.clear();
let mut memory_maps = self.memory_maps.write().await;
memory_maps.clear();
let mut stats = self.load_stats.write().await;
*stats = LoadStats {
total_loads: 0,
cache_hits: 0,
cache_misses: 0,
avg_load_time_ms: 0.0,
cache_hit_rate: 0.0,
};
info!("✅ 缓存清理完成");
Ok(())
}
pub async fn precompile(&self, wasm_bytes: &[u8]) -> Result<Arc<Module>> {
let start_time: _ = std::time::Instant::now();
let module: _ = Module::from_binary(self.engine.as_ref(), wasm_bytes)
.context("预编译模块失败")?;
let compile_time: _ = start_time.elapsed().as_secs_f64() * 1000.0;
debug!("⚡ 预编译完成 (耗时: {:.2}ms)", compile_time);
Ok(Arc::new(Mutex::new(module)),)
}
}
impl Default for WasmZeroCopyLoader {
fn default() -> Self {
Self::new(100, true).expect("初始化 WasmZeroCopyLoader 失败")
}
}