use crate::{
Identifier, MarketStream, SnapshotFetcher,
error::DataError,
exchange::Connector,
instrument::InstrumentData,
process_buffered_events,
subscriber::{Subscribed, Subscriber},
subscription::{Subscription, SubscriptionKind},
transformer::ExchangeTransformer,
};
use futures::Stream;
use rustrade_integration::{
Transformer,
protocol::{
StreamParser,
websocket::{WebSocketSerdeParser, WsError, WsMessage},
},
stream::ExchangeStream,
};
use std::{
pin::Pin,
task::{Context, Poll},
};
use tokio::sync::mpsc;
pub trait SharedTransport:
Stream<Item = Result<WsMessage, WsError>> + Unpin + Send + sealed::Sealed
{
type Attached: Send;
fn take_attached(&mut self) -> Self::Attached;
}
pub(crate) mod sealed {
pub trait Sealed {}
}
pub trait AttachedTransformer<Attached> {
fn attached(&mut self, attached: Attached);
}
impl<T> AttachedTransformer<()> for T {
fn attached(&mut self, (): ()) {}
}
pub struct SharedStream<Transport, StreamTransformer>
where
Transport: SharedTransport,
StreamTransformer: Transformer,
WebSocketSerdeParser: StreamParser<StreamTransformer::Input>,
{
inner: ExchangeStream<WebSocketSerdeParser, Transport, StreamTransformer>,
}
impl<Transport, StreamTransformer> std::fmt::Debug for SharedStream<Transport, StreamTransformer>
where
Transport: SharedTransport,
StreamTransformer: Transformer,
WebSocketSerdeParser: StreamParser<StreamTransformer::Input>,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SharedStream")
.field("buffered", &self.inner.buffer.len())
.finish_non_exhaustive()
}
}
impl<Transport, StreamTransformer> Stream for SharedStream<Transport, StreamTransformer>
where
Transport: SharedTransport,
StreamTransformer: Transformer,
WebSocketSerdeParser: StreamParser<StreamTransformer::Input>,
ExchangeStream<WebSocketSerdeParser, Transport, StreamTransformer>:
Stream<Item = Result<StreamTransformer::Output, StreamTransformer::Error>> + Unpin,
{
type Item = Result<StreamTransformer::Output, StreamTransformer::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.get_mut().inner).poll_next(cx)
}
}
impl<Exchange, Instrument, Kind, Transport, StreamTransformer>
MarketStream<Exchange, Instrument, Kind> for SharedStream<Transport, StreamTransformer>
where
Exchange: Connector + Send + Sync,
Exchange::Subscriber: Subscriber<Transport = Transport>,
Transport: SharedTransport,
Instrument: InstrumentData,
Kind: SubscriptionKind + Send + Sync,
Kind::Event: Send + Sync,
StreamTransformer: ExchangeTransformer<Exchange, Instrument::Key, Kind>
+ AttachedTransformer<Transport::Attached>
+ Send,
WebSocketSerdeParser:
StreamParser<StreamTransformer::Input, Message = WsMessage, Error = WsError>,
{
async fn init<SnapFetcher>(
subscriber: &Exchange::Subscriber,
subscriptions: &[Subscription<Exchange, Instrument, Kind>],
) -> Result<Self, DataError>
where
SnapFetcher: SnapshotFetcher<Exchange, Kind>,
Subscription<Exchange, Instrument, Kind>:
Identifier<Exchange::Channel> + Identifier<Exchange::Market>,
{
let Subscribed {
transport: mut attachment,
map: instrument_map,
buffered_websocket_events,
} = subscriber.subscribe(subscriptions).await?;
let initial_snapshots = SnapFetcher::fetch_snapshots(subscriptions).await?;
let (ws_sink_tx, _unread) = mpsc::unbounded_channel();
let mut transformer =
StreamTransformer::init(instrument_map, &initial_snapshots, ws_sink_tx).await?;
transformer.attached(attachment.take_attached());
let mut processed = process_buffered_events::<WebSocketSerdeParser, _>(
&mut transformer,
buffered_websocket_events,
);
processed.extend(initial_snapshots.into_iter().map(Ok));
Ok(Self {
inner: ExchangeStream::new(attachment, transformer, processed),
})
}
}