Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101use std::time::Duration;
use axum_tws::{Limits, Message, WebSocket};use bytes::Bytes;use futures::{SinkExt, Stream, StreamExt};use tokio::time::{MissedTickBehavior, interval, timeout};use tracing::{debug, warn};
const PING_INTERVAL: Duration = Duration::from_secs(30);const CLOSE_TIMEOUT: Duration = Duration::from_secs(1);const CONTROL_FRAME_PAYLOAD_LIMIT: usize = 125;
pub(super) fn control_frame_only_limits() -> Limits { Limits::default().max_payload_len(Some(CONTROL_FRAME_PAYLOAD_LIMIT))}
pub(super) enum WsAction { Send(Message), #[cfg_attr(not(feature = "indexer_stream"), allow(dead_code))] Skip, Close(Option<Message>),}
pub(super) async fn run_socket<S, F, G>( socket: WebSocket, mut events: S, mut to_action: F, send_timeout: Duration, slow_consumer: G,) where S: Stream + Unpin, F: FnMut(S::Item) -> WsAction, G: FnOnce(Duration) -> Option<Message>,{ let (mut sink, mut ws_recv) = socket.split();
let mut ping_timer = interval(PING_INTERVAL); ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); ping_timer.tick().await;
let mut slow_consumer = Some(slow_consumer);
loop { tokio::select! { inbound = ws_recv.next() => match inbound { Some(Ok(m)) if m.is_close() => break, Some(Ok(m)) if m.is_text() || m.is_binary() => { debug!("client sent unsolicited data frame, closing"); break; } Some(Ok(m)) if m.is_ping() => { if let Err(err) = sink.send(Message::pong(m.into_payload())).await { warn!(err = %err, "ws pong send error"); break; } } Some(Ok(_)) => {} Some(Err(e)) => { warn!(err = %e, "ws recv error"); break; } None => break, }, evt = events.next() => match evt { Some(item) => match to_action(item) { WsAction::Skip => {} WsAction::Send(msg) => match timeout(send_timeout, sink.send(msg)).await { Ok(Ok(())) => {} Ok(Err(err)) => { warn!(err = %err, "ws send error"); break; } Err(_) => { if let Some(final_msg) = slow_consumer.take().and_then(|f| f(send_timeout)) { let _ = timeout(CLOSE_TIMEOUT, sink.send(final_msg)).await; } break; } }, WsAction::Close(final_msg) => { if let Some(m) = final_msg { let _ = timeout(send_timeout, sink.send(m)).await; } break; } }, None => break, }, _ = ping_timer.tick() => { if let Err(err) = sink.send(Message::ping(Bytes::new())).await { warn!(err = %err, "ws ping send error"); break; } } } } let _ = timeout(CLOSE_TIMEOUT, sink.close()).await;}