WebSockets & SSE

use axum::{extract::ws::{Message, WebSocket, WebSocketUpgrade}, response::IntoResponse, routing::get};
 
async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_socket)
}
 
async fn handle_socket(mut socket: WebSocket) {
    while let Some(Ok(msg)) = socket.recv().await {
        if let Message::Text(text) = msg {
            let _ = socket.send(Message::Text(format!("echo: {text}"))).await;
        }
    }
}

use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    response::sse::{Event, KeepAlive, Sse},
    routing::get,
    Router,
};
use futures_util::stream::{self, Stream};
use std::{convert::Infallible, time::Duration};
use tokio_stream::StreamExt as _;
 
async fn ws_upgrade(ws: WebSocketUpgrade) -> impl axum::response::IntoResponse {
    ws.on_upgrade(handle_ws)
}
 
async fn handle_ws(mut socket: WebSocket) {
    while let Some(Ok(msg)) = socket.recv().await {
        if let Message::Text(t) = msg {
            if socket.send(Message::Text(format!("echo: {t}"))).await.is_err() {
                break;
            }
        }
    }
}
 
async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let stream = stream::iter(0..5).map(|i| {
        Ok(Event::default().data(format!("tick {i}")))
    });
    Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
}
 
#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/ws", get(ws_upgrade))
        .route("/events", get(sse_handler));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

// Broadcast channel fan-out to many WebSocket clients
let (tx, _) = tokio::sync::broadcast::channel::<String>(100);