use std::sync::Arc;
use std::time::Duration;
use tensor_wasm_mem::wasm_memory::TensorWasmMemoryCreator;
use tokio::task::JoinHandle;
use wasmtime::{
Config, Enabled, Engine, InstanceAllocationStrategy, PoolingAllocationConfig, Strategy,
};
const DEFAULT_EPOCH_TICK: Duration = Duration::from_millis(10);
const MAX_WASM_STACK_BYTES: usize = 1_048_576;
pub(crate) const TABLE_ENTRY_BYTES: u64 = 16;
#[derive(Debug, Clone, Default)]
pub enum MemoryBackend {
#[default]
UnifiedBuffer,
PoolingMpk {
max_memories: u32,
memory_bytes: usize,
},
}
#[derive(Debug, Clone)]
pub struct EngineConfig {
pub max_memory_bytes: usize,
pub epoch_tick: Duration,
pub strategy: Strategy,
pub component_model: bool,
pub backend: MemoryBackend,
pub max_module_cache_entries: usize,
pub max_instances: Option<usize>,
pub max_module_bytes: usize,
pub max_concurrent_compiles: Option<usize>,
pub auto_offload: bool,
pub auto_offload_detector: Option<tensor_wasm_jit::detector::DetectorConfig>,
pub max_instances_per_tenant: Option<usize>,
}
impl EngineConfig {
pub fn effective_memory_cap(&self) -> usize {
match self.backend {
MemoryBackend::UnifiedBuffer => self.max_memory_bytes,
MemoryBackend::PoolingMpk { memory_bytes, .. } => {
self.max_memory_bytes.min(memory_bytes)
}
}
}
}
impl Default for EngineConfig {
fn default() -> Self {
Self {
max_memory_bytes: 256 * 1024 * 1024,
epoch_tick: DEFAULT_EPOCH_TICK,
strategy: Strategy::Cranelift,
component_model: true,
backend: MemoryBackend::default(),
max_module_cache_entries: 1024,
max_instances: Some(10_000),
max_module_bytes: crate::executor::MAX_MODULE_BYTES,
max_concurrent_compiles: None,
auto_offload: false,
auto_offload_detector: None,
max_instances_per_tenant: None,
}
}
}
pub struct TensorWasmEngine {
engine: Engine,
ticker_handle: Option<JoinHandle<()>>,
config: EngineConfig,
}
impl TensorWasmEngine {
pub const EPOCH_TICK: Duration = DEFAULT_EPOCH_TICK;
pub fn new() -> Result<Self, wasmtime::Error> {
Self::with_config(EngineConfig::default())
}
pub fn with_config(cfg: EngineConfig) -> Result<Self, wasmtime::Error> {
let mut wt_cfg = Config::new();
wt_cfg.epoch_interruption(true);
wt_cfg.consume_fuel(false);
wt_cfg.wasm_component_model(cfg.component_model);
wt_cfg.strategy(cfg.strategy);
wt_cfg.max_wasm_stack(MAX_WASM_STACK_BYTES);
wt_cfg.wasm_memory64(false);
wt_cfg.wasm_multi_memory(false);
wt_cfg.wasm_relaxed_simd(false);
wt_cfg.wasm_tail_call(false);
wt_cfg.wasm_simd(true);
wt_cfg.wasm_bulk_memory(true);
wt_cfg.wasm_multi_value(true);
match cfg.backend {
MemoryBackend::UnifiedBuffer => {
let memory_creator = Arc::new(TensorWasmMemoryCreator::default());
wt_cfg.with_host_memory(memory_creator);
wt_cfg.guard_before_linear_memory(false);
wt_cfg.memory_init_cow(false);
wt_cfg.memory_reservation(0);
wt_cfg.memory_guard_size(0);
}
MemoryBackend::PoolingMpk {
max_memories,
memory_bytes,
} => {
let effective = cfg.max_memory_bytes.min(memory_bytes);
let mut pooling = PoolingAllocationConfig::default();
pooling.total_memories(max_memories);
pooling.max_memory_size(effective);
pooling.total_tables(max_memories);
let table_elems: usize = (effective as u64 / TABLE_ENTRY_BYTES)
.try_into()
.unwrap_or(usize::MAX);
pooling.table_elements(table_elems);
pooling.memory_protection_keys(Enabled::Auto);
wt_cfg.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling));
}
}
let engine = Engine::new(&wt_cfg)?;
let mut this = Self {
engine,
ticker_handle: None,
config: cfg,
};
if tokio::runtime::Handle::try_current().is_ok() {
this.spawn_epoch_ticker();
}
Ok(this)
}
pub fn inner(&self) -> &Engine {
&self.engine
}
pub fn config(&self) -> &EngineConfig {
&self.config
}
pub fn spawn_epoch_ticker(&mut self) {
if self.ticker_handle.is_some() {
return;
}
let engine = self.engine.clone();
let tick = self.config.epoch_tick;
let handle = tokio::spawn(async move {
loop {
tokio::time::sleep(tick).await;
engine.increment_epoch();
}
});
self.ticker_handle = Some(handle);
}
pub fn stop_epoch_ticker(&mut self) {
if let Some(h) = self.ticker_handle.take() {
h.abort();
}
}
pub fn is_epoch_ticker_running(&self) -> bool {
self.ticker_handle
.as_ref()
.is_some_and(|h| !h.is_finished())
}
pub fn tick(&self) {
self.engine.increment_epoch();
}
}
impl Drop for TensorWasmEngine {
fn drop(&mut self) {
self.stop_epoch_ticker();
}
}
impl Default for TensorWasmEngine {
fn default() -> Self {
Self::new().expect("default TensorWasmEngine construction")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn engine_constructs() {
let _engine = TensorWasmEngine::new().expect("construct");
}
#[tokio::test]
async fn ticker_is_idempotent() {
let mut engine = TensorWasmEngine::new().unwrap();
engine.spawn_epoch_ticker();
engine.spawn_epoch_ticker();
engine.stop_epoch_ticker();
engine.stop_epoch_ticker();
}
#[tokio::test]
async fn manual_tick() {
let engine = TensorWasmEngine::new().unwrap();
engine.tick();
engine.tick();
}
#[test]
fn default_config_values() {
let c = EngineConfig::default();
assert_eq!(c.max_memory_bytes, 256 * 1024 * 1024);
assert_eq!(c.epoch_tick, Duration::from_millis(10));
assert!(c.component_model);
assert!(matches!(c.backend, MemoryBackend::UnifiedBuffer));
}
#[test]
fn effective_memory_cap_unified_is_max_memory_bytes() {
let cfg = EngineConfig {
max_memory_bytes: 128 * 1024 * 1024,
backend: MemoryBackend::UnifiedBuffer,
..EngineConfig::default()
};
assert_eq!(cfg.effective_memory_cap(), 128 * 1024 * 1024);
}
#[test]
fn effective_memory_cap_pooling_takes_minimum() {
let cfg = EngineConfig {
max_memory_bytes: 256 * 1024 * 1024,
backend: MemoryBackend::PoolingMpk {
max_memories: 8,
memory_bytes: 64 * 1024 * 1024,
},
..EngineConfig::default()
};
assert_eq!(cfg.effective_memory_cap(), 64 * 1024 * 1024);
let cfg = EngineConfig {
max_memory_bytes: 16 * 1024 * 1024,
backend: MemoryBackend::PoolingMpk {
max_memories: 8,
memory_bytes: 64 * 1024 * 1024,
},
..EngineConfig::default()
};
assert_eq!(cfg.effective_memory_cap(), 16 * 1024 * 1024);
}
#[tokio::test]
async fn engine_constructs_with_unified_backend() {
let cfg = EngineConfig {
backend: MemoryBackend::UnifiedBuffer,
..EngineConfig::default()
};
let engine = TensorWasmEngine::with_config(cfg);
assert!(
engine.is_ok(),
"engine should construct: {:?}",
engine.err()
);
}
#[tokio::test]
async fn engine_constructs_with_pooling_mpk_backend() {
let cfg = EngineConfig {
backend: MemoryBackend::PoolingMpk {
max_memories: 32,
memory_bytes: 64 * 1024,
},
..EngineConfig::default()
};
let engine = TensorWasmEngine::with_config(cfg);
assert!(
engine.is_ok(),
"engine should construct: {:?}",
engine.err()
);
}
}