tensor_wasm_exec/executor.rs
1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Craton Software Company
3
4//! [`TensorWasmExecutor`] — async executor for TensorWasm Wasm instances.
5//!
6//! Owns a shared [`TensorWasmEngine`] and a registry of live [`TensorWasmInstance`]s
7//! keyed by [`InstanceId`]. Exposes the trio of operations
8//! [`TensorWasmExecutor::spawn_instance`], [`TensorWasmExecutor::call_export`], and
9//! [`TensorWasmExecutor::terminate`] — all async, all driven from the calling
10//! Tokio runtime.
11
12use std::num::NonZeroUsize;
13use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
14use std::sync::{Arc, OnceLock};
15use std::time::{Duration, Instant};
16
17use dashmap::{mapref::entry::Entry, DashMap};
18use lru::LruCache;
19use tensor_wasm_core::metrics::TensorWasmMetrics;
20use tensor_wasm_core::types::{InstanceId, TenantId};
21use tensor_wasm_jit::cache::KernelCache;
22use tensor_wasm_jit::rewrite::{rewrite_wasm, RewriteOptions};
23use thiserror::Error;
24use tokio::sync::{Mutex, Semaphore};
25use tracing::{debug, info, instrument, warn};
26use wasmtime::{ExternType, Module, ResourceLimiter, Store, UpdateDeadline, Val};
27
28use crate::engine::TensorWasmEngine;
29use crate::instance::{InstanceState, TensorWasmInstance};
30use crate::instance_pool::{InstancePool, ModuleHash};
31use crate::jit_dispatch::add_jit_dispatch_to_linker;
32use tensor_wasm_wasi_gpu::scheduler::add_scheduler_to_linker;
33use tensor_wasm_wasi_gpu::streaming::{
34 add_input_to_linker, add_streaming_to_linker, InputContext, StreamingContext,
35};
36
37/// Convert a wall-clock [`Duration`] into a number of epoch ticks suitable
38/// for [`wasmtime::Store::set_epoch_deadline`].
39///
40/// Rounds up so a sub-tick remainder still terminates, with a floor of 1 so
41/// `Duration::from_nanos(1)` does not silently become "never trip". Clamps
42/// at [`u64::MAX`] if a caller supplies a pathologically long deadline. The
43/// `tick` parameter is the engine's `epoch_tick` cadence; a zero-or-less tick
44/// is treated as 1 ms to avoid division-by-zero on a malformed config.
45///
46/// # Granularity
47///
48/// Both the deadline and the tick are measured in **whole milliseconds**
49/// (`Duration::as_millis`). Sub-millisecond resolution is therefore clamped:
50/// a deadline of `Duration::from_micros(500)` and one of
51/// `Duration::from_micros(1)` both round up to a single tick. This is
52/// intentional — the epoch interrupt itself only fires on the background
53/// ticker's cadence (10 ms by default), so sub-ms precision in the deadline
54/// would be illusory anyway. Callers needing finer-grained interruption must
55/// shorten [`crate::engine::EngineConfig::epoch_tick`], not the deadline.
56///
57/// The result is clamped to [`MAX_EPOCH_DEADLINE_TICKS`] (not `u64::MAX`)
58/// because `set_epoch_deadline` is *relative*: wasmtime computes
59/// `current_epoch + ticks`, so a `u64::MAX` result overflows once the
60/// background ticker has advanced the epoch.
61fn duration_to_epoch_ticks(d: Duration, tick: Duration) -> u64 {
62 let d_ms = d.as_millis();
63 let t_ms = tick.as_millis().max(1);
64 let ticks_u128 = d_ms.div_ceil(t_ms).max(1);
65 u64::try_from(ticks_u128)
66 .unwrap_or(MAX_EPOCH_DEADLINE_TICKS)
67 .min(MAX_EPOCH_DEADLINE_TICKS)
68}
69
70/// Install the COOPERATIVE epoch-deadline scheme on `store`, arming the first
71/// deadline `first_ticks` ticks out.
72///
73/// This replaces the historical trap-only configuration (`set_epoch_deadline`
74/// alone, which leaves the store in wasmtime's default
75/// [`Store::epoch_deadline_trap`](wasmtime::Store::epoch_deadline_trap) mode).
76/// Under trap mode a compute-bound guest never returns `Pending`: the
77/// `call_async` poll blocks the worker until the epoch traps, so an outer
78/// `timeout` / future-drop cannot cancel it cooperatively and a single-thread
79/// runtime can be wedged by a spinning guest. Cooperative yielding fixes both.
80///
81/// The scheme is two-level:
82///
83/// * **Cooperative yield (between yields the guest is cancellable).** The
84/// store is configured via
85/// [`Store::epoch_deadline_callback`](wasmtime::Store::epoch_deadline_callback).
86/// Each time the epoch deadline trips, the callback returns
87/// [`UpdateDeadline::Yield`] with [`COOPERATIVE_YIELD_TICKS`], so the guest
88/// yields `Pending` to the async executor and the deadline re-arms for
89/// another window. While parked at that yield, the surrounding
90/// `call_async` future can be dropped — which terminates the guest fiber
91/// promptly (wasmtime turns a cancel-while-yielded into a trap on the
92/// unwinding fiber). This is exactly the property the
93/// `orphan_cleanup_on_drop` test needs.
94///
95/// * **Hard deadline (still TRAPs a runaway guest).** Before yielding, the
96/// callback consults the per-store
97/// [`InstanceState::hard_deadline`](crate::instance::InstanceState::hard_deadline).
98/// Once that absolute wall-clock instant has elapsed the callback returns
99/// `Err(..)` instead of yielding, which wasmtime converts into a trap —
100/// bit-for-bit the same termination the old `set_epoch_deadline` trap
101/// produced. The executor sets `hard_deadline` to the per-call deadline
102/// (and, during instantiation, to a value additionally clamped by
103/// [`MAX_START_FN_DURATION`]), so the public timeout contract is preserved:
104/// `call_export_with_args` still classifies the resulting error as
105/// [`ExecError::Timeout`] because the wall clock has crossed `deadline`.
106///
107/// The error message intentionally matches the spirit of a wasmtime epoch
108/// trap ("epoch deadline reached"); the executor never surfaces this string
109/// to callers — it is reclassified into the typed [`ExecError::Timeout`] (or,
110/// for the no-deadline start-fn cap, propagated as `ExecError::Wasmtime`,
111/// exactly as a trap did before).
112fn arm_cooperative_epoch(store: &mut Store<InstanceState>, first_ticks: u64) {
113 // Arm the FIRST deadline relative to the current epoch. Subsequent
114 // re-arms are driven by the value the callback returns.
115 store.set_epoch_deadline(first_ticks);
116 store.epoch_deadline_callback(|ctx| {
117 if ctx.data().hard_deadline_elapsed() {
118 // HARD deadline crossed: trap, terminating the guest. Mirrors the
119 // historical `epoch_deadline_trap` behaviour. The executor's
120 // timeout classification keys off the wall clock, not this string.
121 Err(wasmtime::Error::msg("epoch deadline reached"))
122 } else {
123 // Within the hard deadline: yield `Pending` so the guest is
124 // cancellable, then re-arm for another cooperative window.
125 Ok(UpdateDeadline::Yield(COOPERATIVE_YIELD_TICKS))
126 }
127 });
128}
129
130/// Overflow-safe sentinel for "effectively no deadline", in epoch ticks.
131///
132/// [`wasmtime::Store::set_epoch_deadline`] takes a deadline **relative** to
133/// the current epoch — internally `current_epoch + ticks`. A `u64::MAX`
134/// sentinel therefore overflows the moment the background epoch ticker has
135/// advanced the counter past 0: in a debug build that is an `attempt to add
136/// with overflow` panic inside wasmtime; in release (overflow checks off) it
137/// wraps to a deadline *in the past*, so the guest is interrupted
138/// immediately. `u64::MAX / 2` is effectively infinite (≈9.2e18 ticks —
139/// millions of years at any realistic cadence) while guaranteeing
140/// `current_epoch + MAX_EPOCH_DEADLINE_TICKS` cannot overflow for the life of
141/// the process.
142const MAX_EPOCH_DEADLINE_TICKS: u64 = u64::MAX / 2;
143
144/// Number of epoch ticks between cooperative yields under the async-yield
145/// epoch scheme.
146///
147/// The executor configures every store with
148/// [`wasmtime::Store::epoch_deadline_callback`] (see
149/// [`TensorWasmExecutor::arm_cooperative_epoch`]) so that — instead of trapping
150/// the instant the epoch deadline trips — a compute-bound guest *yields*
151/// `Pending` back to the async runtime and the deadline is re-armed for
152/// another `COOPERATIVE_YIELD_TICKS` window. Yielding is what makes a
153/// runaway guest *cancellable*: while it is parked at a yield point the
154/// surrounding `call_async` future can be dropped (by an outer `timeout`
155/// or task cancellation), which terminates the guest promptly instead of
156/// blocking the worker until a trap fires. The HARD deadline is still
157/// enforced — on each yield the callback first checks the per-store
158/// [`InstanceState::hard_deadline`](crate::instance::InstanceState::hard_deadline)
159/// and traps once it has elapsed (see `arm_cooperative_epoch`).
160///
161/// One tick (the engine's `epoch_tick`, 10 ms by default) keeps yields
162/// frequent enough that future-drop cancellation is observed within a
163/// single tick of cadence, while staying coarse enough that a well-behaved
164/// guest doing real work is not penalised with excessive yield churn. The
165/// background epoch ticker ([`crate::engine::TensorWasmEngine::spawn_epoch_ticker`])
166/// must be running for the deadline to advance at all — the same
167/// pre-existing requirement as the trap scheme, enforced by the
168/// [`ExecError::EpochTickerNotRunning`] admission check.
169const COOPERATIVE_YIELD_TICKS: u64 = 1;
170
171/// Hard upper bound on how long a Wasm module's `start` function (and any
172/// other code that runs inside [`wasmtime::Instance::new_async`]) is allowed
173/// to execute before the epoch interrupt trips.
174///
175/// Without this cap a `SpawnConfig { deadline: None, .. }` would set the
176/// per-store epoch deadline to `u64::MAX`, which means an infinite-loop
177/// start function would burn forever inside `Instance::new_async`. Because
178/// the instance is not registered with the executor until that call
179/// returns, [`TensorWasmExecutor::terminate`] cannot reach it — the only
180/// thing that can interrupt the loop is the epoch deadline. 30 seconds is
181/// generous for legitimate start functions (which typically just call out
182/// to a few initialisers) while still bounding the worst case.
183pub const MAX_START_FN_DURATION: Duration = Duration::from_secs(30);
184
185/// Hard upper bound on the byte length of a Wasm module the executor will
186/// accept for compilation. Modules above this size are rejected with
187/// [`ExecError::ModuleTooLarge`] *before* `Module::from_binary` runs —
188/// pathological code-section blow-ups can otherwise force Cranelift to
189/// burn arbitrary CPU on adversarial input. 64 MiB is comfortably above
190/// any legitimate ML kernel module we've seen (single-digit MiB is
191/// typical), while keeping the Cranelift worst case bounded.
192///
193/// This constant is the floor: embedders can tighten via
194/// [`EngineConfig::max_module_bytes`](crate::engine::EngineConfig::max_module_bytes).
195pub const MAX_MODULE_BYTES: usize = 64 * 1024 * 1024;
196
197/// Errors raised by the executor.
198#[derive(Debug, Error)]
199pub enum ExecError {
200 /// Wasmtime returned an error during compile / instantiate / call.
201 ///
202 /// The full wasmtime error chain (including any inner trap, backtrace,
203 /// or compile-error context) is preserved via `#[from]` (which thiserror
204 /// also wires as `#[source]`) so callers converting to
205 /// [`tensor_wasm_core::error::TensorWasmError`] do not lose detail.
206 #[error("wasmtime error")]
207 Wasmtime(#[from] wasmtime::Error),
208 /// Looked up an instance that does not exist (or has terminated).
209 #[error("no such instance: {0}")]
210 NotFound(InstanceId),
211 /// Looked up an export that the instance does not provide.
212 #[error("instance has no export `{0}`")]
213 MissingExport(String),
214 /// The instance ran past its deadline before the call completed.
215 ///
216 /// Carries the offending [`InstanceId`] plus the real `elapsed_ms` /
217 /// `deadline_ms` figures captured at the time the epoch interrupt
218 /// fired. Surfaced as [`tensor_wasm_core::error::TensorWasmError::KernelTimeout`]
219 /// with the same numbers on the conversion boundary.
220 ///
221 /// Tuple-shaped (rather than a struct variant) so existing match arms
222 /// like `ExecError::Timeout(_)` keep compiling.
223 #[error("{0}")]
224 Timeout(TimeoutContext),
225 /// The module declares — via an exported or imported
226 /// [`wasmtime::ExternType::Memory`] — an initial or maximum linear
227 /// memory size that exceeds `EngineConfig::max_memory_bytes`.
228 ///
229 /// Surfaced *before* `Instance::new_async` because Wasmtime's
230 /// [`wasmtime::ResourceLimiter::memory_growing`] only fires on
231 /// `memory.grow`, not on the initial allocation. A guest declaring
232 /// `(memory 1 65536)` would otherwise force a 4 GiB allocation at
233 /// instantiation. Maps to
234 /// [`tensor_wasm_core::error::TensorWasmError::MemoryExhausted`].
235 #[error("module-declared linear memory {requested_bytes} bytes exceeds engine cap {limit_bytes} bytes")]
236 ModuleMemoryTooLarge {
237 /// Bytes the module asked for (initial or declared maximum,
238 /// whichever tripped the check first).
239 requested_bytes: u64,
240 /// Configured engine-wide per-instance cap in bytes.
241 limit_bytes: u64,
242 },
243 /// The executor refused to admit a new instance because the
244 /// engine-wide live-instance cap
245 /// ([`crate::engine::EngineConfig::max_instances`]) is already
246 /// saturated.
247 ///
248 /// Surfaced from [`TensorWasmExecutor::spawn_instance`] *before* any
249 /// compile / instantiate work; the failed spawn never consumes a
250 /// slot in the registry. Mapped to
251 /// [`tensor_wasm_core::error::TensorWasmError::MemoryExhausted`] on
252 /// the conversion boundary (the API layer surfaces it as 503).
253 /// The executor refused to admit a new instance because an
254 /// admission-control ceiling is already saturated. Surfaced for two
255 /// rejections that the API layer treats identically (503, retryable):
256 ///
257 /// * the engine-wide live-instance ceiling
258 /// ([`crate::engine::EngineConfig::max_instances`], exec S-10); or
259 /// * the spawning tenant's per-tenant fairness cap
260 /// ([`crate::engine::EngineConfig::max_instances_per_tenant`]), even
261 /// though the engine-wide ceiling may still have headroom. One tenant
262 /// hitting its per-tenant cap is refused here WITHOUT affecting any
263 /// other tenant's ability to spawn.
264 ///
265 /// Surfaced from [`TensorWasmExecutor::spawn_instance`] *before* any
266 /// compile / instantiate work; the failed spawn never consumes a slot
267 /// (engine-wide or per-tenant). Mapped to
268 /// [`tensor_wasm_core::error::TensorWasmError::MemoryExhausted`] on the
269 /// conversion boundary (the API layer surfaces it as 503).
270 ///
271 /// This is the *engine-wide* ceiling: the shared `max_instances` budget is
272 /// saturated and the spawn would succeed once aggregate load drops. The
273 /// remedy is retry-with-backoff and it is no single tenant's fault, so the
274 /// API layer maps it to 503. The per-tenant fairness cap is a distinct
275 /// condition — see [`ExecError::TenantCapacityExhausted`].
276 #[error("instance capacity exhausted: {active} active, limit {limit}")]
277 CapacityExhausted {
278 /// Engine-wide live-instance count observed at the rejection point
279 /// (post-increment, so `active > limit`).
280 active: usize,
281 /// Configured engine-wide `max_instances` ceiling that was exceeded.
282 limit: usize,
283 },
284 /// A single tenant exceeded its per-tenant fairness cap
285 /// ([`crate::engine::EngineConfig::max_instances_per_tenant`]) while the
286 /// shared engine-wide budget still had room. This is semantically
287 /// distinct from [`ExecError::CapacityExhausted`]: the offending tenant is
288 /// over *its own* quota, so the corrective action is for that tenant to
289 /// reduce concurrency — not to wait for global load to drop. The API layer
290 /// maps it to 429 (`tenant_capacity_exhausted`) rather than 503, giving
291 /// callers a quota-specific retry signal. Other tenants are unaffected.
292 #[error("tenant {tenant} instance capacity exhausted: {active} active, limit {limit}")]
293 TenantCapacityExhausted {
294 /// The tenant whose per-tenant cap was hit (the spawn's owner).
295 tenant: TenantId,
296 /// That tenant's live-instance count at the rejection point
297 /// (post-increment, so `active > limit`).
298 active: usize,
299 /// The configured per-tenant `max_instances_per_tenant` ceiling.
300 limit: usize,
301 },
302 /// The submitted Wasm module is larger than the configured
303 /// pre-compile size cap
304 /// ([`crate::engine::EngineConfig::max_module_bytes`], floored at
305 /// [`MAX_MODULE_BYTES`]).
306 ///
307 /// Rejected before `Module::from_binary` runs so a pathological
308 /// code section cannot force Cranelift to burn CPU on adversarial
309 /// input. Mapped to
310 /// [`tensor_wasm_core::error::TensorWasmError::MemoryExhausted`] on
311 /// the conversion boundary (the API layer surfaces it as 503).
312 #[error("wasm module byte length {len} exceeds cap {max}")]
313 ModuleTooLarge {
314 /// Length of the rejected wasm blob, in bytes.
315 len: usize,
316 /// Configured per-executor cap, in bytes.
317 max: usize,
318 },
319 /// `spawn_instance` was called with a deadline configured but the
320 /// engine's epoch ticker is not running. Without the ticker, the
321 /// per-store epoch counter never advances, so neither the per-call
322 /// deadline nor [`MAX_START_FN_DURATION`] can fire — a runaway
323 /// guest would wedge the worker thread until it returned of its
324 /// own accord. Refuse the spawn instead of silently dropping the
325 /// deadline contract.
326 #[error("epoch ticker not running — refusing spawn with deadline; call `engine.spawn_epoch_ticker()` first")]
327 EpochTickerNotRunning,
328}
329
330/// Payload for [`ExecError::Timeout`]. Carries the real elapsed and deadline
331/// figures captured when the epoch interrupt fired so the error mapping
332/// layer can surface them through [`tensor_wasm_core::error::TensorWasmError::KernelTimeout`].
333#[derive(Debug, Clone, Copy)]
334pub struct TimeoutContext {
335 /// Instance that exceeded its deadline.
336 pub id: InstanceId,
337 /// Wall-clock milliseconds the call took before being interrupted.
338 pub elapsed_ms: u64,
339 /// Configured per-call deadline in milliseconds.
340 pub deadline_ms: u64,
341}
342
343impl std::fmt::Display for TimeoutContext {
344 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345 write!(
346 f,
347 "instance {} exceeded deadline (elapsed {} ms, deadline {} ms)",
348 self.id, self.elapsed_ms, self.deadline_ms,
349 )
350 }
351}
352
353impl From<ExecError> for tensor_wasm_core::error::TensorWasmError {
354 fn from(e: ExecError) -> Self {
355 use tensor_wasm_core::error::TensorWasmError;
356 match e {
357 ExecError::Wasmtime(err) => {
358 // Distinguish runtime traps from compile/instantiate errors.
359 // wasmtime wraps runtime traps as `wasmtime::Trap` inside the
360 // anyhow error; compile/parse failures do NOT. We classify
361 // accordingly so the unified `TensorWasmError` carries the right
362 // variant (and `is_retryable` / `kind` reflect it).
363 //
364 // SECURITY (exec S-9): the full wasmtime error chain
365 // (`format!("{err:#}")`) walks every `#[source]` link and
366 // routinely surfaces host pointer addresses, host file paths,
367 // and internal stack-frame names — none of which are safe to
368 // hand back to an untrusted caller. We therefore log the full
369 // chain server-side and return a stable opaque string in the
370 // payload so external observers cannot fingerprint the host.
371 let is_trap = err.downcast_ref::<wasmtime::Trap>().is_some();
372 tracing::error!(
373 target: "tensor_wasm_exec::executor",
374 error = ?err,
375 error_chain = %format!("{err:#}"),
376 is_trap,
377 "wasmtime trap",
378 );
379 if is_trap {
380 TensorWasmError::WasmTrap("wasm trap".into())
381 } else {
382 TensorWasmError::WasmCompile("wasm compile failed".into())
383 }
384 }
385 ExecError::NotFound(id) => {
386 TensorWasmError::Serialization(format!("instance not found: {id}").into())
387 }
388 ExecError::MissingExport(name) => {
389 TensorWasmError::Serialization(format!("instance missing export: {name}").into())
390 }
391 ExecError::Timeout(ctx) => TensorWasmError::KernelTimeout {
392 elapsed_ms: ctx.elapsed_ms,
393 deadline_ms: ctx.deadline_ms,
394 },
395 ExecError::ModuleMemoryTooLarge {
396 requested_bytes,
397 limit_bytes,
398 } => TensorWasmError::MemoryExhausted {
399 requested: requested_bytes,
400 limit: limit_bytes,
401 },
402 ExecError::CapacityExhausted { active, limit } => TensorWasmError::MemoryExhausted {
403 requested: active as u64,
404 limit: limit as u64,
405 },
406 // The per-tenant fairness cap collapses to the same resource-
407 // exhaustion shape on the native `TensorWasmError` boundary (the
408 // tenant distinction is preserved on the richer API-layer mapping,
409 // not here). Tenant id is dropped — `TensorWasmError` has no field
410 // for it and it is logged at the rejection site.
411 ExecError::TenantCapacityExhausted { active, limit, .. } => {
412 TensorWasmError::MemoryExhausted {
413 requested: active as u64,
414 limit: limit as u64,
415 }
416 }
417 ExecError::ModuleTooLarge { len, max } => TensorWasmError::MemoryExhausted {
418 requested: len as u64,
419 limit: max as u64,
420 },
421 ExecError::EpochTickerNotRunning => {
422 // Surface as a compile-class failure: the spawn never
423 // got off the ground, no guest code executed, and the
424 // remedy is operational (start the ticker) rather than
425 // anything the caller can retry.
426 TensorWasmError::WasmCompile("epoch ticker not running".into())
427 }
428 }
429 }
430}
431
432/// Configuration passed to [`TensorWasmExecutor::spawn_instance`].
433#[derive(Debug, Clone)]
434pub struct SpawnConfig {
435 /// Owning tenant.
436 pub tenant_id: TenantId,
437 /// Optional per-call deadline.
438 pub deadline: Option<Duration>,
439 /// Arguments forwarded to the first [`TensorWasmExecutor::call_export_with_args`]
440 /// invocation against this instance.
441 ///
442 /// Callers that drive a single spawn-then-call flow (CLI `run`, API
443 /// `/invoke`) populate this field so the caller's argument list survives
444 /// the trip across crate boundaries without a parallel `CallConfig`.
445 /// Multi-call flows should ignore this field and pass arguments directly
446 /// to each `call_export_with_args` invocation.
447 pub args: Vec<WasmArg>,
448 /// Optional streaming context (roadmap feature #2). When `Some`,
449 /// [`TensorWasmExecutor::spawn_instance`] builds a wasmtime
450 /// [`wasmtime::Linker`] that wires the `wasi:tensor/host` host
451 /// functions (`emit-chunk`, `flush`) against this context — guest
452 /// emits land on the matching `mpsc::Receiver<Vec<u8>>` the
453 /// gateway is draining into the SSE / chunked HTTP response.
454 ///
455 /// `None` means streaming is disabled: the spawn uses the
456 /// historical empty-imports `Instance::new_async` path, and a
457 /// guest that imports `wasi:tensor/host` will fail to link.
458 /// `/invoke` (the synchronous route) takes the `None` path; only
459 /// `/invoke-stream` opts in.
460 pub streaming: Option<StreamingContext>,
461 /// Bytes staged for the guest to pull via the `wasi:tensor/host`
462 /// pull-model input channel (`input-len` / `read-input`).
463 ///
464 /// Empty by default (the historical behaviour: the guest has no
465 /// input channel). When non-empty,
466 /// [`TensorWasmExecutor::spawn_instance`] installs an
467 /// [`InputContext`] on
468 /// the per-instance state and registers the `input-len` /
469 /// `read-input` host functions on the spawn linker — so a guest can
470 /// copy these bytes into its own linear memory at the start of the
471 /// invocation. The OpenAI completions shim sets this to the assembled
472 /// prompt bytes.
473 pub input: Vec<u8>,
474}
475
476impl SpawnConfig {
477 /// Construct with just a tenant and no deadline.
478 pub fn for_tenant(tenant_id: TenantId) -> Self {
479 Self {
480 tenant_id,
481 deadline: None,
482 args: Vec::new(),
483 streaming: None,
484 input: Vec::new(),
485 }
486 }
487
488 /// Add a deadline relative to "now at spawn time".
489 pub fn with_deadline(mut self, deadline: Duration) -> Self {
490 self.deadline = Some(deadline);
491 self
492 }
493
494 /// Attach an argument list for the upcoming call. See [`SpawnConfig::args`].
495 pub fn with_args(mut self, args: Vec<WasmArg>) -> Self {
496 self.args = args;
497 self
498 }
499
500 /// Attach a streaming context (roadmap feature #2). See
501 /// [`SpawnConfig::streaming`].
502 ///
503 /// Builder method; pairs with [`Self::for_tenant`] /
504 /// [`Self::with_deadline`]. The API gateway's `/invoke-stream`
505 /// route constructs an `mpsc::channel`, wraps the sender in a
506 /// `StreamingContext` via [`StreamingContext::with_channel`], and
507 /// passes it here so the guest's `wasi:tensor/host.emit-chunk`
508 /// calls land on the matching receiver — which the gateway
509 /// concurrently drains into the SSE / chunked response body.
510 pub fn with_streaming(mut self, ctx: StreamingContext) -> Self {
511 self.streaming = Some(ctx);
512 self
513 }
514
515 /// Stage `input` bytes for the guest to pull via the
516 /// `wasi:tensor/host` input channel (`input-len` / `read-input`). See
517 /// [`SpawnConfig::input`].
518 ///
519 /// Builder method; pairs with [`Self::for_tenant`] /
520 /// [`Self::with_deadline`] / [`Self::with_streaming`]. An empty slice
521 /// is a no-op (the default) — the guest then observes `input-len() ==
522 /// 0`. The OpenAI completions shim passes the assembled prompt bytes
523 /// here so the guest receives the prompt.
524 pub fn with_input(mut self, input: Vec<u8>) -> Self {
525 self.input = input;
526 self
527 }
528}
529
530/// A typed Wasm value supplied to [`TensorWasmExecutor::call_export_with_args`].
531///
532/// Mirrors the four core wasm value types (`i32`, `i64`, `f32`, `f64`). Held
533/// `Copy` so callers can clone an argument list cheaply when retrying. Marked
534/// `#[non_exhaustive]` so additional value types (e.g. `v128`, reference
535/// types) can be added in a future minor release without breaking the
536/// match-arm count on downstream code.
537#[derive(Debug, Clone, Copy, PartialEq)]
538#[non_exhaustive]
539pub enum WasmArg {
540 /// A 32-bit signed integer argument.
541 I32(i32),
542 /// A 64-bit signed integer argument.
543 I64(i64),
544 /// A 32-bit IEEE-754 float argument.
545 F32(f32),
546 /// A 64-bit IEEE-754 float argument.
547 F64(f64),
548}
549
550impl WasmArg {
551 /// Convert a [`serde_json::Value`] into the closest-fitting [`WasmArg`]
552 /// variant.
553 ///
554 /// Integer literals that fit in `i32` become [`WasmArg::I32`]; larger
555 /// integers become [`WasmArg::I64`]; non-integer numerics become
556 /// [`WasmArg::F64`]. Any non-numeric value is rejected with an error
557 /// string suitable for forwarding into a user-facing CLI / HTTP error.
558 /// `f32` cannot be selected from JSON unambiguously — callers needing a
559 /// 32-bit float should construct [`WasmArg::F32`] directly.
560 pub fn from_json(v: &serde_json::Value) -> Result<Self, &'static str> {
561 match v {
562 serde_json::Value::Number(n) => {
563 if let Some(i) = n.as_i64() {
564 if let Ok(i32v) = i32::try_from(i) {
565 Ok(WasmArg::I32(i32v))
566 } else {
567 Ok(WasmArg::I64(i))
568 }
569 } else if let Some(f) = n.as_f64() {
570 Ok(WasmArg::F64(f))
571 } else {
572 Err("unsupported number")
573 }
574 }
575 _ => Err("unsupported arg type — only numbers"),
576 }
577 }
578
579 /// Convert a [`WasmArg`] into the wasmtime [`Val`] expected by
580 /// `Func::call_async`. `f32`/`f64` are stored as bit patterns per the
581 /// wasmtime ABI.
582 pub fn into_val(self) -> wasmtime::Val {
583 match self {
584 WasmArg::I32(v) => wasmtime::Val::I32(v),
585 WasmArg::I64(v) => wasmtime::Val::I64(v),
586 WasmArg::F32(v) => wasmtime::Val::F32(v.to_bits()),
587 WasmArg::F64(v) => wasmtime::Val::F64(v.to_bits()),
588 }
589 }
590}
591
592/// Render a wasmtime [`Val`] as the closest-fitting [`serde_json::Value`].
593///
594/// `i32`/`i64` become JSON numbers (integer); `f32`/`f64` become JSON
595/// numbers (floating-point); other value types — `v128`, references —
596/// degrade to a JSON `null` so callers see a stable shape rather than a
597/// runtime error. Used by [`TensorWasmExecutor::call_export_with_args`]
598/// to project the wasmtime result slice into a JSON array.
599fn val_to_json(v: &Val) -> serde_json::Value {
600 match v {
601 Val::I32(n) => serde_json::json!(*n),
602 Val::I64(n) => serde_json::json!(*n),
603 Val::F32(bits) => serde_json::json!(f32::from_bits(*bits)),
604 Val::F64(bits) => serde_json::json!(f64::from_bits(*bits)),
605 // Unsupported value types fall through as JSON null rather than
606 // erroring — keeps the response shape predictable for callers that
607 // only ever return numeric scalars (the common case for B5.6).
608 // The fallthrough is observable: a guest returning a `v128` or a
609 // reference type silently became `null` before, masking a genuine
610 // signature mismatch. Emit a debug event so operators can see when
611 // a non-numeric return is being lossily projected.
612 other => {
613 // `Val::ty()` needs a store handle we don't thread here, so
614 // describe the variant directly — enough for an operator to
615 // tell a `v128` apart from a reference type.
616 let kind = match other {
617 Val::V128(_) => "v128",
618 Val::FuncRef(_) => "funcref",
619 Val::ExternRef(_) => "externref",
620 Val::AnyRef(_) => "anyref",
621 _ => "unknown",
622 };
623 debug!(
624 target: "tensor_wasm_exec::executor",
625 val_kind = kind,
626 "non-numeric wasm return value mapped to JSON null (v128 / reference types are not representable)",
627 );
628 serde_json::Value::Null
629 }
630 }
631}
632
633/// Per-store [`ResourceLimiter`] that caps linear-memory growth at the
634/// engine-configured `max_memory_bytes`.
635///
636/// One instance is attached to each [`Store`] via [`Store::limiter`].
637/// Constructing it is cheap (a single `usize` plus the cached limit). The
638/// `engine_max` field is duplicated from [`crate::engine::EngineConfig::max_memory_bytes`]
639/// so the limiter does not need to re-borrow the engine during the hot
640/// `memory.grow` path.
641#[derive(Debug)]
642pub struct TensorWasmResourceLimiter {
643 /// Per-instance hard cap on linear memory. Mirrored from the engine config.
644 engine_max: usize,
645}
646
647impl TensorWasmResourceLimiter {
648 /// Construct a limiter that denies any growth past `engine_max` bytes.
649 pub fn new(engine_max: usize) -> Self {
650 Self { engine_max }
651 }
652}
653
654impl ResourceLimiter for TensorWasmResourceLimiter {
655 fn memory_growing(
656 &mut self,
657 _current: usize,
658 desired: usize,
659 maximum: Option<usize>,
660 ) -> wasmtime::Result<bool> {
661 // Reject if the requested size exceeds either the engine-wide cap
662 // or the module's own declared maximum. Returning `Ok(false)` causes
663 // `memory.grow` to return -1 in guest land, mirroring wasmtime's
664 // `StoreLimits` convention. Hard traps would not surface a stable
665 // error to the host; instead the executor maps subsequent OOM
666 // behaviour via `ExecError::Wasmtime`.
667 if desired > self.engine_max {
668 return Ok(false);
669 }
670 if let Some(m) = maximum {
671 if desired > m {
672 return Ok(false);
673 }
674 }
675 Ok(true)
676 }
677
678 fn table_growing(
679 &mut self,
680 _current: usize,
681 desired: usize,
682 maximum: Option<usize>,
683 ) -> wasmtime::Result<bool> {
684 // Cap table growth proportionally to the per-instance memory budget.
685 // Each table entry costs ~16 bytes of host memory on wasmtime (a
686 // tagged pointer plus type-index slot). Without this cap a guest
687 // could `table.grow` up to u32::MAX entries (~64 GiB of host RAM at
688 // 16 B/entry), bypassing the `memory_growing` cap entirely.
689 //
690 // Using `engine_max` (the linear-memory byte cap) as the table-byte
691 // budget keeps the policy a single dial: a tenant gets at most
692 // `engine_max` bytes of *either* linear memory *or* table backing
693 // store. That's loose (allows ~engine_max bytes for each) but it
694 // bounds the worst case from u32::MAX entries down to engine_max/16
695 // entries — the qualitative DoS vector closes.
696 // Shared with the pooling allocator's `table_elements` derivation in
697 // `engine.rs` (MED finding) so both backends budget tables against
698 // the same per-instance byte ceiling.
699 let bytes_needed = (desired as u64).saturating_mul(crate::engine::TABLE_ENTRY_BYTES);
700 if bytes_needed > self.engine_max as u64 {
701 return Ok(false);
702 }
703 if let Some(m) = maximum {
704 if desired > m {
705 return Ok(false);
706 }
707 }
708 Ok(true)
709 }
710}
711
712/// The async executor.
713#[derive(Clone)]
714pub struct TensorWasmExecutor {
715 engine: Arc<TensorWasmEngine>,
716 instances: Arc<DashMap<InstanceId, Arc<Mutex<TensorWasmInstance>>>>,
717 next_instance_id: Arc<AtomicU64>,
718 /// Per-engine compiled-module cache keyed by the full 256-bit BLAKE3
719 /// digest of the wasm bytes. Avoids re-running Cranelift on every
720 /// `spawn_instance` for repeat tenants. Hash is computed with `blake3`
721 /// (SIMD-accelerated; ~5x faster than SipHash on multi-MiB wasm
722 /// modules). The full 32-byte digest is used as the key — truncating
723 /// to 8 bytes would expose a ~2⁻³² birthday-collision window across
724 /// tenants (cross-tenant module-cache poisoning) that an attacker
725 /// crafting modules with colliding prefixes could exploit at scale.
726 ///
727 /// Bounded with LRU eviction (cap from
728 /// [`crate::engine::EngineConfig::max_module_cache_entries`], default
729 /// 1024) — closes exec S-5 where an unbounded `DashMap` let a
730 /// misbehaving tenant pin arbitrarily many compiled modules. The
731 /// guard is a `parking_lot::Mutex` rather than a `DashMap` because
732 /// `lru::LruCache` is not concurrency-safe (every `get` mutates the
733 /// recency list).
734 module_cache: Arc<parking_lot::Mutex<LruCache<[u8; 32], Module>>>,
735 /// Live-instance counter, used to enforce
736 /// [`crate::engine::EngineConfig::max_instances`] (exec S-10).
737 /// Atomically bumped *before* compile/instantiate in `spawn_instance`
738 /// (with rollback on failure) and decremented in `terminate`. We keep
739 /// this separate from `instances.len()` so the admission decision
740 /// commits in a single CAS rather than racing against in-flight
741 /// spawns that have already passed the check but not yet inserted.
742 instance_count: Arc<AtomicUsize>,
743 /// Per-tenant live-instance counts, used to enforce
744 /// [`crate::engine::EngineConfig::max_instances_per_tenant`] (fairness
745 /// bound). Keyed by the spawning [`TenantId`]; the value is the number
746 /// of slots that tenant currently holds. Bumped (with rollback on a
747 /// failed spawn) alongside the engine-wide `instance_count` in
748 /// [`Self::charge_instance_slot`] and decremented in [`Self::terminate`]
749 /// / [`Self::release_instance_slot`]. Empty entries are pruned on
750 /// decrement so a tenant that churns instances does not leak map keys.
751 /// Independent of `instance_count` — the engine-wide cap and the
752 /// per-tenant cap are checked in the same admission step but tracked
753 /// separately so one tenant hitting its cap never perturbs another's
754 /// accounting.
755 tenant_counts: Arc<DashMap<TenantId, usize>>,
756 /// Optional metrics handle. When `Some`, spawn/terminate operations
757 /// increment the corresponding Prometheus counters / gauges.
758 metrics: Option<TensorWasmMetrics>,
759 /// One-shot guard for the "epoch ticker not running" operator warning.
760 ///
761 /// Initialised lazily on first observation of a missing ticker; the
762 /// inner [`AtomicBool`] flips to `true` once the warning has fired so
763 /// subsequent spawns on the same executor stay quiet (the warning is
764 /// load-bearing for operators, but at 1 line per spawn it would flood
765 /// the log). Scoped to the executor (and therefore the engine) so
766 /// distinct engines in the same process each get their own warning.
767 ticker_warned: Arc<OnceLock<AtomicBool>>,
768 /// Optional pre-instantiated instance pool (roadmap feature #5).
769 ///
770 /// v0.3.6 scaffold: when set, the pool is wired through to embedders
771 /// but its [`crate::instance_pool::InstancePool::acquire`] path falls
772 /// through to [`Self::spawn_instance`] on every call. v0.4 lands the
773 /// channel-driven warm-instance draw on top of this same field —
774 /// callers wiring it up today get forward-compatible plumbing for
775 /// free.
776 pool: Option<Arc<InstancePool>>,
777 /// Optional JIT kernel cache. When set, `instantiate_detached` registers
778 /// the `tensor-wasm:jit/host` `dispatch`/`alloc`/`free` imports
779 /// (`jit_dispatch::add_jit_dispatch_to_linker`) on every spawn's linker
780 /// so auto-offloaded guests can reach the cached kernels. `None` (the
781 /// default) leaves the JIT surface unlinked — a guest importing it then
782 /// fails to link, matching the historical behaviour for embedders that
783 /// have not opted in. The cache is intentionally per-executor and
784 /// cross-tenant (lookups are tenant-scoped inside the dispatch closure
785 /// via `CacheKey::for_tenant`); see `jit_dispatch.rs`.
786 jit_cache: Option<Arc<KernelCache>>,
787 /// Bounds the number of concurrent `Module::from_binary` compiles on the
788 /// Tokio blocking pool (MEDIUM finding). A permit is acquired around the
789 /// `spawn_blocking` compile in [`Self::compile_module_cached`]; the
790 /// permit is released as soon as the compile resolves (cache hits never
791 /// touch the semaphore). Capacity comes from
792 /// [`crate::engine::EngineConfig::max_concurrent_compiles`], defaulting
793 /// to [`std::thread::available_parallelism`] (floored at 1). Independent
794 /// of `max_instances` — that caps live instances, this caps in-flight
795 /// Cranelift work.
796 compile_semaphore: Arc<Semaphore>,
797}
798
799/// Resolve the concurrent-compile permit count from the engine config,
800/// falling back to the host parallelism (floored at 1) when the operator
801/// left [`crate::engine::EngineConfig::max_concurrent_compiles`] unset.
802/// A configured value of 0 is coerced to 1 (a 0-permit semaphore would
803/// deadlock every compile).
804fn resolve_compile_permits(requested: Option<usize>) -> usize {
805 match requested {
806 Some(0) | None => std::thread::available_parallelism()
807 .map(|n| n.get())
808 .unwrap_or(1)
809 .max(1),
810 Some(n) => n,
811 }
812 .max(1)
813}
814
815/// Resolve a non-zero LRU cache capacity from a possibly-zero config
816/// value. We coerce 0 to 1 because `LruCache::new(NonZeroUsize)` requires
817/// a non-zero capacity, and operators who set the knob to 0 most plausibly
818/// meant "as small as possible" rather than "panic on construction".
819fn lru_cap(requested: usize) -> NonZeroUsize {
820 NonZeroUsize::new(requested).unwrap_or_else(|| NonZeroUsize::new(1).expect("1 is non-zero"))
821}
822
823/// RAII guard that rolls back a successful `instance_count.fetch_add`
824/// if the spawn path drops it without committing. `commit()` defuses
825/// the rollback once the instance has been inserted into the registry;
826/// every other exit path (`?`, panic during `Instance::new_async`,
827/// store-construction failure) leaves the guard alive and triggers a
828/// decrement on drop.
829///
830/// Without this guard, exec S-10 admission control would leak a slot
831/// for every failed spawn — a misbehaving tenant could trip an
832/// always-failing instantiation in a loop and exhaust the cap with
833/// zero live instances.
834struct InstanceSlotGuard {
835 counter: Arc<AtomicUsize>,
836 /// Per-tenant rollback handle. `Some` carries the per-tenant count map
837 /// plus the tenant whose count was bumped alongside the engine-wide
838 /// `counter`; on a non-committed drop both the engine-wide and the
839 /// per-tenant count are rolled back in lockstep. `None` means no
840 /// per-tenant charge was made (per-tenant cap disabled, or this guard
841 /// only re-protects the engine-wide count — e.g. the post-charge
842 /// register guard in `spawn_instance`).
843 tenant_rollback: Option<(Arc<DashMap<TenantId, usize>>, TenantId)>,
844 committed: bool,
845}
846
847impl InstanceSlotGuard {
848 fn new(counter: Arc<AtomicUsize>) -> Self {
849 Self {
850 counter,
851 tenant_rollback: None,
852 committed: false,
853 }
854 }
855
856 /// Construct a guard that, on a non-committed drop, rolls back BOTH the
857 /// engine-wide count and the `tenant`'s per-tenant count. Used by
858 /// [`TensorWasmExecutor::charge_instance_slot`] when a per-tenant cap is
859 /// configured so a failed spawn never leaks either counter.
860 fn with_tenant(
861 counter: Arc<AtomicUsize>,
862 tenant_counts: Arc<DashMap<TenantId, usize>>,
863 tenant: TenantId,
864 ) -> Self {
865 Self {
866 counter,
867 tenant_rollback: Some((tenant_counts, tenant)),
868 committed: false,
869 }
870 }
871
872 fn commit(mut self) {
873 self.committed = true;
874 }
875}
876
877impl Drop for InstanceSlotGuard {
878 fn drop(&mut self) {
879 if !self.committed {
880 // Relaxed is fine here: the matching `fetch_add` used AcqRel
881 // for admission ordering; the rollback only undoes a count
882 // that no other thread depends on observing.
883 self.counter.fetch_sub(1, Ordering::Relaxed);
884 if let Some((tenant_counts, tenant)) = &self.tenant_rollback {
885 decrement_tenant_count(tenant_counts, *tenant);
886 }
887 }
888 }
889}
890
891/// Decrement (and prune-on-zero) a tenant's per-tenant live-instance count.
892/// Shared by [`InstanceSlotGuard`]'s rollback, [`TensorWasmExecutor::terminate`],
893/// and [`TensorWasmExecutor::release_instance_slot`] so the prune-empty-entry
894/// policy lives in exactly one place. A `None`/zero entry is a no-op (a
895/// double-decrement cannot drive the count negative).
896fn decrement_tenant_count(tenant_counts: &DashMap<TenantId, usize>, tenant: TenantId) {
897 if let Entry::Occupied(mut e) = tenant_counts.entry(tenant) {
898 let v = e.get_mut();
899 *v = v.saturating_sub(1);
900 if *v == 0 {
901 e.remove();
902 }
903 }
904}
905
906/// Walk every exported and imported [`ExternType::Memory`] in `module` and
907/// reject the spawn if either the initial (`minimum`) or the declared
908/// `maximum` size, expressed in bytes via the memory type's own
909/// `page_size()`, exceeds `cap_bytes`.
910///
911/// Returns [`ExecError::ModuleMemoryTooLarge`] on the first offending
912/// memory found. The check runs against the compiled [`Module`] before
913/// `Instance::new_async`, so a rejected module is never instantiated and
914/// no host allocation is attempted on its behalf.
915fn check_module_memory_within_cap(module: &Module, cap_bytes: usize) -> Result<(), ExecError> {
916 let cap_u64 = cap_bytes as u64;
917 let check = |mt: &wasmtime::MemoryType| -> Result<(), ExecError> {
918 let page_size = mt.page_size();
919 // `minimum()` is in pages; multiply with overflow-safe saturating
920 // arithmetic so a pathological declaration cannot wrap on cast.
921 let min_pages = mt.minimum();
922 let min_bytes = min_pages.saturating_mul(page_size);
923 if min_bytes > cap_u64 {
924 return Err(ExecError::ModuleMemoryTooLarge {
925 requested_bytes: min_bytes,
926 limit_bytes: cap_u64,
927 });
928 }
929 if let Some(max_pages) = mt.maximum() {
930 let max_bytes = max_pages.saturating_mul(page_size);
931 if max_bytes > cap_u64 {
932 return Err(ExecError::ModuleMemoryTooLarge {
933 requested_bytes: max_bytes,
934 limit_bytes: cap_u64,
935 });
936 }
937 }
938 Ok(())
939 };
940 for ex in module.exports() {
941 if let ExternType::Memory(mt) = ex.ty() {
942 check(&mt)?;
943 }
944 }
945 for im in module.imports() {
946 if let ExternType::Memory(mt) = im.ty() {
947 check(&mt)?;
948 }
949 }
950 Ok(())
951}
952
953/// Pre-compile sibling of [`check_module_memory_within_cap`] that walks the
954/// module's declared linear-memory types directly from the raw Wasm bytes via
955/// [`wasmparser`], *before* [`Module::new`] runs.
956///
957/// `check_module_memory_within_cap` needs a compiled [`Module`], but under the
958/// pooling allocator a module declaring an oversized *initial* memory is
959/// rejected by wasmtime at compile time — so that post-compile walk never runs
960/// for the initial-size case and the caller sees an opaque
961/// [`ExecError::Wasmtime`] instead of the structured
962/// [`ExecError::ModuleMemoryTooLarge`]. Running the identical cap arithmetic on
963/// the raw bytes up-front makes the rejection backend-independent and preserves
964/// the precise requested-byte figure regardless of the configured allocator.
965///
966/// Parse/validation errors are deliberately *not* surfaced here: a malformed
967/// module is left for [`Module::new`] to reject with its richer diagnostics —
968/// this pass only ever refines the memory-cap case, never swallows other
969/// failures.
970fn check_raw_module_memory_within_cap(wasm: &[u8], cap_bytes: usize) -> Result<(), ExecError> {
971 use wasmparser::{Parser, Payload, TypeRef};
972 let cap_u64 = cap_bytes as u64;
973 let check = |mt: &wasmparser::MemoryType| -> Result<(), ExecError> {
974 // Wasm linear-memory sizes are page counts; the default page size is
975 // 64 KiB, overridable per-memory by the custom-page-sizes proposal
976 // (`page_size_log2`). Saturating arithmetic so a pathological
977 // declaration cannot wrap on the multiply.
978 let page_size: u64 = 1u64 << u64::from(mt.page_size_log2.unwrap_or(16));
979 let min_bytes = mt.initial.saturating_mul(page_size);
980 if min_bytes > cap_u64 {
981 return Err(ExecError::ModuleMemoryTooLarge {
982 requested_bytes: min_bytes,
983 limit_bytes: cap_u64,
984 });
985 }
986 if let Some(max_pages) = mt.maximum {
987 let max_bytes = max_pages.saturating_mul(page_size);
988 if max_bytes > cap_u64 {
989 return Err(ExecError::ModuleMemoryTooLarge {
990 requested_bytes: max_bytes,
991 limit_bytes: cap_u64,
992 });
993 }
994 }
995 Ok(())
996 };
997 for payload in Parser::new(0).parse_all(wasm) {
998 // Bail to `Module::new` on any parse hiccup — we only refine here.
999 let Ok(payload) = payload else { return Ok(()) };
1000 match payload {
1001 Payload::MemorySection(reader) => {
1002 for mem in reader {
1003 let Ok(mt) = mem else { return Ok(()) };
1004 check(&mt)?;
1005 }
1006 }
1007 Payload::ImportSection(reader) => {
1008 for im in reader {
1009 let Ok(im) = im else { return Ok(()) };
1010 if let TypeRef::Memory(mt) = im.ty {
1011 check(&mt)?;
1012 }
1013 }
1014 }
1015 _ => {}
1016 }
1017 }
1018 Ok(())
1019}
1020
1021/// Best-effort refinement of a pooling-allocator instantiation error into a
1022/// typed [`ExecError::ModuleMemoryTooLarge`] (MED finding).
1023///
1024/// The pooling allocator surfaces a memory-slot sizing failure as an opaque
1025/// [`wasmtime::Error`]; its message chain contains the allocator's
1026/// memory-size signature (wasmtime phrases these as "memory ... exceeds the
1027/// limit" / "memory minimum size of N pages exceeds ..."). When we recognise
1028/// that signature we re-tag the error as `ModuleMemoryTooLarge` so callers
1029/// see a stable `MemoryExhausted` on the conversion boundary instead of a
1030/// generic compile error. Anything we do not recognise is returned verbatim
1031/// as [`ExecError::Wasmtime`] — we only ever *refine*, never swallow.
1032fn classify_instantiation_error(err: wasmtime::Error, cap_bytes: usize) -> ExecError {
1033 let chain = format!("{err:#}").to_ascii_lowercase();
1034 // LOW finding (fragile error classification): this is *substring* matching
1035 // on English error text because wasmtime does not expose a structured
1036 // error type for a pooling-allocator memory-slot sizing refusal. It is
1037 // therefore WASMTIME-VERSION-COUPLED — a wasmtime upgrade can silently
1038 // reword these phrasings and quietly stop the refinement. That degradation
1039 // is SAFE BY CONSTRUCTION: this function only ever *refines* a recognised
1040 // error into the typed `ModuleMemoryTooLarge`; every unrecognised error
1041 // falls through to the `else` branch below and is returned verbatim as
1042 // `ExecError::Wasmtime(err)` — the original error is never swallowed or
1043 // dropped, only (best-effort) re-tagged. When bumping wasmtime, re-verify
1044 // the phrasings below against the pooling `memory_pool` error messages.
1045 // Conservative match: require both a "memory" mention and an
1046 // exceeds/limit phrasing so unrelated traps/link errors are not
1047 // misclassified. The pooling allocator's memory-size refusals all
1048 // contain one of these limit phrasings.
1049 // These phrasings track wasmtime 25's pooling memory_pool errors:
1050 // * "memory index N has a minimum byte size of M which exceeds the
1051 // limit of L bytes" (per-slot size refusal)
1052 // * "maximum memory size of 0x… bytes exceeds the configured maximum
1053 // size" (pool construction / sizing refusal)
1054 let is_memory_size_failure = chain.contains("memory")
1055 && (chain.contains("exceeds the limit") || chain.contains("exceeds the configured"));
1056 if is_memory_size_failure {
1057 ExecError::ModuleMemoryTooLarge {
1058 // We do not know the exact requested figure here (the allocator
1059 // does not expose it structurally), so report the cap as both
1060 // bounds — the typed variant's value is the classification, and
1061 // the full opaque chain is already logged server-side on the
1062 // conversion boundary.
1063 requested_bytes: cap_bytes as u64,
1064 limit_bytes: cap_bytes as u64,
1065 }
1066 } else {
1067 ExecError::Wasmtime(err)
1068 }
1069}
1070
1071impl TensorWasmExecutor {
1072 /// Construct an executor over the given shared engine.
1073 pub fn new(engine: Arc<TensorWasmEngine>) -> Self {
1074 let cap = lru_cap(engine.config().max_module_cache_entries);
1075 let permits = resolve_compile_permits(engine.config().max_concurrent_compiles);
1076 Self {
1077 engine,
1078 instances: Arc::new(DashMap::new()),
1079 next_instance_id: Arc::new(AtomicU64::new(1)),
1080 module_cache: Arc::new(parking_lot::Mutex::new(LruCache::new(cap))),
1081 instance_count: Arc::new(AtomicUsize::new(0)),
1082 tenant_counts: Arc::new(DashMap::new()),
1083 metrics: None,
1084 ticker_warned: Arc::new(OnceLock::new()),
1085 pool: None,
1086 jit_cache: None,
1087 compile_semaphore: Arc::new(Semaphore::new(permits)),
1088 }
1089 }
1090
1091 /// Construct an executor that publishes spawn/terminate events to the
1092 /// supplied [`TensorWasmMetrics`] registry. Metric handles are cheaply cloneable;
1093 /// pass a clone of the process-wide registry.
1094 pub fn with_metrics(engine: Arc<TensorWasmEngine>, metrics: TensorWasmMetrics) -> Self {
1095 let cap = lru_cap(engine.config().max_module_cache_entries);
1096 let permits = resolve_compile_permits(engine.config().max_concurrent_compiles);
1097 Self {
1098 engine,
1099 instances: Arc::new(DashMap::new()),
1100 next_instance_id: Arc::new(AtomicU64::new(1)),
1101 module_cache: Arc::new(parking_lot::Mutex::new(LruCache::new(cap))),
1102 instance_count: Arc::new(AtomicUsize::new(0)),
1103 tenant_counts: Arc::new(DashMap::new()),
1104 metrics: Some(metrics),
1105 ticker_warned: Arc::new(OnceLock::new()),
1106 pool: None,
1107 jit_cache: None,
1108 compile_semaphore: Arc::new(Semaphore::new(permits)),
1109 }
1110 }
1111
1112 /// Attach an [`InstancePool`] to this executor and return the
1113 /// modified executor.
1114 ///
1115 /// Builder method; pairs with [`Self::new`] / [`Self::with_metrics`].
1116 /// The pool itself is opt-in — embedders that do not call this method
1117 /// see the v0.3.5 behaviour (every `spawn_instance` does a fresh
1118 /// compile/instantiate). Calling it today wires up the v0.3.6
1119 /// scaffold; the same call gets the v0.4 warm-pool path for free
1120 /// once that lands.
1121 pub fn with_instance_pool(mut self, pool: Arc<InstancePool>) -> Self {
1122 self.pool = Some(pool);
1123 self
1124 }
1125
1126 /// Attach a JIT [`KernelCache`] to this executor and return the modified
1127 /// executor.
1128 ///
1129 /// Builder method; pairs with [`Self::new`] / [`Self::with_metrics`].
1130 /// When set, every instantiation registers the `tensor-wasm:jit/host`
1131 /// `dispatch`/`alloc`/`free` imports (see
1132 /// [`crate::jit_dispatch::add_jit_dispatch_to_linker`]) against this
1133 /// cache, so auto-offloaded guests link successfully. Embedders that do
1134 /// not call this leave the JIT surface unlinked — a guest importing it
1135 /// then fails to link, matching the pre-wiring behaviour.
1136 pub fn with_jit_cache(mut self, cache: Arc<KernelCache>) -> Self {
1137 self.jit_cache = Some(cache);
1138 self
1139 }
1140
1141 /// Borrow the attached JIT [`KernelCache`], if any. Returns `None` for
1142 /// executors constructed without [`Self::with_jit_cache`].
1143 pub fn jit_cache(&self) -> Option<&Arc<KernelCache>> {
1144 self.jit_cache.as_ref()
1145 }
1146
1147 /// Borrow the attached [`InstancePool`], if any. Returns `None` for
1148 /// executors constructed without [`Self::with_instance_pool`].
1149 pub fn instance_pool(&self) -> Option<&Arc<InstancePool>> {
1150 self.pool.as_ref()
1151 }
1152
1153 /// Borrow the underlying engine.
1154 pub fn engine(&self) -> &TensorWasmEngine {
1155 &self.engine
1156 }
1157
1158 /// Number of instances currently present in the **registry** — i.e.
1159 /// those that have been spawned and `register_pooled_instance`'d but
1160 /// not yet `terminate`'d. This is the size of the `instances` `DashMap`.
1161 ///
1162 /// Contrast with [`Self::instances_len`], which reports the
1163 /// **admission** count (the atomic counter that
1164 /// [`crate::engine::EngineConfig::max_instances`] is enforced against).
1165 /// The two can briefly diverge: `instances_len` is bumped *before*
1166 /// compile/instantiate in `spawn_instance` (with rollback on failure)
1167 /// and a pool can hold an admitted-but-unregistered instance in a warm
1168 /// channel, so `instances_len() >= live_count()` always holds, with
1169 /// equality once all in-flight spawns have either registered or rolled
1170 /// back. Use `live_count()` when you mean "how many handles can
1171 /// `call_export` resolve right now"; use `instances_len()` when you mean
1172 /// "how close are we to the admission cap".
1173 pub fn live_count(&self) -> usize {
1174 self.instances.len()
1175 }
1176
1177 /// Number of compiled modules retained in the per-executor cache.
1178 /// Exposed for tests and operators that want to confirm cache reuse.
1179 pub fn cached_module_count(&self) -> usize {
1180 self.module_cache.lock().len()
1181 }
1182
1183 /// Current number of entries held by the bounded LRU module cache.
1184 /// Alias for [`Self::cached_module_count`] under the name used by the
1185 /// exec S-5 admission-control bound work; both delegate to the same
1186 /// underlying length so callers can pick whichever reads better at
1187 /// the call site.
1188 pub fn module_cache_len(&self) -> usize {
1189 self.module_cache.lock().len()
1190 }
1191
1192 /// Current per-tenant **admission** count for `tenant`, or `0` if the
1193 /// tenant holds no live slots. This is the counter the per-tenant
1194 /// fairness cap
1195 /// ([`crate::engine::EngineConfig::max_instances_per_tenant`]) is
1196 /// enforced against in `spawn_instance`. Exposed for tests and operators
1197 /// that want to confirm a tenant's footprint; the count only moves when
1198 /// the per-tenant cap is configured (it is otherwise left at 0 and the
1199 /// map stays empty).
1200 pub fn tenant_instance_count(&self, tenant: TenantId) -> usize {
1201 self.tenant_counts
1202 .get(&tenant)
1203 .map(|e| *e.value())
1204 .unwrap_or(0)
1205 }
1206
1207 /// Current **admission** count, sampled atomically. This is the counter
1208 /// the admission check in `spawn_instance` consults to decide whether a
1209 /// new instance fits under
1210 /// [`crate::engine::EngineConfig::max_instances`].
1211 ///
1212 /// This is NOT the registry size — see [`Self::live_count`] for the
1213 /// distinction. This counter includes instances that have been admitted
1214 /// but are not yet (or are no longer) in the registry: an in-flight
1215 /// `spawn_instance` between the admission bump and the registry insert,
1216 /// and pool-held warm instances that were detached from the registry but
1217 /// still occupy a slot. `instances_len() >= live_count()` always holds.
1218 pub fn instances_len(&self) -> usize {
1219 self.instance_count.load(Ordering::Acquire)
1220 }
1221
1222 /// Generate a fresh, vacant [`InstanceId`].
1223 ///
1224 /// `next_instance_id` is an `AtomicU64` widened to a `u128` on insert.
1225 /// At 1 instance per nanosecond it would take ~584 years to wrap, but
1226 /// we still defend against collisions: if the freshly-allocated id is
1227 /// already occupied (post-wrap or external reservation), bump and
1228 /// retry. A `warn!` event fires on every collision so operators see
1229 /// it long before the registry corrupts.
1230 fn allocate_instance_id(&self) -> InstanceId {
1231 loop {
1232 let raw = self.next_instance_id.fetch_add(1, Ordering::Relaxed);
1233 let id = InstanceId(u128::from(raw));
1234 if !self.instances.contains_key(&id) {
1235 return id;
1236 }
1237 warn!(
1238 target: "tensor_wasm_exec::executor",
1239 raw,
1240 "instance id collision detected; retrying with next sequence value",
1241 );
1242 }
1243 }
1244
1245 /// Compile `wasm` via wasmtime, caching the result so repeat calls with
1246 /// the same bytes return without re-running Cranelift. Cache key is the
1247 /// full 32-byte BLAKE3 digest of the wasm bytes — stable across runs
1248 /// and platforms, and ~5x faster than SipHash on the multi-MiB modules
1249 /// we actually compile. Using the full digest (rather than truncating
1250 /// to 8 bytes) closes a cross-tenant cache-poisoning vector: at 8
1251 /// bytes, a 65k-module corpus has a ~2⁻³² collision chance per pair,
1252 /// which an attacker crafting prefix-colliding modules can amplify.
1253 ///
1254 /// The actual `Module::from_binary` call runs inside
1255 /// [`tokio::task::spawn_blocking`]: Cranelift compile is CPU-bound and
1256 /// can exceed 100 ms on multi-MiB modules — running it on a Tokio
1257 /// worker thread blocks every other I/O task multiplexed onto that
1258 /// worker. Offloading to the blocking pool keeps the reactor responsive.
1259 /// The byte-length cap above runs synchronously before the offload so
1260 /// an oversized blob fails fast without entering the blocking pool.
1261 ///
1262 /// Returns the compiled [`Module`] alongside the 32-byte BLAKE3 digest
1263 /// it was keyed under, so callers (e.g. [`Self::build_pooled_instance`])
1264 /// that also need the digest for the pool key do not re-hash the bytes
1265 /// (PERF: the hash was previously computed here AND in
1266 /// `build_pooled_instance`).
1267 async fn compile_module_cached(&self, wasm: &[u8]) -> Result<(Module, ModuleHash), ExecError> {
1268 // Pre-compile size cap (exec hardening). Reject pathologically
1269 // large blobs *before* hashing or handing them to Cranelift —
1270 // a wasm with a malicious code section can otherwise force
1271 // arbitrary compile-time CPU. The configured cap is floored at
1272 // `MAX_MODULE_BYTES` upstream in `EngineConfig`, but we still
1273 // observe the configured value here so a stricter operator
1274 // policy wins.
1275 let cap = self.engine.config().max_module_bytes;
1276 if wasm.len() > cap {
1277 return Err(ExecError::ModuleTooLarge {
1278 len: wasm.len(),
1279 max: cap,
1280 });
1281 }
1282 let digest = blake3::hash(wasm);
1283 // BLAKE3 outputs a fixed 32-byte digest; use it whole as the cache key.
1284 let key: [u8; 32] = *digest.as_bytes();
1285 // Scoped lock for the get: releasing the mutex before the
1286 // potentially-expensive `Module::from_binary` call below is what
1287 // lets concurrent spawns of *different* modules compile in
1288 // parallel. The cost is that two spawns of the *same* fresh
1289 // module may both compile it — but the second one's `put` simply
1290 // overwrites the first, no correctness hazard.
1291 if let Some(m) = self.module_cache.lock().get(&key).cloned() {
1292 return Ok((m, key));
1293 }
1294 // Cranelift compile is CPU-bound — offload to the blocking
1295 // pool. We clone the wasmtime `Engine` (cheap `Arc`-shaped
1296 // internally) and the wasm bytes so the closure is fully
1297 // owning. `spawn_blocking` returns a `JoinError` which we
1298 // surface as a wasmtime error: a panic inside Cranelift is
1299 // not something a caller can usefully distinguish from a
1300 // parse failure, and either way the spawn must be aborted.
1301 let engine = self.engine.inner().clone();
1302 let bytes = wasm.to_vec();
1303 // Bound concurrent Cranelift compiles (MEDIUM finding). The permit is
1304 // held only for the duration of the `spawn_blocking` compile and
1305 // dropped immediately after — cache hits above never reach here, so
1306 // repeat tenants do not contend for permits. `acquire_owned` cannot
1307 // fail unless the semaphore is closed, which we never do; surface the
1308 // (unreachable) closed case as a wasmtime error rather than panicking.
1309 let _permit = self
1310 .compile_semaphore
1311 .clone()
1312 .acquire_owned()
1313 .await
1314 .map_err(|_| ExecError::Wasmtime(wasmtime::Error::msg("compile semaphore closed")))?;
1315 let module = tokio::task::spawn_blocking(move || Module::from_binary(&engine, &bytes))
1316 .await
1317 .map_err(|join_err| {
1318 ExecError::Wasmtime(wasmtime::Error::msg(format!(
1319 "wasm compile task failed: {join_err}"
1320 )))
1321 })?
1322 .map_err(ExecError::Wasmtime)?;
1323 drop(_permit);
1324 self.module_cache.lock().put(key, module.clone());
1325 Ok((module, key))
1326 }
1327
1328 /// Internal: charge a live-instance slot, enforcing both the engine-wide
1329 /// [`max_instances`](crate::engine::EngineConfig::max_instances) ceiling
1330 /// and the optional per-tenant
1331 /// [`max_instances_per_tenant`](crate::engine::EngineConfig::max_instances_per_tenant)
1332 /// fairness cap keyed by `tenant`. Returns an [`InstanceSlotGuard`] that
1333 /// rolls BOTH counts back unless `commit()` is called. Used by
1334 /// [`Self::build_pooled_instance`] / [`Self::rebuild_pooled_from_module`]
1335 /// (and transitively [`Self::spawn_instance`]) so the pool's pre-spawn /
1336 /// reset paths share the same admission accounting as the bare spawn
1337 /// path.
1338 ///
1339 /// Charge order is engine-wide first, then per-tenant. If the per-tenant
1340 /// cap rejects, the engine-wide charge is rolled back before returning so
1341 /// a tenant hitting its cap never erodes the shared ceiling.
1342 fn charge_instance_slot(&self, tenant: TenantId) -> Result<InstanceSlotGuard, ExecError> {
1343 // Engine-wide ceiling (exec S-10).
1344 if let Some(max) = self.engine.config().max_instances {
1345 let new_count = self.instance_count.fetch_add(1, Ordering::AcqRel) + 1;
1346 if new_count > max {
1347 self.instance_count.fetch_sub(1, Ordering::Relaxed);
1348 return Err(ExecError::CapacityExhausted {
1349 active: new_count,
1350 limit: max,
1351 });
1352 }
1353 } else {
1354 self.instance_count.fetch_add(1, Ordering::AcqRel);
1355 }
1356
1357 // Per-tenant fairness cap. When unset, the engine-wide charge above
1358 // is the whole story and the guard carries no per-tenant rollback.
1359 let Some(per_tenant_max) = self.engine.config().max_instances_per_tenant else {
1360 return Ok(InstanceSlotGuard::new(self.instance_count.clone()));
1361 };
1362
1363 // Charge the per-tenant count under the DashMap entry lock so the
1364 // read-modify-write is atomic against concurrent spawns of the SAME
1365 // tenant (different tenants take different shards / entries and never
1366 // serialise against each other). On overflow, roll the engine-wide
1367 // charge back so the rejected spawn consumes neither counter.
1368 let mut entry = self.tenant_counts.entry(tenant).or_insert(0);
1369 let new_tenant_count = *entry + 1;
1370 if new_tenant_count > per_tenant_max {
1371 // Drop the entry guard before mutating other state; the count was
1372 // never incremented so there is nothing to roll back on the
1373 // per-tenant side.
1374 drop(entry);
1375 self.instance_count.fetch_sub(1, Ordering::Relaxed);
1376 // Distinct from the engine-wide `CapacityExhausted` above: this is
1377 // a per-tenant fairness rejection (the offending tenant is over its
1378 // own quota while the shared budget still has room). It carries the
1379 // tenant id so the API layer can surface a quota-specific 429
1380 // (`tenant_capacity_exhausted`) instead of a generic 503. Still log
1381 // it server-side for operator visibility.
1382 warn!(
1383 target: "tensor_wasm_exec::executor",
1384 tenant = %tenant,
1385 active = new_tenant_count,
1386 limit = per_tenant_max,
1387 "per-tenant instance cap exhausted; refusing spawn (other tenants unaffected)",
1388 );
1389 return Err(ExecError::TenantCapacityExhausted {
1390 tenant,
1391 active: new_tenant_count,
1392 limit: per_tenant_max,
1393 });
1394 }
1395 *entry = new_tenant_count;
1396 drop(entry);
1397 Ok(InstanceSlotGuard::with_tenant(
1398 self.instance_count.clone(),
1399 self.tenant_counts.clone(),
1400 tenant,
1401 ))
1402 }
1403
1404 /// Internal: explicitly release a live-instance slot charged for
1405 /// `tenant`. Used by [`InstancePool`] when an instance held in a warm
1406 /// channel is dropped (channel full on release, reset failed, pool
1407 /// shutdown). Mirrors the slot release that [`Self::terminate`] performs
1408 /// for the registered case — both the engine-wide and the per-tenant
1409 /// count are decremented — but does not touch the registry, since
1410 /// pooled-but-not-handed-out instances were never registered. The
1411 /// `tenant` is the same one the matching
1412 /// [`Self::charge_instance_slot`] was keyed under (the pool always knows
1413 /// it via the channel's `(tenant, module_hash)` key).
1414 pub(crate) fn release_instance_slot(&self, tenant: TenantId) {
1415 self.instance_count.fetch_sub(1, Ordering::AcqRel);
1416 decrement_tenant_count(&self.tenant_counts, tenant);
1417 }
1418
1419 /// Internal: compile + instantiate a Wasm module without registering
1420 /// the result in the executor registry. The admission slot is charged
1421 /// (and never rolled back on the success path) so the caller —
1422 /// [`InstancePool`] — can hold the instance in a warm channel and
1423 /// account it against the per-engine live-instance cap.
1424 ///
1425 /// Returns the detached [`TensorWasmInstance`] plus the compiled
1426 /// [`Module`] (cached by [`Self::compile_module_cached`], so it is
1427 /// nearly free to keep around) plus the wasm BLAKE3 digest. The pool
1428 /// uses the digest as half of its `(tenant_id, module_hash)` channel
1429 /// key, and the [`Module`] for the cheap reset-on-release path
1430 /// (re-instantiate from the cached compile, no Cranelift work).
1431 ///
1432 /// This is the shared implementation under both [`Self::spawn_instance`]
1433 /// (which registers immediately) and [`InstancePool::acquire`] /
1434 /// [`InstancePool::release`] (which hold the instance detached in a
1435 /// channel). Every deadline / ticker / module-cap check from
1436 /// [`Self::spawn_instance`] is preserved verbatim — the only behavioural
1437 /// difference is the missing `instances.insert` at the end.
1438 pub(crate) async fn build_pooled_instance(
1439 &self,
1440 cfg: &SpawnConfig,
1441 wasm: &[u8],
1442 ) -> Result<(TensorWasmInstance, Module, ModuleHash), ExecError> {
1443 // Pre-compile memory-cap walk (exec-S-2 / mem-H5): reject a module
1444 // whose *declared* linear memory exceeds the engine cap with a
1445 // structured `ModuleMemoryTooLarge` BEFORE `Module::new` (and before
1446 // any instance slot is charged). Under the pooling allocator an
1447 // oversized *initial* memory is otherwise refused by wasmtime at
1448 // compile time with an opaque "memory ... exceeds the limit" error
1449 // that masks the typed variant the post-compile
1450 // `check_module_memory_within_cap` already produces for the
1451 // declared-maximum and imported cases. Checking the raw bytes here
1452 // makes the rejection — and its precise requested-byte figure —
1453 // backend-independent (on-demand vs. pooling).
1454 check_raw_module_memory_within_cap(wasm, self.engine.config().effective_memory_cap())?;
1455 let slot_guard = self.charge_instance_slot(cfg.tenant_id)?;
1456 // Compile (and cache) the module first; the cap check lives
1457 // inside `compile_module_cached` so an oversized blob fails
1458 // before the digest computation matters. The digest is computed
1459 // once inside `compile_module_cached` (as the cache key) and
1460 // returned here so the pool key path reuses it rather than
1461 // re-hashing the wasm bytes (PERF: previously hashed twice).
1462 let (module, module_hash) = self.compile_module_cached(wasm).await?;
1463 let inst = self.instantiate_detached(cfg, &module).await?;
1464 // A wasmtime instance was successfully created — count it
1465 // against the monotonic spawn counter exactly once per genuine
1466 // instantiation. The `active_instances` gauge moves only at
1467 // registry insert / detach time (see `register_pooled_instance`
1468 // and `detach_pooled_instance`).
1469 if let Some(m) = &self.metrics {
1470 m.instance_spawns_total().inc();
1471 }
1472 // Slot stays charged — defuse rollback so the pool's caller can
1473 // either register the instance (commit it for real) or release
1474 // the slot explicitly via [`Self::release_instance_slot`].
1475 slot_guard.commit();
1476 Ok((inst, module, module_hash))
1477 }
1478
1479 /// Internal: build a detached instance from an already-cached [`Module`].
1480 /// Used by the pool's reset path: drop the spent instance, re-instantiate
1481 /// from the same compiled module (skipping the Cranelift step entirely),
1482 /// and stash the fresh instance back in the warm channel.
1483 ///
1484 /// The slot is charged on success and the caller decides whether to
1485 /// release it ([`Self::release_instance_slot`]) or register it
1486 /// ([`Self::register_pooled_instance`]).
1487 pub(crate) async fn rebuild_pooled_from_module(
1488 &self,
1489 cfg: &SpawnConfig,
1490 module: &Module,
1491 ) -> Result<TensorWasmInstance, ExecError> {
1492 let slot_guard = self.charge_instance_slot(cfg.tenant_id)?;
1493 let inst = self.instantiate_detached(cfg, module).await?;
1494 if let Some(m) = &self.metrics {
1495 m.instance_spawns_total().inc();
1496 }
1497 slot_guard.commit();
1498 Ok(inst)
1499 }
1500
1501 /// Internal: shared instantiation logic. Builds the [`Store`], wires the
1502 /// limiter, arms the start-function epoch deadline, builds a
1503 /// [`wasmtime::Linker`] with every available host surface registered
1504 /// (scheduler always; the `wasi:tensor/host` input channel — `input-len`
1505 /// / `read-input` — always; streaming when [`SpawnConfig::streaming`] is
1506 /// set; JIT dispatch when a [`KernelCache`] is configured via
1507 /// [`Self::with_jit_cache`]), and instantiates against it. Does NOT touch
1508 /// the registry or the admission counter — callers must pair this with
1509 /// [`Self::charge_instance_slot`] / [`Self::register_pooled_instance`].
1510 async fn instantiate_detached(
1511 &self,
1512 cfg: &SpawnConfig,
1513 module: &Module,
1514 ) -> Result<TensorWasmInstance, ExecError> {
1515 // Reconcile the engine-wide cap with the pooling allocator's slot
1516 // size (MED finding): on the pooling backend a module larger than
1517 // the physical slot would fail to instantiate with an opaque
1518 // allocator error, so we cap the limiter AND the pre-instantiation
1519 // module check at `min(max_memory_bytes, pooling memory_bytes)`.
1520 // `effective_memory_cap()` is `max_memory_bytes` on the
1521 // UnifiedBuffer path, preserving prior behaviour there.
1522 let max_memory_bytes = self.engine.config().effective_memory_cap();
1523 let mut state =
1524 InstanceState::new(cfg.tenant_id, InstanceId(0)).with_memory_limit(max_memory_bytes);
1525 if let Some(ref s) = cfg.streaming {
1526 state = state.with_streaming(s.clone());
1527 }
1528 if !cfg.input.is_empty() {
1529 state = state.with_input(InputContext::new(cfg.input.clone()));
1530 }
1531 if let Some(d) = cfg.deadline {
1532 state = state
1533 .with_deadline(Instant::now() + d)
1534 .with_deadline_duration(d);
1535 }
1536 // HARD-deadline instant the cooperative epoch callback traps at while
1537 // `start` (and anything else running inside `instantiate_async`)
1538 // executes. Bounded by BOTH the per-call deadline (if any) and the
1539 // implicit [`MAX_START_FN_DURATION`] cap, mirroring the old
1540 // `start_deadline_ticks = min(epoch_deadline_ticks, max_start_ticks)`
1541 // trap point — except now expressed as a wall-clock instant the
1542 // callback consults on each yield rather than a one-shot trap count.
1543 // Until the callback observes this instant has passed it yields
1544 // cooperatively, so even a runaway `start` function is interruptible
1545 // by the epoch ticker AND traps no later than this instant.
1546 let start_phase_budget = match cfg.deadline {
1547 Some(d) => d.min(MAX_START_FN_DURATION),
1548 None => MAX_START_FN_DURATION,
1549 };
1550 state = state.with_hard_deadline(Instant::now() + start_phase_budget);
1551 let tick = self.engine.config().epoch_tick;
1552 let epoch_deadline_ticks = match cfg.deadline {
1553 Some(d) => duration_to_epoch_ticks(d, tick),
1554 // Overflow-safe "no deadline": `set_epoch_deadline` is relative
1555 // (`current_epoch + ticks`), so `u64::MAX` would overflow once the
1556 // ticker advances. See [`MAX_EPOCH_DEADLINE_TICKS`].
1557 None => MAX_EPOCH_DEADLINE_TICKS,
1558 };
1559 let max_start_ticks = {
1560 let d_ms = MAX_START_FN_DURATION.as_millis();
1561 let t_ms = tick.as_millis().max(1);
1562 let ticks_u128 = d_ms.div_ceil(t_ms).max(1);
1563 u64::try_from(ticks_u128).unwrap_or(u64::MAX)
1564 };
1565 let start_deadline_ticks = epoch_deadline_ticks.min(max_start_ticks);
1566 let deadline_class_applies =
1567 cfg.deadline.is_some() || MAX_START_FN_DURATION > Duration::ZERO;
1568 if deadline_class_applies && !self.engine.is_epoch_ticker_running() {
1569 let flag = self.ticker_warned.get_or_init(|| AtomicBool::new(false));
1570 if !flag.swap(true, Ordering::AcqRel) {
1571 tracing::error!(
1572 target: "tensor_wasm_exec::executor",
1573 "epoch ticker not running — refusing spawn; call `engine.spawn_epoch_ticker()` before serving traffic",
1574 );
1575 }
1576 return Err(ExecError::EpochTickerNotRunning);
1577 }
1578 check_module_memory_within_cap(module, max_memory_bytes)?;
1579 let mut store = Store::new(self.engine.inner(), state);
1580 store.limiter(|state| &mut state.limiter as &mut dyn ResourceLimiter);
1581 // Arm the COOPERATIVE epoch scheme for the start phase. The first
1582 // deadline is armed at the cooperative cadence (never beyond
1583 // `start_deadline_ticks`, which is the latest the hard deadline could
1584 // possibly fall) so the callback gets a chance to run — and yield —
1585 // long before the start-phase budget elapses. At each epoch trip the
1586 // guest yields `Pending` (so an in-flight `instantiate_async` is
1587 // cancellable on future-drop) and the callback traps only once the
1588 // per-store `hard_deadline` (set just above to the start-phase budget)
1589 // has elapsed — preserving the `MAX_START_FN_DURATION` guarantee that
1590 // a runaway `start` cannot burn forever. The relative tick counts now
1591 // govern only the YIELD cadence; the trap point is the wall-clock
1592 // `hard_deadline` instant the callback consults.
1593 arm_cooperative_epoch(
1594 &mut store,
1595 start_deadline_ticks.min(COOPERATIVE_YIELD_TICKS),
1596 );
1597 // HIGH finding fix: build a single `Linker<InstanceState>` and
1598 // register every host surface whose backing machinery is actually
1599 // present, then instantiate against it. Previously only the
1600 // streaming surface was wired (and only on the streaming path),
1601 // leaving the scheduler (`instance.rs` `SchedulerContext`) and
1602 // JIT-dispatch (`jit_dispatch.rs`) imports unlinkable on the real
1603 // spawn path — a guest importing `wasi:scheduler/host` or
1604 // `tensor-wasm:jit/host` failed to link and that machinery was dead.
1605 //
1606 // All three surfaces now coexist on the same linker:
1607 // - scheduler: always registered. The per-store `SchedulerContext`
1608 // is constructed unconditionally in `InstanceState::new`
1609 // (`unbounded()` default, real budget when a deadline is set), so
1610 // the surface is always safe to expose — a guest that doesn't
1611 // import it pays nothing.
1612 // - streaming: registered when `cfg.streaming.is_some()`
1613 // (preserves prior behaviour; the per-store context is
1614 // `disabled()` otherwise and the surface is only meaningful for
1615 // the `/invoke-stream` route).
1616 // - jit: registered when a `KernelCache` is configured on the
1617 // executor (`with_jit_cache`); the cache is the cross-tenant
1618 // backing store the dispatch closure consults (lookups are
1619 // tenant-scoped internally).
1620 //
1621 // The scheduler surface is always available (its per-store context is
1622 // constructed unconditionally), so we always go through the linker
1623 // path now. A guest with zero imports still instantiates fine — a
1624 // linker with extra registered imports does not force the guest to
1625 // import them. (Wasmtime only errors on *missing* imports, never on
1626 // *unused* registered ones.) This replaces the old
1627 // `Instance::new_async(.., &[])` empty-imports branch.
1628 let mut linker: wasmtime::Linker<InstanceState> =
1629 wasmtime::Linker::new(self.engine.inner());
1630 // Scheduler host functions (`wasi:scheduler/host@0.1.0`). The getter
1631 // borrows the per-store `SchedulerContext`, so two instances sharing
1632 // a linker never cross-talk.
1633 add_scheduler_to_linker(&mut linker, |state: &InstanceState| state.scheduler())
1634 .map_err(ExecError::Wasmtime)?;
1635 // Guest-input pull channel (`wasi:tensor/host` `input-len` /
1636 // `read-input`): always registered. The per-store `InputContext`
1637 // is constructed unconditionally in `InstanceState::new`
1638 // (`empty()` default, populated when `cfg.input` is non-empty),
1639 // so the surface is always safe to expose — a guest that doesn't
1640 // import it pays nothing, and one that does sees `input-len() ==
1641 // 0` when no prompt was staged. Mirrors the always-on scheduler
1642 // surface above. Registered before streaming so both
1643 // `wasi:tensor/host` host-fn families coexist on one linker.
1644 add_input_to_linker(&mut linker).map_err(ExecError::Wasmtime)?;
1645 if cfg.streaming.is_some() {
1646 add_streaming_to_linker(&mut linker).map_err(ExecError::Wasmtime)?;
1647 }
1648 if let Some(cache) = &self.jit_cache {
1649 // `add_jit_dispatch_to_linker` is generic over the store payload
1650 // via `JitArenaProvider + TenantContext`, both of which
1651 // `InstanceState` implements (see `instance.rs` /
1652 // `jit_dispatch.rs`). The arena lives per-store; the cache is the
1653 // shared cross-tenant backing handle.
1654 add_jit_dispatch_to_linker(&mut linker, cache.clone()).map_err(ExecError::Wasmtime)?;
1655 }
1656 let instance = match linker.instantiate_async(&mut store, module).await {
1657 Ok(inst) => inst,
1658 Err(err) => {
1659 // MED finding: the pre-instantiation `check_module_memory_within_cap`
1660 // already rejects modules whose *declared* memory exceeds the
1661 // reconciled cap with a typed `ModuleMemoryTooLarge`. The
1662 // pooling allocator can still refuse instantiation for a
1663 // memory-sizing reason the static check cannot see (e.g. the
1664 // slot byte-size ceiling), surfacing it as an opaque
1665 // wasmtime error. Where the error chain carries the
1666 // allocator's memory-size signature, re-classify it as the
1667 // typed `ModuleMemoryTooLarge` so callers get a stable
1668 // `MemoryExhausted` rather than an opaque compile error.
1669 // This is best-effort string inspection — it only *refines*
1670 // the error; anything unrecognised still flows through as
1671 // `ExecError::Wasmtime` unchanged.
1672 return Err(classify_instantiation_error(err, max_memory_bytes));
1673 }
1674 };
1675 // Start phase complete: re-arm the cooperative epoch deadline for the
1676 // post-instantiation window and swap the hard-deadline instant from
1677 // the start-phase budget over to the per-call deadline.
1678 //
1679 // * Deadline spawn: trap at the configured per-call `deadline`
1680 // instant (`now + d`, seeded above). A later `call_export`
1681 // re-arms both this instant and the ticks for that call's window.
1682 // * Deadline-less spawn: clear the hard deadline so the callback
1683 // yields cooperatively FOREVER (matching the old
1684 // `MAX_EPOCH_DEADLINE_TICKS` "effectively no deadline" semantics) —
1685 // but now the guest is CANCELLABLE on future-drop between yields,
1686 // which the old trap-at-`u64::MAX/2` scheme could not offer.
1687 //
1688 // The relative tick count armed here is only the YIELD cadence (never
1689 // beyond `epoch_deadline_ticks`, the latest a configured deadline
1690 // could fall); the trap decision is driven entirely by the wall-clock
1691 // `hard_deadline` the callback consults. Arming at the cooperative
1692 // cadence — rather than the old `epoch_deadline_ticks` (which for a
1693 // deadline-less spawn was the near-infinite `MAX_EPOCH_DEADLINE_TICKS`
1694 // sentinel) — is what makes even a deadline-less compute-bound guest
1695 // yield, and therefore stay cancellable, instead of running an entire
1696 // epoch sentinel's worth of ticks before its first yield point.
1697 store.set_epoch_deadline(epoch_deadline_ticks.min(COOPERATIVE_YIELD_TICKS));
1698 let post_start_hard_deadline = match cfg.deadline {
1699 Some(_) => store.data().deadline,
1700 None => None,
1701 };
1702 store.data_mut().hard_deadline = post_start_hard_deadline;
1703 Ok(TensorWasmInstance::new(store, instance))
1704 }
1705
1706 /// Internal: register a previously-detached [`TensorWasmInstance`] in
1707 /// the executor registry, allocating a fresh [`InstanceId`]. The slot
1708 /// is assumed to already be charged (via
1709 /// [`Self::charge_instance_slot`] or its [`InstanceSlotGuard::commit`]);
1710 /// this method does NOT bump `instance_count`.
1711 ///
1712 /// Used by [`InstancePool::acquire`] to surface a warm-channel
1713 /// instance through the standard [`InstanceId`] handle that
1714 /// [`Self::call_export_with_args`] consumes.
1715 pub(crate) fn register_pooled_instance(
1716 &self,
1717 mut inst: TensorWasmInstance,
1718 ) -> Result<InstanceId, ExecError> {
1719 let id = self.allocate_instance_id();
1720 // Overwrite the placeholder InstanceId(0) baked in at
1721 // instantiation time with the freshly-allocated registry id so
1722 // host imports observing `caller.data().instance_id` see the
1723 // same value the caller holds.
1724 inst.store.data_mut().instance_id = id;
1725 match self.instances.entry(id) {
1726 Entry::Vacant(v) => {
1727 v.insert(Arc::new(Mutex::new(inst)));
1728 }
1729 Entry::Occupied(_) => {
1730 warn!(
1731 target: "tensor_wasm_exec::executor",
1732 %id,
1733 "instance id race after allocation (pool register); this is a serious bug",
1734 );
1735 return Err(ExecError::Wasmtime(wasmtime::Error::msg(
1736 "instance id collision after allocation",
1737 )));
1738 }
1739 }
1740 if let Some(m) = &self.metrics {
1741 // Only the gauge moves here — `instance_spawns_total` is
1742 // incremented at instantiate time inside
1743 // [`Self::build_pooled_instance`] / [`Self::rebuild_pooled_from_module`]
1744 // so the monotonic counter measures genuine wasmtime
1745 // instantiations, not registry insertions (the pool can
1746 // re-register the same underlying instance after a reset).
1747 m.active_instances().inc();
1748 }
1749 Ok(id)
1750 }
1751
1752 /// Internal: remove an instance from the registry, returning the
1753 /// underlying [`TensorWasmInstance`] WITHOUT decrementing
1754 /// `instance_count`. The slot remains charged so the caller —
1755 /// [`InstancePool::release`] — can keep the instance alive in its
1756 /// warm channel without churning the admission counter.
1757 ///
1758 /// Returns `None` if the id is unknown (already terminated / never
1759 /// registered).
1760 pub(crate) async fn detach_pooled_instance(
1761 &self,
1762 id: InstanceId,
1763 ) -> Option<TensorWasmInstance> {
1764 let (_, handle) = self.instances.remove(&id)?;
1765 // Decrement the "active_instances" gauge: this is the moment the
1766 // instance leaves the externally-visible registry. The
1767 // `instance_spawns_total` counter is intentionally not paired
1768 // with a decrement (it is monotonic), and `instance_terminations_total`
1769 // is reserved for genuine terminate calls, not pool detach.
1770 if let Some(m) = &self.metrics {
1771 m.active_instances().dec();
1772 }
1773 // Capture the owning tenant id BEFORE we attempt `try_unwrap`, while
1774 // we can still cheaply read it. We need it on the try_unwrap-failure
1775 // branch below to decrement the per-tenant fairness counter (the
1776 // engine-wide slot release alone used to leave the per-tenant count
1777 // leaked). Only read it when a per-tenant cap is configured —
1778 // otherwise `tenant_counts` is never populated and we skip the
1779 // per-instance lock entirely, exactly mirroring `terminate`'s
1780 // discipline. The lock here is uncontended on the common path (we hold
1781 // the only handle after the `DashMap::remove` above); a racing
1782 // in-flight call would briefly hold it, which is correct — we want the
1783 // post-call tenant. We drop the lock guard before `try_unwrap` so we
1784 // are not ourselves holding a strong borrow that would defeat it.
1785 let owning_tenant = if self.engine.config().max_instances_per_tenant.is_some() {
1786 Some(handle.lock().await.tenant_id())
1787 } else {
1788 None
1789 };
1790 // Unwrap the Arc<Mutex<_>>: we need the sole strong reference to move
1791 // the instance out into the warm channel.
1792 //
1793 // LOW finding (detach try_unwrap race): the previous comment claimed
1794 // "the registry never hands out Arc clones" — that is NOT true.
1795 // `call_export_with_args` clones the value handle
1796 // (`.value().clone()`) and holds it across its `call_async` await, so
1797 // a `detach` racing an in-flight call on the SAME id observes an
1798 // outstanding strong reference and `try_unwrap` fails. The
1799 // `DashMap::remove` above prevents NEW clones (the entry is gone), but
1800 // an already-cloned handle from a call that started before the remove
1801 // can still be live. This is rare (the pool detaches idle instances)
1802 // but reachable, so the failure branch below must be correct, not
1803 // merely "should not happen".
1804 match Arc::try_unwrap(handle) {
1805 Ok(mutex) => Some(mutex.into_inner()),
1806 Err(_arc) => {
1807 // A concurrent `call_export_with_args` still holds a clone of
1808 // this handle. We drop our reference and skip the pool path —
1809 // returning `None` is the safe behaviour (no use-after-detach;
1810 // the racing call finishes and drops the last Arc, freeing the
1811 // instance). But we already removed the entry from the
1812 // registry, so the engine-wide admission slot it occupied
1813 // would otherwise leak permanently (no `terminate` /
1814 // `release_instance_slot` will ever run for this id). The
1815 // racing in-flight call is the pool path's *non-terminating*
1816 // `call_export_with_args` (see `invoke`): when it finishes it
1817 // merely drops its Arc clone — it never calls `terminate` /
1818 // `release_instance_slot`, and the registry entry is already
1819 // gone, so NO other code path will ever decrement this slot.
1820 // We therefore must release BOTH counters here, exactly
1821 // mirroring the `release_instance_slot` the `Some` path's
1822 // caller (`InstancePool::release`) would have run. This is the
1823 // sole decrement for this slot, so it cannot double-count:
1824 // - engine-wide slot: released here (was previously the only
1825 // release, hence correct);
1826 // - per-tenant slot: now released here too (using the tenant
1827 // id captured up-front above), closing the unbounded
1828 // per-tenant leak that the old "lesser, bounded evil"
1829 // comment described. The decrement is skipped when no
1830 // per-tenant cap is configured (`owning_tenant` is `None`),
1831 // matching `terminate` / `release_instance_slot`.
1832 // The happy path above intentionally leaves the slot charged
1833 // (the pool owns it and releases via `release_instance_slot`).
1834 self.instance_count.fetch_sub(1, Ordering::AcqRel);
1835 if let Some(tenant) = owning_tenant {
1836 decrement_tenant_count(&self.tenant_counts, tenant);
1837 }
1838 warn!(
1839 target: "tensor_wasm_exec::executor",
1840 %id,
1841 "detach_pooled_instance: outstanding Arc reference (racing in-flight call); \
1842 instance not pooled, both engine-wide and per-tenant slots released",
1843 );
1844 None
1845 }
1846 }
1847 }
1848
1849 /// Internal: when [`EngineConfig::auto_offload`](crate::engine::EngineConfig::auto_offload)
1850 /// is enabled, consult the analyser and rewrite offload-candidate
1851 /// function bodies into JIT-dispatch trampolines, returning the
1852 /// rewritten bytes. On every fallback condition — flag disabled, no JIT
1853 /// cache attached (the trampoline's `tensor-wasm:jit/host` imports would
1854 /// be unlinkable), analysis error, rewrite error, or a rewrite that
1855 /// swapped nothing — the original `wasm` slice is borrowed unchanged.
1856 /// This is the activation point for the swap the `auto_offload` module
1857 /// documents as consultation-only: with the flag off it stays exactly
1858 /// that.
1859 ///
1860 /// Never returns an error: a failure to analyse or rewrite must not fail
1861 /// a spawn, so every error path logs and falls back to the original
1862 /// module.
1863 fn maybe_rewrite_for_offload<'w>(
1864 &self,
1865 cfg: &SpawnConfig,
1866 wasm: &'w [u8],
1867 ) -> std::borrow::Cow<'w, [u8]> {
1868 use std::borrow::Cow;
1869
1870 if !self.engine.config().auto_offload {
1871 return Cow::Borrowed(wasm);
1872 }
1873 // The rewritten module imports `tensor-wasm:jit/host` (dispatch /
1874 // alloc / free). Those imports only link when a `KernelCache` is
1875 // attached (`with_jit_cache`) — without one the rewritten module
1876 // would fail to instantiate, which is strictly worse than running
1877 // the original on the CPU. Skip the rewrite and fall back.
1878 let Some(cache) = self.jit_cache.as_ref() else {
1879 debug!(
1880 target: "tensor_wasm_exec::executor",
1881 tenant = %cfg.tenant_id,
1882 "auto_offload enabled but no JIT cache attached; falling back to original module",
1883 );
1884 return Cow::Borrowed(wasm);
1885 };
1886
1887 // Resolve the detector thresholds once and use the SAME config for
1888 // the consultation pass and the rewrite so they agree on candidates.
1889 let detector = self
1890 .engine
1891 .config()
1892 .auto_offload_detector
1893 .unwrap_or_default();
1894
1895 // Consultation pass: emit the per-function verdicts (the historical
1896 // consultation-only behaviour) and decide whether any function is
1897 // worth offloading before paying for the rewrite. If analysis
1898 // errors, fall back.
1899 let verdicts = match crate::auto_offload::analyse_with_config(wasm, &detector) {
1900 Ok(v) => v,
1901 Err(e) => {
1902 warn!(
1903 target: "tensor_wasm_exec::executor",
1904 tenant = %cfg.tenant_id,
1905 error = %e,
1906 "auto_offload analysis failed; falling back to original module",
1907 );
1908 return Cow::Borrowed(wasm);
1909 }
1910 };
1911 let any_offload = verdicts.iter().any(|v| {
1912 matches!(
1913 v.verdict,
1914 tensor_wasm_jit::detector::DetectorVerdict::Offload
1915 )
1916 });
1917 if !any_offload {
1918 // No candidate — nothing to rewrite. Borrow the original.
1919 return Cow::Borrowed(wasm);
1920 }
1921
1922 // Rewrite. Thread the spawning tenant so the cache pre-population is
1923 // keyed under the tenant the runtime dispatch looks up against
1924 // (cache keys are tenant-scoped), and the resolved detector so the
1925 // rewrite swaps exactly the functions the consultation flagged.
1926 let opts = RewriteOptions {
1927 tenant_id: cfg.tenant_id,
1928 detector,
1929 ..RewriteOptions::default()
1930 };
1931 match rewrite_wasm(wasm, &opts, cache) {
1932 Ok(outcome) if !outcome.offloaded_functions.is_empty() => {
1933 info!(
1934 target: "tensor_wasm_exec::executor",
1935 tenant = %cfg.tenant_id,
1936 offloaded = outcome.offloaded_functions.len(),
1937 total_defined = outcome.total_defined_functions,
1938 "auto_offload rewrite applied; instantiating trampoline-augmented module",
1939 );
1940 Cow::Owned(outcome.rewritten_wasm)
1941 }
1942 Ok(_) => {
1943 // The detector flagged a candidate but the rewriter declined
1944 // every swap (e.g. unsupported signature / lowering refusal).
1945 // Nothing changed — borrow the original to skip a needless
1946 // recompile of identical-but-reencoded bytes.
1947 debug!(
1948 target: "tensor_wasm_exec::executor",
1949 tenant = %cfg.tenant_id,
1950 "auto_offload rewrite swapped no functions; using original module",
1951 );
1952 Cow::Borrowed(wasm)
1953 }
1954 Err(e) => {
1955 warn!(
1956 target: "tensor_wasm_exec::executor",
1957 tenant = %cfg.tenant_id,
1958 error = %e,
1959 "auto_offload rewrite failed; falling back to original module",
1960 );
1961 Cow::Borrowed(wasm)
1962 }
1963 }
1964 }
1965
1966 /// Compile + instantiate a Wasm module. Returns the assigned [`InstanceId`].
1967 ///
1968 /// # Deadline / ticker contract
1969 ///
1970 /// If a [`SpawnConfig::deadline`] is set — or if the implicit
1971 /// [`MAX_START_FN_DURATION`] cap would otherwise apply (which it
1972 /// always does, since every spawn runs `Instance::new_async`) —
1973 /// the engine's epoch ticker MUST be running. Without it the
1974 /// per-store epoch counter never advances, so neither the per-call
1975 /// deadline nor the start-function cap can fire, and a runaway
1976 /// guest would wedge the worker thread until it returned of its
1977 /// own accord. We refuse the spawn with
1978 /// [`ExecError::EpochTickerNotRunning`] instead of silently
1979 /// dropping the deadline contract; operators must call
1980 /// `engine.spawn_epoch_ticker()` (typically inside a Tokio runtime
1981 /// at startup) before serving traffic. The engine constructor
1982 /// auto-spawns the ticker when invoked from inside a runtime, so
1983 /// this only trips for sync-startup setups that forget the
1984 /// explicit call.
1985 #[instrument(skip(self, wasm), fields(tenant = %cfg.tenant_id, instance_id = tracing::field::Empty))]
1986 pub async fn spawn_instance(
1987 &self,
1988 cfg: SpawnConfig,
1989 wasm: &[u8],
1990 ) -> Result<InstanceId, ExecError> {
1991 // Auto-offload (opt-in via `EngineConfig::auto_offload`). When
1992 // enabled, consult the analyser and — if any function is flagged —
1993 // rewrite the module's offload-candidate bodies into JIT-dispatch
1994 // trampolines, then instantiate the rewritten module instead of the
1995 // original. Any failure (analysis error, rewrite error, no JIT cache
1996 // to link the trampoline imports, or a rewrite that swaps nothing)
1997 // falls back to the original bytes — enabling the flag can never
1998 // fail a spawn that would otherwise have succeeded. Returns a `Cow`
1999 // so the disabled / fallback paths borrow the caller's slice with
2000 // zero copies.
2001 let effective_wasm = self.maybe_rewrite_for_offload(&cfg, wasm);
2002 // Refactored to share the detached compile+instantiate path
2003 // (`build_pooled_instance`) with [`InstancePool`]. The semantics
2004 // are byte-for-byte preserved: admission control runs first
2005 // (with rollback on failure), then compile / instantiate /
2006 // register. The split lets the pool reuse the heavy work
2007 // without the registry insert when holding warm instances in a
2008 // channel.
2009 let (inst, _module, _module_hash) = self
2010 .build_pooled_instance(&cfg, effective_wasm.as_ref())
2011 .await?;
2012 // `build_pooled_instance` returns with the slot committed (charged),
2013 // so a failure between here and the registry insert must release
2014 // the slot explicitly. Wrap it in a `defer`-style guard so any
2015 // `?` from `register_pooled_instance` does not leak the count.
2016 // `build_pooled_instance` committed BOTH the engine-wide and (when a
2017 // per-tenant cap is configured) the per-tenant count, so this
2018 // re-protect guard must roll back both on a failed register. Use the
2019 // tenant-aware constructor only when the per-tenant cap is active so
2020 // we don't decrement a per-tenant entry that was never charged.
2021 let slot_guard = if self.engine.config().max_instances_per_tenant.is_some() {
2022 InstanceSlotGuard::with_tenant(
2023 self.instance_count.clone(),
2024 self.tenant_counts.clone(),
2025 cfg.tenant_id,
2026 )
2027 } else {
2028 InstanceSlotGuard::new(self.instance_count.clone())
2029 };
2030 let id = self.register_pooled_instance(inst)?;
2031 // Successful register: defuse the rollback so the slot stays
2032 // charged until `terminate`. (Without this defuse the guard's
2033 // Drop would decrement the count we just committed.)
2034 slot_guard.commit();
2035 tracing::Span::current().record("instance_id", tracing::field::display(id));
2036 info!(target: "tensor_wasm_exec::executor", tenant = %cfg.tenant_id, instance = %id, "instance spawned");
2037 Ok(id)
2038 }
2039
2040 /// Invoke `export` with no arguments and no return value.
2041 ///
2042 /// This is the minimal signature needed for the 100-instance integration
2043 /// test. Richer signatures arrive in S17 (HTTP API) and S18 (CLI).
2044 ///
2045 /// # Concurrency note
2046 ///
2047 /// The per-instance mutex is held across the inner `call_async` await
2048 /// point. Concurrent calls into the **same instance** therefore serialise
2049 /// — this matches wasmtime's `Store`-is-not-`Sync` contract. Concurrent
2050 /// calls into **different instances** run in parallel. If you need
2051 /// pipelined invocation on a single instance, spawn additional instances
2052 /// over the same module bytes (the executor's module cache makes that
2053 /// nearly free).
2054 ///
2055 /// If the executor's engine has not had `spawn_epoch_ticker` called and
2056 /// the instance was spawned with a deadline, this call will run until
2057 /// the wasm returns of its own accord — the deadline cannot fire without
2058 /// the ticker. A warning is logged the first time this combination is
2059 /// observed per call.
2060 #[deprecated(
2061 since = "0.3.7",
2062 note = "use `call_export_with_args` with an empty `&[]` for the same semantics; \
2063 v0.4 removes this shim. See `docs/MIGRATING-FROM-WASMTIME-WASMER.md` § \"Typed exports\"."
2064 )]
2065 #[instrument(skip(self), fields(instance = %id, export = %export))]
2066 pub async fn call_export(&self, id: InstanceId, export: &str) -> Result<(), ExecError> {
2067 // Back-compat wrapper: most callers (the bench loop, the executor's
2068 // own tests, the orphan-cleanup integration test) only need the
2069 // `() -> ()` signature and explicitly assert success via `.unwrap()`
2070 // or `?`. Threading the new `Result<serde_json::Value, _>` shape
2071 // through every call site would be churn for no behavioural gain;
2072 // instead we keep the unit-typed surface here and delegate to the
2073 // shared implementation. The result value is discarded — callers
2074 // wanting it should use [`Self::call_export_with_args`] directly.
2075 self.call_export_with_args(id, export, &[])
2076 .await
2077 .map(|_| ())
2078 }
2079
2080 /// Invoke `export` with the supplied `args` (which may be empty) and
2081 /// return the export's result list, serialised as a JSON array.
2082 ///
2083 /// This is the general entry point for guest-export invocation; the
2084 /// `()`-shaped [`Self::call_export`] is a thin wrapper that discards
2085 /// the result. The choice of `serde_json::Value` for the return type
2086 /// keeps the executor's public API free of any tensor-wasm-api types
2087 /// while still giving the HTTP transport a structured payload to
2088 /// forward verbatim.
2089 ///
2090 /// When `args` is empty the implementation uses the typed
2091 /// `func.typed::<(), ()>()` fast path, matching the historical
2092 /// behaviour and keeping every existing `() -> ()` export call
2093 /// branch-for-branch identical. With a non-empty `args` slice the
2094 /// dynamic `func.call_async(&[Val], &mut [Val])` path runs instead;
2095 /// the result slice is sized from the export's declared result arity
2096 /// at runtime, so an export returning `(i32, i32)` produces a JSON
2097 /// array with two numbers.
2098 #[instrument(skip(self, args), fields(instance = %id, export = %export, args_len = args.len()))]
2099 pub async fn call_export_with_args(
2100 &self,
2101 id: InstanceId,
2102 export: &str,
2103 args: &[WasmArg],
2104 ) -> Result<serde_json::Value, ExecError> {
2105 let handle = self
2106 .instances
2107 .get(&id)
2108 .ok_or(ExecError::NotFound(id))?
2109 .value()
2110 .clone();
2111 let mut guard = handle.lock().await;
2112 // Re-arm the deadline at the start of each call.
2113 //
2114 // The fields on `InstanceState` work in concert: `deadline_duration`
2115 // is the configured per-call budget (immutable for the life of the
2116 // instance), and `deadline` is the absolute `Instant` we expect to
2117 // not cross. At spawn time we seeded `deadline = now + d`, but if a
2118 // caller invokes `call_export` twice with delay in between, the
2119 // second call would inherit an already-elapsed deadline — and the
2120 // wasmtime epoch counter set at spawn would already be consumed.
2121 // That used to surface as `Timeout { elapsed_ms: 0, deadline_ms: 0 }`
2122 // because the wasmtime trap fired before any real work happened and
2123 // the legacy `deadline_at.saturating_duration_since(started_at)`
2124 // returned zero. Re-arming here gives every call its own honest
2125 // window (and honest numbers if it does time out).
2126 let call_start = Instant::now();
2127 let configured_deadline = guard.store.data().deadline_duration;
2128 if let Some(d) = configured_deadline {
2129 let new_deadline = call_start + d;
2130 guard.store.data_mut().deadline = Some(new_deadline);
2131 // Re-arm the HARD-deadline instant the cooperative epoch callback
2132 // traps at, in lockstep with `deadline`, so THIS call's window
2133 // (not a previous call's, nor the spawn-time start-phase budget)
2134 // governs when a runaway guest traps. The callback installed at
2135 // spawn (`arm_cooperative_epoch`) persists across calls; only the
2136 // instant it consults and the relative tick count are re-armed.
2137 guard.store.data_mut().hard_deadline = Some(new_deadline);
2138 let tick = self.engine.config().epoch_tick;
2139 // Re-arm at the cooperative cadence (never beyond this call's
2140 // deadline window). The cooperative epoch callback installed at
2141 // spawn yields `Pending` at each trip and traps only once
2142 // `hard_deadline` (re-armed to `new_deadline` above) elapses — so
2143 // the relative ticks here govern only the yield cadence, while the
2144 // wall-clock `hard_deadline` governs the trap. Arming at the
2145 // cooperative cadence keeps a compute-bound guest yielding (hence
2146 // cancellable on future-drop) throughout the call.
2147 let call_ticks = duration_to_epoch_ticks(d, tick).min(COOPERATIVE_YIELD_TICKS);
2148 guard.store.set_epoch_deadline(call_ticks);
2149 }
2150 // Re-arm the cooperative-scheduler context's wall-clock origin
2151 // in lockstep with the deadline above. Without this, a guest
2152 // that calls `wasi:scheduler/host.yield()` on a back-to-back
2153 // call would see the elapsed time from the *previous* call
2154 // charged against its budget, and the first yield would
2155 // immediately return STOP. The re-arm is a no-op for spawns
2156 // without a configured deadline (the unbounded context
2157 // ignores `started_at`).
2158 guard.store.data_mut().rearm_scheduler();
2159 let deadline_at = guard.store.data().deadline;
2160 let configured_deadline_ms = configured_deadline
2161 .map(|d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
2162 .unwrap_or(0);
2163 let wasmtime_instance = *guard.wasmtime_instance();
2164 let func = wasmtime_instance
2165 .get_func(&mut guard.store, export)
2166 .ok_or_else(|| ExecError::MissingExport(export.to_string()))?;
2167
2168 // Branch on the export's full signature. The typed fast path is
2169 // reserved for genuine `() -> ()` exports — no args AND no results
2170 // — so we don't disturb the behaviour every existing void-export
2171 // test relies on (and so the dynamic-call overhead stays off the
2172 // bench path). A no-arg export that nonetheless RETURNS a value
2173 // (e.g. `tick: () -> i32`) must NOT take this branch: a
2174 // `typed::<(), ()>` view of it is rejected by wasmtime with
2175 // "type mismatch with results: expected 0 types, found N". Every
2176 // other shape — args present, results present, or both — takes the
2177 // dynamic `Func::call_async` path with `&[Val]` IO buffers (params
2178 // may legitimately be empty); the result vec is pre-sized to the
2179 // export's declared result arity so wasmtime can write into it.
2180 let func_ty = func.ty(&guard.store);
2181 let call_outcome = if args.is_empty() && func_ty.results().len() == 0 {
2182 match func.typed::<(), ()>(&guard.store) {
2183 Ok(typed) => typed
2184 .call_async(&mut guard.store, ())
2185 .await
2186 .map(|()| serde_json::Value::Array(Vec::new())),
2187 Err(e) => Err(e),
2188 }
2189 } else {
2190 let params: Vec<Val> = args.iter().copied().map(WasmArg::into_val).collect();
2191 // `Val::I32(0)` is just a placeholder — wasmtime overwrites
2192 // every slot before returning. The element count must match
2193 // the export's declared result arity exactly or wasmtime
2194 // returns an error.
2195 let mut results: Vec<Val> = vec![Val::I32(0); func_ty.results().len()];
2196 match func
2197 .call_async(&mut guard.store, ¶ms, &mut results)
2198 .await
2199 {
2200 Ok(()) => {
2201 let json: Vec<serde_json::Value> = results.iter().map(val_to_json).collect();
2202 Ok(serde_json::Value::Array(json))
2203 }
2204 Err(e) => Err(e),
2205 }
2206 };
2207
2208 match call_outcome {
2209 Ok(value) => Ok(value),
2210 Err(err) => {
2211 // If we had a deadline AND the wall clock has tripped past
2212 // it, classify the failure as Timeout with real numbers.
2213 // Otherwise propagate the raw wasmtime error.
2214 let elapsed = call_start.elapsed();
2215 let past_deadline = deadline_at.map(|d| Instant::now() >= d).unwrap_or(false);
2216 if past_deadline {
2217 Err(ExecError::Timeout(TimeoutContext {
2218 id,
2219 elapsed_ms: u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX),
2220 deadline_ms: configured_deadline_ms,
2221 }))
2222 } else {
2223 Err(ExecError::Wasmtime(err))
2224 }
2225 }
2226 }
2227 }
2228
2229 /// Drop the instance, releasing its resources.
2230 #[instrument(skip(self), fields(instance = %id))]
2231 pub async fn terminate(&self, id: InstanceId) -> Result<(), ExecError> {
2232 match self.instances.remove(&id) {
2233 Some((_, handle)) => {
2234 // Release the admission slot reserved at spawn time
2235 // (exec S-10). The decrement only runs on successful
2236 // removal — a `NotFound` terminate must not free a
2237 // slot it never charged, or a tenant could double-
2238 // terminate to inflate their effective cap.
2239 self.instance_count.fetch_sub(1, Ordering::AcqRel);
2240 // Mirror the engine-wide release on the per-tenant count
2241 // (fairness cap). Only taken when a per-tenant cap is
2242 // configured — otherwise `tenant_counts` is never populated
2243 // and we skip the per-instance lock entirely, preserving the
2244 // historical lock-free terminate path. When the cap IS active,
2245 // read the owning tenant off the now-removed instance and
2246 // decrement its slot, pruning the map entry at zero. Locking
2247 // the per-instance mutex to read the tenant is uncontended on
2248 // the common path (the registry held the only handle); a
2249 // racing in-flight call would briefly hold it, which is
2250 // correct — we want the post-call tenant.
2251 if self.engine.config().max_instances_per_tenant.is_some() {
2252 let tenant = handle.lock().await.tenant_id();
2253 decrement_tenant_count(&self.tenant_counts, tenant);
2254 }
2255 if let Some(m) = &self.metrics {
2256 m.instance_terminations_total().inc();
2257 m.active_instances().dec();
2258 }
2259 debug!(target: "tensor_wasm_exec::executor", instance = %id, "instance terminated");
2260 Ok(())
2261 }
2262 None => Err(ExecError::NotFound(id)),
2263 }
2264 }
2265
2266 /// Call an export, then unconditionally terminate the instance —
2267 /// even if the returned future is dropped mid-await (api S-20 +
2268 /// orphan-instance cleanup).
2269 ///
2270 /// The previous flow (`call_export` + explicit `terminate` from the
2271 /// caller) leaks the instance into `instances` when the caller's
2272 /// future is dropped by an outer cancellation (e.g. tower's
2273 /// `TimeoutLayer` firing). The leaked entry holds the wasmtime
2274 /// `Store` and counts against `max_instances`, but is unreachable
2275 /// by id (the caller lost the handle). This wrapper installs an
2276 /// `AutoTerminateGuard`: on the normal exit paths it disarms the
2277 /// guard and calls `terminate` via the async API; on a Future-drop
2278 /// the guard's `Drop` synchronously removes the registry entry
2279 /// (and frees the admission slot) so the leak window closes at
2280 /// the await boundary.
2281 ///
2282 /// Known limitation: this wrapper cannot stop CPU work that is
2283 /// already running inside Wasmtime's blocking compile path
2284 /// (`Instance::new_async` invokes Cranelift, which does not expose
2285 /// a cancellation hook). For now, the cap on
2286 /// [`crate::engine::EngineConfig::max_module_cache_entries`]
2287 /// limits the worst case to one compile per unique module.
2288 /// Per-store epoch cancellation interrupts wasm execution at the
2289 /// next epoch tick, which is what closes the actual run-time
2290 /// window.
2291 #[deprecated(
2292 since = "0.3.7",
2293 note = "use `call_export_with_args_then_terminate` with an empty `&[]` for the same semantics; \
2294 v0.4 removes this shim. See `docs/MIGRATING-FROM-WASMTIME-WASMER.md` § \"Typed exports\"."
2295 )]
2296 pub async fn call_export_then_terminate(
2297 &self,
2298 id: InstanceId,
2299 export: &str,
2300 ) -> Result<(), ExecError> {
2301 // Unit-typed back-compat surface, mirrors [`Self::call_export`].
2302 self.call_export_with_args_then_terminate(id, export, &[])
2303 .await
2304 .map(|_| ())
2305 }
2306
2307 /// High-level "invoke" entry point: spawn (or draw from the warm pool),
2308 /// call the export, return the result, and clean up the instance —
2309 /// all in a single async call.
2310 ///
2311 /// When an [`InstancePool`] is attached via [`Self::with_instance_pool`],
2312 /// this routes through [`InstancePool::acquire`] / [`InstancePool::release`]
2313 /// so the per-(tenant, module-hash) warm channel absorbs the
2314 /// compile+instantiate cost on the hot path (T37). Without a pool
2315 /// attached, behaviour is identical to
2316 /// [`Self::spawn_instance`] + [`Self::call_export_with_args_then_terminate`]
2317 /// — every existing caller of that pair sees no semantic change.
2318 ///
2319 /// Mirrors the routes layer's invoke / invoke-stream / invoke-async
2320 /// pattern: a single `(wasm, cfg, export, args)` shot, no caller-side
2321 /// id juggling. T34's streaming and T36's back-pressure /
2322 /// deadline-near semantics propagate unchanged — the
2323 /// [`SpawnConfig::streaming`] context flows through `acquire`, and
2324 /// per-call deadline re-arming happens inside `call_export_with_args`
2325 /// regardless of which path produced the instance.
2326 ///
2327 /// Pool-side invariant: streaming spawns are NEVER recycled, even
2328 /// when a pool is attached — the streaming context is one-shot
2329 /// (the gateway's channel receiver is drained for the duration of
2330 /// the SSE response and then dropped). [`InstancePool::release`]
2331 /// drops the instance instead of returning it to the warm channel
2332 /// when `SpawnConfig::streaming.is_some()`.
2333 pub async fn invoke(
2334 &self,
2335 cfg: SpawnConfig,
2336 wasm: &[u8],
2337 export: &str,
2338 args: &[WasmArg],
2339 ) -> Result<serde_json::Value, ExecError> {
2340 if let Some(pool) = self.pool.clone() {
2341 // Pool path (T37): acquire a warm (or freshly-spawned)
2342 // instance, call, then release. Release routes through the
2343 // pool's reset path — on success a fresh replacement
2344 // instance returns to the warm channel; on streaming opt-in
2345 // or reset failure the slot is released without
2346 // replenishment.
2347 let pooled = pool.acquire(self, wasm, cfg.clone()).await?;
2348 // Capture the id by value so the subsequent `release`
2349 // (which consumes `pooled`) and the call below see the
2350 // same handle.
2351 let id = pooled.id();
2352 let result = self.call_export_with_args(id, export, args).await;
2353 // Release regardless of call outcome — a guest trap should
2354 // not poison the warm pool (the pool's reset re-instantiates
2355 // from the cached module, so post-trap state is irrelevant).
2356 // Streaming spawns are dropped (never recycled) inside
2357 // `release`.
2358 pool.release(self, pooled, &cfg).await;
2359 result
2360 } else {
2361 // No pool attached — preserve the historical spawn + call
2362 // + terminate flow verbatim, including the auto-terminate
2363 // drop guard (api S-20).
2364 let id = self.spawn_instance(cfg, wasm).await?;
2365 self.call_export_with_args_then_terminate(id, export, args)
2366 .await
2367 }
2368 }
2369
2370 /// Argument-aware sibling of [`Self::call_export_then_terminate`].
2371 ///
2372 /// Identical lifecycle / drop-guard semantics — auto-terminates on
2373 /// success, failure, and Future-drop — but routes through
2374 /// [`Self::call_export_with_args`] so callers receive the export's
2375 /// result list as a JSON array.
2376 pub async fn call_export_with_args_then_terminate(
2377 &self,
2378 id: InstanceId,
2379 export: &str,
2380 args: &[WasmArg],
2381 ) -> Result<serde_json::Value, ExecError> {
2382 // exec H2: capture the per-tenant rollback handle BEFORE the wrapped
2383 // call so the guard's `Drop` (which cannot await) can mirror the
2384 // per-tenant decrement that the async `terminate` path performs by
2385 // reading the tenant off the instance. Only populated when a
2386 // per-tenant cap is configured — otherwise `tenant_counts` is never
2387 // charged and we leave it `None` so the guard skips the per-tenant
2388 // decrement entirely (matching `terminate`, which also skips the
2389 // per-instance lock in that case). We read the owning tenant off the
2390 // already-registered instance here (we *can* await at construction
2391 // time); the cancellation race the guard defends against only opens
2392 // once we enter the `call_export_with_args` await below.
2393 let tenant_rollback = if self.engine.config().max_instances_per_tenant.is_some() {
2394 // Read the owning tenant off the already-registered instance under
2395 // its per-instance lock (uncontended on the common path). If the
2396 // id is unknown the call below will return `NotFound` and there is
2397 // nothing to roll back, so a `None` here is harmless.
2398 match self.instances.get(&id).map(|h| h.value().clone()) {
2399 Some(handle) => {
2400 let tenant = handle.lock().await.tenant_id();
2401 Some((Arc::clone(&self.tenant_counts), tenant))
2402 }
2403 None => None,
2404 }
2405 } else {
2406 None
2407 };
2408 let guard = AutoTerminateGuard {
2409 instances: Arc::clone(&self.instances),
2410 instance_count: Arc::clone(&self.instance_count),
2411 metrics: self.metrics.clone(),
2412 id,
2413 tenant_rollback,
2414 // Re-arm on construction; only the success/error path below
2415 // is allowed to disarm.
2416 armed: true,
2417 };
2418 let result = self.call_export_with_args(id, export, args).await;
2419 // Disarm BEFORE the async terminate so a panic in `terminate`
2420 // does not double-fire. Both paths still remove the instance
2421 // exactly once: the guard via the sync DashMap::remove if it
2422 // is armed, the explicit terminate via the same DashMap::remove
2423 // when the guard is disarmed.
2424 let mut guard = guard;
2425 guard.armed = false;
2426 let _ = self.terminate(id).await; // ignore NotFound on success-after-cancel races
2427 result
2428 }
2429}
2430
2431/// RAII drop-guard that synchronously removes an instance from the
2432/// registry if it is still armed when dropped. See
2433/// [`TensorWasmExecutor::call_export_then_terminate`] for the threat
2434/// model that motivates the design.
2435///
2436/// Holds `Arc` clones of the registry and the admission counter so
2437/// the guard can run without borrowing the executor — which matters
2438/// because the original `&self` reference is consumed by the
2439/// `call_export` await this guard wraps.
2440struct AutoTerminateGuard {
2441 instances: Arc<DashMap<InstanceId, Arc<Mutex<TensorWasmInstance>>>>,
2442 instance_count: Arc<AtomicUsize>,
2443 metrics: Option<TensorWasmMetrics>,
2444 id: InstanceId,
2445 /// Per-tenant rollback handle (exec H2). `Some((tenant_counts, tenant))`
2446 /// only when a per-tenant fairness cap is configured — in that case the
2447 /// async `terminate` path decrements the per-tenant count too, so a
2448 /// cancelled-future drop MUST mirror it or the tenant leaks a slot
2449 /// permanently (lockout once the cap is reached). The owning `TenantId`
2450 /// is captured up front (before the wrapped call) because `Drop` cannot
2451 /// await to read it off the instance the way `terminate` does. `None`
2452 /// when the cap is disabled — `tenant_counts` is never populated then and
2453 /// the guard skips the per-tenant decrement entirely.
2454 tenant_rollback: Option<(Arc<DashMap<TenantId, usize>>, TenantId)>,
2455 armed: bool,
2456}
2457
2458impl Drop for AutoTerminateGuard {
2459 fn drop(&mut self) {
2460 if !self.armed {
2461 return;
2462 }
2463 // Sync remove: we cannot await in Drop. The async `terminate`
2464 // method does exactly the same work plus a debug! log, so this
2465 // is a faithful sync mirror.
2466 if self.instances.remove(&self.id).is_some() {
2467 self.instance_count.fetch_sub(1, Ordering::AcqRel);
2468 // exec H2: mirror the per-tenant decrement the async `terminate`
2469 // path performs. Without this a cancelled/dropped invoke future
2470 // released the engine-wide slot but leaked the per-tenant slot,
2471 // so a tenant that accumulated cancelled invokes would hit its
2472 // `max_instances_per_tenant` cap with zero live instances and be
2473 // locked out permanently. Same saturating / prune-on-zero
2474 // semantics as `terminate` (both go through
2475 // `decrement_tenant_count`).
2476 if let Some((tenant_counts, tenant)) = &self.tenant_rollback {
2477 decrement_tenant_count(tenant_counts, *tenant);
2478 }
2479 if let Some(m) = &self.metrics {
2480 m.instance_terminations_total().inc();
2481 m.active_instances().dec();
2482 }
2483 tracing::warn!(
2484 target: "tensor_wasm_exec::executor",
2485 instance = %self.id,
2486 "instance auto-terminated by drop-guard (handler future cancelled \
2487 mid-call_export; see api S-20)"
2488 );
2489 }
2490 }
2491}
2492
2493#[cfg(test)]
2494mod tests {
2495 use super::*;
2496
2497 fn trivial_wasm() -> Vec<u8> {
2498 wat::parse_str(r#"(module (func (export "noop")))"#).unwrap()
2499 }
2500
2501 #[tokio::test]
2502 async fn spawn_then_terminate() {
2503 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2504 let exec = TensorWasmExecutor::new(engine);
2505 let id = exec
2506 .spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
2507 .await
2508 .unwrap();
2509 assert_eq!(exec.live_count(), 1);
2510 exec.call_export_with_args(id, "noop", &[]).await.unwrap();
2511 exec.terminate(id).await.unwrap();
2512 assert_eq!(exec.live_count(), 0);
2513 }
2514
2515 #[tokio::test]
2516 async fn missing_export() {
2517 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2518 let exec = TensorWasmExecutor::new(engine);
2519 let id = exec
2520 .spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
2521 .await
2522 .unwrap();
2523 let err = exec
2524 .call_export_with_args(id, "does_not_exist", &[])
2525 .await
2526 .unwrap_err();
2527 assert!(matches!(err, ExecError::MissingExport(_)));
2528 }
2529
2530 /// Back-compat smoke test: the `#[deprecated]` `call_export` shim must
2531 /// still drive a guest to completion until v0.4 removes it. Tagged
2532 /// `#[allow(deprecated)]` because exercising the deprecated path is
2533 /// the *point* of this test — without it the test would emit the
2534 /// migration warning we ship to external callers.
2535 ///
2536 /// The `let _f = TensorWasmExecutor::call_export;` line at the bottom
2537 /// is a static proof-of-presence: if the `#[deprecated]` attribute
2538 /// were ever dropped again (as it was in merge `66af7db`), the
2539 /// `#[allow(deprecated)]` would become an "unnecessary attribute"
2540 /// warning under `-D unused`, signalling regression. The same logic
2541 /// applies to `call_export_then_terminate`.
2542 #[tokio::test]
2543 #[allow(deprecated)]
2544 async fn legacy_call_export_still_works() {
2545 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2546 let exec = TensorWasmExecutor::new(engine);
2547 let id = exec
2548 .spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
2549 .await
2550 .unwrap();
2551 exec.call_export(id, "noop").await.unwrap();
2552 exec.terminate(id).await.unwrap();
2553
2554 // Static proof-of-presence: both deprecated shims must be
2555 // resolvable as `fn` items. If a future refactor removes the
2556 // `#[deprecated]` attribute, the `#[allow(deprecated)]` above
2557 // is flagged as unused and CI breaks — exactly the regression
2558 // guard we want.
2559 let _f = TensorWasmExecutor::call_export;
2560 let _g = TensorWasmExecutor::call_export_then_terminate;
2561 }
2562
2563 #[tokio::test]
2564 async fn terminate_unknown() {
2565 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2566 let exec = TensorWasmExecutor::new(engine);
2567 let err = exec.terminate(InstanceId(999)).await.unwrap_err();
2568 assert!(matches!(err, ExecError::NotFound(_)));
2569 }
2570
2571 #[tokio::test]
2572 async fn metrics_increment_on_spawn_and_terminate() {
2573 use tensor_wasm_core::metrics::TensorWasmMetrics;
2574 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2575 let metrics = TensorWasmMetrics::new();
2576 let exec = TensorWasmExecutor::with_metrics(engine, metrics.clone());
2577 let id = exec
2578 .spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &trivial_wasm())
2579 .await
2580 .unwrap();
2581 let text = metrics.encode_text();
2582 assert!(
2583 text.contains("tensor_wasm_instance_spawns_total 1"),
2584 "got:\n{text}"
2585 );
2586 assert!(
2587 text.contains("tensor_wasm_active_instances 1"),
2588 "got:\n{text}"
2589 );
2590 exec.terminate(id).await.unwrap();
2591 let text = metrics.encode_text();
2592 assert!(
2593 text.contains("tensor_wasm_instance_terminations_total 1"),
2594 "got:\n{text}"
2595 );
2596 assert!(
2597 text.contains("tensor_wasm_active_instances 0"),
2598 "got:\n{text}"
2599 );
2600 }
2601
2602 #[test]
2603 fn exec_error_converts_to_tensor_wasm_error() {
2604 use tensor_wasm_core::error::TensorWasmError;
2605 let e = ExecError::NotFound(InstanceId(99));
2606 let b: TensorWasmError = e.into();
2607 assert!(matches!(b, TensorWasmError::Serialization(_)));
2608 // `TensorWasmError`'s Display is deliberately sanitised and does NOT
2609 // echo the inner string (it can carry host paths / pointers from
2610 // third-party crates). The raw context is reachable via `inner()`.
2611 assert!(
2612 b.inner().unwrap_or("").contains("instance not found"),
2613 "inner: {:?}",
2614 b.inner()
2615 );
2616
2617 let e = ExecError::Timeout(TimeoutContext {
2618 id: InstanceId(1),
2619 elapsed_ms: 150,
2620 deadline_ms: 100,
2621 });
2622 let b: TensorWasmError = e.into();
2623 match b {
2624 TensorWasmError::KernelTimeout {
2625 elapsed_ms,
2626 deadline_ms,
2627 } => {
2628 assert_eq!(elapsed_ms, 150);
2629 assert_eq!(deadline_ms, 100);
2630 }
2631 other => panic!("expected KernelTimeout, got {other:?}"),
2632 }
2633 }
2634
2635 #[tokio::test]
2636 async fn module_cache_reuses_compilation() {
2637 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2638 let exec = TensorWasmExecutor::new(engine);
2639 let wasm = trivial_wasm();
2640 let _id1 = exec
2641 .spawn_instance(SpawnConfig::for_tenant(TenantId(1)), &wasm)
2642 .await
2643 .unwrap();
2644 let _id2 = exec
2645 .spawn_instance(SpawnConfig::for_tenant(TenantId(2)), &wasm)
2646 .await
2647 .unwrap();
2648 // Both spawns hit the same wasm bytes — the cache should hold one entry.
2649 assert_eq!(exec.cached_module_count(), 1);
2650 }
2651
2652 #[test]
2653 fn classify_instantiation_error_refines_memory_size_failure() {
2654 // A wasmtime error whose chain carries the pooling allocator's
2655 // memory-size signature must be re-tagged as ModuleMemoryTooLarge
2656 // (MED finding) so callers see a typed MemoryExhausted.
2657 let err = wasmtime::Error::msg("memory index 0 has a minimum that exceeds the limit");
2658 let classified = classify_instantiation_error(err, 64 * 1024 * 1024);
2659 match classified {
2660 ExecError::ModuleMemoryTooLarge {
2661 requested_bytes,
2662 limit_bytes,
2663 } => {
2664 assert_eq!(limit_bytes, 64 * 1024 * 1024);
2665 assert_eq!(requested_bytes, 64 * 1024 * 1024);
2666 }
2667 other => panic!("expected ModuleMemoryTooLarge, got {other:?}"),
2668 }
2669 }
2670
2671 #[test]
2672 fn classify_instantiation_error_passes_through_unrelated() {
2673 // An unrelated error (e.g. an import-link failure) must NOT be
2674 // misclassified — it flows through verbatim as ExecError::Wasmtime.
2675 let err = wasmtime::Error::msg("unknown import: `env::foo` has not been defined");
2676 let classified = classify_instantiation_error(err, 64 * 1024 * 1024);
2677 assert!(
2678 matches!(classified, ExecError::Wasmtime(_)),
2679 "unrelated errors must not be re-tagged, got {classified:?}",
2680 );
2681 }
2682
2683 #[test]
2684 fn resource_limiter_allows_under_cap() {
2685 let mut lim = TensorWasmResourceLimiter::new(2 * 1024 * 1024);
2686 assert!(lim.memory_growing(0, 1024 * 1024, None).unwrap());
2687 }
2688
2689 #[test]
2690 fn resource_limiter_rejects_over_cap() {
2691 let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
2692 assert!(!lim.memory_growing(0, 2 * 1024 * 1024, None).unwrap());
2693 }
2694
2695 #[test]
2696 fn resource_limiter_respects_module_maximum() {
2697 let mut lim = TensorWasmResourceLimiter::new(usize::MAX);
2698 // Even if the engine cap is unbounded, the module's declared max wins.
2699 assert!(!lim.memory_growing(0, 4096, Some(2048)).unwrap());
2700 }
2701
2702 #[test]
2703 fn resource_limiter_rejects_huge_table_growth() {
2704 // 1 MiB engine cap → at 16 B/entry that's ~65k table entries max.
2705 // u32::MAX (~4.3 billion entries × 16 B = ~64 GiB) must be denied.
2706 let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
2707 assert!(!lim.table_growing(0, usize::MAX, None).unwrap());
2708 }
2709
2710 #[test]
2711 fn resource_limiter_allows_modest_table_growth() {
2712 // 1 MiB engine cap → ~65k entries should still fit.
2713 let mut lim = TensorWasmResourceLimiter::new(1024 * 1024);
2714 assert!(lim.table_growing(0, 1024, None).unwrap());
2715 }
2716
2717 #[test]
2718 fn resource_limiter_table_respects_module_maximum() {
2719 // Even with an unbounded engine cap, the module's declared table max wins.
2720 let mut lim = TensorWasmResourceLimiter::new(usize::MAX);
2721 assert!(!lim.table_growing(0, 4096, Some(2048)).unwrap());
2722 }
2723
2724 /// exec H2 regression: a cancelled/dropped invoke future must release the
2725 /// per-tenant slot, not just the engine-wide one. Before the fix,
2726 /// `AutoTerminateGuard::drop` decremented `instance_count` but left
2727 /// `tenant_counts` untouched, so a tenant that accumulated cancelled
2728 /// invokes would hit `max_instances_per_tenant` with zero live instances
2729 /// and be locked out permanently.
2730 ///
2731 /// We exercise the guard directly (the struct + its `Drop` are the unit
2732 /// the fix lives in) rather than racing a real future-cancellation, which
2733 /// is non-deterministic: we charge a per-tenant slot exactly the way
2734 /// `spawn_instance` does, build an *armed* guard with the per-tenant
2735 /// rollback populated the way `call_export_with_args_then_terminate` now
2736 /// does, drop it, and assert BOTH counters returned to zero.
2737 #[tokio::test]
2738 async fn dropped_invoke_releases_per_tenant_slot() {
2739 use crate::engine::EngineConfig;
2740 let cfg = EngineConfig {
2741 max_instances_per_tenant: Some(1),
2742 ..EngineConfig::default()
2743 };
2744 let engine = Arc::new(TensorWasmEngine::with_config(cfg).unwrap());
2745 let exec = TensorWasmExecutor::new(engine);
2746 let tenant = TenantId(7);
2747
2748 // Spawn one instance: charges engine-wide + per-tenant to 1, and (with
2749 // the cap == 1) the tenant is now at its ceiling.
2750 let id = exec
2751 .spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
2752 .await
2753 .unwrap();
2754 assert_eq!(exec.tenant_instance_count(tenant), 1);
2755 assert_eq!(exec.instances_len(), 1);
2756
2757 // Build the guard exactly as `call_export_with_args_then_terminate`
2758 // does: armed, with the per-tenant rollback captured up front. Then
2759 // drop it to simulate the cancelled-future path (the guard stays
2760 // armed because the normal disarm-then-terminate sequence never ran).
2761 {
2762 let _guard = AutoTerminateGuard {
2763 instances: Arc::clone(&exec.instances),
2764 instance_count: Arc::clone(&exec.instance_count),
2765 metrics: exec.metrics.clone(),
2766 id,
2767 tenant_rollback: Some((Arc::clone(&exec.tenant_counts), tenant)),
2768 armed: true,
2769 };
2770 // dropped here
2771 }
2772
2773 // Both counters must be back to zero — the tenant is no longer locked
2774 // out and the engine-wide slot was freed.
2775 assert_eq!(
2776 exec.tenant_instance_count(tenant),
2777 0,
2778 "per-tenant slot leaked on dropped invoke (exec H2)"
2779 );
2780 assert_eq!(exec.instances_len(), 0, "engine-wide slot leaked");
2781 assert_eq!(exec.live_count(), 0, "registry entry leaked");
2782
2783 // The tenant must be able to spawn again (proves no lockout).
2784 let id2 = exec
2785 .spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
2786 .await
2787 .expect("tenant should not be locked out after a cancelled invoke");
2788 exec.terminate(id2).await.unwrap();
2789 assert_eq!(exec.tenant_instance_count(tenant), 0);
2790 }
2791
2792 /// exec H2 companion: with NO per-tenant cap configured, the guard's
2793 /// `tenant_rollback` is `None` and the drop path must not touch
2794 /// `tenant_counts` (it is never populated) — only the engine-wide slot is
2795 /// released, matching the async `terminate` path which also skips the
2796 /// per-tenant decrement when the cap is unset.
2797 #[tokio::test]
2798 async fn dropped_invoke_no_tenant_cap_releases_engine_slot_only() {
2799 let engine = Arc::new(TensorWasmEngine::new().unwrap());
2800 let exec = TensorWasmExecutor::new(engine);
2801 let tenant = TenantId(3);
2802 let id = exec
2803 .spawn_instance(SpawnConfig::for_tenant(tenant), &trivial_wasm())
2804 .await
2805 .unwrap();
2806 assert_eq!(exec.instances_len(), 1);
2807 // No per-tenant cap => tenant_counts never charged.
2808 assert_eq!(exec.tenant_instance_count(tenant), 0);
2809 {
2810 let _guard = AutoTerminateGuard {
2811 instances: Arc::clone(&exec.instances),
2812 instance_count: Arc::clone(&exec.instance_count),
2813 metrics: exec.metrics.clone(),
2814 id,
2815 tenant_rollback: None,
2816 armed: true,
2817 };
2818 }
2819 assert_eq!(exec.instances_len(), 0, "engine-wide slot leaked");
2820 assert_eq!(exec.tenant_instance_count(tenant), 0);
2821 }
2822}