diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 0aae7adcc..9bc796579 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -88,6 +88,7 @@ smg-auth = "1.0.0" wfaas = "1.0.0" data-connector = "1.0.0" smg-mcp = "1.0.0" +smg-wasm = "1.0.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } rustls-pemfile = "2.2" openssl = "0.10.73" diff --git a/sgl-model-gateway/src/middleware.rs b/sgl-model-gateway/src/middleware.rs index 224854264..ccade365f 100644 --- a/sgl-model-gateway/src/middleware.rs +++ b/sgl-model-gateway/src/middleware.rs @@ -38,7 +38,7 @@ use crate::{ module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, spec::{ apply_modify_action_to_headers, build_wasm_headers_from_axum_headers, - sgl::model_gateway::middleware_types::{ + smg::gateway::middleware_types::{ Action, Request as WasmRequest, Response as WasmResponse, }, }, diff --git a/sgl-model-gateway/src/wasm/README.md b/sgl-model-gateway/src/wasm/README.md deleted file mode 100644 index 7bd898fe4..000000000 --- a/sgl-model-gateway/src/wasm/README.md +++ /dev/null @@ -1,227 +0,0 @@ -# WebAssembly (WASM) Extensibility for sgl-model-gateway - -This module provides WebAssembly-based extensibility for sgl-model-gateway, enabling dynamic, safe, and portable middleware execution without requiring router restarts or recompilation. - -## Overview - -The WASM module allows you to extend sgl-model-gateway functionality by deploying WebAssembly components that can: - -- **Intercept requests/responses** at various lifecycle points (OnRequest, OnResponse) -- **Modify HTTP headers and bodies** before/after processing -- **Reject requests** with custom status codes -- **Execute custom logic** in a sandboxed, isolated environment - -## Architecture - -### Components - -The WASM module consists of several key components: - -``` -src/wasm/ -├── module.rs # Data structures (metadata, types, attach points) -├── module_manager.rs # Module lifecycle management (add/remove/list) -├── runtime.rs # WASM execution engine and thread pool -├── route.rs # HTTP API endpoints for module management -├── spec.rs # WASM interface types bindings and type conversions -├── types.rs # Generic input/output types -├── errors.rs # Error definitions -├── config.rs # Runtime configuration -└── interface/ # WebAssembly Interface Types definitions -``` - -### Execution Flow - -``` -1. HTTP Request arrives at router - ↓ -2. Middleware chain checks for WASM modules attached to OnRequest - ↓ -3. For each module: - a. Module manager retrieves pre-loaded WASM bytes - b. Runtime executes component in isolated worker thread - c. Component processes request via WASM type interface - d. Returns Action (Continue/Reject/Modify) - ↓ -4. If Continue: proceed to next middleware/upstream - If Reject: return error response immediately - If Modify: apply changes (headers, body, status) - ↓ -5. After upstream response: - - Modules attached to OnResponse process response - - Apply modifications - ↓ -6. Return final response to client -``` - -### WebAssembly Interface Types - -The module uses the WebAssembly Component Model with WASM interface type for type-safe communication between host and WASM components: - -- **Request Processing**: `middleware-on-request::on-request(req: Request) -> Action` -- **Response Processing**: `middleware-on-response::on-response(resp: Response) -> Action` -- **Actions**: `Continue`, `Reject(status)`, or `Modify(modify-action)` - -See [`interface/`](./interface/) for the complete interface definition. - -## Usage - -### Prerequisites - -- sgl-model-gateway compiled with WASM support -- Rust toolchain (for building WASM components) -- `wasm32-wasip2` target: `rustup target add wasm32-wasip2` -- `wasm-tools`: `cargo install wasm-tools` - -### Starting the Router - -Enable WASM support when starting the router: - -```bash -./sgl-model-gateway --enable-wasm --worker-urls=http://0.0.0.0:30000 --port=3000 -``` - -### Deploying a WASM Module - -Use the `/wasm` POST endpoint to deploy modules: - -```bash -curl -X POST http://localhost:3000/wasm \ - -H "Content-Type: application/json" \ - -d '{ - "modules": [{ - "name": "my-middleware", - "file_path": "/path/to/my-component.component.wasm", - "module_type": "Middleware", - "attach_points": [{"Middleware": "OnRequest"}] - }] - }' -``` - -### Managing Modules - -**List all modules:** -```bash -curl http://localhost:3000/wasm -``` - -**Remove a module:** -```bash -curl -X DELETE http://localhost:3000/wasm/{module-uuid} -``` - -### Module Configuration - -Each module requires: - -- **name**: Unique identifier for the module -- **file_path**: Absolute path to the WASM component file -- **module_type**: Currently supports `"Middleware"` -- **attach_points**: List of attachment points, e.g., `[{"Middleware": "OnRequest"}]` - -Supported attachment points: -- `{"Middleware": "OnRequest"}` - Execute before forwarding to upstream -- `{"Middleware": "OnResponse"}` - Execute after receiving upstream response -- `{"Middleware": "OnError"}` - Not yet implemented - -## Examples - -See [`examples/wasm/`](../../examples/wasm/) for complete examples: - -1. **[wasm-guest-auth](../../examples/wasm/wasm-guest-auth/)** - API key authentication middleware -2. **[wasm-guest-logging](../../examples/wasm/wasm-guest-logging/)** - Request tracking and status code conversion -3. **[wasm-guest-ratelimit](../../examples/wasm/wasm-guest-ratelimit/)** - Rate limiting middleware - -Each example includes: -- Complete source code -- Build instructions -- Deployment examples -- Testing guidelines - -## Security and Resource Management - -### Sandboxing - -WASM modules run in isolated environments provided by wasmtime, preventing: -- Direct system access -- Memory corruption of the host process -- Unauthorized network access -- File system access (unless explicitly granted via WASI) - -### Resource Limits - -Runtime configuration allows setting limits: - -```rust -WasmRuntimeConfig { - max_memory_pages: 1024, // 64MB limit - max_execution_time_ms: 1000, // 1 second timeout - max_stack_size: 1024 * 1024, // 1MB stack - thread_pool_size: 4, // Worker threads - module_cache_size: 10, // Cached modules per worker -} -``` - -### Error Handling - -- Failed module executions are logged and don't crash the router -- Invalid WASM components are rejected during load time -- Metrics track execution success/failure rates - -## Metrics - -The module exposes execution metrics via the `/wasm` GET endpoint: - -```json -{ - "modules": [...], - "metrics": { - "total_executions": 1000, - "successful_executions": 995, - "failed_executions": 5, - "total_execution_time_ms": 50000, - "max_execution_time_ms": 150, - "average_execution_time_ms": 50.0 - } -} -``` - -## Development - -### Building WASM Components - -WASM components must be built using the Component Model. For Rust: - -```bash -# 1. Build as WASM module -cargo build --target wasm32-wasip2 --release - -# 2. Wrap into component format -wasm-tools component new target/wasm32-wasip2/release/my_module.wasm \ - -o my_module.component.wasm -``` - -### WASM Interface Type - -Define your component using the WASM interface from `interface/spec.*`: - -```rust -wit_bindgen::generate!({ - path: "../../../src/wasm/interface", - world: "sgl-model-gateway", -}); - -use exports::sgl::model_gateway::middleware_on_request::Guest as OnRequestGuest; -use sgl::model_gateway::middleware_types::{Request, Action}; - -struct Middleware; - -impl OnRequestGuest for Middleware { - fn on_request(req: Request) -> Action { - // Your logic here - Action::Continue - } -} - -export!(Middleware); -``` diff --git a/sgl-model-gateway/src/wasm/config.rs b/sgl-model-gateway/src/wasm/config.rs deleted file mode 100644 index 37d95388a..000000000 --- a/sgl-model-gateway/src/wasm/config.rs +++ /dev/null @@ -1,353 +0,0 @@ -//! WASM Runtime Configuration -//! -//! Defines configuration parameters for the WASM runtime, -//! including memory limits, execution timeouts, and thread pool settings. - -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Deserialize, Serialize)] -pub struct WasmRuntimeConfig { - /// Maximum memory size in pages (64KB per page) - pub max_memory_pages: u32, - /// Maximum execution time in milliseconds - pub max_execution_time_ms: u64, - /// Maximum stack size in bytes - pub max_stack_size: usize, - /// Number of worker threads in the pool - pub thread_pool_size: usize, - /// Maximum number of modules to cache per worker - pub module_cache_size: usize, - /// Maximum HTTP body size in bytes for middleware request/response processing - pub max_body_size: usize, -} - -impl Default for WasmRuntimeConfig { - fn default() -> Self { - let default_thread_pool_size = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(4) - .clamp(1, 4); - - Self { - max_memory_pages: 1024, // 64MB - max_execution_time_ms: 1000, // 1 seconds - max_stack_size: 1024 * 1024, // 1MB - thread_pool_size: default_thread_pool_size, // based on cpu count and capped - module_cache_size: 10, // Cache up to 10 modules per worker - max_body_size: 10 * 1024 * 1024, // 10MB - } - } -} - -impl WasmRuntimeConfig { - /// Validate the configuration parameters - pub fn validate(&self) -> Result<(), String> { - // Validate max_memory_pages - if self.max_memory_pages == 0 { - return Err("max_memory_pages cannot be 0".to_string()); - } - if self.max_memory_pages > 65536 { - return Err("max_memory_pages cannot exceed 65536 (4GB)".to_string()); - } - - // Validate max_execution_time_ms - if self.max_execution_time_ms == 0 { - return Err("max_execution_time_ms cannot be 0".to_string()); - } - if self.max_execution_time_ms > 300000 { - return Err("max_execution_time_ms cannot exceed 300000ms (5 minutes)".to_string()); - } - - // Validate max_stack_size - if self.max_stack_size == 0 { - return Err("max_stack_size cannot be 0".to_string()); - } - if self.max_stack_size < 64 * 1024 { - return Err("max_stack_size must be at least 64KB".to_string()); - } - if self.max_stack_size > 16 * 1024 * 1024 { - return Err("max_stack_size cannot exceed 16MB".to_string()); - } - - // Validate thread_pool_size - if self.thread_pool_size == 0 { - return Err("thread_pool_size cannot be 0".to_string()); - } - if self.thread_pool_size > 128 { - return Err("thread_pool_size cannot exceed 128".to_string()); - } - - // Validate module_cache_size - if self.module_cache_size == 0 { - return Err("module_cache_size cannot be 0".to_string()); - } - if self.module_cache_size > 1000 { - return Err("module_cache_size cannot exceed 1000".to_string()); - } - - // Validate max_body_size - if self.max_body_size == 0 { - return Err("max_body_size cannot be 0".to_string()); - } - if self.max_body_size > 100 * 1024 * 1024 { - return Err("max_body_size cannot exceed 100MB".to_string()); - } - - Ok(()) - } - - /// Create a new config with validation - pub fn new( - max_memory_pages: u32, - max_execution_time_ms: u64, - max_stack_size: usize, - thread_pool_size: usize, - module_cache_size: usize, - max_body_size: usize, - ) -> Result { - let config = Self { - max_memory_pages, - max_execution_time_ms, - max_stack_size, - thread_pool_size, - module_cache_size, - max_body_size, - }; - config.validate()?; - Ok(config) - } - - /// Get the total memory size in bytes - pub fn get_total_memory_bytes(&self) -> u64 { - self.max_memory_pages as u64 * 64 * 1024 // 64KB per page - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_default_config_validation() { - let config = WasmRuntimeConfig::default(); - assert!(config.validate().is_ok()); - } - - #[test] - fn test_config_new_with_validation() { - let config = WasmRuntimeConfig::new(1024, 1000, 1024 * 1024, 2, 10, 10 * 1024 * 1024); - assert!(config.is_ok()); - } - - #[test] - fn test_validation_max_memory_pages_zero() { - let config = WasmRuntimeConfig { - max_memory_pages: 0, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result.unwrap_err().contains("max_memory_pages cannot be 0")); - } - - #[test] - fn test_validation_max_memory_pages_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 65537, // Exceeds 4GB limit - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_memory_pages cannot exceed 65536")); - } - - #[test] - fn test_validation_max_execution_time_zero() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 0, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_execution_time_ms cannot be 0")); - } - - #[test] - fn test_validation_max_execution_time_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 300001, // Exceeds 5 minutes - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_execution_time_ms cannot exceed 300000ms")); - } - - #[test] - fn test_validation_max_stack_size_too_small() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 32 * 1024, // Less than 64KB - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_stack_size must be at least 64KB")); - } - - #[test] - fn test_validation_max_stack_size_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 17 * 1024 * 1024, // Exceeds 16MB - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_stack_size cannot exceed 16MB")); - } - - #[test] - fn test_validation_thread_pool_size_zero() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 0, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result.unwrap_err().contains("thread_pool_size cannot be 0")); - } - - #[test] - fn test_validation_thread_pool_size_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 129, // Exceeds 128 - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("thread_pool_size cannot exceed 128")); - } - - #[test] - fn test_validation_module_cache_size_zero() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 0, - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("module_cache_size cannot be 0")); - } - - #[test] - fn test_validation_module_cache_size_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 1001, // Exceeds 1000 - max_body_size: 10 * 1024 * 1024, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("module_cache_size cannot exceed 1000")); - } - - #[test] - fn test_get_total_memory_bytes() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 10 * 1024 * 1024, - }; - // 1024 pages * 64KB = 64MB - assert_eq!(config.get_total_memory_bytes(), 64 * 1024 * 1024); - } - - #[test] - fn test_validation_max_body_size_zero() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 0, - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result.unwrap_err().contains("max_body_size cannot be 0")); - } - - #[test] - fn test_validation_max_body_size_too_large() { - let config = WasmRuntimeConfig { - max_memory_pages: 1024, - max_execution_time_ms: 1000, - max_stack_size: 1024 * 1024, - thread_pool_size: 2, - module_cache_size: 10, - max_body_size: 101 * 1024 * 1024, // Exceeds 100MB - }; - let result = config.validate(); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .contains("max_body_size cannot exceed 100MB")); - } -} diff --git a/sgl-model-gateway/src/wasm/errors.rs b/sgl-model-gateway/src/wasm/errors.rs deleted file mode 100644 index 561ee75b1..000000000 --- a/sgl-model-gateway/src/wasm/errors.rs +++ /dev/null @@ -1,117 +0,0 @@ -//! WASM Error Types -//! -//! Defines comprehensive error types for the WASM subsystem, -//! including module, manager, and runtime errors. - -use std::fmt; - -use thiserror::Error; - -pub type Result = std::result::Result; - -/// SHA256 hash wrapper for display purposes -#[derive(Debug, Clone, Copy)] -pub struct Sha256Hash(pub [u8; 32]); - -impl fmt::Display for Sha256Hash { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let hex_string: String = self.0.iter().map(|b| format!("{:02x}", b)).collect(); - write!(f, "{}", hex_string) - } -} - -impl From<[u8; 32]> for Sha256Hash { - fn from(hash: [u8; 32]) -> Self { - Sha256Hash(hash) - } -} - -#[derive(Debug, Error)] -pub enum WasmError { - #[error(transparent)] - Module(#[from] WasmModuleError), - - #[error(transparent)] - Manager(#[from] WasmManagerError), - - #[error(transparent)] - Runtime(#[from] WasmRuntimeError), - - #[error(transparent)] - Io(#[from] std::io::Error), - - #[error("{0}")] - Other(String), -} - -#[derive(Debug, Error)] -pub enum WasmModuleError { - #[error("invalid module descriptor: {0}")] - InvalidDescriptor(String), - - #[error("module with same sha256 already exists: {0}")] - DuplicateSha256(Sha256Hash), - - #[error("module not found: {0}")] - NotFound(uuid::Uuid), - - #[error("failed to read module file: {0}")] - FileRead(String), - - #[error("validation failed: {0}")] - ValidationFailed(String), - - #[error("attach point missing: {0}")] - AttachPointMissing(String), - - #[error("invalid function for attach point: {0}")] - AttachPointFunctionInvalid(String), -} - -#[derive(Debug, Error)] -pub enum WasmManagerError { - #[error("failed to acquire lock: {0}")] - LockFailed(String), - - #[error("module add failed: {0}")] - ModuleAddFailed(String), - - #[error("module remove failed: {0}")] - ModuleRemoveFailed(String), - - #[error("runtime unavailable")] - RuntimeUnavailable, - - #[error("execution failed: {0}")] - ExecutionFailed(String), - - #[error("module {0} not found")] - ModuleNotFound(uuid::Uuid), -} - -#[derive(Debug, Error)] -pub enum WasmRuntimeError { - #[error("failed to create engine: {0}")] - EngineCreateFailed(String), - - #[error("failed to compile module: {0}")] - CompileFailed(String), - - #[error("failed to create instance: {0}")] - InstanceCreateFailed(String), - - #[error("function not found: {0}")] - FunctionNotFound(String), - - #[error("execution timeout after {0}ms")] - Timeout(u64), - - #[error("execution failed: {0}")] - CallFailed(String), -} - -impl From for WasmError { - fn from(value: wasmtime::Error) -> Self { - WasmError::Runtime(WasmRuntimeError::CallFailed(value.to_string())) - } -} diff --git a/sgl-model-gateway/src/wasm/interface/spec.wit b/sgl-model-gateway/src/wasm/interface/spec.wit deleted file mode 100644 index 7750a32da..000000000 --- a/sgl-model-gateway/src/wasm/interface/spec.wit +++ /dev/null @@ -1,54 +0,0 @@ -package sgl:model-gateway; - -interface middleware-types { - record header { name: string, value: string } - - // onRequest - record request { - method: string, - path: string, - query: string, - headers: list
, - body: list, - request-id: string, - now-epoch-ms: u64, - } - - // onResponse - record response { - status: u16, - headers: list
, - body: list, - } - - // modify action - record modify-action { - status: option, - headers-set: list
, - headers-add: list
, - headers-remove: list, - body-replace: option>, - } - - // return actions - variant action { - continue, - reject(u16), // status code - modify(modify-action), - } -} - -interface middleware-on-request { - use middleware-types.{request, action}; - on-request: func(req: request) -> action; -} - -interface middleware-on-response { - use middleware-types.{response, action}; - on-response: func(resp: response) -> action; -} - -world sgl-model-gateway { - export middleware-on-request; - export middleware-on-response; -} diff --git a/sgl-model-gateway/src/wasm/mod.rs b/sgl-model-gateway/src/wasm/mod.rs deleted file mode 100644 index fa8d55aac..000000000 --- a/sgl-model-gateway/src/wasm/mod.rs +++ /dev/null @@ -1,13 +0,0 @@ -//! WebAssembly (WASM) module support for sgl-model-gateway -//! -//! This module provides WASM component execution capabilities using the WebAssembly Component Model. -//! It supports middleware execution at various attach points (OnRequest, OnResponse) with async support. - -pub mod config; -pub mod errors; -pub mod module; -pub mod module_manager; -pub mod route; -pub mod runtime; -pub mod spec; -pub mod types; diff --git a/sgl-model-gateway/src/wasm/module.rs b/sgl-model-gateway/src/wasm/module.rs deleted file mode 100644 index ff49ef07c..000000000 --- a/sgl-model-gateway/src/wasm/module.rs +++ /dev/null @@ -1,209 +0,0 @@ -//! WASM Module Data Structures and Types -//! -//! This module defines the core data structures for managing WebAssembly components: -//! - Module metadata (UUID, name, file path, hash, timestamps, metrics) -//! - Module types and attachment points (Middleware hooks: OnRequest, OnResponse, OnError) -//! - API request/response types for module management -//! - Execution metrics and statistics -//! -//! The module provides custom serialization for: -//! - SHA256 hashes (hex string representation) -//! - Timestamps (ISO 8601 format for JSON output) - -use serde::{Deserialize, Serialize, Serializer}; -use uuid::Uuid; - -/// Serialize [u8; 32] as hex string -fn serialize_sha256_hash(hash: &[u8; 32], serializer: S) -> Result -where - S: Serializer, -{ - let hex_string = hash - .iter() - .map(|b| format!("{:02x}", b)) - .collect::(); - serializer.serialize_str(&hex_string) -} - -/// Deserialize hex string to [u8; 32] -fn deserialize_sha256_hash<'de, D>(deserializer: D) -> Result<[u8; 32], D::Error> -where - D: serde::Deserializer<'de>, -{ - use serde::Deserialize; - let hex_string = String::deserialize(deserializer)?; - - // Parse hex string to bytes - if hex_string.len() != 64 { - return Err(serde::de::Error::custom(format!( - "Invalid SHA256 hash length: expected 64 hex characters, got {}", - hex_string.len() - ))); - } - - let mut hash = [0u8; 32]; - for (i, chunk) in hex_string.as_bytes().chunks(2).enumerate() { - if chunk.len() != 2 { - return Err(serde::de::Error::custom("Invalid hex string format")); - } - let byte_str = std::str::from_utf8(chunk) - .map_err(|e| serde::de::Error::custom(format!("Invalid UTF-8: {}", e)))?; - hash[i] = u8::from_str_radix(byte_str, 16) - .map_err(|e| serde::de::Error::custom(format!("Invalid hex digit: {}", e)))?; - } - - Ok(hash) -} - -/// Serialize u64 timestamp (nanoseconds since epoch) as ISO 8601 string -fn serialize_timestamp(timestamp: &u64, serializer: S) -> Result -where - S: Serializer, -{ - use chrono::{DateTime, Utc}; - - // Convert nanoseconds to seconds and remaining nanoseconds - let secs = (*timestamp / 1_000_000_000) as i64; - let nanos = (*timestamp % 1_000_000_000) as u32; - - match DateTime::::from_timestamp(secs, nanos) { - Some(dt) => { - let s = dt.to_rfc3339_opts(chrono::SecondsFormat::Nanos, true); - serializer.serialize_str(&s) - } - None => { - // Fallback: format manually if timestamp is out of range - let s = format!("{}", timestamp); - serializer.serialize_str(&s) - } - } -} - -/// Deserialize ISO 8601 string to u64 timestamp (nanoseconds since epoch) -fn deserialize_timestamp<'de, D>(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - use chrono::{DateTime, Utc}; - use serde::Deserialize; - - let timestamp_str = String::deserialize(deserializer)?; - - // Try to parse as ISO 8601 datetime (RFC 3339) - match DateTime::parse_from_rfc3339(×tamp_str) { - Ok(dt) => { - // Convert to UTC and then to nanoseconds since epoch - let dt_utc = dt.with_timezone(&Utc); - let secs = dt_utc.timestamp(); - let nanos = dt_utc.timestamp_subsec_nanos(); - Ok((secs as u64) * 1_000_000_000 + (nanos as u64)) - } - Err(_) => { - // Fallback: try to parse as u64 directly - timestamp_str - .parse::() - .map_err(|e| serde::de::Error::custom(format!("Invalid timestamp format: {}", e))) - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModule { - // unique identifier for the module - pub module_uuid: Uuid, - pub module_meta: WasmModuleMeta, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum WasmModuleAddResult { - Success(Uuid), - Error(String), -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModuleDescriptor { - pub name: String, - pub file_path: String, - pub module_type: WasmModuleType, - pub attach_points: Vec, - pub add_result: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModuleMeta { - // module name provided by the user - pub name: String, - // path to the module file - pub file_path: String, - // sha256 hash of the module file - #[serde( - serialize_with = "serialize_sha256_hash", - deserialize_with = "deserialize_sha256_hash" - )] - pub sha256_hash: [u8; 32], - // size of the module file in bytes - pub size_bytes: u64, - // timestamp of when the module was created (nanoseconds since epoch) - #[serde( - serialize_with = "serialize_timestamp", - deserialize_with = "deserialize_timestamp" - )] - pub created_at: u64, - // timestamp of when the module was last accessed (nanoseconds since epoch) - #[serde( - serialize_with = "serialize_timestamp", - deserialize_with = "deserialize_timestamp" - )] - pub last_accessed_at: u64, - // number of times the module was accessed - pub access_count: u64, - // attach points for the module - pub attach_points: Vec, - // Pre-loaded WASM component bytes (loaded into memory for faster execution) - #[serde(skip)] - pub wasm_bytes: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)] -pub enum WasmModuleType { - Middleware, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)] -pub enum MiddlewareAttachPoint { - OnRequest, - OnResponse, - OnError, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)] -pub enum WasmModuleAttachPoint { - Middleware(MiddlewareAttachPoint), -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModuleAddRequest { - pub modules: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModuleAddResponse { - pub modules: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmModuleListResponse { - pub modules: Vec, - pub metrics: WasmMetrics, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WasmMetrics { - pub total_executions: u64, - pub successful_executions: u64, - pub failed_executions: u64, - pub total_execution_time_ms: u64, - pub max_execution_time_ms: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub average_execution_time_ms: Option, -} diff --git a/sgl-model-gateway/src/wasm/module_manager.rs b/sgl-model-gateway/src/wasm/module_manager.rs deleted file mode 100644 index ad5454346..000000000 --- a/sgl-model-gateway/src/wasm/module_manager.rs +++ /dev/null @@ -1,267 +0,0 @@ -//! WASM Module Manager - -use std::{ - collections::HashMap, - sync::{ - atomic::{AtomicU64, Ordering}, - Arc, RwLock, - }, -}; - -use uuid::Uuid; - -use crate::wasm::{ - config::WasmRuntimeConfig, - errors::{Result, WasmError, WasmManagerError, WasmModuleError}, - module::{WasmModule, WasmModuleAttachPoint}, - runtime::WasmRuntime, - types::{WasmComponentInput, WasmComponentOutput}, -}; - -pub struct WasmModuleManager { - modules: Arc>>, - runtime: Arc, - // Metrics - total_executions: AtomicU64, - successful_executions: AtomicU64, - failed_executions: AtomicU64, - total_execution_time_ms: AtomicU64, - max_execution_time_ms: AtomicU64, -} - -impl WasmModuleManager { - pub fn new(config: WasmRuntimeConfig) -> Result { - let runtime = Arc::new(WasmRuntime::new(config)?); - Ok(Self { - modules: Arc::new(RwLock::new(HashMap::new())), - runtime, - total_executions: AtomicU64::new(0), - successful_executions: AtomicU64::new(0), - failed_executions: AtomicU64::new(0), - total_execution_time_ms: AtomicU64::new(0), - max_execution_time_ms: AtomicU64::new(0), - }) - } - - pub fn with_default_config() -> Result { - Self::new(WasmRuntimeConfig::default()) - } - - /// Register a module (for workflow steps) - pub(crate) fn register_module_internal(&self, module: WasmModule) -> Result<()> { - let mut modules = self - .modules - .write() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - modules.insert(module.module_uuid, module); - Ok(()) - } - - /// Remove a module (for workflow steps) - pub(crate) fn remove_module_internal(&self, module_uuid: Uuid) -> Result<()> { - let mut modules = self - .modules - .write() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - if !modules.contains_key(&module_uuid) { - return Err(WasmManagerError::ModuleNotFound(module_uuid).into()); - } - modules.remove(&module_uuid); - Ok(()) - } - - pub(crate) fn check_duplicate_sha256_hash(&self, sha256_hash: &[u8; 32]) -> Result<()> { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - if modules - .values() - .any(|module: &WasmModule| module.module_meta.sha256_hash == *sha256_hash) - { - return Err(WasmModuleError::DuplicateSha256((*sha256_hash).into()).into()); - } - Ok(()) - } - - pub fn get_all_modules(&self) -> Result> { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - Ok(modules.values().cloned().collect()) - } - - pub fn get_module(&self, module_uuid: Uuid) -> Result> { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - Ok(modules.get(&module_uuid).cloned()) - } - - pub fn get_modules(&self) -> Result> { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - Ok(modules.values().cloned().collect()) - } - - /// get modules by attach point - pub fn get_modules_by_attach_point( - &self, - attach_point: WasmModuleAttachPoint, - ) -> Result> { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - Ok(modules - .values() - .filter(|module| module.module_meta.attach_points.contains(&attach_point)) - .cloned() - .collect()) - } - - pub fn get_runtime(&self) -> &Arc { - &self.runtime - } - - /// Get the configured maximum body size for HTTP request/response processing - pub fn get_max_body_size(&self) -> usize { - self.runtime.get_config().max_body_size - } - - /// Execute WASM module using WebAssembly component model based on attach_point - pub async fn execute_module_interface( - &self, - module_uuid: Uuid, - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - ) -> Result { - let start_time = std::time::Instant::now(); - - // First, get the WASM bytes with a read lock (faster) - let (wasm_bytes, wasm_hash) = { - let modules = self - .modules - .read() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - let module = modules - .get(&module_uuid) - .ok_or_else(|| WasmError::from(WasmManagerError::ModuleNotFound(module_uuid)))?; - - // Clone the pre-loaded WASM bytes (already in memory, no file I/O) - ( - module.module_meta.wasm_bytes.clone(), - module.module_meta.sha256_hash, - ) - }; - - { - let mut modules = self - .modules - .write() - .map_err(|e| WasmManagerError::LockFailed(e.to_string()))?; - if let Some(module) = modules.get_mut(&module_uuid) { - // SystemTime::duration_since only fails if the system time is before UNIX_EPOCH, - // which should never happen in normal operation. If it does, use current time as fallback. - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_else(|_| { - // Fallback to a reasonable timestamp if system time is invalid - // This should never occur in practice, but provides a safe fallback - std::time::Duration::from_nanos(0) - }) - .as_nanos() as u64; - module.module_meta.last_accessed_at = now; - module.module_meta.access_count += 1; - } - } - - let result = self - .runtime - .execute_component_async(wasm_bytes, wasm_hash, attach_point, input) - .await; - - // Record metrics - let execution_time_ms = start_time.elapsed().as_millis() as u64; - self.total_executions.fetch_add(1, Ordering::Relaxed); - self.total_execution_time_ms - .fetch_add(execution_time_ms, Ordering::Relaxed); - // Update max execution time - self.max_execution_time_ms - .fetch_max(execution_time_ms, Ordering::Relaxed); - - if result.is_ok() { - self.successful_executions.fetch_add(1, Ordering::Relaxed); - } else { - self.failed_executions.fetch_add(1, Ordering::Relaxed); - } - - result - } - - /// Execute WASM module using WebAssembly component model (sync version) - pub fn execute_module_interface_sync( - &self, - module_uuid: Uuid, - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - ) -> Result { - let handle = tokio::runtime::Handle::current(); - handle.block_on(self.execute_module_interface(module_uuid, attach_point, input)) - } - - /// Get current metrics - pub fn get_metrics(&self) -> (u64, u64, u64, u64, u64) { - ( - self.total_executions.load(Ordering::Relaxed), - self.successful_executions.load(Ordering::Relaxed), - self.failed_executions.load(Ordering::Relaxed), - self.total_execution_time_ms.load(Ordering::Relaxed), - self.max_execution_time_ms.load(Ordering::Relaxed), - ) - } - - /// Execute a WASM module for a given attach point - /// Returns the Action if successful, or None if execution failed - /// - /// This is a convenience method that wraps execute_module_interface and handles - /// error logging automatically. - pub async fn execute_module_for_attach_point( - &self, - module: &WasmModule, - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - ) -> Option { - use tracing::error; - - let action_result = self - .execute_module_interface(module.module_uuid, attach_point, input) - .await; - - match action_result { - Ok(output) => match output { - WasmComponentOutput::MiddlewareAction(action) => Some(action), - }, - Err(e) => { - error!( - "Failed to execute WASM module {}: {}", - module.module_meta.name, e - ); - None - } - } - } -} - -impl Default for WasmModuleManager { - fn default() -> Self { - // with_default_config() should always succeed with default configuration. - // If it fails, it indicates a critical system configuration error. - Self::with_default_config() - .expect("Failed to create WasmModuleManager with default config. This should never happen with valid default configuration.") - } -} diff --git a/sgl-model-gateway/src/wasm/route.rs b/sgl-model-gateway/src/wasm/route.rs deleted file mode 100644 index 6d9b18d68..000000000 --- a/sgl-model-gateway/src/wasm/route.rs +++ /dev/null @@ -1,228 +0,0 @@ -//! WASM HTTP API Routes -//! -//! Provides REST API endpoints for managing WASM modules: -//! - POST /wasm - Add modules -//! - DELETE /wasm/:uuid - Remove a module -//! - GET /wasm - List all modules with metrics - -use std::{sync::Arc, time::Duration}; - -use axum::{ - extract::{Json, Path, State}, - http::StatusCode, - response::{IntoResponse, Response}, -}; -use uuid::Uuid; - -use crate::{ - core::{job_queue::Job, steps::WasmModuleConfigRequest}, - server::AppState, - wasm::module::{ - WasmMetrics, WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAddResult, - WasmModuleListResponse, - }, -}; - -/// Wait for job completion by polling job status -/// Returns the job result message if successful -async fn wait_for_job_completion( - job_queue: &crate::core::job_queue::JobQueue, - status_key: &str, - timeout_duration: Duration, -) -> Result { - let start = std::time::Instant::now(); - let mut poll_interval = Duration::from_millis(100); - let max_poll_interval = Duration::from_millis(2000); - let poll_backoff = Duration::from_millis(200); - - loop { - if start.elapsed() > timeout_duration { - return Err(format!("Job timeout after {}s", timeout_duration.as_secs())); - } - - if let Some(job_status) = job_queue.get_status(status_key) { - match job_status.status.as_str() { - "pending" | "processing" => { - tokio::time::sleep(poll_interval).await; - poll_interval = (poll_interval + poll_backoff).min(max_poll_interval); - continue; - } - "failed" => { - let error_msg = job_status - .message - .unwrap_or_else(|| "Unknown error".to_string()); - job_queue.remove_status(status_key); - return Err(error_msg); - } - _ => { - // Should not happen, but handle gracefully - job_queue.remove_status(status_key); - return Err("Unexpected job status".to_string()); - } - } - } else { - // Job completed successfully (status was removed by record_job_completion) - // We need to get the result from the job execution - // Since job queue removes status on success, we can't get the result here - // We'll need to query the wasm manager to find the module by name - // For now, return a success message and let caller extract UUID from manager - return Ok("Job completed successfully".to_string()); - } - } -} - -pub async fn add_wasm_module( - State(state): State>, - Json(config): Json, -) -> Response { - let Some(_) = state.context.wasm_manager.as_ref() else { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - }; - - let Some(job_queue) = state.context.worker_job_queue.get() else { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - }; - - let mut status = StatusCode::OK; - let mut modules = config.modules.clone(); - - for module in modules.iter_mut() { - let wasm_config = WasmModuleConfigRequest { - descriptor: module.clone(), - }; - - let job = Job::AddWasmModule { - config: Box::new(wasm_config), - }; - - let worker_url = job.worker_url().to_string(); - - // Submit job to queue - match job_queue.submit(job).await { - Ok(_) => { - // Wait for job completion (timeout: 5 minutes) - let timeout = Duration::from_secs(300); - match wait_for_job_completion(job_queue, &worker_url, timeout).await { - Ok(_) => { - // Job completed successfully, but we need to get the UUID - // Since job queue removes status on success, we need to query - // the workflow engine or wasm manager to get the UUID - // For now, let's try to get it from the wasm manager by name - if let Some(wasm_manager) = state.context.wasm_manager.as_ref() { - if let Ok(all_modules) = wasm_manager.get_modules() { - if let Some(registered_module) = all_modules - .iter() - .find(|m| m.module_meta.name == module.name) - { - module.add_result = Some(WasmModuleAddResult::Success( - registered_module.module_uuid, - )); - } else { - module.add_result = Some(WasmModuleAddResult::Error( - "Module registered but UUID not found".to_string(), - )); - status = StatusCode::BAD_REQUEST; - } - } else { - module.add_result = Some(WasmModuleAddResult::Error( - "Failed to query registered modules".to_string(), - )); - status = StatusCode::BAD_REQUEST; - } - } else { - module.add_result = Some(WasmModuleAddResult::Error( - "WASM manager not available".to_string(), - )); - status = StatusCode::BAD_REQUEST; - } - } - Err(e) => { - module.add_result = Some(WasmModuleAddResult::Error(e)); - status = StatusCode::BAD_REQUEST; - } - } - } - Err(e) => { - module.add_result = Some(WasmModuleAddResult::Error(format!( - "Failed to submit job: {}", - e - ))); - status = StatusCode::BAD_REQUEST; - } - } - } - - let response = WasmModuleAddResponse { modules }; - (status, Json(response)).into_response() -} - -pub async fn remove_wasm_module( - State(state): State>, - Path(module_uuid_str): Path, -) -> Response { - let Ok(module_uuid) = Uuid::parse_str(&module_uuid_str) else { - return StatusCode::BAD_REQUEST.into_response(); - }; - - let Some(_) = state.context.wasm_manager.as_ref() else { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - }; - - let Some(job_queue) = state.context.worker_job_queue.get() else { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - }; - - use crate::core::steps::WasmModuleRemovalRequest; - - let removal_request = WasmModuleRemovalRequest::new(module_uuid); - - let job = Job::RemoveWasmModule { - request: Box::new(removal_request), - }; - - let worker_url = job.worker_url().to_string(); - - // Submit job to queue - match job_queue.submit(job).await { - Ok(_) => { - // Wait for job completion (timeout: 1 minute) - let timeout = Duration::from_secs(60); - match wait_for_job_completion(job_queue, &worker_url, timeout).await { - Ok(_) => (StatusCode::OK, "Module removed successfully").into_response(), - Err(e) => (StatusCode::BAD_REQUEST, e).into_response(), - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Failed to submit job: {}", e), - ) - .into_response(), - } -} - -pub async fn list_wasm_modules(State(state): State>) -> Response { - let Some(wasm_manager) = state.context.wasm_manager.as_ref() else { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - }; - let modules = wasm_manager.get_modules(); - if let Ok(modules) = modules { - let (total, success, failed, total_time_ms, max_time_ms) = wasm_manager.get_metrics(); - let average_execution_time_ms = if total > 0 { - Some(total_time_ms as f64 / total as f64) - } else { - None - }; - let metrics = WasmMetrics { - total_executions: total, - successful_executions: success, - failed_executions: failed, - total_execution_time_ms: total_time_ms, - max_execution_time_ms: max_time_ms, - average_execution_time_ms, - }; - let response = WasmModuleListResponse { modules, metrics }; - (StatusCode::OK, Json(response)).into_response() - } else { - StatusCode::INTERNAL_SERVER_ERROR.into_response() - } -} diff --git a/sgl-model-gateway/src/wasm/runtime.rs b/sgl-model-gateway/src/wasm/runtime.rs deleted file mode 100644 index 9b34e36cc..000000000 --- a/sgl-model-gateway/src/wasm/runtime.rs +++ /dev/null @@ -1,615 +0,0 @@ -//! WASM Runtime -//! -//! Manages WASM component execution using wasmtime with async support. -//! Provides a thread pool for concurrent WASM execution and metrics tracking. - -use std::{ - num::NonZeroUsize, - sync::{ - atomic::{AtomicU64, Ordering}, - Arc, - }, - time::Duration, -}; - -use lru::LruCache; -use tokio::sync::oneshot; -use tracing::{debug, error}; -use wasmtime::{ - component::{Component, Linker, ResourceTable}, - Config, Engine, InstanceAllocationStrategy, PoolingAllocationConfig, Store, StoreLimitsBuilder, - UpdateDeadline, -}; -use wasmtime_wasi::WasiCtx; - -/// Epoch increment interval in milliseconds. -/// Epochs are used for cooperative timeout enforcement in WASM execution. -/// A smaller interval gives finer-grained timeout control but slightly more overhead. -const EPOCH_INTERVAL_MS: u64 = 100; - -use crate::wasm::{ - config::WasmRuntimeConfig, - errors::{Result, WasmError, WasmRuntimeError}, - module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, - spec::SglModelGateway, - types::{WasiState, WasmComponentInput, WasmComponentOutput}, -}; - -pub struct WasmRuntime { - config: WasmRuntimeConfig, - thread_pool: Arc, - // Metrics - total_executions: AtomicU64, - successful_executions: AtomicU64, - failed_executions: AtomicU64, - total_execution_time_ms: AtomicU64, - max_execution_time_ms: AtomicU64, -} - -pub struct WasmThreadPool { - sender: async_channel::Sender, - receiver: async_channel::Receiver, - workers: Vec>, - // Metrics - total_tasks: AtomicU64, - completed_tasks: AtomicU64, - failed_tasks: AtomicU64, -} - -pub enum WasmTask { - ExecuteComponent { - wasm_bytes: Vec, - wasm_hash: [u8; 32], - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - response: oneshot::Sender>, - }, -} - -impl WasmRuntime { - pub fn new(config: WasmRuntimeConfig) -> Result { - let thread_pool = Arc::new(WasmThreadPool::new(config.clone())?); - - Ok(Self { - config, - thread_pool, - total_executions: AtomicU64::new(0), - successful_executions: AtomicU64::new(0), - failed_executions: AtomicU64::new(0), - total_execution_time_ms: AtomicU64::new(0), - max_execution_time_ms: AtomicU64::new(0), - }) - } - - pub fn with_default_config() -> Result { - Self::new(WasmRuntimeConfig::default()) - } - - pub fn get_config(&self) -> &WasmRuntimeConfig { - &self.config - } - - /// get available cpu count and max recommended cpu count - pub fn get_cpu_info() -> (usize, usize) { - let cpu_count = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(4); - let max_recommended = cpu_count.max(1); - (cpu_count, max_recommended) - } - - /// get current thread pool status - pub fn get_thread_pool_info(&self) -> (usize, usize) { - let (_cpu_count, max_recommended) = Self::get_cpu_info(); - let current_workers = self.thread_pool.workers.len(); - (current_workers, max_recommended) - } - - /// Execute WASM component using WASM interface based on attach_point - pub async fn execute_component_async( - &self, - wasm_bytes: Vec, - wasm_hash: [u8; 32], - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - ) -> Result { - let start_time = std::time::Instant::now(); - let (response_tx, response_rx) = oneshot::channel(); - - let task = WasmTask::ExecuteComponent { - wasm_bytes, - wasm_hash, - attach_point, - input, - response: response_tx, - }; - - self.thread_pool.sender.send(task).await.map_err(|e| { - WasmRuntimeError::CallFailed(format!("Failed to send task to thread pool: {}", e)) - })?; - - let result = response_rx.await.map_err(|e| { - WasmRuntimeError::CallFailed(format!( - "Failed to receive response from thread pool: {}", - e - )) - })?; - - let execution_time_ms = start_time.elapsed().as_millis() as u64; - self.total_executions.fetch_add(1, Ordering::Relaxed); - self.total_execution_time_ms - .fetch_add(execution_time_ms, Ordering::Relaxed); - // Update max execution time - self.max_execution_time_ms - .fetch_max(execution_time_ms, Ordering::Relaxed); - - if result.is_ok() { - self.successful_executions.fetch_add(1, Ordering::Relaxed); - } else { - self.failed_executions.fetch_add(1, Ordering::Relaxed); - } - - result - } - - /// Get current metrics - pub fn get_metrics(&self) -> (u64, u64, u64, u64, u64) { - ( - self.total_executions.load(Ordering::Relaxed), - self.successful_executions.load(Ordering::Relaxed), - self.failed_executions.load(Ordering::Relaxed), - self.total_execution_time_ms.load(Ordering::Relaxed), - self.max_execution_time_ms.load(Ordering::Relaxed), - ) - } -} - -/// Maps a wasmtime error to a WasmError, detecting epoch interruption (timeout) traps. -fn map_wasm_error(e: wasmtime::Error, timeout_ms: u64) -> WasmError { - // Use proper trap code detection instead of brittle string matching. - // Wasmtime uses Trap::Interrupt for epoch-based interruptions. - if e.downcast_ref::() == Some(&wasmtime::Trap::Interrupt) { - WasmError::from(WasmRuntimeError::Timeout(timeout_ms)) - } else { - WasmError::from(WasmRuntimeError::CallFailed(e.to_string())) - } -} - -impl WasmThreadPool { - pub fn new(config: WasmRuntimeConfig) -> Result { - let (sender, receiver) = async_channel::unbounded(); - - let mut workers = Vec::new(); - // set thread pool size based on cpu count - let max_workers = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(4) - .max(1); - let num_workers = config.thread_pool_size.clamp(1, max_workers); - - debug!( - target: "smg::wasm::runtime", - "Initializing WASM runtime with {} workers", - num_workers - ); - - for worker_id in 0..num_workers { - let receiver = receiver.clone(); - let config = config.clone(); - - let worker = std::thread::spawn(move || { - // create independent tokio runtime for this thread - let rt = match tokio::runtime::Runtime::new() { - Ok(rt) => rt, - Err(e) => { - error!( - target: "smg::wasm::runtime", - worker_id = worker_id, - "Failed to create tokio runtime: {}", - e - ); - return; - } - }; - - rt.block_on(async { - Self::worker_loop(worker_id, receiver, config).await; - }); - }); - - workers.push(worker); - } - - Ok(Self { - sender, - receiver, - workers, - total_tasks: AtomicU64::new(0), - completed_tasks: AtomicU64::new(0), - failed_tasks: AtomicU64::new(0), - }) - } - - /// Get current thread pool metrics - pub fn get_metrics(&self) -> (u64, u64, u64) { - ( - self.total_tasks.load(Ordering::Relaxed), - self.completed_tasks.load(Ordering::Relaxed), - self.failed_tasks.load(Ordering::Relaxed), - ) - } - - async fn worker_loop( - worker_id: usize, - receiver: async_channel::Receiver, - config: WasmRuntimeConfig, - ) { - debug!( - target: "smg::wasm::runtime", - worker_id = worker_id, - thread_id = ?std::thread::current().id(), - "Worker started" - ); - - let mut pool_config = PoolingAllocationConfig::default(); - let max_memory_bytes = (config.max_memory_pages as usize) * 65536; - - // Since this thread handles tasks sequentially, we don't need a large pool per thread. - // A pool size of 20 allows for efficient reuse without hogging memory. - pool_config.total_core_instances(20); - pool_config.max_memory_size(max_memory_bytes); - pool_config.max_component_instance_size(max_memory_bytes); - pool_config.max_tables_per_component(5); - - let mut wasmtime_config = Config::new(); - wasmtime_config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool_config)); - - wasmtime_config.async_stack_size(config.max_stack_size); - wasmtime_config.async_support(true); - wasmtime_config.wasm_component_model(true); // Enable component model - wasmtime_config.epoch_interruption(true); // Enable epoch-based timeout interruption - - let engine = match Engine::new(&wasmtime_config) { - Ok(engine) => engine, - Err(e) => { - error!( - target: "smg::wasm::runtime", - worker_id = worker_id, - "Failed to create engine: {}", - e - ); - return; - } - }; - - let cache_capacity = - NonZeroUsize::new(config.module_cache_size).unwrap_or(NonZeroUsize::new(10).unwrap()); - let mut component_cache: LruCache<[u8; 32], Component> = LruCache::new(cache_capacity); - - // Start epoch incrementer for timeout enforcement. - // The engine's epoch counter is incremented periodically, and each Store - // can set a deadline (number of epochs). When the deadline is reached, - // WASM execution is interrupted with a trap. - let engine_for_epoch = engine.clone(); - let epoch_handle = tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_millis(EPOCH_INTERVAL_MS)); - loop { - interval.tick().await; - engine_for_epoch.increment_epoch(); - } - }); - - debug!( - target: "smg::wasm::runtime", - worker_id = worker_id, - epoch_interval_ms = EPOCH_INTERVAL_MS, - "Epoch incrementer started for timeout enforcement" - ); - - loop { - let task = match receiver.recv().await { - Ok(task) => task, - Err(_) => { - debug!( - target: "smg::wasm::runtime", - worker_id = worker_id, - "Worker shutting down" - ); - epoch_handle.abort(); // Stop the epoch incrementer - break; // channel closed, exit loop - } - }; - - match task { - WasmTask::ExecuteComponent { - wasm_bytes, - wasm_hash, - attach_point, - input, - response, - } => { - let result = Self::execute_component_in_worker( - &engine, - &mut component_cache, // Pass the cache - wasm_bytes, - wasm_hash, - attach_point, - input, - &config, - ) - .await; - - let _ = response.send(result); - } - } - } - } - - async fn execute_component_in_worker( - engine: &Engine, - cache: &mut LruCache<[u8; 32], Component>, // cache argument - wasm_bytes: Vec, - wasm_hash: [u8; 32], - attach_point: WasmModuleAttachPoint, - input: WasmComponentInput, - config: &WasmRuntimeConfig, - ) -> Result { - // Compile component from bytes OR retrieve from cache - // Note: The WASM file must be in component format (not plain WASM module) - let component = if let Some(comp) = cache.get(&wasm_hash) { - comp.clone() // Component is just a handle (cheap clone) - } else { - // Compile new component - let comp = Component::new(engine, &wasm_bytes).map_err(|e| { - WasmError::Runtime(WasmRuntimeError::CompileFailed(format!( - "failed to parse WebAssembly component: {}. \ - Hint: The WASM file must be in component format. \ - If you're using wit-bindgen, use 'wasm-tools component new' to wrap the WASM module into a component.", - e - ))) - })?; - - cache.push(wasm_hash, comp.clone()); - comp - }; - - 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. - let deadline_epochs = (config.max_execution_time_ms / EPOCH_INTERVAL_MS).max(1); - store.set_epoch_deadline(deadline_epochs); - - // Configure what happens when the deadline is reached during async yields - store.epoch_deadline_callback(|_store| Ok(UpdateDeadline::Yield(1))); - - let output = match attach_point { - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest) => { - let request = match input { - WasmComponentInput::MiddlewareRequest(req) => req, - _ => { - return Err(WasmError::from(WasmRuntimeError::CallFailed( - "Expected MiddlewareRequest input for OnRequest attach point" - .to_string(), - ))); - } - }; - - // Instantiate component (must use async instantiation when async support is enabled) - let bindings = SglModelGateway::instantiate_async(&mut store, &component, &linker) - .await - .map_err(|e| { - WasmError::from(WasmRuntimeError::InstanceCreateFailed(e.to_string())) - })?; - - // Call on-request (async call when async support is enabled) - let action_result = bindings - .sgl_model_gateway_middleware_on_request() - .call_on_request(&mut store, &request) - .await - .map_err(|e| map_wasm_error(e, config.max_execution_time_ms))?; - - WasmComponentOutput::MiddlewareAction(action_result) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse) => { - // Extract Response input - let response = match input { - WasmComponentInput::MiddlewareResponse(resp) => resp, - _ => { - return Err(WasmError::from(WasmRuntimeError::CallFailed( - "Expected MiddlewareResponse input for OnResponse attach point" - .to_string(), - ))); - } - }; - - // Instantiate component (must use async instantiation when async support is enabled) - let bindings = SglModelGateway::instantiate_async(&mut store, &component, &linker) - .await - .map_err(|e| { - WasmError::from(WasmRuntimeError::InstanceCreateFailed(e.to_string())) - })?; - - // Call on-response (async call when async support is enabled) - let action_result = bindings - .sgl_model_gateway_middleware_on_response() - .call_on_response(&mut store, &response) - .await - .map_err(|e| map_wasm_error(e, config.max_execution_time_ms))?; - - WasmComponentOutput::MiddlewareAction(action_result) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnError) => { - return Err(WasmError::from(WasmRuntimeError::CallFailed( - "OnError attach point not yet implemented".to_string(), - ))); - } - }; - - Ok(output) - } -} - -impl Drop for WasmThreadPool { - fn drop(&mut self) { - // close sender and receiver - self.sender.close(); - self.receiver.close(); - - // wait for all workers to complete - for worker in self.workers.drain(..) { - let _ = worker.join(); - } - } -} - -#[cfg(test)] -mod tests { - use std::{num::NonZeroUsize, time::Instant}; - - use lru::LruCache; - - use super::*; - use crate::wasm::config::WasmRuntimeConfig; - - #[test] - fn test_get_cpu_info() { - let (cpu_count, max_recommended) = WasmRuntime::get_cpu_info(); - assert!(cpu_count > 0); - assert!(max_recommended > 0); - assert!(max_recommended >= cpu_count); - } - - #[test] - fn test_config_default_values() { - let config = WasmRuntimeConfig::default(); - - assert_eq!(config.max_memory_pages, 1024); - assert_eq!(config.max_execution_time_ms, 1000); - assert_eq!(config.max_stack_size, 1024 * 1024); - assert!(config.thread_pool_size > 0); - assert_eq!(config.module_cache_size, 10); - } - - #[test] - fn test_config_clone() { - let config = WasmRuntimeConfig::default(); - let cloned_config = config.clone(); - - assert_eq!(config.max_memory_pages, cloned_config.max_memory_pages); - assert_eq!( - config.max_execution_time_ms, - cloned_config.max_execution_time_ms - ); - assert_eq!(config.max_stack_size, cloned_config.max_stack_size); - assert_eq!(config.thread_pool_size, cloned_config.thread_pool_size); - assert_eq!(config.module_cache_size, cloned_config.module_cache_size); - } - #[test] - fn test_wasm_instantiation_performance_threshold() -> Result<()> { - // A simple WASM module forcing memory allocation - const WASM_WAT: &str = r#" - (module - (memory (export "memory") 1) - (func (export "run") (param i32 i32) (result i32) - local.get 0 - local.get 1 - i32.add) - ) - "#; - - let iterations = 1000; - - // Scenario A: Baseline (No Pool, No Cache) - let engine_standard = Engine::default(); - let start_standard = Instant::now(); - for _ in 0..iterations { - // Simulate compilation + instantiation overhead - let module = wasmtime::Module::new(&engine_standard, WASM_WAT).unwrap(); - let mut store = Store::new(&engine_standard, ()); - let instance = wasmtime::Instance::new(&mut store, &module, &[]).unwrap(); - let run_func = instance - .get_typed_func::<(i32, i32), i32>(&mut store, "run") - .unwrap(); - let _ = run_func.call(&mut store, (10, 20)).unwrap(); - } - let duration_standard = start_standard.elapsed(); - - // --- Scenario B: Optimized (Pool + Cache) - let mut pool_config = PoolingAllocationConfig::default(); - - pool_config.total_core_instances(100); - - let mut config = Config::new(); - config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool_config)); - - let engine_pooled = Engine::new(&config).unwrap(); - - // Setup LRU Cache - let cache_capacity = NonZeroUsize::new(100).unwrap(); - let mut cache: LruCache, wasmtime::Module> = LruCache::new(cache_capacity); - - // Pre-warm cache (simulating the "cached" state) - let key = WASM_WAT.as_bytes().to_vec(); - let module_compiled = wasmtime::Module::new(&engine_pooled, WASM_WAT).unwrap(); - cache.push(key.clone(), module_compiled); - - let start_pooled = Instant::now(); - for _ in 0..iterations { - let module = cache.get(&key).unwrap().clone(); - let mut store = Store::new(&engine_pooled, ()); - let instance = wasmtime::Instance::new(&mut store, &module, &[]).unwrap(); - let run_func = instance - .get_typed_func::<(i32, i32), i32>(&mut store, "run") - .unwrap(); - let _ = run_func.call(&mut store, (10, 20)).unwrap(); - } - let duration_pooled = start_pooled.elapsed(); - - // Verify Speedup - let standard_secs = duration_standard.as_secs_f64(); - let pooled_secs = duration_pooled.as_secs_f64(); - - if pooled_secs > 0.0 { - let speedup = standard_secs / pooled_secs; - - assert!( - speedup > 5.0, - "Optimization regression: Pooling+Caching was only {:.2}x faster", - speedup - ); - } - - Ok(()) - } -} diff --git a/sgl-model-gateway/src/wasm/spec.rs b/sgl-model-gateway/src/wasm/spec.rs deleted file mode 100644 index 263002b1f..000000000 --- a/sgl-model-gateway/src/wasm/spec.rs +++ /dev/null @@ -1,60 +0,0 @@ -//! WebAssembly Interface Bindings and Type Conversions -//! -//! Contains wasmtime component bindings generated from interface definitions, -//! and helper functions to convert between Axum HTTP types and interface types. - -use axum::http::{header, HeaderMap, HeaderValue}; - -wasmtime::component::bindgen!({ - path: "src/wasm/interface", - world: "sgl-model-gateway", - imports: { default: async | trappable }, - exports: { default: async }, -}); - -/// Build WebAssembly headers from Axum HeaderMap -pub fn build_wasm_headers_from_axum_headers( - headers: &HeaderMap, -) -> Vec { - let mut wasm_headers = Vec::new(); - for (name, value) in headers.iter() { - if let Ok(value_str) = value.to_str() { - wasm_headers.push(sgl::model_gateway::middleware_types::Header { - name: name.as_str().to_string(), - value: value_str.to_string(), - }); - } - } - wasm_headers -} - -/// Apply ModifyAction header modifications to Axum HeaderMap -pub fn apply_modify_action_to_headers( - headers: &mut HeaderMap, - modify: &sgl::model_gateway::middleware_types::ModifyAction, -) { - // Apply headers_set - for header_mod in &modify.headers_set { - if let (Ok(name), Ok(value)) = ( - header_mod.name.parse::(), - header_mod.value.parse::(), - ) { - headers.insert(name, value); - } - } - // Apply headers_add - for header_mod in &modify.headers_add { - if let (Ok(name), Ok(value)) = ( - header_mod.name.parse::(), - header_mod.value.parse::(), - ) { - headers.append(name, value); - } - } - // Apply headers_remove - for name_str in &modify.headers_remove { - if let Ok(name) = name_str.parse::() { - headers.remove(name); - } - } -} diff --git a/sgl-model-gateway/src/wasm/types.rs b/sgl-model-gateway/src/wasm/types.rs deleted file mode 100644 index 1e243311c..000000000 --- a/sgl-model-gateway/src/wasm/types.rs +++ /dev/null @@ -1,102 +0,0 @@ -//! WASM Component Type System -//! -//! Provides generic input/output types for WASM component execution -//! based on attach points. - -use wasmtime::{component::ResourceTable, StoreLimits}; -use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView}; - -use crate::wasm::{ - module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, - spec::sgl::model_gateway::middleware_types, -}; - -/// Generic input type for WASM component execution -/// -/// This enum represents all possible input types that can be passed -/// to a WASM component, determined by the attach_point. -#[derive(Debug, Clone)] -pub enum WasmComponentInput { - /// Middleware OnRequest input - MiddlewareRequest(middleware_types::Request), - /// Middleware OnResponse input - MiddlewareResponse(middleware_types::Response), -} - -/// Generic output type from WASM component execution -/// -/// This enum represents all possible output types that can be returned -/// from a WASM component, determined by the attach_point. -#[derive(Debug, Clone)] -pub enum WasmComponentOutput { - /// Middleware Action output - MiddlewareAction(middleware_types::Action), -} - -impl WasmComponentInput { - /// Create input based on attach_point and raw data - /// - /// This helper function validates that the attach_point matches - /// the expected input type. - pub fn from_attach_point(attach_point: &WasmModuleAttachPoint) -> Result { - match attach_point { - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest) => { - // OnRequest expects a Request type, but we can't construct it here - // The caller should use MiddlewareRequest variant directly - Err("OnRequest requires MiddlewareRequest input. Use WasmComponentInput::MiddlewareRequest directly.".to_string()) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse) => { - // OnResponse expects a Response type - Err("OnResponse requires MiddlewareResponse input. Use WasmComponentInput::MiddlewareResponse directly.".to_string()) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnError) => { - Err("OnError attach point not yet implemented".to_string()) - } - } - } - - /// Get the expected attach_point for this input type - pub fn expected_attach_point(&self) -> WasmModuleAttachPoint { - match self { - WasmComponentInput::MiddlewareRequest(_) => { - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest) - } - WasmComponentInput::MiddlewareResponse(_) => { - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse) - } - } - } -} - -impl WasmComponentOutput { - /// Get the attach_point that produced this output type - pub fn from_attach_point(attach_point: &WasmModuleAttachPoint) -> Result { - match attach_point { - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest) => { - // This would be set after execution - Err("Cannot create output before execution".to_string()) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse) => { - Err("Cannot create output before execution".to_string()) - } - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnError) => { - Err("OnError attach point not yet implemented".to_string()) - } - } - } -} - -pub struct WasiState { - pub ctx: WasiCtx, - pub table: ResourceTable, - pub limits: StoreLimits, -} - -impl WasiView for WasiState { - fn ctx(&mut self) -> WasiCtxView<'_> { - WasiCtxView { - ctx: &mut self.ctx, - table: &mut self.table, - } - } -} diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index e397d5a20..008c53791 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -744,7 +744,7 @@ async fn test_wasm_module_execution() { // Execute the module use smg::wasm::{ - spec::sgl::model_gateway::middleware_types, + spec::smg::gateway::middleware_types, types::{WasmComponentInput, WasmComponentOutput}, };