use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use dashmap::{mapref::entry::Entry, DashMap};
use lru::LruCache;
use tensor_wasm_core::metrics::TensorWasmMetrics;
use tensor_wasm_core::types::{InstanceId, TenantId};
use tensor_wasm_jit::cache::KernelCache;
use tensor_wasm_jit::rewrite::{rewrite_wasm, RewriteOptions};
use thiserror::Error;
use tokio::sync::{Mutex, Semaphore};
use tracing::{debug, info, instrument, warn};
use wasmtime::{ExternType, Module, ResourceLimiter, Store, UpdateDeadline, Val};
use crate::engine::TensorWasmEngine;
use crate::instance::{InstanceState, TensorWasmInstance};
use crate::instance_pool::{InstancePool, ModuleHash};
use crate::jit_dispatch::add_jit_dispatch_to_linker;
use tensor_wasm_wasi_gpu::scheduler::add_scheduler_to_linker;
use tensor_wasm_wasi_gpu::streaming::{
add_input_to_linker, add_streaming_to_linker, InputContext, StreamingContext,
};
fn duration_to_epoch_ticks(d: Duration, tick: Duration) -> u64 {
let d_ms = d.as_millis();
let t_ms = tick.as_millis().max(1);
let ticks_u128 = d_ms.div_ceil(t_ms).max(1);
u64::try_from(ticks_u128)
.unwrap_or(MAX_EPOCH_DEADLINE_TICKS)
.min(MAX_EPOCH_DEADLINE_TICKS)
}
fn arm_cooperative_epoch(store: &mut Store<InstanceState>, first_ticks: u64) {
store.set_epoch_deadline(first_ticks);
store.epoch_deadline_callback(|ctx| {
if ctx.data().hard_deadline_elapsed() {
Err(wasmtime::Error::msg("epoch deadline reached"))
} else {
Ok(UpdateDeadline::Yield(COOPERATIVE_YIELD_TICKS))
}
});
}
const MAX_EPOCH_DEADLINE_TICKS: u64 = u64::MAX / 2;
const COOPERATIVE_YIELD_TICKS: u64 = 1;
pub const MAX_START_FN_DURATION: Duration = Duration::from_secs(30);
pub const MAX_MODULE_BYTES: usize = 64 * 1024 * 1024;
#[derive(Debug, Error)]
pub enum ExecError {
#[error("wasmtime error")]
Wasmtime(#[from] wasmtime::Error),
#[error("no such instance: {0}")]
NotFound(InstanceId),
#[error("instance has no export `{0}`")]
MissingExport(String),
#[error("{0}")]
Timeout(TimeoutContext),
#[error("module-declared linear memory {requested_bytes} bytes exceeds engine cap {limit_bytes} bytes")]
ModuleMemoryTooLarge {
requested_bytes: u64,
limit_bytes: u64,
},
#[error("instance capacity exhausted: {active} active, limit {limit}")]
CapacityExhausted {
active: usize,
limit: usize,
},
#[error("tenant {tenant} instance capacity exhausted: {active} active, limit {limit}")]
TenantCapacityExhausted {
tenant: TenantId,
active: usize,
limit: usize,
},
#[error("wasm module byte length {len} exceeds cap {max}")]
ModuleTooLarge {
len: usize,
max: usize,
},
#[error("epoch ticker not running — refusing spawn with deadline; call `engine.spawn_epoch_ticker()` first")]
EpochTickerNotRunning,
}
#[derive(Debug, Clone, Copy)]
pub struct TimeoutContext {
pub id: InstanceId,
pub elapsed_ms: u64,
pub deadline_ms: u64,
}
impl std::fmt::Display for TimeoutContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"instance {} exceeded deadline (elapsed {} ms, deadline {} ms)",
self.id, self.elapsed_ms, self.deadline_ms,
)
}
}
impl From<ExecError> for tensor_wasm_core::error::TensorWasmError {
fn from(e: ExecError) -> Self {
use tensor_wasm_core::error::TensorWasmError;
match e {
ExecError::Wasmtime(err) => {
let is_trap = err.downcast_ref::<wasmtime::Trap>().is_some();
tracing::error!(
target: "tensor_wasm_exec::executor",
error = ?err,
error_chain = %format!("{err:#}"),
is_trap,
"wasmtime trap",
);
if is_trap {
TensorWasmError::WasmTrap("wasm trap".into())
} else {
TensorWasmError::WasmCompile("wasm compile failed".into())
}
}
ExecError::NotFound(id) => {
TensorWasmError::Serialization(format!("instance not found: {id}").into())
}
ExecError::MissingExport(name) => {
TensorWasmError::Serialization(format!("instance missing export: {name}").into())
}
ExecError::Timeout(ctx) => TensorWasmError::KernelTimeout {
elapsed_ms: ctx.elapsed_ms,
deadline_ms: ctx.deadline_ms,
},
ExecError::ModuleMemoryTooLarge {
requested_bytes,
limit_bytes,
} => TensorWasmError::MemoryExhausted {
requested: requested_bytes,
limit: limit_bytes,
},
ExecError::CapacityExhausted { active, limit } => TensorWasmError::MemoryExhausted {
requested: active as u64,
limit: limit as u64,
},
ExecError::TenantCapacityExhausted { active, limit, .. } => {
TensorWasmError::MemoryExhausted {
requested: active as u64,
limit: limit as u64,
}
}
ExecError::ModuleTooLarge { len, max } => TensorWasmError::MemoryExhausted {
requested: len as u64,
limit: max as u64,
},
ExecError::EpochTickerNotRunning => {
TensorWasmError::WasmCompile("epoch ticker not running".into())
}
}
}
}
#[derive(Debug, Clone)]
pub struct SpawnConfig {
pub tenant_id: TenantId,
pub deadline: Option<Duration>,
pub args: Vec<WasmArg>,
pub streaming: Option<StreamingContext>,
pub input: Vec<u8>,
}
impl SpawnConfig {
pub fn for_tenant(tenant_id: TenantId) -> Self {
Self {
tenant_id,
deadline: None,
args: Vec::new(),
streaming: None,
input: Vec::new(),
}
}
pub fn with_deadline(mut self, deadline: Duration) -> Self {
self.deadline = Some(deadline);
self
}
pub fn with_args(mut self, args: Vec<WasmArg>) -> Self {
self.args = args;
self
}
pub fn with_streaming(mut self, ctx: StreamingContext) -> Self {
self.streaming = Some(ctx);
self
}
pub fn with_input(mut self, input: Vec<u8>) -> Self {
self.input = input;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
pub enum WasmArg {
I32(i32),
I64(i64),
F32(f32),
F64(f64),
}
impl WasmArg {
pub fn from_json(v: &serde_json::Value) -> Result<Self, &'static str> {
match v {
serde_json::Value::Number(n) => {
if let Some(i) = n.as_i64() {
if let Ok(i32v) = i32::try_from(i) {
Ok(WasmArg::I32(i32v))
} else {
Ok(WasmArg::I64(i))
}
} else if let Some(f) = n.as_f64() {
Ok(WasmArg::F64(f))
} else {
Err("unsupported number")
}
}
_ => Err("unsupported arg type — only numbers"),
}
}
pub fn into_val(self) -> wasmtime::Val {
match self {
WasmArg::I32(v) => wasmtime::Val::I32(v),
WasmArg::I64(v) => wasmtime::Val::I64(v),
WasmArg::F32(v) => wasmtime::Val::F32(v.to_bits()),
WasmArg::F64(v) => wasmtime::Val::F64(v.to_bits()),
}
}
}
fn val_to_json(v: &Val) -> serde_json::Value {
match v {
Val::I32(n) => serde_json::json!(*n),
Val::I64(n) => serde_json::json!(*n),
Val::F32(bits) => serde_json::json!(f32::from_bits(*bits)),
Val::F64(bits) => serde_json::json!(f64::from_bits(*bits)),
other => {
let kind = match other {
Val::V128(_) => "v128",
Val::FuncRef(_) => "funcref",
Val::ExternRef(_) => "externref",
Val::AnyRef(_) => "anyref",
_ => "unknown",
};
debug!(
target: "tensor_wasm_exec::executor",
val_kind = kind,
"non-numeric wasm return value mapped to JSON null (v128 / reference types are not representable)",
);
serde_json::Value::Null
}
}
}
#[derive(Debug)]
pub struct TensorWasmResourceLimiter {
engine_max: usize,
}
impl TensorWasmResourceLimiter {
pub fn new(engine_max: usize) -> Self {
Self { engine_max }
}
}
impl ResourceLimiter for TensorWasmResourceLimiter {
fn memory_growing(
&mut self,
_current: usize,
desired: usize,
maximum: Option<usize>,
) -> wasmtime::Result<bool> {
if desired > self.engine_max {
return Ok(false);
}
if let Some(m) = maximum {
if desired > m {
return Ok(false);
}
}
Ok(true)
}
fn table_growing(
&mut self,
_current: usize,
desired: usize,
maximum: Option<usize>,
) -> wasmtime::Result<bool> {
let bytes_needed = (desired as u64).saturating_mul(crate::engine::TABLE_ENTRY_BYTES);
if bytes_needed > self.engine_max as u64 {
return Ok(false);
}
if let Some(m) = maximum {
if desired > m {
return Ok(false);
}
}
Ok(true)
}
}
#[derive(Clone)]
pub struct TensorWasmExecutor {
engine: Arc<TensorWasmEngine>,
instances: Arc<DashMap<InstanceId, Arc<Mutex<TensorWasmInstance>>>>,
next_instance_id: Arc<AtomicU64>,
module_cache: Arc<parking_lot::Mutex<LruCache<[u8; 32], Module>>>,
instance_count: Arc<AtomicUsize>,
tenant_counts: Arc<DashMap<TenantId, usize>>,
metrics: Option<TensorWasmMetrics>,
ticker_warned: Arc<OnceLock<AtomicBool>>,
pool: Option<Arc<InstancePool>>,
jit_cache: Option<Arc<KernelCache>>,
compile_semaphore: Arc<Semaphore>,
}
fn resolve_compile_permits(requested: Option<usize>) -> usize {
match requested {
Some(0) | None => std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
.max(1),
Some(n) => n,
}
.max(1)
}
fn lru_cap(requested: usize) -> NonZeroUsize {
NonZeroUsize::new(requested).unwrap_or_else(|| NonZeroUsize::new(1).expect("1 is non-zero"))
}
struct InstanceSlotGuard {
counter: Arc<AtomicUsize>,
tenant_rollback: Option<(Arc<DashMap<TenantId, usize>>, TenantId)>,
committed: bool,
}
impl InstanceSlotGuard {
fn new(counter: Arc<AtomicUsize>) -> Self {
Self {
counter,
tenant_rollback: None,
committed: false,
}
}
fn with_tenant(
counter: Arc<AtomicUsize>,
tenant_counts: Arc<DashMap<TenantId, usize>>,
tenant: TenantId,
) -> Self {
Self {
counter,
tenant_rollback: Some((tenant_counts, tenant)),
committed: false,
}
}
fn commit(mut self) {
self.committed = true;
}
}
impl Drop for InstanceSlotGuard {
fn drop(&mut self) {
if !self.committed {
self.counter.fetch_sub(1, Ordering::Relaxed);
if let Some((tenant_counts, tenant)) = &self.tenant_rollback {
decrement_tenant_count(tenant_counts, *tenant);
}
}
}
}
fn decrement_tenant_count(tenant_counts: &DashMap<TenantId, usize>, tenant: TenantId) {
if let Entry::Occupied(mut e) = tenant_counts.entry(tenant) {
let v = e.get_mut();
*v = v.saturating_sub(1);
if *v == 0 {
e.remove();
}
}
}
fn check_module_memory_within_cap(module: &Module, cap_bytes: usize) -> Result<(), ExecError> {
let cap_u64 = cap_bytes as u64;
let check = |mt: &wasmtime::MemoryType| -> Result<(), ExecError> {
let page_size = mt.page_size();
let min_pages = mt.minimum();
let min_bytes = min_pages.saturating_mul(page_size);
if min_bytes > cap_u64 {
return Err(ExecError::ModuleMemoryTooLarge {
requested_bytes: min_bytes,
limit_bytes: cap_u64,
});
}
if let Some(max_pages) = mt.maximum() {
let max_bytes = max_pages.saturating_mul(page_size);
if max_bytes > cap_u64 {
return Err(ExecError::ModuleMemoryTooLarge {
requested_bytes: max_bytes,
limit_bytes: cap_u64,
});
}
}
Ok(())
};
for ex in module.exports() {
if let ExternType::Memory(mt) = ex.ty() {
check(&mt)?;
}
}
for im in module.imports() {
if let ExternType::Memory(mt) = im.ty() {
check(&mt)?;
}
}
Ok(())
}
fn check_raw_module_memory_within_cap(wasm: &[u8], cap_bytes: usize) -> Result<(), ExecError> {
use wasmparser::{Parser, Payload, TypeRef};
let cap_u64 = cap_bytes as u64;
let check = |mt: &wasmparser::MemoryType| -> Result<(), ExecError> {
let page_size: u64 = 1u64 << u64::from(mt.page_size_log2.unwrap_or(16));
let min_bytes = mt.initial.saturating_mul(page_size);
if min_bytes > cap_u64 {
return Err(ExecError::ModuleMemoryTooLarge {
requested_bytes: min_bytes,
limit_bytes: cap_u64,
});
}
if let Some(max_pages) = mt.maximum {
let max_bytes = max_pages.saturating_mul(page_size);
if max_bytes > cap_u64 {
return Err(ExecError::ModuleMemoryTooLarge {
requested_bytes: max_bytes,
limit_bytes: cap_u64,
});
}
}
Ok(())
};
for payload in Parser::new(0).parse_all(wasm) {
let Ok(payload) = payload else { return Ok(()) };
match payload {
Payload::MemorySection(reader) => {
for mem in reader {
let Ok(mt) = mem else { return Ok(()) };
check(&mt)?;
}
}
Payload::ImportSection(reader) => {
for im in reader {
let Ok(im) = im else { return Ok(()) };
if let TypeRef::Memory(mt) = im.ty {
check(&mt)?;
}
}
}
_ => {}
}
}
Ok(())
}
fn classify_instantiation_error(err: wasmtime::Error, cap_bytes: usize) -> ExecError {
let chain = format!("{err:#}").to_ascii_lowercase();
let is_memory_size_failure = chain.contains("memory")
&& (chain.contains("exceeds the limit") || chain.contains("exceeds the configured"));
if is_memory_size_failure {
ExecError::ModuleMemoryTooLarge {
requested_bytes: cap_bytes as u64,
limit_bytes: cap_bytes as u64,
}
} else {
ExecError::Wasmtime(err)
}
}
impl TensorWasmExecutor {
pub fn new(engine: Arc<TensorWasmEngine>) -> Self {
let cap = lru_cap(engine.config().max_module_cache_entries);
let permits = resolve_compile_permits(engine.config().max_concurrent_compiles);
Self {
engine,
instances: Arc::new(DashMap::new()),
next_instance_id: Arc::new(AtomicU64::new(1)),
module_cache: Arc::new(parking_lot::Mutex::new(LruCache::new(cap))),
instance_count: Arc::new(AtomicUsize::new(0)),
tenant_counts: Arc::new(DashMap::new()),
metrics: None,
ticker_warned: Arc::new(OnceLock::new()),
pool: None,
jit_cache: None,
compile_semaphore: Arc::new(Semaphore::new(permits)),
}
}
pub fn with_metrics(engine: Arc<TensorWasmEngine>, metrics: TensorWasmMetrics) -> Self {
let cap = lru_cap(engine.config().max_module_cache_entries);
let permits = resolve_compile_permits(engine.config().max_concurrent_compiles);
Self {
engine,
instances: Arc::new(DashMap::new()),
next_instance_id: Arc::new(AtomicU64::new(1)),
module_cache: Arc::new(parking_lot::Mutex::new(LruCache::new(cap))),
instance_count: Arc::new(AtomicUsize::new(0)),
tenant_counts: Arc::new(DashMap::new()),
metrics: Some(metrics),
ticker_warned: Arc::new(OnceLock::new()),
pool: None,
jit_cache: None,
compile_semaphore: Arc::new(Semaphore::new(permits)),
}
}
pub fn with_instance_pool(mut self, pool: Arc<InstancePool>) -> Self {
self.pool = Some(pool);
self
}
pub fn with_jit_cache(mut self, cache: Arc<KernelCache>) -> Self {
self.jit_cache = Some(cache);
self
}
pub fn jit_cache(&self) -> Option<&Arc<KernelCache>> {
self.jit_cache.as_ref()
}
pub fn instance_pool(&self) -> Option<&Arc<InstancePool>> {
self.pool.as_ref()
}
pub fn engine(&self) -> &TensorWasmEngine {
&self.engine
}
pub fn live_count(&self) -> usize {
self.instances.len()
}
pub fn cached_module_count(&self) -> usize {
self.module_cache.lock().len()
}
pub fn module_cache_len(&self) -> usize {
self.module_cache.lock().len()
}
pub fn tenant_instance_count(&self, tenant: TenantId) -> usize {
self.tenant_counts
.get(&tenant)
.map(|e| *e.value())
.unwrap_or(0)
}
pub fn instances_len(&self) -> usize {
self.instance_count.load(Ordering::Acquire)
}
fn allocate_instance_id(&self) -> InstanceId {
loop {
let raw = self.next_instance_id.fetch_add(1, Ordering::Relaxed);
let id = InstanceId(u128::from(raw));
if !self.instances.contains_key(&id) {
return id;
}
warn!(
target: "tensor_wasm_exec::executor",
raw,
"instance id collision detected; retrying with next sequence value",
);
}
}
async fn compile_module_cached(&self, wasm: &[u8]) -> Result<(Module, ModuleHash), ExecError> {
let cap = self.engine.config().max_module_bytes;
if wasm.len() > cap {
return Err(ExecError::ModuleTooLarge {
len: wasm.len(),
max: cap,
});
}
let digest = blake3::hash(wasm);
let key: [u8; 32] = *digest.as_bytes();
if let Some(m) = self.module_cache.lock().get(&key).cloned() {
return Ok((m, key));
}
let engine = self.engine.inner().clone();
let bytes = wasm.to_vec();
let _permit = self
.compile_semaphore
.clone()
.acquire_owned()
.await
.map_err(|_| ExecError::Wasmtime(wasmtime::Error::msg("compile semaphore closed")))?;
let module = tokio::task::spawn_blocking(move || Module::from_binary(&engine, &bytes))
.await
.map_err(|join_err| {
ExecError::Wasmtime(wasmtime::Error::msg(format!(
"wasm compile task failed: {join_err}"
)))
})?
.map_err(ExecError::Wasmtime)?;
drop(_permit);
self.module_cache.lock().put(key, module.clone());
Ok((module, key))
}
fn charge_instance_slot(&self, tenant: TenantId) -> Result<InstanceSlotGuard, ExecError> {
if let Some(max) = self.engine.config().max_instances {
let new_count = self.instance_count.fetch_add(1, Ordering::AcqRel) + 1;
if new_count > max {
self.instance_count.fetch_sub(1, Ordering::Relaxed);
return Err(ExecError::CapacityExhausted {
active: new_count,
limit: max,
});
}
} else {
self.instance_count.fetch_add(1, Ordering::AcqRel);
}
let Some(per_tenant_max) = self.engine.config().max_instances_per_tenant else {
return Ok(InstanceSlotGuard::new(self.instance_count.clone()));
};
let mut entry = self.tenant_counts.entry(tenant).or_insert(0);
let new_tenant_count = *entry + 1;
if new_tenant_count > per_tenant_max {
drop(entry);
self.instance_count.fetch_sub(1, Ordering::Relaxed);
warn!(
target: "tensor_wasm_exec::executor",
tenant = %tenant,
active = new_tenant_count,
limit = per_tenant_max,
"per-tenant instance cap exhausted; refusing spawn (other tenants unaffected)",
);
return Err(ExecError::TenantCapacityExhausted {
tenant,
active: new_tenant_count,
limit: per_tenant_max,
});
}
*entry = new_tenant_count;
drop(entry);
Ok(InstanceSlotGuard::with_tenant(
self.instance_count.clone(),
self.tenant_counts.clone(),
tenant,
))
}
pub(crate) fn release_instance_slot(&self, tenant: TenantId) {
self.instance_count.fetch_sub(1, Ordering::AcqRel);
decrement_tenant_count(&self.tenant_counts, tenant);
}
pub(crate) async fn build_pooled_instance(
&self,
cfg: &SpawnConfig,
wasm: &[u8],
) -> Result<(TensorWasmInstance, Module, ModuleHash), ExecError> {
check_raw_module_memory_within_cap(wasm, self.engine.config().effective_memory_cap())?;
let slot_guard = self.charge_instance_slot(cfg.tenant_id)?;
let (module, module_hash) = self.compile_module_cached(wasm).await?;
let inst = self.instantiate_detached(cfg, &module).await?;
if let Some(m) = &self.metrics {
m.instance_spawns_total().inc();
}
slot_guard.commit();
Ok((inst, module, module_hash))
}
pub(crate) async fn rebuild_pooled_from_module(
&self,
cfg: &SpawnConfig,
module: &Module,
) -> Result<TensorWasmInstance, ExecError> {
let slot_guard = self.charge_instance_slot(cfg.tenant_id)?;
let inst = self.instantiate_detached(cfg, module).await?;
if let Some(m) = &self.metrics {
m.instance_spawns_total().inc();
}
slot_guard.commit();
Ok(inst)
}
async fn instantiate_detached(
&self,
cfg: &SpawnConfig,
module: &Module,
) -> Result<TensorWasmInstance, ExecError> {
let max_memory_bytes = self.engine.config().effective_memory_cap();
let mut state =
InstanceState::new(cfg.tenant_id, InstanceId(0)).with_memory_limit(max_memory_bytes);
if let Some(ref s) = cfg.streaming {
state = state.with_streaming(s.clone());
}
if !cfg.input.is_empty() {
state = state.with_input(InputContext::new(cfg.input.clone()));
}
if let Some(d) = cfg.deadline {
state = state
.with_deadline(Instant::now() + d)
.with_deadline_duration(d);
}
let start_phase_budget = match cfg.deadline {
Some(d) => d.min(MAX_START_FN_DURATION),
None => MAX_START_FN_DURATION,
};
state = state.with_hard_deadline(Instant::now() + start_phase_budget);
let tick = self.engine.config().epoch_tick;
let epoch_deadline_ticks = match cfg.deadline {
Some(d) => duration_to_epoch_ticks(d, tick),
None => MAX_EPOCH_DEADLINE_TICKS,
};
let max_start_ticks = {
let d_ms = MAX_START_FN_DURATION.as_millis();
let t_ms = tick.as_millis().max(1);
let ticks_u128 = d_ms.div_ceil(t_ms).max(1);
u64::try_from(ticks_u128).unwrap_or(u64::MAX)
};
let start_deadline_ticks = epoch_deadline_ticks.min(max_start_ticks);
let deadline_class_applies =
cfg.deadline.is_some() || MAX_START_FN_DURATION > Duration::ZERO;
if deadline_class_applies && !self.engine.is_epoch_ticker_running() {
let flag = self.ticker_warned.get_or_init(|| AtomicBool::new(false));
if !flag.swap(true, Ordering::AcqRel) {
tracing::error!(
target: "tensor_wasm_exec::executor",
"epoch ticker not running — refusing spawn; call `engine.spawn_epoch_ticker()` before serving traffic",
);
}
return Err(ExecError::EpochTickerNotRunning);
}
check_module_memory_within_cap(module, max_memory_bytes)?;
let mut store = Store::new(self.engine.inner(), state);
store.limiter(|state| &mut state.limiter as &mut dyn ResourceLimiter);
arm_cooperative_epoch(
&mut store,
start_deadline_ticks.min(COOPERATIVE_YIELD_TICKS),
);
let mut linker: wasmtime::Linker<InstanceState> =
wasmtime::Linker::new(self.engine.inner());
add_scheduler_to_linker(&mut linker, |state: &InstanceState| state.scheduler())
.map_err(ExecError::Wasmtime)?;
add_input_to_linker(&mut linker).map_err(ExecError::Wasmtime)?;
if cfg.streaming.is_some() {
add_streaming_to_linker(&mut linker).map_err(ExecError::Wasmtime)?;
}
if let Some(cache) = &self.jit_cache {
add_jit_dispatch_to_linker(&mut linker, cache.clone()).map_err(ExecError::Wasmtime)?;
}
let instance = match linker.instantiate_async(&mut store, module).await {
Ok(inst) => inst,
Err(err) => {
return Err(classify_instantiation_error(err, max_memory_bytes));
}
};
store.set_epoch_deadline(epoch_deadline_ticks.min(COOPERATIVE_YIELD_TICKS));
let post_start_hard_deadline = match cfg.deadline {
Some(_) => store.data().deadline,
None => None,
};
store.data_mut().hard_deadline = post_start_hard_deadline;
Ok(TensorWasmInstance::new(store, instance))
}
pub(crate) fn register_pooled_instance(
&self,
mut inst: TensorWasmInstance,
) -> Result<InstanceId, ExecError> {
let id = self.allocate_instance_id();
inst.store.data_mut().instance_id = id;
match self.instances.entry(id) {
Entry::Vacant(v) => {
v.insert(Arc::new(Mutex::new(inst)));
}
Entry::Occupied(_) => {
warn!(
target: "tensor_wasm_exec::executor",
%id,
"instance id race after allocation (pool register); this is a serious bug",
);
return Err(ExecError::Wasmtime(wasmtime::Error::msg(
"instance id collision after allocation",
)));
}
}
if let Some(m) = &self.metrics {
m.active_instances().inc();
}
Ok(id)
}
pub(crate) async fn detach_pooled_instance(
&self,
id: InstanceId,
) -> Option<TensorWasmInstance> {
let (_, handle) = self.instances.remove(&id)?;
if let Some(m) = &self.metrics {
m.active_instances().dec();
}
let owning_tenant = if self.engine.config().max_instances_per_tenant.is_some() {
Some(handle.lock().await.tenant_id())
} else {
None
};
match Arc::try_unwrap(handle) {
Ok(mutex) => Some(mutex.into_inner()),
Err(_arc) => {
self.instance_count.fetch_sub(1, Ordering::AcqRel);
if let Some(tenant) = owning_tenant {
decrement_tenant_count(&self.tenant_counts, tenant);
}
warn!(
target: "tensor_wasm_exec::executor",
%id,
"detach_pooled_instance: outstanding Arc reference (racing in-flight call); \
instance not pooled, both engine-wide and per-tenant slots released",
);
None
}
}
}
fn maybe_rewrite_for_offload<'w>(
&self,
cfg: &SpawnConfig,
wasm: &'w [u8],
) -> std::borrow::Cow<'w, [u8]> {
use std::borrow::Cow;
if !self.engine.config().auto_offload {
return Cow::Borrowed(wasm);
}
let Some(cache) = self.jit_cache.as_ref() else {
debug!(
target: "tensor_wasm_exec::executor",
tenant = %cfg.tenant_id,
"auto_offload enabled but no JIT cache attached; falling back to original module",
);
return Cow::Borrowed(wasm);
};
let detector = self
.engine
.config()
.auto_offload_detector
.unwrap_or_default();
let verdicts = match crate::auto_offload::analyse_with_config(wasm, &detector) {
Ok(v) => v,
Err(e) => {
warn!(
target: "tensor_wasm_exec::executor",
tenant = %cfg.tenant_id,
error = %e,
"auto_offload analysis failed; falling back to original module",
);
return Cow::Borrowed(wasm);
}
};
let any_offload = verdicts.iter().any(|v| {
matches!(
v.verdict,
tensor_wasm_jit::detector::DetectorVerdict::Offload
)
});
if !any_offload {
return Cow::Borrowed(wasm);
}
let opts = RewriteOptions {
tenant_id: cfg.tenant_id,
detector,
..RewriteOptions::default()
};
match rewrite_wasm(wasm, &opts, cache) {
Ok(outcome) if !outcome.offloaded_functions.is_empty() => {
info!(
target: "tensor_wasm_exec::executor",
tenant = %cfg.tenant_id,
offloaded = outcome.offloaded_functions.len(),
total_defined = outcome.total_defined_functions,
"auto_offload rewrite applied; instantiating trampoline-augmented module",
);
Cow::Owned(outcome.rewritten_wasm)
}
Ok(_) => {
debug!(
target: "tensor_wasm_exec::executor",
tenant = %cfg.tenant_id,
"auto_offload rewrite swapped no functions; using original module",
);
Cow::Borrowed(wasm)
}
Err(e) => {
warn!(
target: "tensor_wasm_exec::executor",
tenant = %cfg.tenant_id,
error = %e,
"auto_offload rewrite failed; falling back to original module",
);
Cow::Borrowed(wasm)
}
}
}
#[instrument(skip(self, wasm), fields(tenant = %cfg.tenant_id, instance_id = tracing::field::Empty))]
pub async fn spawn_instance(
&self,
cfg: SpawnConfig,
wasm: &[u8],
) -> Result<InstanceId, ExecError> {
let effective_wasm = self.maybe_rewrite_for_offload(&cfg, wasm);
let (inst, _module, _module_hash) = self
.build_pooled_instance(&cfg, effective_wasm.as_ref())
.await?;
let slot_guard = if self.engine.config().max_instances_per_tenant.is_some() {
InstanceSlotGuard::with_tenant(
self.instance_count.clone(),
self.tenant_counts.clone(),
cfg.tenant_id,
)
} else {
InstanceSlotGuard::new(self.instance_count.clone())
};
let id = self.register_pooled_instance(inst)?;
slot_guard.commit();
tracing::Span::current().record("instance_id", tracing::field::display(id));
info!(target: "tensor_wasm_exec::executor", tenant = %cfg.tenant_id, instance = %id, "instance spawned");
Ok(id)
}
#[deprecated(
since = "0.3.7",
note = "use `call_export_with_args` with an empty `&[]` for the same semantics; \
v0.4 removes this shim. See `docs/MIGRATING-FROM-WASMTIME-WASMER.md` § \"Typed exports\"."
)]
#[instrument(skip(self), fields(instance = %id, export = %export))]
pub async fn call_export(&self, id: InstanceId, export: &str) -> Result<(), ExecError> {
self.call_export_with_args(id, export, &[])
.await
.map(|_| ())
}
#[instrument(skip(self, args), fields(instance = %id, export = %export, args_len = args.len()))]
pub async fn call_export_with_args(
&self,
id: InstanceId,
export: &str,
args: &[WasmArg],
) -> Result<serde_json::Value, ExecError> {
let handle = self
.instances
.get(&id)
.ok_or(ExecError::NotFound(id))?
.value()
.clone();
let mut guard = handle.lock().await;
let call_start = Instant::now();
let configured_deadline = guard.store.data().deadline_duration;
if let Some(d) = configured_deadline {
let new_deadline = call_start + d;
guard.store.data_mut().deadline = Some(new_deadline);
guard.store.data_mut().hard_deadline = Some(new_deadline);
let tick = self.engine.config().epoch_tick;
let call_ticks = duration_to_epoch_ticks(d, tick).min(COOPERATIVE_YIELD_TICKS);
guard.store.set_epoch_deadline(call_ticks);
}
guard.store.data_mut().rearm_scheduler();
let deadline_at = guard.store.data().deadline;
let configured_deadline_ms = configured_deadline
.map(|d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
.unwrap_or(0);
let wasmtime_instance = *guard.wasmtime_instance();
let func = wasmtime_instance
.get_func(&mut guard.store, export)
.ok_or_else(|| ExecError::MissingExport(export.to_string()))?;
let func_ty = func.ty(&guard.store);
let call_outcome = if args.is_empty() && func_ty.results().len() == 0 {
match func.typed::<(), ()>(&guard.store) {
Ok(typed) => typed
.call_async(&mut guard.store, ())
.await
.map(|()| serde_json::Value::Array(Vec::new())),
Err(e) => Err(e),
}
} else {
let params: Vec<Val> = args.iter().copied().map(WasmArg::into_val).collect();
let mut results: Vec<Val> = vec![Val::I32(0); func_ty.results().len()];
match func
.call_async(&mut guard.store, ¶ms, &mut results)
.await
{
Ok(()) => {
let json: Vec<serde_json::Value> = results.iter().map(val_to_json).collect();
Ok(serde_json::Value::Array(json))
}
Err(e) => Err(e),
}
};
match call_outcome {
Ok(value) => Ok(value),
Err(err) => {
let elapsed = call_start.elapsed();
let past_deadline = deadline_at.map(|d| Instant::now() >= d).unwrap_or(false);
if past_deadline {
Err(ExecError::Timeout(TimeoutContext {
id,
elapsed_ms: u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX),
deadline_ms: configured_deadline_ms,
}))
} else {
Err(ExecError::Wasmtime(err))
}
}
}
}
#[instrument(skip(self), fields(instance = %id))]
pub async fn terminate(&self, id: InstanceId) -> Result<(), ExecError> {
match self.instances.remove(&id) {
Some((_, handle)) => {
self.instance_count.fetch_sub(1, Ordering::AcqRel);
if self.engine.config().max_instances_per_tenant.is_some() {
let tenant = handle.lock().await.tenant_id();
decrement_tenant_count(&self.tenant_counts, tenant);
}
if let Some(m) = &self.metrics {
m.instance_terminations_total().inc();
m.active_instances().dec();
}
debug!(target: "tensor_wasm_exec::executor", instance = %id, "instance terminated");
Ok(())
}
None => Err(ExecError::NotFound(id)),
}
}
#[deprecated(
since = "0.3.7",
note = "use `call_export_with_args_then_terminate` with an empty `&[]` for the same semantics; \
v0.4 removes this shim. See `docs/MIGRATING-FROM-WASMTIME-WASMER.md` § \"Typed exports\"."
)]
pub async fn call_export_then_terminate(
&self,
id: InstanceId,
export: &str,
) -> Result<(), ExecError> {
self.call_export_with_args_then_terminate(id, export, &[])
.await
.map(|_| ())
}
pub async fn invoke(
&self,
cfg: SpawnConfig,
wasm: &[u8],
export: &str,
args: &[WasmArg],
) -> Result<serde_json::Value, ExecError> {
if let Some(pool) = self.pool.clone() {
let pooled = pool.acquire(self, wasm, cfg.clone()).await?;
let id = pooled.id();
let result = self.call_export_with_args(id, export, args).await;
pool.release(self, pooled, &cfg).await;
result
} else {
let id = self.spawn_instance(cfg, wasm).await?;
self.call_export_with_args_then_terminate(id, export, args)
.await
}
}
pub async fn call_export_with_args_then_terminate(
&self,
id: InstanceId,
export: &str,
args: &[WasmArg],
) -> Result<serde_json::Value, ExecError> {
let tenant_rollback = if self.engine.config().max_instances_per_tenant.is_some() {
match self.instances.get(&id).map(|h| h.value().clone()) {
Some(handle) => {
let tenant = handle.lock().await.tenant_id();
Some((Arc::clone(&self.tenant_counts), tenant))
}
None => None,
}
} else {
None
};
let guard = AutoTerminateGuard {
instances: Arc::clone(&self.instances),
instance_count: Arc::clone(&self.instance_count),
metrics: self.metrics.clone(),
id,
tenant_rollback,
armed: true,
};
let result = self.call_export_with_args(id, export, args).await;
let mut guard = guard;
guard.armed = false;
let _ = self.terminate(id).await; result
}
}
struct AutoTerminateGuard {
instances: Arc<DashMap<InstanceId, Arc<Mutex<TensorWasmInstance>>>>,
instance_count: Arc<AtomicUsize>,
metrics: Option<TensorWasmMetrics>,
id: InstanceId,
tenant_rollback: Option<(Arc<DashMap<TenantId, usize>>, TenantId)>,
armed: bool,
}
impl Drop for AutoTerminateGuard {
fn drop(&mut self) {
if !self.armed {
return;
}
if self.instances.remove(&self.id).is_some() {
self.instance_count.fetch_sub(1, Ordering::AcqRel);
if let Some((tenant_counts, tenant)) = &self.tenant_rollback {
decrement_tenant_count(tenant_counts, *tenant);
}
if let Some(m) = &self.metrics {
m.instance_terminations_total().inc();
m.active_instances().dec();
}
tracing::warn!(
target: "tensor_wasm_exec::executor",
instance = %self.id,
"instance auto-terminated by drop-guard (handler future cancelled \
mid-call_export; see api S-20)"
);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn trivial_wasm() -> Vec<u8> {
wat::parse_str(r#"(module (func (export "noop")))"#).unwrap()
}
#[tokio::test]
async fn spawn_then_terminate() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let id = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
.await
.unwrap();
assert_eq!(exec.live_count(), 1);
exec.call_export_with_args(id, "noop", &[]).await.unwrap();
exec.terminate(id).await.unwrap();
assert_eq!(exec.live_count(), 0);
}
#[tokio::test]
async fn missing_export() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let id = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
.await
.unwrap();
let err = exec
.call_export_with_args(id, "does_not_exist", &[])
.await
.unwrap_err();
assert!(matches!(err, ExecError::MissingExport(_)));
}
#[tokio::test]
#[allow(deprecated)]
async fn legacy_call_export_still_works() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let id = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
.await
.unwrap();
exec.call_export(id, "noop").await.unwrap();
exec.terminate(id).await.unwrap();
let _f = TensorWasmExecutor::call_export;
let _g = TensorWasmExecutor::call_export_then_terminate;
}
#[tokio::test]
async fn terminate_unknown() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let err = exec.terminate(InstanceId(999)).await.unwrap_err();
assert!(matches!(err, ExecError::NotFound(_)));
}
#[tokio::test]
async fn metrics_increment_on_spawn_and_terminate() {
use tensor_wasm_core::metrics::TensorWasmMetrics;
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let metrics = TensorWasmMetrics::new();
let exec = TensorWasmExecutor::with_metrics(engine, metrics.clone());
let id = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
.await
.unwrap();
let text = metrics.encode_text();
assert!(
text.contains("tensor_wasm_instance_spawns_total 1"),
"got:\n{text}"
);
assert!(
text.contains("tensor_wasm_active_instances 1"),
"got:\n{text}"
);
exec.terminate(id).await.unwrap();
let text = metrics.encode_text();
assert!(
text.contains("tensor_wasm_instance_terminations_total 1"),
"got:\n{text}"
);
assert!(
text.contains("tensor_wasm_active_instances 0"),
"got:\n{text}"
);
}
#[test]
fn exec_error_converts_to_tensor_wasm_error() {
use tensor_wasm_core::error::TensorWasmError;
let e = ExecError::NotFound(InstanceId(99));
let b: TensorWasmError = e.into();
assert!(matches!(b, TensorWasmError::Serialization(_)));
assert!(
b.inner().unwrap_or("").contains("instance not found"),
"inner: {:?}",
b.inner()
);
let e = ExecError::Timeout(TimeoutContext {
id: InstanceId(1),
elapsed_ms: 150,
deadline_ms: 100,
});
let b: TensorWasmError = e.into();
match b {
TensorWasmError::KernelTimeout {
elapsed_ms,
deadline_ms,
} => {
assert_eq!(elapsed_ms, 150);
assert_eq!(deadline_ms, 100);
}
other => panic!("expected KernelTimeout, got {other:?}"),
}
}
#[tokio::test]
async fn module_cache_reuses_compilation() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let wasm = trivial_wasm();
let _id1 = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &wasm)
.await
.unwrap();
let _id2 = exec
.spawn_instance(SpawnConfig::for_tenant(TenantId(2)), &wasm)
.await
.unwrap();
assert_eq!(exec.cached_module_count(), 1);
}
#[test]
fn classify_instantiation_error_refines_memory_size_failure() {
let err = wasmtime::Error::msg("memory index 0 has a minimum that exceeds the limit");
let classified = classify_instantiation_error(err, 64 * 1024 * 1024);
match classified {
ExecError::ModuleMemoryTooLarge {
requested_bytes,
limit_bytes,
} => {
assert_eq!(limit_bytes, 64 * 1024 * 1024);
assert_eq!(requested_bytes, 64 * 1024 * 1024);
}
other => panic!("expected ModuleMemoryTooLarge, got {other:?}"),
}
}
#[test]
fn classify_instantiation_error_passes_through_unrelated() {
let err = wasmtime::Error::msg("unknown import: `env::foo` has not been defined");
let classified = classify_instantiation_error(err, 64 * 1024 * 1024);
assert!(
matches!(classified, ExecError::Wasmtime(_)),
"unrelated errors must not be re-tagged, got {classified:?}",
);
}
#[test]
fn resource_limiter_allows_under_cap() {
let mut lim = TensorWasmResourceLimiter::new(2 * 1024 * 1024);
assert!(lim.memory_growing(0, 1024 * 1024, None).unwrap());
}
#[test]
fn resource_limiter_rejects_over_cap() {
let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
assert!(!lim.memory_growing(0, 2 * 1024 * 1024, None).unwrap());
}
#[test]
fn resource_limiter_respects_module_maximum() {
let mut lim = TensorWasmResourceLimiter::new(usize::MAX);
assert!(!lim.memory_growing(0, 4096, Some(2048)).unwrap());
}
#[test]
fn resource_limiter_rejects_huge_table_growth() {
let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
assert!(!lim.table_growing(0, usize::MAX, None).unwrap());
}
#[test]
fn resource_limiter_allows_modest_table_growth() {
let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
assert!(lim.table_growing(0, 1024, None).unwrap());
}
#[test]
fn resource_limiter_table_respects_module_maximum() {
let mut lim = TensorWasmResourceLimiter::new(usize::MAX);
assert!(!lim.table_growing(0, 4096, Some(2048)).unwrap());
}
#[tokio::test]
async fn dropped_invoke_releases_per_tenant_slot() {
use crate::engine::EngineConfig;
let cfg = EngineConfig {
max_instances_per_tenant: Some(1),
..EngineConfig::default()
};
let engine = Arc::new(TensorWasmEngine::with_config(cfg).unwrap());
let exec = TensorWasmExecutor::new(engine);
let tenant = TenantId(7);
let id = exec
.spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
.await
.unwrap();
assert_eq!(exec.tenant_instance_count(tenant), 1);
assert_eq!(exec.instances_len(), 1);
{
let _guard = AutoTerminateGuard {
instances: Arc::clone(&exec.instances),
instance_count: Arc::clone(&exec.instance_count),
metrics: exec.metrics.clone(),
id,
tenant_rollback: Some((Arc::clone(&exec.tenant_counts), tenant)),
armed: true,
};
}
assert_eq!(
exec.tenant_instance_count(tenant),
0,
"per-tenant slot leaked on dropped invoke (exec H2)"
);
assert_eq!(exec.instances_len(), 0, "engine-wide slot leaked");
assert_eq!(exec.live_count(), 0, "registry entry leaked");
let id2 = exec
.spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
.await
.expect("tenant should not be locked out after a cancelled invoke");
exec.terminate(id2).await.unwrap();
assert_eq!(exec.tenant_instance_count(tenant), 0);
}
#[tokio::test]
async fn dropped_invoke_no_tenant_cap_releases_engine_slot_only() {
let engine = Arc::new(TensorWasmEngine::new().unwrap());
let exec = TensorWasmExecutor::new(engine);
let tenant = TenantId(3);
let id = exec
.spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
.await
.unwrap();
assert_eq!(exec.instances_len(), 1);
assert_eq!(exec.tenant_instance_count(tenant), 0);
{
let _guard = AutoTerminateGuard {
instances: Arc::clone(&exec.instances),
instance_count: Arc::clone(&exec.instance_count),
metrics: exec.metrics.clone(),
id,
tenant_rollback: None,
armed: true,
};
}
assert_eq!(exec.instances_len(), 0, "engine-wide slot leaked");
assert_eq!(exec.tenant_instance_count(tenant), 0);
}
}