Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970use crate::control::Hydrant;use axum::Router;use axum::routing::get;use axum::{ extract::{Query, State}, response::IntoResponse,};use axum_tws::{Message, WebSocket, WebSocketUpgrade};use serde::Deserialize;use tracing::error;
use super::ws::{WsAction, control_frame_only_limits, run_socket};
pub fn router() -> Router<Hydrant> { Router::new().route("/", get(handle_stream))}
#[derive(Deserialize)]pub struct StreamQuery { pub cursor: Option<u64>,}
pub async fn handle_stream( State(hydrant): State<Hydrant>, Query(query): Query<StreamQuery>, ws: WebSocketUpgrade,) -> impl IntoResponse { ws.limits(control_frame_only_limits()) .on_upgrade(move |socket| handle_socket(socket, hydrant, query))}
async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { let send_timeout = hydrant.stream_send_timeout(); let events = hydrant.subscribe(query.cursor); run_socket( socket, events, |item| match item { Ok(evt) => match serde_json::to_string(&evt) { Ok(json) => WsAction::Send(Message::text(json)), Err(e) => { error!(err = %e, "failed to serialize event"); WsAction::Skip } }, Err(err) => { let json = serde_json::json!({ "type": "error", "error": err.code(), "message": err.to_string(), }); WsAction::Close(Some(Message::text(json.to_string()))) } }, send_timeout, |timeout_dur| { let json = serde_json::json!({ "type": "error", "error": "ConsumerTooSlow", "message": format!( "stream socket send blocked for at least {} seconds", timeout_dur.as_secs() ), }); Some(Message::text(json.to_string())) }, ) .await;}