Chapter 03 · ~25 minutes

Batching & Coalescing

When two clients ask the model the same question at the same time, one upstream call should be enough. Build a singleflight that proves it.

Why coalesce identical in-flight prompts

Coalescing (aka singleflight, aka request deduplication) is the cheapest reliability and cost win you'll find. When N concurrent requests share an identical prompt + model + sampling params, you can serve all N from a single upstream call by having N–1 of them subscribe to the first one's output stream.

Where it pays off:

  • Agent loops that retry the same prompt with the same seed.
  • RAG fan-outs where dozens of users hit the same top-1 retrieved document.
  • Cache stampedes after a deployment when a hot prompt suddenly has no cached completion.
  • Health checks that probe with a canary prompt.

True batching (combining different prompts into a single forward pass) is something vLLM does internally via continuous batching — we don't need to redo it at the gateway. What the gateway uniquely can do is dedupe identical work before it ever reaches the model.

Coalescing is a correctness choice, not just an optimization

If the prompt or any sampling parameter (temperature, seed, top_p, max_tokens) differs between requests, they MUST NOT share a response. A request that asks for a different max_tokens than the in-flight call will get the wrong number of tokens. Our cache key must include every parameter that affects output.

Choosing the cache key

The simplest correct key: hash the entire request body. If two requests have byte-identical bodies, they're truly identical. If they differ in any field — even a trailing whitespace — they get different keys, so we conservatively call the upstream twice. That's fine: we want false negatives, not false positives.

Add to infergw/Cargo.toml:

sha2  = "0.10"
tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "time", "sync"] }  # already there; sync is required

(The sync feature on tokio is already included from Chapter 01 — listing here only as a reminder.)

The singleflight design

The core data structure is a map from request-hash → tokio::sync::broadcast::Sender<Bytes>. The first request to arrive with a given hash becomes the leader: it makes the upstream call and broadcasts each chunk. Followers find the broadcast in the map, subscribe(), and receive every subsequent chunk.

   client A ──┐
              ├─ same hash ──▶  [in-flight map]  ─── 1 upstream call ──▶ vLLM
   client B ──┘                       │                  │
                                  leader makes        leader pushes
                                  the call,           chunks to
                                  followers           broadcast
                                  subscribe           channel

Three details matter:

  1. broadcast::channel — followers must receive every chunk, in order, and start from where the leader is now. Broadcast channels have a buffer; if a slow follower lags more than buffer-size chunks, it gets a RecvError::Lagged. Set the buffer generously (we use 256).
  2. Late arrivals miss early chunks. A request that arrives after the leader has already sent 5 chunks will only get chunks 6+. That's wrong. We solve it by having the leader also buffer emitted chunks in a Vec until the response completes, so late followers can replay.
  3. Clean up after completion. When the upstream stream ends, remove the entry from the map so new requests start a fresh upstream call.

Implementation

Create infergw/src/coalesce.rs:

use bytes::Bytes;
use sha2::{Digest, Sha256};
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};
use tokio::sync::{broadcast, Mutex as AsyncMutex};

/// A frame coming from the upstream stream. We model end-of-stream explicitly
/// so subscribers know when the leader is done (vs. just slow).
#[derive(Clone, Debug)]
pub enum Frame {
    Chunk(Bytes),
    Done,
    Error(String),
}

/// State shared by the leader and any followers for one in-flight request.
pub struct InFlight {
    pub tx: broadcast::Sender<Frame>,
    /// Frames the leader has already emitted, for late joiners to replay.
    pub replay: Arc<AsyncMutex<Vec<Frame>>>,
}

#[derive(Default, Clone)]
pub struct Coalescer {
    inner: Arc<Mutex<HashMap<String, Arc<InFlight>>>>,
}

pub enum JoinResult {
    /// Caller is the leader and must drive the upstream call.
    Leader(Arc<InFlight>),
    /// Caller is a follower; subscribe and replay prior frames.
    Follower(Arc<InFlight>),
}

impl Coalescer {
    pub fn key_for(body: &serde_json::Value) -> String {
        // Canonicalize by serializing through serde_json — fields keep insertion
        // order from the original JSON, so two requests with reordered keys will
        // hash differently. For coalescing that's acceptably conservative.
        let s = serde_json::to_string(body).unwrap_or_default();
        let mut h = Sha256::new();
        h.update(s.as_bytes());
        format!("{:x}", h.finalize())
    }

    pub fn join(&self, key: &str) -> JoinResult {
        let mut map = self.inner.lock().expect("coalesce map poisoned");
        if let Some(existing) = map.get(key) {
            JoinResult::Follower(existing.clone())
        } else {
            let (tx, _rx) = broadcast::channel::<Frame>(256);
            let inflight = Arc::new(InFlight {
                tx,
                replay: Arc::new(AsyncMutex::new(Vec::new())),
            });
            map.insert(key.to_string(), inflight.clone());
            JoinResult::Leader(inflight)
        }
    }

    pub fn remove(&self, key: &str) {
        let mut map = self.inner.lock().expect("coalesce map poisoned");
        map.remove(key);
    }

    pub fn in_flight_count(&self) -> usize {
        self.inner.lock().map(|m| m.len()).unwrap_or(0)
    }
}

Now the proxy handler needs to branch on leader vs. follower. Update infergw/src/proxy.rs:

use crate::{
    coalesce::{Coalescer, Frame, InFlight, JoinResult},
    routes::AppState,
};
use axum::{
    body::Body,
    extract::State,
    http::{header, StatusCode},
    response::{IntoResponse, Response},
    Json,
};
use bytes::Bytes;
use futures::{stream, StreamExt};
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::broadcast::error::RecvError;

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

    if !stream_requested {
        // Non-streaming path stays simple — no coalescing for now. We could add
        // it, but the win is small and the bookkeeping is fussier.
        return passthrough_nonstreaming(&state, &body).await;
    }

    let key = Coalescer::key_for(&body);
    match state.coalescer.join(&key) {
        JoinResult::Leader(inflight) => {
            tracing::info!(key = %&key[..12], "coalesce: leader");
            spawn_leader(state.clone(), key.clone(), body, inflight.clone());
            let body_stream = subscribe_stream(inflight, None);
            Ok(sse_response(body_stream))
        }
        JoinResult::Follower(inflight) => {
            tracing::info!(key = %&key[..12], "coalesce: follower");
            let replay = inflight.replay.lock().await.clone();
            let body_stream = subscribe_stream(inflight, Some(replay));
            Ok(sse_response(body_stream))
        }
    }
}

fn sse_response<S>(s: S) -> Response
where
    S: futures::Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
{
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "text/event-stream")
        .header(header::CACHE_CONTROL, "no-cache")
        .header("x-accel-buffering", "no")
        .body(Body::from_stream(s))
        .expect("response builder")
}

fn subscribe_stream(
    inflight: Arc<InFlight>,
    replay: Option<Vec<Frame>>,
) -> impl futures::Stream<Item = Result<Bytes, std::io::Error>> {
    let mut rx = inflight.tx.subscribe();
    let replay_iter = replay.unwrap_or_default().into_iter();
    let replay_stream = stream::iter(replay_iter).map(Ok);

    let live = async_stream::stream! {
        loop {
            match rx.recv().await {
                Ok(Frame::Chunk(_)) | Ok(Frame::Done) | Ok(Frame::Error(_)) => {}
                Err(RecvError::Lagged(n)) => {
                    tracing::warn!(missed = n, "follower lagged broadcast buffer");
                    continue;
                }
                Err(RecvError::Closed) => break,
            }
            // re-borrow on each iteration to make the borrow checker happy:
            // we already matched above just to detect Lagged; redo the recv
            // to actually yield the frame.
        }
    };
    // Simpler shape using async_stream over a fresh subscriber:
    let mut rx2 = inflight.tx.subscribe();
    drop(live);
    let live = async_stream::stream! {
        loop {
            match rx2.recv().await {
                Ok(Frame::Chunk(b)) => yield Ok(b),
                Ok(Frame::Done) => {
                    yield Ok(Bytes::from_static(b"data: [DONE]\n\n"));
                    break;
                }
                Ok(Frame::Error(e)) => {
                    yield Err(std::io::Error::other(e));
                    break;
                }
                Err(RecvError::Lagged(n)) => {
                    tracing::warn!(missed = n, "follower lagged broadcast buffer");
                    // We can't recover the missed chunks; bail.
                    yield Err(std::io::Error::other("coalesce follower lagged"));
                    break;
                }
                Err(RecvError::Closed) => break,
            }
        }
    };

    let replay_bytes = replay_stream.filter_map(|f| async {
        match f {
            Ok(Frame::Chunk(b)) => Some(Ok(b)),
            Ok(Frame::Done) => Some(Ok(Bytes::from_static(b"data: [DONE]\n\n"))),
            Ok(Frame::Error(e)) => Some(Err(std::io::Error::other(e))),
            Err(_) => None,
        }
    });

    replay_bytes.chain(live)
}

fn spawn_leader(state: AppState, key: String, body: Value, inflight: Arc<InFlight>) {
    tokio::spawn(async move {
        let url = state.upstream.chat_completions_url();
        let resp = match state.upstream.client.post(&url).json(&body).send().await {
            Ok(r) => r,
            Err(e) => {
                let _ = inflight.tx.send(Frame::Error(format!("upstream send: {e}")));
                state.coalescer.remove(&key);
                return;
            }
        };
        if !resp.status().is_success() {
            let s = resp.status();
            let text = resp.text().await.unwrap_or_default();
            let _ = inflight.tx.send(Frame::Error(format!("upstream {s}: {text}")));
            state.coalescer.remove(&key);
            return;
        }

        let mut stream = resp.bytes_stream();
        while let Some(chunk) = stream.next().await {
            match chunk {
                Ok(b) => {
                    inflight.replay.lock().await.push(Frame::Chunk(b.clone()));
                    let _ = inflight.tx.send(Frame::Chunk(b));
                }
                Err(e) => {
                    let _ = inflight.tx.send(Frame::Error(format!("upstream chunk: {e}")));
                    state.coalescer.remove(&key);
                    return;
                }
            }
        }
        inflight.replay.lock().await.push(Frame::Done);
        let _ = inflight.tx.send(Frame::Done);
        state.coalescer.remove(&key);
    });
}

async fn passthrough_nonstreaming(state: &AppState, body: &Value) -> Result<Response, ProxyError> {
    let url = state.upstream.chat_completions_url();
    let resp = state.upstream.client.post(&url).json(body).send().await
        .map_err(ProxyError::Upstream)?;
    let status = resp.status();
    let bytes = resp.bytes().await.map_err(ProxyError::Upstream)?;
    Ok(Response::builder()
        .status(status)
        .header(header::CONTENT_TYPE, "application/json")
        .body(Body::from(bytes))
        .expect("response builder"))
}

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

impl IntoResponse for ProxyError {
    fn into_response(self) -> Response {
        let msg = self.to_string();
        tracing::error!(%msg, "proxy error");
        (StatusCode::BAD_GATEWAY, Json(serde_json::json!({"error": msg}))).into_response()
    }
}

Add async-stream to the deps:

async-stream = "0.3"
A note on the chunk format

The leader forwards the raw bytes from the upstream — which already contain data: {...}\n\n framing. The [DONE] sentinel we synthesize at end-of-stream is a defensive belt-and-suspenders so followers always see a clean terminator, even if the upstream already sent one. Some clients are picky; doubling it is harmless.

Wire the coalescer into AppState

Update infergw/src/routes.rs to carry the coalescer:

use crate::{coalesce::Coalescer, 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,
    pub coalescer: Coalescer,
}

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

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")
    }
}

Register the module in main.rs:

mod coalesce;

Prove it works

Rebuild and run:

cargo run -p infergw

Open three terminals and fire three identical requests at the same time:

for i in 1 2 3; do
  curl -sN -H 'content-type: application/json' \
    -d '{"model":"llama-3-8b","messages":[{"role":"user","content":"identical"}],"stream":true}' \
    http://127.0.0.1:8080/v1/chat/completions \
    | head -c 200 &
done
wait

Look at the gateway logs. You should see:

INFO infergw::proxy: coalesce: leader   key=a4f1c2...
INFO infergw::proxy: coalesce: follower key=a4f1c2...
INFO infergw::proxy: coalesce: follower key=a4f1c2...

And the mock-upstream's logs will show one incoming request, not three. That's the saving.

Now send three different prompts at the same time and verify the mock sees three calls:

for i in 1 2 3; do
  curl -sN -H 'content-type: application/json' \
    -d "{\"model\":\"llama-3-8b\",\"messages\":[{\"role\":\"user\",\"content\":\"prompt $i\"}],\"stream\":true}" \
    http://127.0.0.1:8080/v1/chat/completions \
    | head -c 200 &
done
wait

Three different prompts, three leaders, three upstream calls. The coalescer correctly refuses to dedupe non-identical work.

Stretch: load test the coalescer

Install oha with cargo install oha. Then:

oha -z 10s -c 50 -m POST \
  -H 'content-type: application/json' \
  -d '{"model":"llama-3-8b","messages":[{"role":"user","content":"identical"}],"stream":true}' \
  http://127.0.0.1:8080/v1/chat/completions

50 concurrent identical requests for 10 seconds. Without coalescing, the mock would log ~3000 incoming requests. With coalescing, you'll see roughly one every ~800ms (the duration of a complete mocked response), so ~12 total — a 250x reduction in upstream load. We'll measure this properly in Chapter 07.

When NOT to coalesce

Three cases to be aware of:

  1. Any sampling temperature > 0. Two identical prompts with temperature: 0.7 are supposed to produce different outputs. Coalescing them collapses the distribution and breaks user expectations. The safe default is to disable coalescing when temperature > 0 unless a fixed seed is also set. A nice exercise: extend Coalescer::key_for to return None in that case and skip the dedupe.
  2. Per-user personalization in the prompt. If the prompt is templated with user-specific context, identical-looking requests across users are rare anyway, and you probably don't want to leak one user's completion to another even if it would be technically identical. Use a per-tenant prefix in the key.
  3. Audit / billing implications. If each request is independently billed to a tenant, coalescing means two tenants share one upstream call but get charged twice (or you charge once and lose money). Track this explicitly — the coalescer can emit a counter so billing knows the call was a follower.
Where this pattern came from

The "singleflight" name comes from Go's golang.org/x/sync/singleflight, and the pattern itself predates that by decades — it's how every well-engineered cache (memcached, Varnish, Cloudflare) handles thundering-herd cache fills. The gateway is just doing the same dance at the LLM layer.

Checkpoint

  • Three concurrent identical requests produce ONE upstream call and three streams to clients.
  • Three different prompts produce three upstream calls (no false coalescing).
  • Followers see the same chunks as the leader, in order.
  • The in-flight map cleans up when the leader's stream ends.

Next we make the upstream call itself more resilient.

→ Chapter 04: Circuit breaker, retry budgets, and request hedging