use std::{
env::current_dir,
path::{Path, PathBuf},
};
use anyhow::{Context, Error};
use parking_lot::Mutex;
use swc_common::{
collections::AHashMap,
sync::{Lazy, OnceCell},
};
use wasmer::{Module, Store};
#[cfg(all(not(target_arch = "wasm32"), feature = "filesystem_cache"))]
use wasmer_cache::{Cache as WasmerCache, FileSystemCache, Hash};
#[cfg(all(not(feature = "filesystem_cache"), not(feature = "memory_cache")))]
compile_error!("Plugin_runner should enable either filesystem, or memory cache");
#[cfg(all(feature = "filesystem_cache", feature = "memory_cache"))]
compile_error!(
"Only one cache feature should be enabled. If you enabled filesystem_cache, it activates its \
memory cache as well."
);
const MODULE_SERIALIZATION_VERSION: &str = "v3";
pub static PLUGIN_MODULE_CACHE: Lazy<PluginModuleCache> = Lazy::new(Default::default);
#[cfg(feature = "filesystem_cache")]
#[derive(Default)]
pub struct CacheInner {
fs_cache: Option<FileSystemCache>,
loaded_module_bytes: AHashMap<PathBuf, Module>,
}
#[cfg(feature = "memory_cache")]
#[derive(Default)]
pub struct CacheInner {
loaded_module_bytes: AHashMap<PathBuf, Vec<u8>>,
}
#[derive(Default)]
pub struct PluginModuleCache {
inner: OnceCell<Mutex<CacheInner>>,
instantiation_lock: Mutex<()>,
}
#[cfg(feature = "filesystem_cache")]
#[tracing::instrument(level = "info", skip_all)]
fn create_filesystem_cache(filesystem_cache_root: &Option<String>) -> Option<FileSystemCache> {
let mut root_path = if let Some(root) = filesystem_cache_root {
Some(PathBuf::from(root))
} else if let Ok(cwd) = current_dir() {
Some(cwd.join(".swc"))
} else {
None
};
if let Some(root_path) = &mut root_path {
root_path.push("plugins");
root_path.push(MODULE_SERIALIZATION_VERSION);
return FileSystemCache::new(&root_path).ok();
}
None
}
#[cfg(feature = "filesystem_cache")]
pub fn init_plugin_module_cache_once(filesystem_cache_root: &Option<String>) {
PLUGIN_MODULE_CACHE.inner.get_or_init(|| {
Mutex::new(CacheInner {
fs_cache: create_filesystem_cache(filesystem_cache_root),
loaded_module_bytes: Default::default(),
})
});
}
#[cfg(feature = "memory_cache")]
pub fn init_plugin_module_cache_once() {
PLUGIN_MODULE_CACHE.inner.get_or_init(|| {
Mutex::new(CacheInner {
loaded_module_bytes: Default::default(),
})
});
}
impl PluginModuleCache {
pub fn new() -> Self {
PluginModuleCache {
inner: OnceCell::from(Mutex::new(Default::default())),
instantiation_lock: Mutex::new(()),
}
}
#[cfg(feature = "filesystem_cache")]
#[tracing::instrument(level = "info", skip_all)]
pub fn load_module(&self, binary_path: &Path) -> Result<Module, Error> {
let binary_path = binary_path.to_path_buf();
let mut inner_cache = self.inner.get().expect("Cache should be available").lock();
let in_memory_module = inner_cache.loaded_module_bytes.get(&binary_path);
if let Some(module) = in_memory_module {
return Ok(module.clone());
}
let module_bytes =
std::fs::read(&binary_path).context("Cannot read plugin from specified path")?;
let module_bytes_hash = Hash::generate(&module_bytes);
let wasmer_store = Store::default();
let load_cold_wasm_bytes = || {
let span = tracing::span!(
tracing::Level::INFO,
"load_cold_wasm_bytes",
plugin_module = binary_path.to_str()
);
let span_guard = span.enter();
let _lock = self.instantiation_lock.lock();
let ret =
Module::new(&wasmer_store, module_bytes).context("Cannot compile plugin binary");
drop(span_guard);
ret
};
let module = if let Some(fs_cache) = &mut inner_cache.fs_cache {
let load_result = unsafe { fs_cache.load(&wasmer_store, module_bytes_hash) };
if let Ok(module) = load_result {
module
} else {
let cold_bytes = load_cold_wasm_bytes()?;
fs_cache.store(module_bytes_hash, &cold_bytes)?;
cold_bytes
}
} else {
load_cold_wasm_bytes()?
};
inner_cache
.loaded_module_bytes
.insert(binary_path, module.clone());
Ok(module)
}
#[cfg(feature = "memory_cache")]
#[tracing::instrument(level = "info", skip_all)]
pub fn load_module(&self, binary_path: &Path) -> Result<Module, Error> {
let binary_path = binary_path.to_path_buf();
let mut inner_cache = self.inner.get().expect("Cache should be available").lock();
let in_memory_module_bytes = inner_cache
.loaded_module_bytes
.get(&binary_path)
.ok_or_else(|| {
anyhow::anyhow!("Could not locate plugin binary {}", binary_path.display())
})?;
let wasmer_store = Store::default();
let module = Module::new(&wasmer_store, in_memory_module_bytes)?;
Ok(module)
}
#[cfg(feature = "memory_cache")]
#[tracing::instrument(level = "info", skip_all)]
pub fn store_once(&self, module_name: &str, module_bytes: Vec<u8>) {
let binary_path = PathBuf::from(module_name);
let mut inner_cache = self.inner.get().expect("Cache should be available").lock();
if !inner_cache.loaded_module_bytes.contains_key(&binary_path) {
inner_cache
.loaded_module_bytes
.insert(binary_path, module_bytes);
}
}
}