diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index c986ec881..1c178fc49 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -149,6 +149,10 @@ tonic-v12 = { version = "0.12.3", package = "tonic" } serial_test = "3.0" rsa = { version = "0.9", features = ["sha2"] } +[[bench]] +name = "wasm_middleware_latency" +harness = false +path = "benches/wasm_middleware_latency.rs" [[bench]] name = "request_processing" harness = false diff --git a/sgl-model-gateway/benches/wasm_middleware_latency.rs b/sgl-model-gateway/benches/wasm_middleware_latency.rs new file mode 100644 index 000000000..9f33cb951 --- /dev/null +++ b/sgl-model-gateway/benches/wasm_middleware_latency.rs @@ -0,0 +1,108 @@ +use std::sync::Arc; + +use axum::{ + body::Body, + http::{HeaderMap, Request, Response, StatusCode}, + middleware, + response::IntoResponse, +}; +use criterion::{criterion_group, criterion_main, Criterion}; +use http_body_util::BodyExt; +use smg::{ + app_context::AppContext, config::RouterConfig, middleware::wasm_middleware, + protocols::chat::ChatCompletionRequest, routers::RouterTrait, server::AppState, +}; +use tokio::runtime::Runtime; +use tower::{Layer, Service}; + +#[derive(Debug)] +struct MockRouter; + +#[async_trait::async_trait] +impl RouterTrait for MockRouter { + fn as_any(&self) -> &dyn std::any::Any { + self + } + async fn route_chat( + &self, + _headers: Option<&HeaderMap>, + _body: &ChatCompletionRequest, + _model_id: Option<&str>, + ) -> Response { + StatusCode::OK.into_response() + } + fn router_type(&self) -> &'static str { + "mock" + } +} + +/// Mock service that simulates a streaming response with a 500ms delay. +async fn mock_next_streaming(_req: Request) -> Response { + let (tx, rx) = tokio::sync::mpsc::channel(16); + + tokio::spawn(async move { + // Send first chunk immediately + let _ = tx + .send(Ok::<_, std::io::Error>(bytes::Bytes::from("chunk 1 "))) + .await; + // Simulate generation delay + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + // Send final chunk + let _ = tx + .send(Ok::<_, std::io::Error>(bytes::Bytes::from("chunk 2"))) + .await; + }); + + Response::new(Body::from_stream( + tokio_stream::wrappers::ReceiverStream::new(rx), + )) +} + +fn bench_wasm_middleware_buffering(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + + // Setup AppContext with WASM enabled + let config = RouterConfig::builder().enable_wasm(true).build_unchecked(); + + let context = rt.block_on(AppContext::from_config(config, 30)).unwrap(); + let app_state = Arc::new(AppState { + router: Arc::new(MockRouter), + context: Arc::new(context), + concurrency_queue_tx: None, + router_manager: None, + }); + + c.bench_function("wasm_middleware_pre_fix_latency", |b| { + b.iter(|| { + rt.block_on(async { + let req = Request::builder() + .uri("/v1/chat/completions") + .body(Body::empty()) + .unwrap(); + + // Create the service by applying the middleware layer to the mock streamer + let mut service = + middleware::from_fn_with_state(app_state.clone(), wasm_middleware).layer( + tower::service_fn(|req: Request| async move { + Ok::<_, std::convert::Infallible>(mock_next_streaming(req).await) + }), + ); + + // Explicitly poll the service + let response: Response = + service.call(req).await.expect("Middleware service failed"); + + // Measure how long it takes to receive the FIRST frame + let mut body = response.into_body(); + let _first_frame = body.frame().await; + }); + }); + }); +} + +criterion_group! { + name = benches; + config = Criterion::default().sample_size(10); + targets = bench_wasm_middleware_buffering +} +criterion_main!(benches); diff --git a/sgl-model-gateway/src/middleware.rs b/sgl-model-gateway/src/middleware.rs index cc4ee546c..876e18a8f 100644 --- a/sgl-model-gateway/src/middleware.rs +++ b/sgl-model-gateway/src/middleware.rs @@ -773,95 +773,91 @@ pub async fn wasm_middleware( } }; - // Extract request body once before processing modules - let method = request.method().clone(); - let uri = request.uri().clone(); - let mut headers = request.headers().clone(); - let max_body_size = wasm_manager.get_max_body_size(); - let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await { - Ok(bytes) => bytes.to_vec(), - Err(e) => { - error!("Failed to read request body: {}", e); - // Create a minimal request with empty body for error recovery - let error_request = Request::builder() - .uri(uri) - .body(Body::empty()) - .unwrap_or_else(|_| Request::new(Body::empty())); - return Ok(next.run(error_request).await); - } - }; - - // Process each OnRequest module - let mut modified_body = body_bytes; - - // Pre-compute strings once before the loop to avoid repeated allocations - let method_str = method.to_string(); - let path_str = uri.path().to_string(); - let query_str = uri.query().unwrap_or("").to_string(); - - for module in modules_on_request { - // Build WebAssembly request from collected data - let wasm_headers = build_wasm_headers_from_axum_headers(&headers); - let wasm_request = WasmRequest { - method: method_str.clone(), - path: path_str.clone(), - query: query_str.clone(), - headers: wasm_headers, - body: modified_body.clone(), - request_id: request_id.clone(), - now_epoch_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_else(|_| { - // Fallback to 0 if system time is before UNIX_EPOCH - // This should never happen in practice, but provides a safe fallback - Duration::from_millis(0) - }) - .as_millis() as u64, + let response = if modules_on_request.is_empty() { + next.run(request).await + } else { + // Extract request body once before processing modules + let method = request.method().clone(); + let uri = request.uri().clone(); + let mut headers = request.headers().clone(); + let max_body_size = wasm_manager.get_max_body_size(); + let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await { + Ok(bytes) => bytes.to_vec(), + Err(e) => { + error!("Failed to read request body: {}", e); + // Create a minimal request with empty body for error recovery + let error_request = Request::builder() + .uri(uri) + .body(Body::empty()) + .unwrap_or_else(|_| Request::new(Body::empty())); + return Ok(next.run(error_request).await); + } }; - // Execute WASM component - let action = match wasm_manager - .execute_module_for_attach_point( - &module, - on_request_attach_point.clone(), - WasmComponentInput::MiddlewareRequest(wasm_request), - ) - .await - { - Some(action) => action, - None => continue, // Continue to next module on error - }; + // Process each OnRequest module + let mut modified_body = body_bytes; - // Process action - match action { - Action::Continue => { - // Continue to next module or request processing - } - Action::Reject(status) => { - // Immediately reject the request - return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST)); - } - Action::Modify(modify) => { - // Apply modifications to headers and body - apply_modify_action_to_headers(&mut headers, &modify); - // Apply body_replace - if let Some(body_bytes) = modify.body_replace { - modified_body = body_bytes; + // Pre-compute strings once before the loop to avoid repeated allocations + let method_str = method.to_string(); + let path_str = uri.path().to_string(); + let query_str = uri.query().unwrap_or("").to_string(); + + for module in modules_on_request { + // Build WebAssembly request from collected data + let wasm_headers = build_wasm_headers_from_axum_headers(&headers); + let wasm_request = WasmRequest { + method: method_str.clone(), + path: path_str.clone(), + query: query_str.clone(), + headers: wasm_headers, + body: modified_body.clone(), + request_id: request_id.clone(), + now_epoch_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_millis(0)) + .as_millis() as u64, + }; + + // Execute WASM component + let action = match wasm_manager + .execute_module_for_attach_point( + &module, + on_request_attach_point.clone(), + WasmComponentInput::MiddlewareRequest(wasm_request), + ) + .await + { + Some(action) => action, + None => continue, // Continue to next module on error + }; + + // Process action + match action { + Action::Continue => {} + Action::Reject(status) => { + return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST)); + } + Action::Modify(modify) => { + // Apply modifications to headers and body + apply_modify_action_to_headers(&mut headers, &modify); + if let Some(body_bytes) = modify.body_replace { + modified_body = body_bytes; + } } } } - } - // Reconstruct request with modifications - let mut final_request = Request::builder() - .method(method) - .uri(uri) - .body(Body::from(modified_body)) - .unwrap_or_else(|_| Request::new(Body::empty())); - *final_request.headers_mut() = headers; + // Reconstruct request with modifications + let mut final_request = Request::builder() + .method(method) + .uri(uri) + .body(Body::from(modified_body)) + .unwrap_or_else(|_| Request::new(Body::empty())); + *final_request.headers_mut() = headers; - // Continue with request processing - let response = next.run(final_request).await; + // Continue with request processing + next.run(final_request).await + }; // ===== OnResponse Phase ===== let on_response_attach_point = @@ -875,10 +871,13 @@ pub async fn wasm_middleware( return Ok(response); } }; - + if modules_on_response.is_empty() { + return Ok(response); + } // Extract response data once before processing modules let mut status = response.status(); let mut headers = response.headers().clone(); + let max_body_size = wasm_manager.get_max_body_size(); let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await { Ok(bytes) => bytes.to_vec(), Err(e) => {