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::{connect as connect_running, Handle, Running, Settings},
DEFAULT_MAX_CONNECTION_AGE,
},
spot::ws::{
error::SpotWsError, event::Update, protocol::SpotStreams, stream::StreamName, StreamPath,
SPOT_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(SpotWsError),
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 SpotWsBuilder {
base_url: String,
timing: Timing,
streams: Vec<StreamName>,
runtime: Option<Arc<dyn Runtime>>,
socket_transport: Option<Arc<dyn SocketTransport>>,
}
impl Default for SpotWsBuilder {
fn default() -> Self {
Self::new()
}
}
impl SpotWsBuilder {
pub fn new() -> Self {
let mut timing = Timing::default();
timing.max_connection_age = Some(DEFAULT_MAX_CONNECTION_AGE);
Self {
base_url: SPOT_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<SupervisedSpotWs, SpotWsError> {
let running = connect_running::<SpotStreams>(Settings {
base_url: self.base_url,
timing: self.timing,
streams: self.streams,
runtime: self.runtime,
socket_transport: self.socket_transport,
})
.await?;
Ok(SupervisedSpotWs { running })
}
}
#[derive(Debug, Clone)]
pub struct MembershipHandle {
handle: Handle<SpotStreams>,
}
impl MembershipHandle {
pub async fn subscribe(
&self,
streams: impl IntoIterator<Item = StreamName>,
) -> Result<(), SpotWsError> {
self.handle.subscribe(streams.into_iter().collect()).await
}
pub async fn unsubscribe(
&self,
streams: impl IntoIterator<Item = StreamName>,
) -> Result<(), SpotWsError> {
self.handle.unsubscribe(streams.into_iter().collect()).await
}
}
pub struct SupervisedSpotWs {
running: Running<SpotStreams>,
}
impl fmt::Debug for SupervisedSpotWs {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SupervisedSpotWs")
.field("running", &self.running.is_running())
.finish_non_exhaustive()
}
}
impl SupervisedSpotWs {
pub fn membership(&self) -> MembershipHandle {
MembershipHandle {
handle: self.running.membership(),
}
}
pub async fn close(self) -> Result<(), SpotWsError> {
self.running.close().await
}
}
impl Stream for SupervisedSpotWs {
type Item = Result<Event, SpotWsError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.running.poll_event(cx)
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
#[test]
fn the_defaults_match_the_measured_cadences() {
let b = SpotWsBuilder::new();
assert_eq!(b.base_url, "wss://stream.binance.com:9443");
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(23 * 3600 + 50 * 60))
);
assert!(b.timing.max_connection_age < Some(Duration::from_secs(24 * 3600)));
assert!(b.streams.is_empty());
}
#[test]
fn the_setters_replace_their_field_and_default_is_new() {
let b = SpotWsBuilder::default()
.base_url("ws://127.0.0.1:9")
.streams([StreamName::AllMiniTickers]);
assert_eq!(b.base_url, "ws://127.0.0.1:9");
assert_eq!(b.streams, [StreamName::AllMiniTickers]);
assert_eq!(
SpotWsBuilder::default().timing.stale_after,
SpotWsBuilder::new().timing.stale_after
);
}
#[test]
fn a_disconnect_reason_reads_as_the_texts_usdm_gives() {
assert_eq!(
DisconnectReason::Closed {
code: Some(1008),
reason: "Too many requests".to_owned()
}
.to_string(),
"closed by the server (1008 Too many requests)"
);
assert_eq!(
DisconnectReason::Closed {
code: None,
reason: String::new()
}
.to_string(),
"closed by the server"
);
let refused = SpotWsError::Refused {
code: 2,
msg: "Invalid request".to_owned(),
};
let text = refused.to_string();
assert_eq!(DisconnectReason::Error(refused).to_string(), text);
assert_eq!(
DisconnectReason::Stale.to_string(),
"no frame within the staleness window"
);
assert_eq!(
DisconnectReason::Rotation.to_string(),
"rotated before the 24-hour cutoff"
);
}
#[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::spot::ws::scripted::{Script, ScriptedServer};
let server = ScriptedServer::start(vec![Script::default()]).await;
let runtime = std::sync::Arc::new(CountingRuntime::default());
let btc = crate::spot::types::Symbol::new("BTCUSDT").unwrap();
let feed = SpotWsBuilder::new()
.base_url(server.url.clone())
.runtime(runtime.clone())
.streams([StreamName::AggTrade(btc)])
.connect()
.await
.expect("connect");
assert_eq!(runtime.spawns.load(std::sync::atomic::Ordering::SeqCst), 3);
feed.close().await.expect("close");
}
}