use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use futures::stream::Stream;
use super::pricing::PriceUpdate;
use super::source::{StreamCommand, StreamSource, run_stream_loop};
use super::subscription::Subscription;
use super::yahoo::YahooStreamSource;
use crate::error::FinanceError;
pub type StreamResult<T> = std::result::Result<T, StreamError>;
#[derive(Debug, Clone)]
pub enum StreamError {
ConnectionFailed(String),
WebSocketError(String),
DecodeError(String),
}
impl std::fmt::Display for StreamError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
StreamError::ConnectionFailed(e) => write!(f, "Connection failed: {}", e),
StreamError::WebSocketError(e) => write!(f, "WebSocket error: {}", e),
StreamError::DecodeError(e) => write!(f, "Decode error: {}", e),
}
}
}
impl std::error::Error for StreamError {}
impl From<StreamError> for FinanceError {
fn from(e: StreamError) -> Self {
FinanceError::ResponseStructureError {
field: "streaming".to_string(),
context: e.to_string(),
}
}
}
const RECONNECT_BACKOFF_SECS: u64 = 3;
const CHANNEL_CAPACITY: usize = 1024;
pub struct PriceStream {
inner: Subscription<PriceUpdate, StreamCommand>,
}
impl PriceStream {
pub async fn subscribe<S, I>(symbols: I) -> StreamResult<Self>
where
S: Into<String>,
I: IntoIterator<Item = S>,
{
Self::subscribe_with_source(
Arc::new(YahooStreamSource),
symbols,
Duration::from_secs(RECONNECT_BACKOFF_SECS),
)
.await
}
pub(crate) async fn subscribe_with_source<S, I>(
source: Arc<dyn StreamSource>,
symbols: I,
retry_delay: Duration,
) -> StreamResult<Self>
where
S: Into<String>,
I: IntoIterator<Item = S>,
{
let initial_symbols: Vec<String> = symbols.into_iter().map(Into::into).collect();
let inner = Subscription::start(
CHANNEL_CAPACITY,
32,
move |broadcast_tx, command_rx| async move {
let _ = run_stream_loop(
source,
initial_symbols,
broadcast_tx,
command_rx,
retry_delay,
)
.await;
},
);
Ok(PriceStream { inner })
}
pub fn resubscribe(&self) -> Self {
PriceStream {
inner: self.inner.resubscribe(),
}
}
pub async fn add_symbols<S, I>(&self, symbols: I)
where
S: Into<String>,
I: IntoIterator<Item = S>,
{
let symbols: Vec<String> = symbols.into_iter().map(Into::into).collect();
self.inner.send(StreamCommand::Subscribe(symbols)).await;
}
pub async fn remove_symbols<S, I>(&self, symbols: I)
where
S: Into<String>,
I: IntoIterator<Item = S>,
{
let symbols: Vec<String> = symbols.into_iter().map(Into::into).collect();
self.inner.send(StreamCommand::Unsubscribe(symbols)).await;
}
pub async fn close(&self) {
self.inner.send(StreamCommand::Close).await;
}
}
impl Stream for PriceStream {
type Item = PriceUpdate;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.inner).poll_next(cx)
}
}
pub struct PriceStreamBuilder {
symbols: Vec<String>,
retry_delay: Duration,
}
impl PriceStreamBuilder {
pub fn new() -> Self {
Self {
symbols: Vec::new(),
retry_delay: Duration::from_secs(RECONNECT_BACKOFF_SECS),
}
}
pub fn symbols<S, I>(mut self, symbols: I) -> Self
where
S: Into<String>,
I: IntoIterator<Item = S>,
{
self.symbols.extend(symbols.into_iter().map(Into::into));
self
}
pub fn retry(mut self, delay: Duration) -> Self {
self.retry_delay = delay;
self
}
pub async fn build(self) -> StreamResult<PriceStream> {
PriceStream::subscribe_with_source(
Arc::new(YahooStreamSource),
self.symbols,
self.retry_delay,
)
.await
}
}
impl Default for PriceStreamBuilder {
fn default() -> Self {
Self::new()
}
}