minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

feature: upstream cancel (#19524)

Browse Source 
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Douglas Yang
2026-05-18 16:20:09 +08:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 4.7
parent 54eb2904a4
commit e5589843a3
13 changed files with 4206 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/tests/common/mock_worker.rs
+528 -23
View File
@@ -20,7 +20,7 @@ use axum::{
};
use futures_util::stream::{self, StreamExt};
use serde_json::json;
use tokio::sync::RwLock;
use tokio::sync::{Notify, RwLock};
use uuid::Uuid;
/// Configuration for mock worker behavior
@@ -158,6 +158,15 @@ async fn should_fail(config: &MockWorkerConfig) -> bool {
rand::random::<f32>() < config.fail_rate
}
/// Pick the HTTP status used when `should_fail` triggers. Defaults to 500
/// for backwards compatibility; tests can override via
/// [`set_fail_status_code`] to exercise 4xx/non-5xx failure paths.
fn fail_status_code(port: u16) -> StatusCode {
get_fail_status_code_for_port(port)
.and_then(|s| StatusCode::from_u16(s).ok())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR)
}
async fn health_handler(State(config): State<Arc<RwLock<MockWorkerConfig>>>) -> Response {
let config = config.read().await;
@@ -303,7 +312,7 @@ async fn generate_handler(
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
fail_status_code(config.port),
[("x-worker-id", worker_id)],
Json(json!({
"error": "Random failure for testing"
@@ -324,6 +333,71 @@ async fn generate_handler(
if is_stream {
let stream_delay = config.response_delay_ms;
if let Some(num_chunks) = get_slow_stream_chunks_for_port(config.port) {
let port = config.port;
let delay_ms = stream_delay;
let error_after = get_stream_error_after_for_port(port);
init_stream_tracking(port, num_chunks);
let (tx, rx) =
tokio::sync::mpsc::channel::<Result<Event, std::io::Error>>(MOCK_STREAM_BUFFER);
tokio::spawn(async move {
let _exit_guard = install_stream_exit_notifier(port);
let timestamp_start = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs_f64();
for i in 0..num_chunks {
if let Some(n) = error_after {
if i == n {
let _ = tx
.send(Err(std::io::Error::other(
"simulated upstream worker crash",
)))
.await;
return;
}
}
if delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
}
let data = json!({
"text": format!("chunk-{} ", i),
"meta_info": {
"prompt_tokens": 10,
"completion_tokens": (i + 1) as u64,
"completion_tokens_wo_jump_forward": (i + 1) as u64,
"input_token_logprobs": null,
"output_token_logprobs": null,
"first_token_latency": delay_ms as f64 / 1000.0,
"time_to_first_token": delay_ms as f64 / 1000.0,
"time_per_output_token": 0.01,
"start_time": timestamp_start,
"finish_reason": null
},
"stage": "mid"
});
if tx
.send(Ok(Event::default().data(data.to_string())))
.await
.is_err()
{
return;
}
record_chunk_sent(port);
}
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
mark_stream_completed(port);
});
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
return (
[("x-worker-id", worker_id)],
Sse::new(stream).keep_alive(KeepAlive::default()),
)
.into_response();
}
// Check if it's a batch request
let is_batch = payload.get("text").and_then(|t| t.as_array()).is_some();
@@ -442,28 +516,99 @@ async fn chat_completions_handler(
if is_stream {
let request_id = format!("chatcmpl-{}", Uuid::new_v4());
let stream = stream::once(async move {
let chunk = json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": timestamp,
"model": "mock-model",
"choices": [{
"index": 0,
"delta": {
"content": "This is a mock chat response."
},
"finish_reason": null
}]
// Check for slow streaming mode (used by upstream cancel tests).
// Reads from the global SLOW_STREAM_CONFIG (set via set_slow_stream_chunks)
// rather than the payload, because the gateway deserializes/re-serializes
// the request body and drops unknown fields.
let slow_chunks = get_slow_stream_chunks_for_port(config.port);
if let Some(num_chunks) = slow_chunks {
let port = config.port;
let delay_ms = config.response_delay_ms;
let error_after = get_stream_error_after_for_port(port);
init_stream_tracking(port, num_chunks);
// Small bounded capacity gives a bit of slack between the producer
// task and the SSE consumer; on receiver drop, send().await
// returns Err and the loop exits regardless of capacity.
let (tx, rx) =
tokio::sync::mpsc::channel::<Result<Event, std::io::Error>>(MOCK_STREAM_BUFFER);
tokio::spawn(async move {
let _exit_guard = install_stream_exit_notifier(port);
for i in 0..num_chunks {
if let Some(n) = error_after {
if i == n {
// Inject a transport-level error to exercise the
// gateway's `Some(Err(_))` arm.
let _ = tx
.send(Err(std::io::Error::other(
"simulated upstream worker crash",
)))
.await;
return;
}
}
if delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
}
let chunk = json!({
"id": &request_id,
"object": "chat.completion.chunk",
"created": timestamp,
"model": "mock-model",
"choices": [{
"index": 0,
"delta": {
"content": format!("chunk-{} ", i)
},
"finish_reason": null
}]
});
if tx
.send(Ok(Event::default().data(chunk.to_string())))
.await
.is_err()
{
// Client disconnected, stream was cancelled
return;
}
record_chunk_sent(port);
}
// Send [DONE]
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
mark_stream_completed(port);
});
Ok::<_, Infallible>(Event::default().data(chunk.to_string()))
})
.chain(stream::once(async { Ok(Event::default().data("[DONE]")) }));
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
} else {
let stream = stream::once(async move {
let chunk = json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": timestamp,
"model": "mock-model",
"choices": [{
"index": 0,
"delta": {
"content": "This is a mock chat response."
},
"finish_reason": null
}]
});
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
Ok::<_, Infallible>(Event::default().data(chunk.to_string()))
})
.chain(stream::once(async { Ok(Event::default().data("[DONE]")) }));
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
}
} else {
Json(json!({
"id": format!("chatcmpl-{}", Uuid::new_v4()),
@@ -793,8 +938,13 @@ async fn responses_handler(
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
} else if has_tools && has_function_output {
// Second turn: emit streaming text response
} else if has_tools
&& has_function_output
&& get_slow_stream_chunks_for_port(config.port).is_none()
{
// Second turn: emit streaming text response.
// If slow-stream is configured, fall through to the slow-stream
// branch below so cancel tests can disconnect mid second-turn.
let rid = request_id.clone();
let msg_id = format!(
"msg_{}",
@@ -941,6 +1091,126 @@ async fn responses_handler(
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
} else if let Some(num_chunks) = get_slow_stream_chunks_for_port(config.port) {
// Slow-stream mode for /responses cancel tests. Mirrors the
// chat-completions slow-stream path so the same set_slow_stream_chunks
// helper drives both endpoints.
let port = config.port;
let delay_ms = config.response_delay_ms;
let error_after = get_stream_error_after_for_port(port);
let rid = request_id.clone();
let msg_id = format!(
"msg_{}",
Uuid::new_v4().to_string().split('-').next().unwrap()
);
init_stream_tracking(port, num_chunks);
let (tx, rx) =
tokio::sync::mpsc::channel::<Result<Event, std::io::Error>>(MOCK_STREAM_BUFFER);
tokio::spawn(async move {
let _exit_guard = install_stream_exit_notifier(port);
// Emit response.created and response.in_progress so the
// gateway's /responses persistence accumulator has the
// structural events it expects.
let created = Event::default().event("response.created").data(
json!({
"type": "response.created",
"response": {
"id": rid.clone(),
"object": "response",
"created_at": timestamp,
"model": "mock-model",
"status": "in_progress"
}
})
.to_string(),
);
if tx.send(Ok(created)).await.is_err() {
return;
}
let in_progress = Event::default().event("response.in_progress").data(
json!({
"type": "response.in_progress",
"response": {
"id": rid.clone(),
"object": "response",
"created_at": timestamp,
"model": "mock-model",
"status": "in_progress"
}
})
.to_string(),
);
if tx.send(Ok(in_progress)).await.is_err() {
return;
}
for i in 0..num_chunks {
if let Some(n) = error_after {
if i == n {
let _ = tx
.send(Err(std::io::Error::other(
"simulated upstream worker crash",
)))
.await;
return;
}
}
if delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
}
let delta = Event::default().event("response.output_text.delta").data(
json!({
"type": "response.output_text.delta",
"output_index": 0,
"content_index": 0,
"item_id": msg_id.clone(),
"delta": format!("chunk-{} ", i)
})
.to_string(),
);
if tx.send(Ok(delta)).await.is_err() {
return;
}
record_chunk_sent(port);
}
let aggregated_text = (0..num_chunks)
.map(|i| format!("chunk-{} ", i))
.collect::<String>();
let completed = Event::default().event("response.completed").data(
json!({
"type": "response.completed",
"response": {
"id": rid,
"object": "response",
"created_at": timestamp,
"model": "mock-model",
"status": "completed",
"output": [{
"id": msg_id,
"type": "message",
"role": "assistant",
"content": [{
"type": "output_text",
"text": aggregated_text
}]
}]
}
})
.to_string(),
);
let _ = tx.send(Ok(completed)).await;
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
mark_stream_completed(port);
});
let stream = tokio_stream::wrappers::ReceiverStream::new(rx);
Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response()
} else {
// Default streaming response
let stream = stream::once(async move {
@@ -1195,6 +1465,241 @@ async fn responses_cancel_handler(
}
}
// --- Slow-stream configuration (for upstream cancel tests) ---
// Configured via a global map keyed by worker port so that tests
// can enable slow streaming WITHOUT relying on the request payload
// (the gateway deserializes/re-serializes the body, dropping unknown fields).
static SLOW_STREAM_CONFIG: OnceLock<Mutex<HashMap<u16, usize>>> = OnceLock::new();
fn get_slow_stream_config() -> &'static Mutex<HashMap<u16, usize>> {
SLOW_STREAM_CONFIG.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Configure a worker (by port) to send `num_chunks` chunks with
/// `response_delay_ms` between each when handling a streaming request.
/// Call this before making the request through the gateway.
pub fn set_slow_stream_chunks(port: u16, num_chunks: usize) {
let mut map = get_slow_stream_config().lock().unwrap();
map.insert(port, num_chunks);
}
/// Clear slow-stream configuration for a worker port.
pub fn clear_slow_stream_chunks(port: u16) {
let mut map = get_slow_stream_config().lock().unwrap();
map.remove(&port);
}
fn get_slow_stream_chunks_for_port(port: u16) -> Option<usize> {
let map = get_slow_stream_config().lock().unwrap();
map.get(&port).copied()
}
// --- Stream error injection (for upstream cancel + error tests) ---
// When set for `port`, the slow-stream producer emits an io::Error to the
// SSE stream after the configured number of successfully-sent chunks.
// reqwest will surface this as a transport error, which exercises the
// gateway's `Some(Err(_))` arm.
static STREAM_ERROR_AFTER: OnceLock<Mutex<HashMap<u16, usize>>> = OnceLock::new();
fn get_stream_error_after_config() -> &'static Mutex<HashMap<u16, usize>> {
STREAM_ERROR_AFTER.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Configure a worker (by port) to abort its SSE stream with an error
/// after sending `n` chunks. Must be combined with
/// [`set_slow_stream_chunks`] to take effect.
pub fn set_stream_error_after_chunks(port: u16, n: usize) {
let mut map = get_stream_error_after_config().lock().unwrap();
map.insert(port, n);
}
/// Clear error-injection configuration for a worker port.
pub fn clear_stream_error_after_chunks(port: u16) {
let mut map = get_stream_error_after_config().lock().unwrap();
map.remove(&port);
}
fn get_stream_error_after_for_port(port: u16) -> Option<usize> {
let map = get_stream_error_after_config().lock().unwrap();
map.get(&port).copied()
}
// --- Failure-status override (for breaker attribution tests) ---
// When set for `port`, `should_fail`-triggered failures return this HTTP
// status instead of the default 500. Lets a test pin breaker semantics for
// the 4xx-from-worker case (the gateway treats 4xx as "not a worker fault")
// without having to fabricate a separate mock worker.
static FAIL_STATUS_CODE: OnceLock<Mutex<HashMap<u16, u16>>> = OnceLock::new();
fn get_fail_status_code_config() -> &'static Mutex<HashMap<u16, u16>> {
FAIL_STATUS_CODE.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Configure a worker (by port) to return `status` when `fail_rate`
/// triggers a failure response, instead of the default 500.
pub fn set_fail_status_code(port: u16, status: u16) {
let mut map = get_fail_status_code_config().lock().unwrap();
map.insert(port, status);
}
/// Clear failure-status override for a worker port.
pub fn clear_fail_status_code(port: u16) {
let mut map = get_fail_status_code_config().lock().unwrap();
map.remove(&port);
}
fn get_fail_status_code_for_port(port: u16) -> Option<u16> {
let map = get_fail_status_code_config().lock().unwrap();
map.get(&port).copied()
}
// --- Stream cancellation tracking (for upstream cancel tests) ---
/// Tracks the state of a streaming response for cancel verification.
#[derive(Clone, Debug)]
pub struct StreamTrackingState {
pub total_chunks: usize,
pub chunks_sent: usize,
pub completed: bool,
}
static STREAM_CANCEL_TRACKER: OnceLock<Mutex<HashMap<u16, StreamTrackingState>>> = OnceLock::new();
fn get_stream_tracker() -> &'static Mutex<HashMap<u16, StreamTrackingState>> {
STREAM_CANCEL_TRACKER.get_or_init(|| Mutex::new(HashMap::new()))
}
// Per-port `Notify` fired when the worker's producer task exits (either
// because its outbound `send().await` failed — i.e. the gateway dropped
// the upstream connection — or because the stream completed naturally).
// Tests await this notification instead of polling counters, so cancel
// assertions don't depend on timing windows.
static STREAM_FINISH_NOTIFIERS: OnceLock<Mutex<HashMap<u16, Arc<Notify>>>> = OnceLock::new();
fn get_stream_finish_notifier_map() -> &'static Mutex<HashMap<u16, Arc<Notify>>> {
STREAM_FINISH_NOTIFIERS.get_or_init(|| Mutex::new(HashMap::new()))
}
fn get_stream_finish_notifier(port: u16) -> Arc<Notify> {
let mut map = get_stream_finish_notifier_map().lock().unwrap();
map.entry(port)
.or_insert_with(|| Arc::new(Notify::new()))
.clone()
}
/// Bound on the per-stream mpsc buffer used by every slow-stream producer
/// task in this mock worker. Tests assert that `chunks_sent` after a cancel
/// grew by at most this many over the pre-drop snapshot, on the theory that
/// anything more means the gateway did not propagate the disconnect
/// upstream.
pub const MOCK_STREAM_BUFFER: usize = 4;
/// RAII guard that fires the per-port finish notifier on drop, so the
/// notification fires whether the producer task exits normally or returns
/// early on `tx.send(...).await.is_err()`.
#[must_use = "StreamExitNotifier must be bound to a local (typically `_exit_guard`) \
and held until the producer task ends — dropping it immediately fires \
the notifier early, causing `wait_for_stream_finish` to return before \
the producer has actually exited"]
pub struct StreamExitNotifier(Arc<Notify>);
impl Drop for StreamExitNotifier {
fn drop(&mut self) {
self.0.notify_one();
}
}
/// Install the exit notifier inside a producer task. Hold the returned
/// guard until the task ends (typically by binding it to `_exit_guard`).
#[must_use = "the returned guard fires the exit notifier on drop; bind it to a local \
(e.g. `let _exit_guard = install_stream_exit_notifier(port);`) so it lives \
for the producer task's lifetime"]
pub fn install_stream_exit_notifier(port: u16) -> StreamExitNotifier {
StreamExitNotifier(get_stream_finish_notifier(port))
}
/// Reset the stream tracker for a given port before starting a new test.
/// Also replaces the finish notifier so any unconsumed permit from a
/// previous test doesn't satisfy this test's wait immediately.
pub fn reset_stream_tracker(port: u16) {
let mut map = get_stream_tracker().lock().unwrap();
map.remove(&port);
let mut nmap = get_stream_finish_notifier_map().lock().unwrap();
nmap.insert(port, Arc::new(Notify::new()));
}
/// Get the stream tracking state for a given port.
pub fn get_stream_tracking_state(port: u16) -> Option<StreamTrackingState> {
let map = get_stream_tracker().lock().unwrap();
map.get(&port).cloned()
}
/// Wait until the worker's producer task for `port` exits — either because
/// the gateway dropped the upstream connection (`send().await` failed) or
/// because the stream finished naturally. Returns the final tracking state.
/// The `timeout` is a safety net for hung tests; a healthy run returns the
/// instant the producer task drops its exit guard.
///
/// **Precondition:** call [`reset_stream_tracker`] before issuing the
/// gateway request whose producer you intend to wait on. The reset
/// installs a fresh `Notify` so a stale permit left by a previous test
/// on the same port can't satisfy this wait immediately.
pub async fn wait_for_stream_finish(
port: u16,
timeout: tokio::time::Duration,
) -> Option<StreamTrackingState> {
let notifier = get_stream_finish_notifier(port);
if tokio::time::timeout(timeout, notifier.notified())
.await
.is_err()
{
// A hung producer would silently look like a successful cancel
// (chunks_sent < total_chunks, completed=false) if we just
// returned what we have. Panic instead so the test fails loudly.
panic!(
"wait_for_stream_finish timed out after {:?} for port {} — \
producer task never fired its exit notifier. Last tracker \
state: {:?}",
timeout,
port,
get_stream_tracking_state(port)
);
}
get_stream_tracking_state(port)
}
// Initialize tracking for a new stream. `map.insert` overwrites any prior
// entry for this port, so callers don't need to reset first; we still expose
// `reset_stream_tracker` so tests can opt into removing the entry entirely.
fn init_stream_tracking(port: u16, total_chunks: usize) {
let mut map = get_stream_tracker().lock().unwrap();
map.insert(
port,
StreamTrackingState {
total_chunks,
chunks_sent: 0,
completed: false,
},
);
}
fn record_chunk_sent(port: u16) {
let mut map = get_stream_tracker().lock().unwrap();
if let Some(state) = map.get_mut(&port) {
state.chunks_sent += 1;
}
}
fn mark_stream_completed(port: u16) {
let mut map = get_stream_tracker().lock().unwrap();
if let Some(state) = map.get_mut(&port) {
state.completed = true;
}
}
// --- Simple in-memory response store per worker port (for tests) ---
static RESP_STORE: OnceLock<Mutex<HashMap<u16, HashSet<String>>>> = OnceLock::new();
sgl-model-gateway/tests/reliability/mod.rs
+1
View File
@@ -4,3 +4,4 @@ pub mod circuit_breaker_test;
pub mod fault_tolerance_test;
pub mod rate_limiting_test;
pub mod retries_test;
pub mod upstream_cancel_test;
sgl-model-gateway/tests/reliability/upstream_cancel_test.rs
+2714
View File
File diff suppressed because it is too large Load Diff