Server-Sent Events & Streaming

use std::convert::Infallible;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use axum::{
    Json, Router,
    extract::{Query, State},
    http::HeaderMap,
    response::sse::{Event, KeepAlive, Sse},
    routing::{get, post},
};
use futures::{Stream, StreamExt, stream};
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;

#[derive(Clone, Serialize, Deserialize)]
struct LoanDue {
    member_id: u64,
    isbn: String,
    due_in_days: u8,
}

/// A notification with its sequence number (the SSE event id).
type Numbered = (u64, LoanDue);

#[derive(Clone)]
struct AppState {
    live: broadcast::Sender<Numbered>,
    history: Arc<Mutex<Vec<Numbered>>>, // a table or Redis stream in production
}

#[derive(Deserialize)]
struct Member {
    member_id: u64, // from the session in the real application
}

fn to_event((id, due): Numbered) -> Result<Event, Infallible> {
    let event = Event::default().event("loan-due").id(id.to_string());
    // json_data only fails if serialisation fails; fall back to plain data
    Ok(event.clone().json_data(&due).unwrap_or_else(|_| event.data("{}")))
}

async fn notifications(
    State(state): State<AppState>,
    Query(Member { member_id }): Query<Member>,
    headers: HeaderMap,
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    // 1. subscribe first, so nothing published during the replay is lost
    let live = BroadcastStream::new(state.live.subscribe())
        .filter_map(|msg| async move { msg.ok() }) // drop "lagged" notices
        .filter(move |(_, due)| std::future::ready(due.member_id == member_id))
        .map(to_event);

    // 2. replay what the client missed while disconnected (EventSource sends Last-Event-ID)
    let last_seen: u64 = headers.get("last-event-id").and_then(|v| v.to_str().ok()).and_then(|v| v.parse().ok()).unwrap_or(0);
    let missed: Vec<_> = state
        .history
        .lock()
        .expect("lock")
        .iter()
        .filter(|(id, due)| *id > last_seen && due.member_id == member_id)
        .cloned()
        .collect();
    let retry = stream::once(async { Ok(Event::default().retry(Duration::from_secs(5)).comment("connected")) });

    let events = retry.chain(stream::iter(missed).map(to_event)).chain(live);
    Sse::new(events).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)).text("keep-alive"))
}

async fn publish(State(state): State<AppState>, Json(due): Json<LoanDue>) -> String {
    let mut history = state.history.lock().expect("lock");
    let id = history.len() as u64 + 1;
    history.push((id, due.clone()));
    let _ = state.live.send((id, due));
    format!("event {id}\n")
}

#[tokio::main]
async fn main() {
    let (live, _) = broadcast::channel(256);
    let state = AppState { live, history: Arc::default() };
    let app = Router::new()
        .route("/events/loans", get(notifications))
        .route("/internal/loan-due", post(publish))
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ curl -s -N 'localhost:3000/events/loans?member_id=42' &
retry: 5000
: connected

$ curl -s -H 'content-type: application/json' -d '{"member_id":42,"isbn":"9780441013593","due_in_days":2}' localhost:3000/internal/loan-due
event 1
event: loan-due
id: 1
data: {"member_id":42,"isbn":"9780441013593","due_in_days":2}

$ curl -s -N -H 'last-event-id: 0' 'localhost:3000/events/loans?member_id=42'   # a reconnect replays event 1

<script>
  const source = new EventSource("/events/loans");          // cookies are sent: same-origin session
  source.addEventListener("loan-due", (e) => {
    const due = JSON.parse(e.data);
    showToast(`"${due.isbn}" is due in ${due.due_in_days} days`);
  });
</script>

use std::convert::Infallible;

use axum::{Router, routing::get, routing::post};
use axum_extra::json_lines::{AsResponse, JsonLines};
use futures::{Stream, StreamExt, stream};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug)]
struct Book {
    isbn: String,
    title: String,
}

// GET: one JSON object per line, produced lazily (e.g. from sqlx's fetch stream)
async fn export() -> JsonLines<impl Stream<Item = Result<Book, Infallible>>, AsResponse> {
    JsonLines::new(stream::iter(1..=3).map(|i| Ok(Book { isbn: format!("97800000000{i:02}"), title: format!("Book {i}") })))
}

// POST: consume an NDJSON upload incrementally, without buffering it all
async fn import(mut books: JsonLines<Book>) -> String {
    let mut count = 0;
    while let Some(book) = books.next().await {
        match book {
            Ok(book) => {
                tracing::info!(?book, "importing");
                count += 1;
            }
            Err(err) => return format!("stopped after {count} books: {err}\n"),
        }
    }
    format!("imported {count} books\n")
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/api/v1/books.ndjson", get(export)).route("/api/v1/imports", post(import));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

$ curl -s localhost:3000/api/v1/books.ndjson
{"isbn":"9780000000001","title":"Book 1"}
{"isbn":"9780000000002","title":"Book 2"}
{"isbn":"9780000000003","title":"Book 3"}
$ curl -s localhost:3000/api/v1/books.ndjson | curl -s --data-binary @- localhost:3000/api/v1/imports
imported 3 books

use axum::{
    Router,
    body::Body,
    extract::State,
    http::{StatusCode, header},
    response::{IntoResponse, Response},
    routing::get,
};

async fn proxy_export(State(client): State<reqwest::Client>) -> Response {
    let upstream = match client.get("http://127.0.0.1:3000/upstream/catalogue.csv").send().await {
        Ok(resp) if resp.status().is_success() => resp,
        Ok(resp) => return (StatusCode::BAD_GATEWAY, format!("upstream answered {}", resp.status())).into_response(),
        Err(_) => return StatusCode::BAD_GATEWAY.into_response(),
    };
    let content_type = upstream.headers().get(header::CONTENT_TYPE).cloned();
    let mut response = Body::from_stream(upstream.bytes_stream()).into_response();
    if let Some(ct) = content_type {
        response.headers_mut().insert(header::CONTENT_TYPE, ct);
    }
    response
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/exports/catalogue.csv", get(proxy_export))
        .route("/upstream/catalogue.csv", get(|| async { ([(header::CONTENT_TYPE, "text/csv")], "isbn,title\n9780441013593,Dune\n") }))
        .with_state(reqwest::Client::new());
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}