barter_data_sniper/streams/builder/
multi.rs

1use super::{ExchangeChannel, StreamBuilder, Streams};
2use crate::{
3    error::DataError,
4    streams::{consumer::MarketStreamResult, reconnect::stream::ReconnectingStream},
5    subscription::SubscriptionKind,
6};
7use barter_instrument_copy::exchange::ExchangeId;
8use futures_util::StreamExt;
9use std::{collections::HashMap, fmt::Debug, future::Future, pin::Pin};
10
11/// Communicative type alias representing the [`Future`] result of a [`StreamBuilder::init`] call
12/// generated whilst executing [`MultiStreamBuilder::add`].
13pub type BuilderInitFuture = Pin<Box<dyn Future<Output = Result<(), DataError>>>>;
14
15/// Builder to configure and initialise a common [`Streams<Output>`](Streams) instance from
16/// multiple [`StreamBuilder<SubscriptionKind>`](StreamBuilder)s.
17#[derive(Default)]
18pub struct MultiStreamBuilder<Output> {
19    pub channels: HashMap<ExchangeId, ExchangeChannel<Output>>,
20    pub futures: Vec<BuilderInitFuture>,
21}
22
23impl<Output> Debug for MultiStreamBuilder<Output>
24where
25    Output: Debug,
26{
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("MultiStreamBuilder<Output>")
29            .field("channels", &self.channels)
30            .field("num_futures", &self.futures.len())
31            .finish()
32    }
33}
34
35impl<Output> MultiStreamBuilder<Output> {
36    /// Construct a new [`Self`].
37    pub fn new() -> Self {
38        Self {
39            channels: HashMap::new(),
40            futures: Vec::new(),
41        }
42    }
43
44    /// Add a [`StreamBuilder<SubscriptionKind>`](StreamBuilder) to the [`MultiStreamBuilder`]. Creates a
45    /// [`Future`] that calls [`StreamBuilder::init`] and maps the [`SubscriptionKind::Event`](SubscriptionKind)
46    /// into a common `Output`.
47    ///
48    /// Note that the created [`Future`] is not awaited until the [`MultiStreamBuilder::init`]
49    /// method is invoked.
50    #[allow(clippy::should_implement_trait)]
51    pub fn add<InstrumentKey, Kind>(mut self, builder: StreamBuilder<InstrumentKey, Kind>) -> Self
52    where
53        Output: From<MarketStreamResult<InstrumentKey, Kind::Event>> + Clone + Send + 'static,
54        InstrumentKey: Debug + Send + 'static,
55        Kind: SubscriptionKind + 'static,
56        Kind::Event: Send,
57    {
58        // Allocate HashMap to hold the exchange_tx<Output> for each StreamBuilder exchange present
59        let mut exchange_txs = HashMap::with_capacity(builder.channels.len());
60
61        // Iterate over each StreamBuilder exchange present
62        for exchange in builder.channels.keys().cloned() {
63            // Insert ExchangeChannel<Output> Entry to Self for each exchange
64            let exchange_tx = self.channels.entry(exchange).or_default().tx.clone();
65
66            // Insert new exchange_tx<Output> into HashMap for each exchange
67            exchange_txs.insert(exchange, exchange_tx);
68        }
69
70        // Init Streams<Kind::Event> & send mapped Outputs to the associated exchange_tx
71        self.futures.push(Box::pin(async move {
72            builder
73                .init()
74                .await?
75                .streams
76                .into_iter()
77                .for_each(|(exchange, exchange_rx)| {
78                    // Remove exchange_tx<Output> from HashMap that's associated with this tuple:
79                    // (ExchangeId, exchange_rx<MarketStreamResult<InstrumentKey, SubscriptionKind::Event>>)
80                    let exchange_tx = exchange_txs
81                        .remove(&exchange)
82                        .expect("all exchange_txs should be present here");
83
84                    // Task to receive MarketStreamResult<SubscriptionKind::Event> and send Outputs via exchange_tx
85                    tokio::spawn(
86                        exchange_rx
87                            .into_stream()
88                            .map(Output::from)
89                            .forward_to(exchange_tx),
90                    );
91                });
92
93            Ok(())
94        }));
95
96        self
97    }
98
99    /// Initialise each [`StreamBuilder<SubscriptionKind>`](StreamBuilder) that was added to the
100    /// [`MultiStreamBuilder`] and map all [`Streams<SubscriptionKind::Event>`](Streams) into a common
101    /// [`Streams<Output>`](Streams).
102    pub async fn init(self) -> Result<Streams<Output>, DataError> {
103        // Await Stream initialisation perpetual and ensure success
104        futures::future::try_join_all(self.futures).await?;
105
106        // Construct Streams<Output> using each ExchangeChannel receiver
107        Ok(Streams {
108            streams: self
109                .channels
110                .into_iter()
111                .map(|(exchange, channel)| (exchange, channel.rx))
112                .collect(),
113        })
114    }
115}
barter_data_sniper/streams/builder/multi.rs

barter_data_sniper/streams/builder/
multi.rs
