Chapter 02 · ~30 minutes

Proxy & Streaming

Forward chat-completion requests to the upstream and stream the SSE response back chunk-by-chunk, with proper backpressure and no buffering.

New dependencies

Append to infergw/Cargo.toml under [dependencies]:

reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] }
futures = "0.3"
bytes   = "1"
Why rustls and not native-tls

Rustls links a pure-Rust TLS stack into the binary. No OpenSSL on the host, no version drift, no surprises when you move to distroless in Chapter 06. The image stays small and the build is reproducible.

One shared reqwest client

reqwest::Client is internally an Arc over a connection pool. Build one at startup, clone freely. Building one per request is the most common rookie mistake — you'll burn through ports and lose all keep-alive benefits.

Create infergw/src/upstream.rs:

use std::time::Duration;

#[derive(Clone)]
pub struct Upstream {
    pub client: reqwest::Client,
    pub base_url: String,
}

impl Upstream {
    pub fn new(base_url: String, timeout: Duration) -> anyhow::Result<Self> {
        let client = reqwest::Client::builder()
            .pool_idle_timeout(Duration::from_secs(60))
            .pool_max_idle_per_host(32)
            // We don't set .timeout(), because that would cap the *whole* SSE
            // stream — we only want a connect timeout for streaming requests.
            .connect_timeout(timeout)
            .http2_prior_knowledge()    // vLLM supports HTTP/2; if your upstream
            .http2_adaptive_window(true) // doesn't, remove these two lines.
            .build()?;
        Ok(Self { client, base_url })
    }

    pub fn chat_completions_url(&self) -> String {
        format!("{}/v1/chat/completions", self.base_url.trim_end_matches('/'))
    }
}
A subtle reqwest footgun

Client::builder().timeout(d) applies to the entire response, including the body. For a streaming SSE response that can run for many seconds, that timeout will kill the stream mid-flight. Use connect_timeout for streaming requests, and reserve full request timeouts for non-streaming calls.

Wire it into AppState. Replace infergw/src/routes.rs with:

use crate::upstream::Upstream;
use axum::{http::StatusCode, response::IntoResponse, routing::{get, post}, Router};
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};

#[derive(Clone)]
pub struct AppState {
    pub ready: Arc<AtomicBool>,
    pub upstream: Upstream,
}

impl AppState {
    pub fn new(upstream: Upstream) -> Self {
        Self {
            ready: Arc::new(AtomicBool::new(true)),
            upstream,
        }
    }
}

pub fn router(state: AppState) -> Router {
    Router::new()
        .route("/healthz", get(healthz))
        .route("/readyz", get(readyz))
        .route("/v1/chat/completions", post(crate::proxy::chat_completions))
        .with_state(state)
}

async fn healthz() -> impl IntoResponse { (StatusCode::OK, "ok") }

async fn readyz(
    axum::extract::State(state): axum::extract::State<AppState>,
) -> impl IntoResponse {
    if state.ready.load(Ordering::Relaxed) {
        (StatusCode::OK, "ready")
    } else {
        (StatusCode::SERVICE_UNAVAILABLE, "draining")
    }
}

The proxy handler

This is the meat of the chapter. The pattern: take the incoming JSON body, forward it to the upstream, get back a byte stream, wrap that in an axum SSE response, hand it to the client. Backpressure is automatic — if the client is slow, the upstream read pauses.

Create infergw/src/proxy.rs:

use crate::routes::AppState;
use axum::{
    body::Body,
    extract::State,
    http::{header, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use bytes::Bytes;
use futures::StreamExt;
use serde_json::Value;

pub async fn chat_completions(
    State(state): State<AppState>,
    Json(body): Json<Value>,
) -> Result<Response, ProxyError> {
    let url = state.upstream.chat_completions_url();
    let stream_requested = body
        .get("stream")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    let resp = state
        .upstream
        .client
        .post(&url)
        .json(&body)
        .send()
        .await
        .map_err(ProxyError::Upstream)?;

    let upstream_status = resp.status();
    if !upstream_status.is_success() {
        let text = resp.text().await.unwrap_or_default();
        tracing::warn!(%upstream_status, body = %text, "upstream non-2xx");
        return Err(ProxyError::UpstreamStatus(upstream_status, text));
    }

    if stream_requested {
        // Forward as SSE byte stream. We don't reparse the SSE framing here;
        // we trust the upstream and let bytes pass through. (Chapter 05 will
        // tap into this stream for token-count metrics.)
        let upstream_body = resp.bytes_stream().map(|chunk| {
            chunk.map_err(std::io::Error::other)
        });

        let response = Response::builder()
            .status(StatusCode::OK)
            .header(header::CONTENT_TYPE, "text/event-stream")
            .header(header::CACHE_CONTROL, "no-cache")
            .header("x-accel-buffering", "no") // nginx: do not buffer
            .body(Body::from_stream(upstream_body))
            .expect("response builder");

        Ok(response)
    } else {
        // Non-streaming: collect and return JSON verbatim.
        let bytes: Bytes = resp.bytes().await.map_err(ProxyError::Upstream)?;
        let response = Response::builder()
            .status(StatusCode::OK)
            .header(header::CONTENT_TYPE, "application/json")
            .body(Body::from(bytes))
            .expect("response builder");
        Ok(response)
    }
}

#[derive(thiserror::Error, Debug)]
pub enum ProxyError {
    #[error("upstream transport error: {0}")]
    Upstream(#[from] reqwest::Error),
    #[error("upstream returned {0}: {1}")]
    UpstreamStatus(StatusCode, String),
}

impl IntoResponse for ProxyError {
    fn into_response(self) -> Response {
        let (status, msg) = match &self {
            ProxyError::Upstream(e) if e.is_connect() => (
                StatusCode::BAD_GATEWAY,
                format!("upstream unreachable: {e}"),
            ),
            ProxyError::Upstream(e) if e.is_timeout() => (
                StatusCode::GATEWAY_TIMEOUT,
                format!("upstream timeout: {e}"),
            ),
            ProxyError::Upstream(e) => (
                StatusCode::BAD_GATEWAY,
                format!("upstream error: {e}"),
            ),
            ProxyError::UpstreamStatus(s, body) => (*s, body.clone()),
        };
        tracing::error!(?status, %msg, "proxy error");
        (status, Json(serde_json::json!({"error": msg}))).into_response()
    }
}
The three things that make this work

1. resp.bytes_stream() returns a Stream<Item = Result<Bytes, reqwest::Error>> — chunks arrive as the upstream sends them. No buffering.

2. Body::from_stream(...) hands that stream to axum, which writes each chunk to the client socket as soon as it lands.

3. x-accel-buffering: no and Cache-Control: no-cache are belt-and-suspenders against nginx / CDNs / proxies that would otherwise buffer the response and ruin streaming. Nobody enjoys debugging "why are my tokens arriving in one burst at the end."

Wire it into main.rs

mod config;
mod proxy;
mod routes;
mod shutdown;
mod telemetry;
mod upstream;

use std::time::Duration;
use tower_http::trace::TraceLayer;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    telemetry::init_tracing();
    let cfg = config::Config::from_env()?;
    tracing::info!(?cfg, "starting infergw");

    let upstream = upstream::Upstream::new(
        cfg.upstream_url.clone(),
        Duration::from_millis(cfg.upstream_timeout_ms),
    )?;
    let state = routes::AppState::new(upstream);
    let app = routes::router(state.clone()).layer(TraceLayer::new_for_http());

    let listener = tokio::net::TcpListener::bind(&cfg.bind_addr).await?;
    tracing::info!(addr = %cfg.bind_addr, "listening");

    let grace = Duration::from_millis(cfg.shutdown_grace_ms);
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown::shutdown_signal(state, grace))
        .await?;

    Ok(())
}

And make sure the mock from Chapter 00 used HTTP/1.1, not HTTP/2 — the .http2_prior_knowledge() call assumes the upstream supports it. The Chapter 00 mock served HTTP/1.1. To match, comment those two lines out in upstream.rs:

// .http2_prior_knowledge()
// .http2_adaptive_window(true)

Real vLLM speaks HTTP/2; flip them back on when you switch to a real upstream.

Run it

Make sure the mock from Chapter 00 is still running on :9000. Then start the gateway:

cargo run -p infergw

From a third terminal:

curl -N -H 'content-type: application/json' \
  -d '{"model":"llama-3-8b","messages":[{"role":"user","content":"Hi there"}],"stream":true}' \
  http://127.0.0.1:8080/v1/chat/completions

You should see exactly the same 19 chunks plus [DONE] as in Chapter 00 — except now they're flowing through your gateway. Watch them arrive one at a time, 40ms apart. That's the streaming working end-to-end.

The gateway logs will show:

INFO tower_http::trace::on_request: started processing request method=POST uri=/v1/chat/completions
INFO tower_http::trace::on_response: finished processing request latency=816 ms status=200

Reading the bytes (optional but eye-opening)

If you want to see exactly what crosses the wire, pipe through xxd or just slow it down with awk:

curl -sN -H 'content-type: application/json' \
  -d '{"model":"llama-3-8b","messages":[{"role":"user","content":"Hi"}],"stream":true}' \
  http://127.0.0.1:8080/v1/chat/completions \
  | awk '{ print systime() " | " $0; fflush(); }'

Each printed line shows the unix timestamp when that chunk arrived at your terminal. Subtract consecutive timestamps and you'll see ~40ms gaps — proof that no layer in between is buffering.

Failure modes worth touching

What happens if the upstream is down?

Stop the mock-upstream (Ctrl-C in its terminal). Now re-run the curl against the gateway:

curl -i -H 'content-type: application/json' \
  -d '{"model":"llama-3-8b","messages":[{"role":"user","content":"Hi"}],"stream":true}' \
  http://127.0.0.1:8080/v1/chat/completions

You'll get a 502 Bad Gateway with a structured error body: {"error":"upstream unreachable: ..."}. That's ProxyError::Upstream doing its job. Chapter 04 will turn this into a circuit-breaker trip instead.

What if the client disconnects mid-stream?

Start a curl, then Ctrl-C it after the first chunk lands. Look at the gateway's logs: the upstream stream is dropped because the response body is dropped — reqwest sees the receiver is gone and closes the connection. No leak, no wasted upstream tokens past that point. This is one of the things you get for free with proper streaming; the alternative (buffer everything, then write to client) keeps the upstream churning even after the client is gone.

What about non-streaming requests?

Drop "stream": true from the body and re-curl. You get a single JSON document. The mock doesn't currently return a non-streaming form — extending it is a 10-minute exercise if you want to handle that path. (For OpenAI / a real vLLM, this Just Works because they handle both modes.)

Checkpoint

  • POST /v1/chat/completions with "stream":true proxies through the gateway and streams chunks to the client.
  • Upstream-down returns a structured 502.
  • Client disconnect cleans up the upstream side.
  • No buffering anywhere — chunks land at the client roughly when the upstream emits them.

This is already a useful gateway. The next four chapters are about making it reliable.

→ Chapter 03: Batching & coalescing