add the rust server ingress tests, guard, and submit modules

rust/sglang-server/src/api_server/guard.rs

@@ -0,0 +1,130 @@
+//! Abort-on-disconnect guard for in-flight requests. Handlers arm a guard per
+//! submitted rid; axum dropping the handler/SSE stream (client disconnected)
+//! drops the guard, which aborts whatever wasn't disarmed (mirrors Python's
+//! `is_disconnected` abort).
+
+use std::collections::HashSet;
+
+use crate::ids::Rid;
+use crate::tokenizer_manager::{AbortSource, Senders};
+
+/// Aborts still-in-flight rids on drop. Each rid is disarmed on natural finish;
+/// whatever remains at drop is aborted.
+pub(super) struct AbortGuard {
+    senders: Senders,
+    /// Rids still in flight. `Rid` carries its own partition key, so there is no
+    /// separate routing value to keep alongside it.
+    ///
+    /// A set, not a `Vec`: `disarm` runs once per request that finishes, and over a
+    /// batch a linear scan makes the guard quadratic in the batch size — measured
+    /// 13.3 ms for a 4096-item batch, more than all of that batch's real transform
+    /// work combined. `Rid`'s identity is its id string, so set membership is the
+    /// same relation `retain` was testing. The cost is two hashes of a ~40-byte
+    /// string on the single-request path (~80 ns against a ~40 µs request), which
+    /// is why the trade is worth making rather than threading slot indices in from
+    /// the batch call sites.
+    rids: HashSet<Rid>,
+}
+
+impl AbortGuard {
+    pub(super) fn new(senders: Senders, rid: Rid) -> Self {
+        Self {
+            senders,
+            rids: HashSet::from([rid]),
+        }
+    }
+
+    /// Guard covering no rids yet — a batch arms each as it's submitted so a
+    /// mid-fan-out disconnect aborts every request already handed to the scheduler.
+    pub(super) fn new_empty(senders: Senders) -> Self {
+        Self {
+            senders,
+            rids: HashSet::new(),
+        }
+    }
+
+    /// Track a request for abort-on-drop.
+    pub(super) fn arm(&mut self, rid: Rid) {
+        self.rids.insert(rid);
+    }
+
+    /// Request finished naturally — don't abort it on drop.
+    pub(super) fn disarm(&mut self, rid: &Rid) {
+        self.rids.remove(rid);
+    }
+}
+
+impl Drop for AbortGuard {
+    fn drop(&mut self) {
+        // Report the abort and nothing more. There is no in-flight rid registry to
+        // release from: `Rid::from_client` makes each client rid internally unique,
+        // so a resubmit of the "same" rid is a different `Rid` and cannot be caught
+        // up in this abort. That removes the ordering hazard split ownership created.
+        //
+        // The lane is unbounded, so this send only fails at shutdown, when the loop
+        // is gone and nothing is generating anyway.
+        for rid in self.rids.drain() {
+            let _ = self.senders.abort.send(AbortSource::Guard(rid));
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn senders_with_abort(abort: flume::Sender<AbortSource>) -> Senders {
+        Senders {
+            tm: flume::unbounded().0,
+            abort,
+            tok: flume::unbounded().0,
+            detok: vec![],
+        }
+    }
+
+    /// A batch guard aborts exactly the rids still armed at drop — the ones whose
+    /// requests never reached a terminal — and leaves the finished ones alone.
+    #[test]
+    fn guard_aborts_only_the_rids_still_armed() {
+        let (abort_tx, abort_rx) = flume::unbounded();
+        let done: Rid = "done".into();
+        let mut guard = AbortGuard::new(senders_with_abort(abort_tx), done.clone());
+        guard.arm("aborted".into());
+        guard.disarm(&done); // finished naturally
+        drop(guard);
+
+        assert!(
+            matches!(abort_rx.try_recv().unwrap(), AbortSource::Guard(r) if r.as_str() == "aborted")
+        );
+        assert!(
+            abort_rx.try_recv().is_err(),
+            "a disarmed rid must not be aborted"
+        );
+    }
+
+    /// An armed guard aborts its rid on drop — exactly the cleanup a busy-skipped
+    /// `/health_generate` probe relies on. It never sees a terminal frame here, so
+    /// dropping the guard is the only path that deregisters its detok sink (via the
+    /// ingress `on_abort`). Regression for the detok-entry leak per health probe.
+    #[test]
+    fn armed_guard_aborts_on_drop() {
+        let (tm_tx, tm_rx) = flume::unbounded();
+        drop(AbortGuard::new(senders_with_abort(tm_tx), "r7".into()));
+        assert!(
+            matches!(tm_rx.try_recv(), Ok(AbortSource::Guard(rid)) if rid.as_str() == "r7"),
+            "armed guard must abort its rid on drop",
+        );
+        assert!(tm_rx.try_recv().is_err(), "exactly one abort");
+    }
+
+    /// A disarmed rid (finished naturally) is not aborted on drop.
+    #[test]
+    fn disarmed_guard_does_not_abort() {
+        let (tm_tx, tm_rx) = flume::unbounded();
+        let id = Rid::from("r9");
+        let mut guard = AbortGuard::new(senders_with_abort(tm_tx), "r9".into());
+        guard.disarm(&id);
+        drop(guard);
+        assert!(tm_rx.try_recv().is_err(), "disarmed rid must not abort");
+    }
+}

rust/sglang-server/src/api_server/submit.rs

@@ -0,0 +1,126 @@
+//! Request submission into the ingress pipeline, shared by every endpoint
+//! module: mint the client-visible rid (uuid hex, Python-parity), build the
+//! `Request`, and hand it to the TM with an egress receiver for the response.
+
+use std::convert::Infallible;
+
+use axum::{
+    Json,
+    http::StatusCode,
+    response::{
+        IntoResponse, Response,
+        sse::{Event, Sse},
+    },
+};
+use tokio::sync::mpsc;
+
+use super::AppState;
+use super::frame::error_value;
+use crate::fsm::RequestState;
+use crate::ids::Rid;
+use crate::message::{EgressItem, EgressSink, Request, RequestKind};
+use crate::tokenizer_manager::TmEvent;
+
+/// Submit one request; returns the rid, its hashed routing key, and the egress
+/// receiver. Every request arrives with its final rid — a generate request from
+/// `into_requests` (or the `HEALTH_CHECK_<uuid>` the health probe sets), a
+/// control request from its constructor — so this only echoes it back.
+pub(super) async fn submit(
+    state: &AppState,
+    kind: RequestKind,
+    // `stream`: the client is reading an SSE stream, so it expects 200 plus an
+    // error frame rather than a 4xx — same rule `pre_submit_error` applies
+    // everywhere else.
+    stream: bool,
+) -> Result<(Rid, mpsc::Receiver<EgressItem>), Response> {
+    let rid = match &kind {
+        // Generate rids are already final: `GenerateBody::into_requests` normalized the
+        // client's, or minted one. Control requests have no client-facing rid.
+        RequestKind::Generate(g) => g.rid.clone(),
+        RequestKind::Control(c) => c.rid().into(),
+    };
+    // Two in-flight requests can name the same client rid, but they cannot share a
+    // `Rid`: `into_requests` built each through `Rid::from_client`, which appends a
+    // uniquifier. So nothing here needs to check for a collision — the detok table
+    // key is unique by construction, and `client_facing` restores what the client
+    // sent for `meta_info.id`.
+    // Async-aware send so a full TM inbox yields (backpressure) instead of parking
+    // a thread; Err only when the inbox is closed (shutdown).
+    let (tx, rx) = mpsc::channel::<EgressItem>(state.egress_buf);
+    let request = Request {
+        rid: rid.clone(),
+        state: RequestState::Received,
+        sink: EgressSink::Local(tx),
+        kind,
+    };
+    match state.senders.tm.send_async(TmEvent::Ingress(request)).await {
+        Ok(()) => Ok((rid, rx)),
+        // `SendError` has a single meaning — the channel is disconnected.
+        Err(_) => {
+            tracing::error!(%rid, "tm inbox closed; request rejected");
+            // Return 503 so the client can retry.
+            Err(pre_submit_error(
+                StatusCode::SERVICE_UNAVAILABLE,
+                "service unavailable",
+                stream,
+            ))
+        }
+    }
+}
+
+/// Shape an error that occurs *before* (or instead of) a successful submit into a
+/// client response. Two parity points with Python's `generate_request`. The body
+/// is the same `{"error": {...}}` object every other path emits — not bare text,
+/// which a client parsing JSON chokes on. And a streaming request gets 200 plus
+/// one SSE error frame and `[DONE]`, not a 4xx: the client has already committed
+/// to reading a stream, and Python answers it inside `stream_results()`.
+pub(super) fn pre_submit_error(code: StatusCode, message: &str, stream: bool) -> Response {
+    let body = error_value(code.as_u16(), message);
+    if !stream {
+        return (code, Json(body)).into_response();
+    }
+    let frames = [body.to_string(), "[DONE]".to_string()];
+    Sse::new(futures::stream::iter(
+        frames.map(|data| Ok::<_, Infallible>(Event::default().data(data))),
+    ))
+    .into_response()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Unary pre-submit errors are a 4xx/5xx with a JSON `{"error":...}` body;
+    /// streaming ones are 200 + an SSE error frame + `[DONE]`, because Python
+    /// answers from inside `stream_results()` once the stream is committed.
+    #[tokio::test]
+    async fn pre_submit_errors_match_python_shape() {
+        let unary = pre_submit_error(StatusCode::BAD_REQUEST, "bad input", false);
+        assert_eq!(unary.status(), StatusCode::BAD_REQUEST);
+        let body = axum::body::to_bytes(unary.into_body(), 64 * 1024)
+            .await
+            .unwrap();
+        let v: serde_json::Value = serde_json::from_slice(&body).expect("JSON body");
+        assert_eq!(v["error"]["message"], "bad input");
+        assert_eq!(v["error"]["code"], 400);
+
+        let streamed = pre_submit_error(StatusCode::BAD_REQUEST, "bad input", true);
+        assert_eq!(
+            streamed.status(),
+            StatusCode::OK,
+            "the stream itself is 200"
+        );
+        let body = axum::body::to_bytes(streamed.into_body(), 64 * 1024)
+            .await
+            .unwrap();
+        let text = String::from_utf8(body.to_vec()).unwrap();
+        assert!(
+            text.contains(r#""code":400"#),
+            "carries the status in-band: {text}"
+        );
+        assert!(
+            text.trim_end().ends_with("data: [DONE]"),
+            "terminated: {text}"
+        );
+    }
+}