diff --git a/sgl-model-gateway/src/wasm/runtime.rs b/sgl-model-gateway/src/wasm/runtime.rs index 58c8ffc74d9b..65ae819c6843 100644 --- a/sgl-model-gateway/src/wasm/runtime.rs +++ b/sgl-model-gateway/src/wasm/runtime.rs @@ -15,7 +15,7 @@ use tokio::sync::oneshot; use tracing::{debug, error, info}; use wasmtime::{ component::{Component, Linker, ResourceTable}, - Config, Engine, Store, UpdateDeadline, + Config, Engine, Store, StoreLimitsBuilder, UpdateDeadline, }; use wasmtime_wasi::WasiCtx; @@ -342,14 +342,34 @@ impl WasmThreadPool { let mut linker = Linker::::new(engine); wasmtime_wasi::p2::add_to_linker_async(&mut linker)?; let mut builder = WasiCtx::builder(); + + // Create memory limits from config. + // Use the config helper to get total bytes, then safely convert to usize. + let memory_limit_bytes = + usize::try_from(config.get_total_memory_bytes()).map_err(|_| { + WasmError::from(WasmRuntimeError::CallFailed( + "Configured WASM memory limit exceeds addressable space on this platform." + .to_string(), + )) + })?; + let limits = StoreLimitsBuilder::new() + .memory_size(memory_limit_bytes) + .trap_on_grow_failure(true) // Trap instead of returning -1 for easier debugging + .build(); + let mut store = Store::new( engine, WasiState { ctx: builder.build(), table: ResourceTable::new(), + limits, }, ); + // Apply resource limits to the store. + // This enforces max_memory_pages by preventing memory.grow beyond the limit. + store.limiter(|state| &mut state.limits); + // Set epoch deadline for timeout enforcement. // The deadline is the number of epoch ticks before execution is interrupted. // With EPOCH_INTERVAL_MS=100ms and max_execution_time_ms=1000ms, deadline=10 epochs. diff --git a/sgl-model-gateway/src/wasm/types.rs b/sgl-model-gateway/src/wasm/types.rs index 261280d5b402..1e243311c775 100644 --- a/sgl-model-gateway/src/wasm/types.rs +++ b/sgl-model-gateway/src/wasm/types.rs @@ -3,7 +3,7 @@ //! Provides generic input/output types for WASM component execution //! based on attach points. -use wasmtime::component::ResourceTable; +use wasmtime::{component::ResourceTable, StoreLimits}; use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView}; use crate::wasm::{ @@ -89,6 +89,7 @@ impl WasmComponentOutput { pub struct WasiState { pub ctx: WasiCtx, pub table: ResourceTable, + pub limits: StoreLimits, } impl WasiView for WasiState {