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 \
  | perl -MTime::HiRes=time -ne 'BEGIN{$|=1} printf "%.3f | %s", time, $_'

Each printed line shows the unix timestamp, to the millisecond, when that line 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

✓ Check yourself

Answer from memory before opening each one.

Stop the mock, then send a streaming request to the gateway with curl -i. What status and body should you get, and what would a request that simply hangs tell you?

Show answer
A 502 with a JSON body whose error starts with upstream unreachable, because a refused connection maps to Bad Gateway. A hang means the upstream accepted the connection but never sent headers. Only the connect timeout is set, so send() would wait without limit. See failure modes.
The shared client sets a connect timeout but no request timeout. A colleague adds .timeout(30s) to be safe. What happens to a response that streams for 40 seconds?

Show answer
The timeout covers the whole response including the body, so the stream is cut at 30 seconds and the client sees a truncated answer. Use a connect timeout, plus a limit on the wait for the first byte if you need one, and keep full-request timeouts for non-streaming calls. See the shared client.
The gateway hands resp.bytes_stream() to Body::from_stream. A client on a slow link reads 1 KB a second. What does that do to the upstream read, and why does it matter?

Show answer
Nothing is buffered between the two sockets, so when the client stops reading, the body is not polled, the upstream read pauses, and TCP flow control pushes back. Memory per request stays bounded. A buffering proxy would absorb the whole response and hold it for the slow client. See the proxy handler.
Through the gateway directly, tokens arrive 40 ms apart. Behind a company nginx they arrive in one burst after about 800 ms. What do you check first?

Show answer
Proxy buffering in the layer in front of the gateway. The x-accel-buffering: no header asks nginx not to buffer, but other proxies and CDNs have their own settings. Compare curl -N direct with curl -N through the proxy to find the hop that holds bytes. See the proxy handler.
The handler passes bytes through without parsing SSE frames. When would you give that up and parse, and what is the trap?

Show answer
Parse when you need to count tokens, spot the [DONE] sentinel, or read errors that arrive mid-stream. The trap is splitting on network chunks: a TCP read can end in the middle of a frame, so you must buffer until the blank line that ends a record. Passthrough is faster and cannot misparse, so take the cost only for a feature that needs it. See reading the bytes.

Takeaway

A streaming proxy gives the upstream byte stream to the response body, so backpressure and cleanup follow from the streams themselves.
Build one reqwest client at startup and clone it. A client per request throws away keep-alive and burns ports.
A client-wide timeout caps the whole body, so streaming calls get a connect timeout instead, and a hung upstream needs its own first-byte limit.
A refused connection becomes a 502 with a JSON error, but only before the first byte is sent. After that the status is already committed.
When the client disconnects, dropping the body drops the upstream stream, so no tokens are generated for nobody.
No-buffering headers protect the stream from intermediaries that would turn it into one burst.

Every client still costs one upstream call. → Batching & coalescing

New dependencies

One shared reqwest client

The proxy handler

Wire it into main.rs

Run it

Reading the bytes (optional but eye-opening)

Failure modes worth touching

Checkpoint

✓ Check yourself

Takeaway

Further reading

Wire it into `main.rs`