Batching & Coalescing

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

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. Its default map type is
        // sorted, so reordered keys and extra whitespace give the same key. Any
        // difference in a value gives a different key, which is the safe direction.
        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)
    }
}

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

async-stream = "0.3"

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

mod coalesce;

cargo run -p infergw

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

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

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