#[cfg(test)]
use std::time::Duration;
use std::{
fmt,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use futures_util::Stream;
use polyoxide_rt::Runtime;
use polyoxide_ws::{SocketTransport, Timing};
use crate::{
shared::ws::{
supervised::{self as shared_supervised, Handle, Running, Settings},
DEFAULT_MAX_CONNECTION_AGE,
},
usdm::ws::{
error::UsdmWsError, event::Update, protocol::UsdmStreams, stream::StreamName, StreamPath,
USDM_WS_BASE,
},
};
#[derive(Debug)]
#[non_exhaustive]
pub enum Event {
Update(Box<Update>),
Disconnected {
path: StreamPath,
reason: DisconnectReason,
},
Reconnected {
path: StreamPath,
},
}
#[derive(Debug)]
#[non_exhaustive]
pub enum DisconnectReason {
Closed {
code: Option<u16>,
reason: String,
},
Error(UsdmWsError),
Stale,
Rotation,
}
impl fmt::Display for DisconnectReason {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Closed { code, reason } => match code {
Some(code) => write!(f, "closed by the server ({code} {reason})"),
None => f.write_str("closed by the server"),
},
Self::Error(err) => write!(f, "{err}"),
Self::Stale => f.write_str("no frame within the staleness window"),
Self::Rotation => f.write_str("rotated before the 24-hour cutoff"),
}
}
}
#[derive(Debug, Clone)]
pub struct UsdmWsBuilder {
base_url: String,
timing: Timing,
streams: Vec<StreamName>,
runtime: Option<Arc<dyn Runtime>>,
socket_transport: Option<Arc<dyn SocketTransport>>,
}
impl Default for UsdmWsBuilder {
fn default() -> Self {
Self::new()
}
}
impl UsdmWsBuilder {
pub fn new() -> Self {
let mut timing = Timing::default();
timing.max_connection_age = Some(DEFAULT_MAX_CONNECTION_AGE);
Self {
base_url: USDM_WS_BASE.to_owned(),
timing,
streams: Vec::new(),
runtime: None,
socket_transport: None,
}
}
pub fn base_url(mut self, url: impl Into<String>) -> Self {
self.base_url = url.into();
self
}
polyoxide_ws::timing_setters! {
ping_interval;
stale_after;
backoff;
connect_timeout;
max_connection_age;
}
polyoxide_ws::transport_setters! {
runtime;
socket_transport;
}
pub fn streams(mut self, streams: impl IntoIterator<Item = StreamName>) -> Self {
self.streams = streams.into_iter().collect();
self
}
pub async fn connect(self) -> Result<SupervisedUsdmWs, UsdmWsError> {
let inner = shared_supervised::connect::<UsdmStreams>(Settings {
base_url: self.base_url,
timing: self.timing,
streams: self.streams,
runtime: self.runtime,
socket_transport: self.socket_transport,
})
.await?;
Ok(SupervisedUsdmWs { inner })
}
}
#[derive(Debug, Clone)]
pub struct MembershipHandle {
inner: Handle<UsdmStreams>,
}
impl MembershipHandle {
pub async fn subscribe(
&self,
streams: impl IntoIterator<Item = StreamName>,
) -> Result<(), UsdmWsError> {
self.inner.subscribe(streams.into_iter().collect()).await
}
pub async fn unsubscribe(
&self,
streams: impl IntoIterator<Item = StreamName>,
) -> Result<(), UsdmWsError> {
self.inner.unsubscribe(streams.into_iter().collect()).await
}
}
pub struct SupervisedUsdmWs {
inner: Running<UsdmStreams>,
}
impl fmt::Debug for SupervisedUsdmWs {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SupervisedUsdmWs")
.field("running", &self.inner.is_running())
.finish_non_exhaustive()
}
}
impl SupervisedUsdmWs {
pub fn membership(&self) -> MembershipHandle {
MembershipHandle {
inner: self.inner.membership(),
}
}
pub async fn close(self) -> Result<(), UsdmWsError> {
self.inner.close().await
}
}
impl Stream for SupervisedUsdmWs {
type Item = Result<Event, UsdmWsError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.inner.poll_event(cx)
}
}
#[cfg(test)]
mod tests {
use polyoxide_ws::Backoff;
use super::*;
#[test]
fn defaults_match_the_documented_cadences() {
let b = UsdmWsBuilder::new();
assert_eq!(b.base_url, USDM_WS_BASE);
assert_eq!(b.timing.ping_interval, Some(Duration::from_secs(20)));
assert_eq!(b.timing.stale_after, Duration::from_secs(30));
assert_eq!(b.timing.initial_backoff, Duration::from_millis(500));
assert_eq!(b.timing.max_backoff, Duration::from_secs(60));
assert_eq!(b.timing.connect_timeout, Duration::from_secs(10));
assert_eq!(
b.timing.max_connection_age,
Some(Duration::from_secs(85_800))
);
assert!(b.timing.max_connection_age < Some(Duration::from_secs(24 * 3600)));
}
#[test]
fn backoff_doubles_to_the_ceiling_and_resets_after_a_working_connection() {
let mut b = Backoff::new(Duration::from_millis(100), Duration::from_millis(350));
assert_eq!(b.take(), Duration::from_millis(100));
assert_eq!(b.take(), Duration::from_millis(200));
assert_eq!(b.take(), Duration::from_millis(350));
b.after_connection_ended(true);
assert_eq!(b.take(), Duration::from_millis(100));
b.after_connection_ended(false);
assert_eq!(b.take(), Duration::from_millis(200));
}
#[test]
fn a_zero_initial_delay_still_grows_and_a_huge_one_saturates() {
let mut b = Backoff::new(Duration::ZERO, Duration::from_secs(1));
assert_eq!(b.take(), Duration::from_millis(1));
assert_eq!(b.take(), Duration::from_millis(2));
let mut b = Backoff::new(Duration::MAX, Duration::MAX);
assert_eq!(b.take(), Duration::MAX);
assert_eq!(b.take(), Duration::MAX);
}
#[derive(Debug, Default)]
struct CountingRuntime {
inner: polyoxide_rt::TokioRuntime,
spawns: std::sync::atomic::AtomicUsize,
}
impl Runtime for CountingRuntime {
fn spawn(&self, task: polyoxide_rt::BoxFuture<()>) {
self.spawns
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
self.inner.spawn(task);
}
fn sleep(&self, duration: Duration) -> polyoxide_rt::BoxFuture<()> {
self.inner.sleep(duration)
}
fn yield_now(&self) -> polyoxide_rt::BoxFuture<()> {
Box::pin(tokio::task::yield_now())
}
fn now(&self) -> polyoxide_rt::Instant {
self.inner.now()
}
}
#[tokio::test]
async fn the_router_and_each_path_are_spawned_on_the_injected_runtime() {
use crate::usdm::ws::scripted::{Script, ScriptedServer};
let server = ScriptedServer::start(vec![Script::default()]).await;
let runtime = std::sync::Arc::new(CountingRuntime::default());
let btc = crate::usdm::types::Symbol::new("BTCUSDT").unwrap();
let feed = UsdmWsBuilder::new()
.base_url(server.url.clone())
.runtime(runtime.clone())
.streams([StreamName::MarkPrice(btc)])
.connect()
.await
.expect("connect");
assert_eq!(runtime.spawns.load(std::sync::atomic::Ordering::SeqCst), 3);
feed.close().await.expect("close");
}
}