use std::{
future::Future,
pin::Pin,
time::{Duration, SystemTime},
vec,
};
use ahash::AHashSet;
use eyre::{Result, bail};
use futures_util::{
FutureExt as _, SinkExt as _, StreamExt as _,
stream::{FuturesUnordered, SplitSink, SplitStream},
};
use jiff::Timestamp;
use reqwest::Url;
use tokio::net::TcpStream;
use tokio_tungstenite::{
MaybeTlsStream, WebSocketStream,
tungstenite::{self, Bytes, Message},
};
use crate::{ConstructAuthError, RetryConfig, UrlError, retry::ExponentialBackoff};
type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
type WsSink = SplitSink<WsStream, Message>;
type WsRead = SplitStream<WsStream>;
type BoxedFu = Pin<Box<dyn Future<Output = FuEvent> + Send + Sync>>;
pub trait WsHandler: std::fmt::Debug {
fn config(&self) -> Result<WsConfig, UrlError> {
Ok(WsConfig::default())
}
#[allow(unused_variables)]
fn handle_auth(&mut self) -> Result<Vec<tungstenite::Message>, WsError> {
Ok(vec![])
}
#[allow(unused_variables)]
fn handle_subscribe(&mut self, topics: AHashSet<Topic>) -> Result<Vec<tungstenite::Message>, WsError>;
fn active_ping(&self) -> Vec<tungstenite::Message> {
vec![]
}
#[allow(unused_variables)]
fn handle_jrpc(&mut self, jrpc: serde_json::Value) -> Result<ResponseOrContent, WsError>;
}
#[derive(Clone, Debug)]
pub enum ResponseOrContent {
Response(Vec<tungstenite::Message>),
Content(ContentEvent),
}
#[derive(Clone, Debug)]
pub struct ContentEvent {
pub data: serde_json::Value,
pub topic: String,
pub time: Timestamp,
pub event_type: String,
}
#[derive(Clone, Debug, Eq)]
pub struct TopicInterpreter<T> {
pub event_name: String,
pub interpret: fn(&serde_json::Value) -> Result<T, WsError>,
}
pub struct WsConnection<H: WsHandler> {
url: Url,
config: WsConfig,
handler: H,
backoff: ExponentialBackoff,
reconnect_after: Option<tokio::time::Instant>,
fu: FuturesUnordered<BoxedFu>,
sink: Option<WsSink>,
outbox: Vec<Message>,
connected_since: Option<SystemTime>,
last_unanswered_communication: Option<SystemTime>,
pending_reconnect: bool,
active_ping_freq: Option<Duration>,
}
impl<H: WsHandler> WsConnection<H> {
#[allow(missing_docs)]
pub fn try_new(url_suffix: &str, handler: H) -> Result<Self, WsError> {
let config = handler.config()?;
let url = match &config.base_url {
Some(base_url) => base_url.join(url_suffix).map_err(UrlError::Parse)?,
None => Url::parse(url_suffix).map_err(UrlError::Parse)?,
};
let backoff = ExponentialBackoff::try_from(&config.reconnect).map_err(|e| WsError::Other(eyre::eyre!("Invalid reconnect backoff configuration: {e}")))?;
let active_ping_freq = config.active_ping_freq;
Ok(Self {
url,
config,
handler,
backoff,
reconnect_after: None,
fu: FuturesUnordered::new(),
sink: None,
outbox: Vec::new(),
connected_since: None,
last_unanswered_communication: None,
pending_reconnect: false,
active_ping_freq,
})
}
pub async fn next(&mut self) -> Result<Vec<ContentEvent>, WsError> {
if let Some(until) = self.reconnect_after.take() {
tokio::time::sleep_until(until).await;
}
if self.pending_reconnect {
self.pending_reconnect = false;
self.reconnect().await?;
}
if let Some(since) = self.connected_since
&& since + self.config.refresh_after < SystemTime::now()
{
tracing::info!("Refreshing connection, as `refresh_after` specified in WsConfig has elapsed ({:?})", self.config.refresh_after);
self.reconnect().await?;
}
if self.connected_since.is_none() {
self.connect().await?;
}
let mut content: Vec<ContentEvent> = Vec::new();
loop {
self.try_flush_outbox();
let timeout = match self.choose_timeout() {
Some(d) => d,
None => {
tracing::error!("Timeout for last unanswered communication ended before `.next()` was called. This likely indicates a clientside implementation error.");
self.reconnect().await?;
continue;
}
};
match tokio::time::timeout(timeout, self.fu.next()).await {
Err(_) => {
if !content.is_empty() {
return Ok(content);
}
if self.last_unanswered_communication.is_some() {
tracing::warn!("Response to a forced communication timed out after {timeout:?}. Reconnecting.");
self.reconnect().await?;
} else {
self.outbox.push(Message::Ping(Bytes::default()));
}
continue;
}
Ok(None) => {
tracing::warn!("FuturesUnordered empty despite permanent reader. Reconnecting.");
self.reconnect().await?;
continue;
}
Ok(Some(FuEvent::Write { sink, result })) => {
self.sink = Some(sink); if let Err(e) = result
&& is_reconnecting(&e)
{
tracing::warn!("Write failed ({e:?}). Reconnecting.");
self.reconnect().await?;
}
continue; }
Ok(Some(FuEvent::PingDue)) => {
self.outbox.extend(self.handler.active_ping());
self.arm_ping();
continue; }
Ok(Some(FuEvent::Read { reader, batch })) => {
if batch.is_empty() {
drop(reader);
if !content.is_empty() {
self.pending_reconnect = true;
return Ok(content);
}
tracing::warn!("tungstenite read EOF from the stream. Reconnecting.");
self.reconnect().await?;
continue;
}
self.arm_reader(reader); self.last_unanswered_communication = None;
let mut terminal = false; for frame in batch {
let __pong_ack = || tracing::trace!("Received app-level pong (active-ping ack)");
match frame {
Ok(Message::Text(text)) => {
let value: serde_json::Value =
serde_json::from_str(&text).expect("API sent invalid JSON, which is completely unexpected. Disappointment is immeasurable and the day is ruined.");
if value["op"] == "pong" || value["ret_msg"] == "pong" {
__pong_ack();
continue;
}
tracing::trace!("{value:#?}");
match self.handler.handle_jrpc(value)? {
ResponseOrContent::Response(messages) => self.outbox.extend(messages),
ResponseOrContent::Content(c) => content.push(c),
}
}
Ok(Message::Ping(_)) => tracing::trace!("Received ping (tungstenite auto-pongs)"),
Ok(Message::Pong(_)) => __pong_ack(),
Ok(Message::Close(maybe_reason)) => {
match maybe_reason {
Some(close_frame) => tracing::info!("Server closed connection; reason: {close_frame:?}"),
None => tracing::info!("Server closed connection; no reason specified."),
}
terminal = true;
break;
}
Ok(Message::Binary(_)) => panic!("Received binary. But exchanges are not smart enough to send this, what is happening"),
Ok(Message::Frame(_)) => unreachable!("Can't get from reading"),
Err(e) if is_reconnecting(&e) => {
tracing::warn!("Reconnecting-class error mid-batch: {e:?}");
terminal = true;
break;
}
Err(e) => {
debug_assert!(!is_reconnecting(&e));
tracing::warn!("Skipping non-fatal polling error: {e:?}");
}
}
}
if terminal {
self.outbox.clear();
if !content.is_empty() {
self.pending_reconnect = true; return Ok(content); }
self.reconnect().await?;
continue;
}
if !content.is_empty() {
return Ok(content);
}
continue; }
}
}
}
fn arm_reader(&mut self, reader: WsRead) {
self.fu.push(Box::pin(read_future(reader)));
}
fn arm_ping(&mut self) {
if let Some(freq) = self.active_ping_freq {
self.fu.push(Box::pin(ping_future(freq)));
}
}
fn try_flush_outbox(&mut self) {
if self.sink.is_none() {
return; }
if self.outbox.is_empty() {
return;
}
let sink = self.sink.take().expect("guarded `is_none` above");
let msgs = std::mem::take(&mut self.outbox);
tracing::debug!("flushing to server: {msgs:#?}");
self.last_unanswered_communication = Some(SystemTime::now());
self.fu.push(Box::pin(write_future(sink, msgs)));
}
fn choose_timeout(&self) -> Option<Duration> {
match self.last_unanswered_communication {
Some(last_unanswered) => match last_unanswered + self.config.response_timeout > SystemTime::now() {
true => Some(self.config.response_timeout),
false => None,
},
None => Some(self.config.message_timeout),
}
}
async fn connect(&mut self) -> Result<(), WsError> {
tracing::info!("Connecting to {}...", self.url);
let (stream, http_resp) = match tokio_tungstenite::connect_async(self.url.as_str()).await {
Ok(result) => result,
Err(e) => {
let delay = self.backoff.next_duration();
if !delay.is_zero() {
tracing::warn!(delay_ms = delay.as_millis(), "Connection failed, backing off before retry.");
self.reconnect_after = Some(tokio::time::Instant::now() + delay);
}
return Err(e.into());
}
};
tracing::debug!("Ws handshake with server: {http_resp:#?}");
let (sink, reader) = stream.split();
self.fu = FuturesUnordered::new();
self.sink = Some(sink);
self.arm_reader(reader);
self.arm_ping(); self.outbox.clear();
self.last_unanswered_communication = None;
self.connected_since = Some(SystemTime::now());
let auth_messages = self.handler.handle_auth()?;
self.outbox.extend(auth_messages);
self.try_flush_outbox();
self.reconnect_after = None;
self.backoff.reset();
Ok(())
}
pub async fn reconnect(&mut self) -> Result<(), WsError> {
self.reconnect_after = None;
if let Some(mut sink) = self.sink.take() {
tracing::info!("Dropping old connection before reconnecting...");
if let Err(e) = sink.send(Message::Close(None)).await {
tracing::debug!("Failed to send Close frame (connection likely already dead): {e}");
}
}
self.fu = FuturesUnordered::new(); self.connected_since = None;
self.connect().await
}
}
#[derive(Clone, Debug)]
pub struct WsConfig {
pub auth: bool,
pub base_url: Option<Url>,
pub reconnect: RetryConfig,
refresh_after: Duration,
message_timeout: Duration,
response_timeout: Duration,
pub topics: AHashSet<String>,
active_ping_freq: Option<Duration>,
}
impl WsConfig {
pub fn set_reconnect(&mut self, reconnect: RetryConfig) {
self.reconnect = reconnect;
}
pub fn set_refresh_after(&mut self, refresh_after: Duration) -> Result<()> {
if refresh_after.is_zero() {
bail!("refresh_after must be greater than 0");
}
self.refresh_after = refresh_after;
Ok(())
}
pub fn set_message_timeout(&mut self, message_timeout: Duration) -> Result<()> {
if message_timeout.is_zero() {
bail!("message_timeout must be greater than 0");
}
self.message_timeout = message_timeout;
Ok(())
}
pub fn set_response_timout(&mut self, response_timeout: Duration) -> Result<()> {
if response_timeout.is_zero() {
bail!("response_timeout must be greater than 0");
}
self.response_timeout = response_timeout;
Ok(())
}
pub fn set_active_ping_freq(&mut self, active_ping_freq: Duration) -> Result<()> {
if active_ping_freq.is_zero() {
bail!("active_ping_freq must be greater than 0");
}
self.active_ping_freq = Some(active_ping_freq);
Ok(())
}
}
#[derive(Debug, miette::Diagnostic, derive_more::Display, thiserror::Error, derive_more::From)]
pub enum WsError {
#[diagnostic(transparent)]
Definition(WsDefinitionError),
#[diagnostic(code(v_exchanges::ws::tungstenite), help("WebSocket protocol error. The connection may need to be reestablished."))]
Tungstenite(tungstenite::Error),
#[diagnostic(transparent)]
Auth(ConstructAuthError),
#[diagnostic(code(v_exchanges::ws::parse), help("Failed to parse WebSocket message. Check if the exchange API has changed."))]
Parse(serde_json::Error),
#[diagnostic(code(v_exchanges::ws::subscription))]
Subscription(String),
#[diagnostic(code(v_exchanges::ws::network), help("Network connection failed. Check your internet connection."))]
NetworkConnection,
#[diagnostic(transparent)]
Url(UrlError),
#[diagnostic(code(v_exchanges::ws::unexpected_event), help("Received an unexpected event from the WebSocket. This may indicate an API change."))]
UnexpectedEvent(serde_json::Value),
#[error(transparent)]
Other(eyre::Report),
}
#[derive(Debug, miette::Diagnostic, derive_more::Display, thiserror::Error)]
pub enum WsDefinitionError {
#[diagnostic(code(v_exchanges::ws::definition::missing_url), help("WebSocket base URL must be configured in WsConfig."))]
MissingUrl,
}
#[derive(Clone, Debug, derive_more::Display, Eq, Hash, PartialEq, serde::Serialize)]
pub enum Topic {
String(String),
Order(serde_json::Value),
}
enum FuEvent {
Read { reader: WsRead, batch: Vec<Result<Message, tungstenite::Error>> },
Write { sink: WsSink, result: Result<(), tungstenite::Error> },
PingDue,
}
async fn read_future(mut reader: WsRead) -> FuEvent {
let mut batch = Vec::new();
if let Some(first) = reader.next().await {
batch.push(first);
while let Some(Some(m)) = reader.next().now_or_never() {
batch.push(m);
}
} FuEvent::Read { reader, batch }
}
async fn write_future(mut sink: WsSink, msgs: Vec<Message>) -> FuEvent {
let mut s = futures_util::stream::iter(msgs).map(Ok);
let result = sink.send_all(&mut s).await;
FuEvent::Write { sink, result }
}
async fn ping_future(freq: Duration) -> FuEvent {
tokio::time::sleep(freq).await;
FuEvent::PingDue
}
fn is_reconnecting(err: &tungstenite::Error) -> bool {
match err {
tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed | tungstenite::Error::Io(_) | tungstenite::Error::Protocol(_) | tungstenite::Error::AttackAttempt => true,
tungstenite::Error::Capacity(_) => false,
tungstenite::Error::Utf8(e) => panic!("received `tungstenite::Error::Utf8` from polling: {e:?}. Exchange is going crazy, aborting"),
tungstenite::Error::WriteBufferFull(_) => unreachable!("can only get from writing"),
tungstenite::Error::Tls(_) | tungstenite::Error::Url(_) | tungstenite::Error::Http(_) | tungstenite::Error::HttpFormat(_) => todo!(),
}
}
impl<T> std::hash::Hash for TopicInterpreter<T> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.event_name.hash(state);
}
}
impl<T> PartialEq for TopicInterpreter<T> {
fn eq(&self, other: &Self) -> bool {
self.event_name == other.event_name
}
}
impl<H: WsHandler> std::fmt::Debug for WsConnection<H> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WsConnection")
.field("url", &self.url)
.field("config", &self.config)
.field("handler", &self.handler)
.field("backoff", &self.backoff)
.field("reconnect_after", &self.reconnect_after)
.field("connected_since", &self.connected_since)
.field("last_unanswered_communication", &self.last_unanswered_communication)
.field("pending_reconnect", &self.pending_reconnect)
.field("outbox_len", &self.outbox.len())
.field("active_ping_freq", &self.active_ping_freq)
.finish_non_exhaustive()
}
}
impl Default for WsConfig {
fn default() -> Self {
Self {
auth: false,
base_url: None,
reconnect: RetryConfig {
max_retries: u32::MAX,
initial_delay_ms: 1_000,
max_delay_ms: 30_000,
backoff_factor: 2.0,
jitter_ms: 500,
immediate_first: false,
max_elapsed_ms: None,
},
refresh_after: Duration::from_hours(12),
message_timeout: Duration::from_mins(16),
response_timeout: Duration::from_mins(2),
topics: AHashSet::new(),
active_ping_freq: None,
}
}
}
#[cfg(test)]
mod tests {
use futures_util::SinkExt as _;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use super::*;
#[derive(Debug)]
struct EchoHandler;
impl WsHandler for EchoHandler {
fn config(&self) -> Result<WsConfig, UrlError> {
let mut c = WsConfig::default();
c.set_message_timeout(Duration::from_millis(200)).expect("non-zero literal");
c.set_response_timout(Duration::from_millis(200)).expect("non-zero literal");
Ok(c)
}
fn handle_subscribe(&mut self, _topics: AHashSet<Topic>) -> Result<Vec<Message>, WsError> {
Ok(vec![])
}
fn handle_jrpc(&mut self, jrpc: serde_json::Value) -> Result<ResponseOrContent, WsError> {
Ok(ResponseOrContent::Content(ContentEvent {
data: jrpc.clone(),
topic: "test".to_owned(),
time: Timestamp::UNIX_EPOCH,
event_type: "test".to_owned(),
}))
}
}
#[derive(Debug)]
struct PingHandler;
impl WsHandler for PingHandler {
fn config(&self) -> Result<WsConfig, UrlError> {
let mut c = WsConfig::default();
c.set_active_ping_freq(Duration::from_millis(80)).expect("non-zero literal");
c.set_message_timeout(Duration::from_secs(5)).expect("non-zero literal");
c.set_response_timout(Duration::from_secs(5)).expect("non-zero literal");
Ok(c)
}
fn active_ping(&self) -> Vec<Message> {
vec![Message::Text("{\"op\":\"ping\"}".into())]
}
fn handle_subscribe(&mut self, _topics: AHashSet<Topic>) -> Result<Vec<Message>, WsError> {
Ok(vec![])
}
fn handle_jrpc(&mut self, jrpc: serde_json::Value) -> Result<ResponseOrContent, WsError> {
Ok(ResponseOrContent::Content(ContentEvent {
data: jrpc,
topic: "test".to_owned(),
time: Timestamp::UNIX_EPOCH,
event_type: "test".to_owned(),
}))
}
}
async fn bind() -> (TcpListener, String) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("loopback bind");
let url = format!("ws://{}", listener.local_addr().unwrap());
(listener, url)
}
#[tokio::test]
async fn drains_whole_buffer_in_one_call() {
const N: usize = 8;
let (listener, url) = bind().await;
let server = async move {
let (tcp, _) = listener.accept().await.expect("accept");
let mut ws = accept_async(tcp).await.expect("handshake");
for i in 0..N {
ws.feed(Message::Text(format!("{{\"n\":{i}}}").into())).await.expect("feed");
}
ws.flush().await.expect("flush");
tokio::time::sleep(Duration::from_secs(3)).await;
drop(ws);
};
let handle = tokio::spawn(server);
let mut conn = WsConnection::try_new(&url, EchoHandler).expect("try_new");
let batch = conn.next().await.expect("next");
assert_eq!(batch.len(), N, "expected all {N} buffered frames drained in a single next()");
handle.abort();
}
#[tokio::test]
async fn pongs_once_and_returns_content_from_same_batch() {
let (listener, url) = bind().await;
let server = async move {
let (tcp, _) = listener.accept().await.expect("accept");
let mut ws = accept_async(tcp).await.expect("handshake");
ws.feed(Message::Ping(Bytes::from_static(b"hi"))).await.expect("feed ping");
ws.feed(Message::Text("{\"n\":1}".into())).await.expect("feed t1");
ws.feed(Message::Text("{\"n\":2}".into())).await.expect("feed t2");
ws.flush().await.expect("flush");
let mut pongs = 0usize;
while let Some(msg) = ws.next().await {
match msg {
Ok(Message::Pong(_)) => pongs += 1,
Ok(Message::Close(_)) | Err(_) => break,
_ => {}
}
}
pongs
};
let handle = tokio::spawn(server);
let mut conn = WsConnection::try_new(&url, EchoHandler).expect("try_new");
let batch = conn.next().await.expect("next");
assert_eq!(batch.len(), 2, "two text frames in the batch become two content events");
let _drove_flush = tokio::time::timeout(Duration::from_millis(150), conn.next()).await;
drop(conn);
let pongs = tokio::time::timeout(Duration::from_secs(1), handle).await.expect("server join timeout").expect("server task");
assert_eq!(pongs, 1, "exactly one Pong (tungstenite's auto-reply) for the single Ping");
}
#[tokio::test]
async fn returns_content_before_close_then_reconnects() {
let (listener, url) = bind().await;
let server = async move {
let (tcp, _) = listener.accept().await.expect("accept 1");
let mut ws = accept_async(tcp).await.expect("handshake 1");
ws.feed(Message::Text("{\"n\":1}".into())).await.expect("feed");
ws.feed(Message::Close(None)).await.expect("feed close");
ws.flush().await.expect("flush");
drop(ws);
let (tcp2, _) = listener.accept().await.expect("accept 2 (reconnect)");
let mut ws2 = accept_async(tcp2).await.expect("handshake 2");
ws2.feed(Message::Text("{\"n\":2}".into())).await.expect("feed 2");
ws2.flush().await.expect("flush 2");
tokio::time::sleep(Duration::from_secs(2)).await;
};
let handle = tokio::spawn(server);
let mut conn = WsConnection::try_new(&url, EchoHandler).expect("try_new");
let first = conn.next().await.expect("next 1");
assert_eq!(first.len(), 1, "content collected before Close is returned, not dropped");
let second = conn.next().await.expect("next 2 after reconnect");
assert_eq!(second.len(), 1, "after reconnect we receive the new connection's frame");
handle.abort();
}
#[tokio::test]
async fn sends_active_ping_on_quiet_connection() {
let (listener, url) = bind().await;
let server = async move {
let (tcp, _) = listener.accept().await.expect("accept");
let mut ws = accept_async(tcp).await.expect("handshake");
loop {
match ws.next().await {
Some(Ok(Message::Text(t))) if t.contains("\"op\":\"ping\"") => return true,
Some(Ok(_)) => continue, _ => return false, }
}
};
let handle = tokio::spawn(server);
let mut conn = WsConnection::try_new(&url, PingHandler).expect("try_new");
let _drove_ping = tokio::time::timeout(Duration::from_millis(400), conn.next()).await;
let saw_ping = tokio::time::timeout(Duration::from_secs(1), handle).await.expect("server join timeout").expect("server task");
assert!(saw_ping, "client must send an app-level `{{\"op\":\"ping\"}}` on a quiet connection");
}
#[tokio::test]
async fn drops_app_level_pong_ack() {
let (listener, url) = bind().await;
let server = async move {
let (tcp, _) = listener.accept().await.expect("accept");
let mut ws = accept_async(tcp).await.expect("handshake");
ws.feed(Message::Text("{\"op\":\"pong\"}".into())).await.expect("feed pong");
ws.feed(Message::Text("{\"n\":1}".into())).await.expect("feed content");
ws.flush().await.expect("flush");
tokio::time::sleep(Duration::from_secs(2)).await;
};
let handle = tokio::spawn(server);
let mut conn = WsConnection::try_new(&url, PingHandler).expect("try_new");
let batch = conn.next().await.expect("next");
assert_eq!(batch.len(), 1, "pong-ack dropped upstream; only the real content frame surfaces");
assert_eq!(batch[0].data["n"], 1, "the surviving event is the content frame, not the pong");
handle.abort();
}
}