chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

remove self managed wasm as it has been replaced with official smg wa… (#17746)

Browse Source
This commit is contained in:
Simo Lin
2026-01-26 01:56:33 -05:00
committed by GitHub
parent 30b3192039
commit ed75136e85
14 changed files with 3 additions and 2247 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-model-gateway/Cargo.toml
View File

@@ -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"

2
sgl-model-gateway/src/middleware.rs
View File

@@ -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,
},
},

227
sgl-model-gateway/src/wasm/README.md
View File

@@ -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);
```

353
sgl-model-gateway/src/wasm/config.rs
View File

@@ -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<Self, String> {
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"));
}
}

117
sgl-model-gateway/src/wasm/errors.rs
View File

@@ -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<T> = std::result::Result<T, WasmError>;
/// 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<wasmtime::Error> for WasmError {
fn from(value: wasmtime::Error) -> Self {
WasmError::Runtime(WasmRuntimeError::CallFailed(value.to_string()))
}
}

54
sgl-model-gateway/src/wasm/interface/spec.wit
View File

@@ -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<header>,
body: list<u8>,
request-id: string,
now-epoch-ms: u64,
}
// onResponse
record response {
status: u16,
headers: list<header>,
body: list<u8>,
}
// modify action
record modify-action {
status: option<u16>,
headers-set: list<header>,
headers-add: list<header>,
headers-remove: list<string>,
body-replace: option<list<u8>>,
}
// 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;
}

13
sgl-model-gateway/src/wasm/mod.rs
View File

@@ -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;

209
sgl-model-gateway/src/wasm/module.rs
View File

@@ -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<S>(hash: &[u8; 32], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let hex_string = hash
.iter()
.map(|b| format!("{:02x}", b))
.collect::<String>();
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<S>(timestamp: &u64, serializer: S) -> Result<S::Ok, S::Error>
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::<Utc>::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<u64, D::Error>
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(&timestamp_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::<u64>()
.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<WasmModuleAttachPoint>,
pub add_result: Option<WasmModuleAddResult>,
}
#[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<WasmModuleAttachPoint>,
// Pre-loaded WASM component bytes (loaded into memory for faster execution)
#[serde(skip)]
pub wasm_bytes: Vec<u8>,
}
#[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<WasmModuleDescriptor>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmModuleAddResponse {
pub modules: Vec<WasmModuleDescriptor>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmModuleListResponse {
pub modules: Vec<WasmModule>,
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<f64>,
}

267
sgl-model-gateway/src/wasm/module_manager.rs
View File

@@ -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<RwLock<HashMap<Uuid, WasmModule>>>,
runtime: Arc<WasmRuntime>,
// 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<Self> {
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> {
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<Vec<WasmModule>> {
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<Option<WasmModule>> {
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<Vec<WasmModule>> {
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<Vec<WasmModule>> {
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<WasmRuntime> {
&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<WasmComponentOutput> {
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<WasmComponentOutput> {
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<crate::wasm::spec::sgl::model_gateway::middleware_types::Action> {
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.")
}
}

228
sgl-model-gateway/src/wasm/route.rs
View File

@@ -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<String, String> {
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<Arc<AppState>>,
Json(config): Json<WasmModuleAddRequest>,
) -> 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<Arc<AppState>>,
Path(module_uuid_str): Path<String>,
) -> 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<Arc<AppState>>) -> 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()
}
}

615
sgl-model-gateway/src/wasm/runtime.rs
View File

@@ -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<WasmThreadPool>,
// 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<WasmTask>,
receiver: async_channel::Receiver<WasmTask>,
workers: Vec<std::thread::JoinHandle<()>>,
// Metrics
total_tasks: AtomicU64,
completed_tasks: AtomicU64,
failed_tasks: AtomicU64,
}
pub enum WasmTask {
ExecuteComponent {
wasm_bytes: Vec<u8>,
wasm_hash: [u8; 32],
attach_point: WasmModuleAttachPoint,
input: WasmComponentInput,
response: oneshot::Sender<Result<WasmComponentOutput>>,
},
}
impl WasmRuntime {
pub fn new(config: WasmRuntimeConfig) -> Result<Self> {
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> {
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<u8>,
wasm_hash: [u8; 32],
attach_point: WasmModuleAttachPoint,
input: WasmComponentInput,
) -> Result<WasmComponentOutput> {
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::<wasmtime::Trap>() == 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<Self> {
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<WasmTask>,
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<u8>,
wasm_hash: [u8; 32],
attach_point: WasmModuleAttachPoint,
input: WasmComponentInput,
config: &WasmRuntimeConfig,
) -> Result<WasmComponentOutput> {
// 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::<WasiState>::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<Vec<u8>, 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(())
}
}

60
sgl-model-gateway/src/wasm/spec.rs
View File

@@ -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<sgl::model_gateway::middleware_types::Header> {
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::HeaderName>(),
header_mod.value.parse::<HeaderValue>(),
) {
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::HeaderName>(),
header_mod.value.parse::<HeaderValue>(),
) {
headers.append(name, value);
}
}
// Apply headers_remove
for name_str in &modify.headers_remove {
if let Ok(name) = name_str.parse::<header::HeaderName>() {
headers.remove(name);
}
}
}

102
sgl-model-gateway/src/wasm/types.rs
View File

@@ -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<Self, String> {
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<Self, String> {
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,
}
}
}

2
sgl-model-gateway/tests/wasm_test.rs
View File

@@ -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},
};