sgl-router: experimental Rust HTTP router for SGLang worker pools (#25851)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-25 15:34:05 +08:00
parent aae04b1241
commit 6e8fe176be
131 changed files with 28623 additions and 55 deletions
@@ -0,0 +1,115 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Tests that the router does not wedge indefinitely when an upstream
+//! worker accepts the TCP connection but never sends response headers.
+//!
+//! Without a configured `.timeout(...)` on the reqwest client, a stalled
+//! backend hangs the axum handler future forever and the test harness
+//! would just timeout. We assert here that the router returns a fast,
+//! clean 502 (`upstream_timeout`) instead.
+
+use axum::body::Body;
+use axum::http::{Request, StatusCode};
+use http_body_util::BodyExt;
+use sgl_router::config::{
+    ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
+    PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
+};
+use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
+use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
+use sgl_router::proxy::Proxy;
+use sgl_router::server::app::build_router;
+use sgl_router::server::app_context::AppContext;
+use sgl_router::tokenizer::TokenizerRegistry;
+use sgl_router::workers::WorkerRegistry;
+use std::sync::Arc;
+use std::time::Duration;
+use tower::ServiceExt;
+
+fn config(_worker_url: &str) -> Config {
+    Config {
+        server: ServerConfig {
+            host: "0".into(),
+            port: 0,
+        },
+        observability: ObservabilityConfig::default(),
+        models: vec![ModelConfig {
+            id: "tiny".into(),
+            tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
+            policy: PolicyKind::RoundRobin,
+            circuit_breaker: None,
+            cache_aware: None,
+        }],
+        discovery: DiscoveryConfig {
+            backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
+                urls: vec!["http://placeholder:0".into()],
+            }),
+        },
+        proxy: ProxyConfig::default(),
+        active_load: ActiveLoadConfig::default(),
+    }
+}
+
+#[tokio::test]
+async fn non_streaming_request_times_out_when_worker_hangs() {
+    // Worker accepts and then sleeps for 5s; router timeout is 200ms.
+    let worker =
+        crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(5)).await;
+    let cfg = config(&worker.url);
+    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
+    let registry = Arc::new(WorkerRegistry::default());
+    let _ = registry.add(WorkerSpec {
+        id: WorkerId("w1".into()),
+        url: worker.url.clone(),
+        mode: WorkerMode::Plain,
+        model_ids: vec![ModelId("tiny".into())],
+        bootstrap_port: None,
+    });
+    let policies = Arc::new(build_policy_registry(&cfg).unwrap());
+    let proxy = Arc::new(Proxy::new(Duration::from_millis(200)).unwrap());
+    let ctx = Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies));
+    let app = build_router(ctx);
+
+    let req = Request::builder()
+        .method("POST")
+        .uri("/v1/chat/completions")
+        .header("content-type", "application/json")
+        .body(Body::from(
+            serde_json::to_vec(&serde_json::json!({
+                "model": "tiny",
+                "messages": [{"role": "user", "content": "hi"}],
+                "stream": false
+            }))
+            .unwrap(),
+        ))
+        .unwrap();
+
+    let started = std::time::Instant::now();
+    // Outer guard so a regression doesn't wedge CI forever.
+    let res = tokio::time::timeout(Duration::from_secs(2), app.oneshot(req))
+        .await
+        .expect("router must return within 2s when proxy timeout is 200ms")
+        .unwrap();
+    let elapsed = started.elapsed();
+    assert!(
+        elapsed < Duration::from_secs(1),
+        "router must short-circuit on upstream timeout; elapsed {elapsed:?}"
+    );
+    assert_eq!(res.status(), StatusCode::BAD_GATEWAY);
+    assert_eq!(
+        res.headers().get("x-router-error-code").unwrap(),
+        "upstream_timeout"
+    );
+    let bytes = res.into_body().collect().await.unwrap().to_bytes();
+    let body_str = String::from_utf8_lossy(&bytes);
+    assert!(
+        body_str.contains("\"code\":\"upstream_timeout\""),
+        "body: {body_str}"
+    );
+    // No leak of worker URL or reqwest source chain to the client.
+    assert!(
+        !body_str.contains(&worker.url),
+        "worker URL must not leak in client-visible body: {body_str}"
+    );
+}