barter_data/streams/builder/
mod.rs

1use super::Streams;
2use crate::{
3    Identifier,
4    error::DataError,
5    exchange::StreamSelector,
6    instrument::InstrumentData,
7    streams::{
8        consumer::{MarketStreamResult, STREAM_RECONNECTION_POLICY, init_market_stream},
9        reconnect::stream::ReconnectingStream,
10    },
11    subscription::{Subscription, SubscriptionKind},
12};
13use barter_instrument::exchange::ExchangeId;
14use barter_integration::{Validator, channel::Channel};
15use std::{
16    collections::HashMap,
17    fmt::{Debug, Display},
18    future::Future,
19    pin::Pin,
20};
21
22/// Defines the [`MultiStreamBuilder`](multi::MultiStreamBuilder) API for ergonomically
23/// initialising a common [`Streams<Output>`](Streams) from multiple
24/// [`StreamBuilder<SubscriptionKind>`](StreamBuilder)s.
25pub mod multi;
26
27/// Defines the [`DynamicStreams`](dynamic::DynamicStreams) API for initialising an arbitrary number
28/// of `MarketStream`s from the [`ExchangeId`] and [`SubKind`](crate::subscription::SubKind) enums, rather than concrete
29/// types.
30pub mod dynamic;
31
32/// Communicative type alias representing the [`Future`] result of a [`Subscription`] validation
33/// call generated whilst executing [`StreamBuilder::subscribe`].
34pub type SubscribeFuture = Pin<Box<dyn Future<Output = Result<(), DataError>>>>;
35
36/// Builder to configure and initialise a [`Streams<MarketEvent<SubscriptionKind::Event>`](Streams) instance
37/// for a specific [`SubscriptionKind`].
38#[derive(Default)]
39pub struct StreamBuilder<InstrumentKey, Kind>
40where
41    Kind: SubscriptionKind,
42{
43    pub channels: HashMap<ExchangeId, Channel<MarketStreamResult<InstrumentKey, Kind::Event>>>,
44    pub futures: Vec<SubscribeFuture>,
45}
46
47impl<InstrumentKey, Kind> Debug for StreamBuilder<InstrumentKey, Kind>
48where
49    InstrumentKey: Debug,
50    Kind: SubscriptionKind,
51{
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("StreamBuilder<InstrumentKey, SubscriptionKind>")
54            .field("channels", &self.channels)
55            .field("num_futures", &self.futures.len())
56            .finish()
57    }
58}
59
60impl<InstrumentKey, Kind> StreamBuilder<InstrumentKey, Kind>
61where
62    Kind: SubscriptionKind,
63{
64    /// Construct a new [`Self`].
65    pub fn new() -> Self {
66        Self {
67            channels: HashMap::new(),
68            futures: Vec::new(),
69        }
70    }
71
72    /// Add a collection of [`Subscription`]s to the [`StreamBuilder`] that will be actioned on
73    /// a distinct [`WebSocket`](barter_integration::protocol::websocket::WebSocket) connection.
74    ///
75    /// Note that [`Subscription`]s are not actioned until the
76    /// [`init()`](StreamBuilder::init()) method is invoked.
77    pub fn subscribe<SubIter, Sub, Exchange, Instrument>(mut self, subscriptions: SubIter) -> Self
78    where
79        SubIter: IntoIterator<Item = Sub>,
80        Sub: Into<Subscription<Exchange, Instrument, Kind>>,
81        Exchange: StreamSelector<Instrument, Kind> + Ord + Send + Sync + 'static,
82        Instrument: InstrumentData<Key = InstrumentKey> + Ord + Display + 'static,
83        Instrument::Key: Debug + Clone + Send + 'static,
84        Kind: Ord + Display + Send + Sync + 'static,
85        Kind::Event: Clone + Send,
86        Subscription<Exchange, Instrument, Kind>:
87            Identifier<Exchange::Channel> + Identifier<Exchange::Market>,
88    {
89        // Construct Vec<Subscriptions> from input SubIter
90        let subscriptions = subscriptions.into_iter().map(Sub::into).collect::<Vec<_>>();
91
92        // Acquire channel Sender to send Market<Kind::Event> from consumer loop to user
93        // '--> Add ExchangeChannel Entry if this Exchange <--> SubscriptionKind combination is new
94        let exchange_tx = self.channels.entry(Exchange::ID).or_default().tx.clone();
95
96        // Add Future that once awaited will yield the Result<(), SocketError> of subscribing
97        self.futures.push(Box::pin(async move {
98            // Validate Subscriptions
99            let mut subscriptions = subscriptions
100                .into_iter()
101                .map(Subscription::validate)
102                .collect::<Result<Vec<_>, _>>()?;
103
104            // Remove duplicate Subscriptions
105            subscriptions.sort();
106            subscriptions.dedup();
107
108            // Initialise a MarketEvent `ReconnectingStream`
109            let stream = init_market_stream(STREAM_RECONNECTION_POLICY, subscriptions).await?;
110
111            // Forward MarketEvents to ExchangeTx
112            tokio::spawn(stream.forward_to(exchange_tx));
113
114            Ok(())
115        }));
116
117        self
118    }
119
120    /// Spawn a [`MarketStreamResult<SubscriptionKind::Event>`](MarketStreamResult) consumer loop
121    /// for each collection of [`Subscription`]s added to [`StreamBuilder`] via the
122    /// [`subscribe()`](StreamBuilder::subscribe()) method.
123    ///
124    /// Each consumer loop distributes consumed [`MarketStreamResult`] to
125    /// the [`Streams`] `HashMap` returned by this method.
126    pub async fn init(
127        self,
128    ) -> Result<Streams<MarketStreamResult<InstrumentKey, Kind::Event>>, DataError> {
129        // Await Stream initialisation perpetual and ensure success
130        futures::future::try_join_all(self.futures).await?;
131
132        // Construct Streams using each ExchangeChannel receiver
133        Ok(Streams {
134            streams: self
135                .channels
136                .into_iter()
137                .map(|(exchange, channel)| (exchange, channel.rx))
138                .collect(),
139        })
140    }
141}