Skip to main content

rustrade_data/streams/
consumer.rs

1use crate::{
2    Identifier, MarketStream,
3    error::DataError,
4    event::MarketEvent,
5    exchange::StreamSelector,
6    instrument::InstrumentData,
7    streams::{
8        reconnect,
9        reconnect::stream::{
10            ReconnectingStream, ReconnectionBackoffPolicy, ReinitFailure, init_reconnecting_stream,
11            terminate_on_error,
12        },
13    },
14    subscription::{Subscription, SubscriptionKind, display_subscriptions_without_exchange},
15};
16use derive_more::Constructor;
17use futures::{Stream, StreamExt, stream::BoxStream};
18use rustrade_instrument::exchange::ExchangeId;
19use serde::{Deserialize, Serialize};
20use std::fmt::Display;
21use tracing::info;
22
23/// Default [`ReconnectionBackoffPolicy`] for a [`reconnecting`](`ReconnectingStream`) [`MarketStream`].
24pub const STREAM_RECONNECTION_POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
25    backoff_ms_initial: 125,
26    backoff_multiplier: 2,
27    backoff_ms_max: 60000,
28};
29
30/// Convenient type alias for a [`MarketEvent`] [`Result`] consumed via a
31/// [`reconnecting`](`ReconnectingStream`) [`MarketStream`].
32pub type MarketStreamResult<InstrumentKey, Kind> =
33    reconnect::Event<ExchangeId, Result<MarketEvent<InstrumentKey, Kind>, DataError>>;
34
35/// Convenient type alias for a [`MarketEvent`] consumed via a
36/// [`reconnecting`](`ReconnectingStream`) [`MarketStream`].
37pub type MarketStreamEvent<InstrumentKey, Kind> =
38    reconnect::Event<ExchangeId, MarketEvent<InstrumentKey, Kind>>;
39
40/// Initialises a [`reconnecting`](`ReconnectingStream`) [`MarketStream`] using a collection of
41/// [`Subscription`]s.
42///
43/// The provided [`ReconnectionBackoffPolicy`] dictates how the exponential backoff scales
44/// between reconnections.
45///
46/// The `subscriber` is cloned into the reconnect closure, so authenticated subscribers
47/// will have their credentials available on reconnection.
48///
49/// # Errors
50/// Returns the first initialisation's error. A batch the venue did not fully acknowledge is
51/// [`DataError::SubscriptionsUnacknowledged`], naming the subscriptions that received no
52/// acknowledgement.
53///
54/// # Re-initialisation
55/// When the stream ends it yields one [`reconnect::Event::Reconnecting`] and re-initialises with
56/// exponential backoff. Each failed attempt yields an `Event::Item(Err(`[`DataError::ReinitFailed`]`))`
57/// counting consecutive failures; retrying continues until an attempt succeeds. Deciding when to
58/// give up, or to drop a subscription the venue no longer accepts, is the caller's.
59pub async fn init_market_stream<Exchange, Instrument, Kind>(
60    policy: ReconnectionBackoffPolicy,
61    subscriber: Exchange::Subscriber,
62    subscriptions: Vec<Subscription<Exchange, Instrument, Kind>>,
63) -> Result<impl Stream<Item = MarketStreamResult<Instrument::Key, Kind::Event>>, DataError>
64where
65    Exchange: StreamSelector<Instrument, Kind> + Send + Sync + 'static,
66    Instrument: InstrumentData + Display + 'static,
67    Kind: SubscriptionKind + Display + Send + Sync + 'static,
68    Subscription<Exchange, Instrument, Kind>:
69        Identifier<Exchange::Channel> + Identifier<Exchange::Market>,
70{
71    // Determine ExchangeId associated with these Subscriptions
72    let exchange = Exchange::ID;
73
74    // Determine StreamKey for use in logging
75    let stream_key = subscriptions
76        .first()
77        .map(|sub| StreamKey::new("market_stream", exchange, Some(sub.kind.as_str())))
78        .ok_or(DataError::SubscriptionsEmpty)?;
79
80    info!(
81        %exchange,
82        subscriptions = %display_subscriptions_without_exchange(&subscriptions),
83        ?policy,
84        ?stream_key,
85        "MarketStream with auto reconnect initialising"
86    );
87
88    // Both the exchange stream and the stream of attempts are boxed before the reconnect
89    // combinators wrap them. Unboxed, every combinator's type name repeats the whole exchange
90    // stream and init closure types, nested ten or more deep, once per exchange and subscription
91    // kind. Those names are almost all of this crate's debug info: they took a full-debug build
92    // of it to several GB in one rustc. Boxed, the combinators are instantiated once per event
93    // type. The cost is one allocation per connection attempt and a virtual call per poll.
94    let attempts = init_reconnecting_stream(move || {
95        let subscriber = subscriber.clone();
96        let subscriptions = subscriptions.clone();
97        async move {
98            Exchange::Stream::init::<Exchange::SnapFetcher>(&subscriber, &subscriptions)
99                .await
100                .map(StreamExt::boxed)
101        }
102    })
103    .await?
104    .boxed();
105
106    Ok(market_stream_events(attempts, policy, stream_key, exchange))
107}
108
109/// Turns a market stream's initialisation attempts into the events its consumer receives.
110///
111/// Each initialised stream's items, ended early by a terminal error, then one `Reconnecting`. A
112/// failed attempt ends no stream, so it yields `ReinitFailed` and no further `Reconnecting`.
113fn market_stream_events<T>(
114    attempts: BoxStream<'static, Result<BoxStream<'static, Result<T, DataError>>, DataError>>,
115    policy: ReconnectionBackoffPolicy,
116    stream_key: StreamKey,
117    exchange: ExchangeId,
118) -> impl Stream<Item = reconnect::Event<ExchangeId, Result<T, DataError>>> {
119    attempts
120        .with_reconnect_backoff_reporting(policy, stream_key)
121        .map(move |initialised| {
122            initialised.map(|stream| terminate_on_error(stream, DataError::is_terminal, stream_key))
123        })
124        .with_reconnection_events_reporting(exchange, |ReinitFailure { attempt, error }| {
125            Err(DataError::ReinitFailed {
126                attempt,
127                error: Box::new(error),
128            })
129        })
130}
131
132#[derive(
133    Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize, Constructor,
134)]
135pub struct StreamKey<Kind = &'static str> {
136    pub stream: &'static str,
137    pub exchange: ExchangeId,
138    pub kind: Option<Kind>,
139}
140
141impl StreamKey {
142    pub fn new_general(stream: &'static str, exchange: ExchangeId) -> Self {
143        Self::new(stream, exchange, None)
144    }
145}
146
147impl std::fmt::Debug for StreamKey {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        match self.kind {
150            None => write!(f, "{}-{}", self.stream, self.exchange),
151            Some(kind) => write!(f, "{}-{}-{}", self.stream, self.exchange, kind),
152        }
153    }
154}
155
156#[cfg(test)]
157#[allow(clippy::unwrap_used)] // Test code: panics on bad input are acceptable
158mod tests {
159    use super::*;
160    use crate::streams::reconnect::Event;
161    use futures::stream;
162
163    const POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
164        backoff_ms_initial: 1,
165        backoff_multiplier: 2,
166        backoff_ms_max: 10,
167    };
168
169    type Item = Result<u8, DataError>;
170
171    fn initialised(items: Vec<Item>) -> Result<BoxStream<'static, Item>, DataError> {
172        Ok(stream::iter(items).boxed())
173    }
174
175    fn unacknowledged() -> DataError {
176        DataError::SubscriptionsUnacknowledged {
177            reason: "WebSocket stream terminated unexpectedly".to_owned(),
178            unacknowledged: vec!["l2Book|XYZ:TSLA".into()],
179        }
180    }
181
182    #[tokio::test(start_paused = true)]
183    async fn each_failed_reinit_is_yielded_and_one_reconnecting_marks_each_disconnect() {
184        let exchange = ExchangeId::HyperliquidPerp;
185        let attempts = stream::iter([
186            initialised(vec![Ok(1)]),
187            Err(unacknowledged()),
188            Err(unacknowledged()),
189            initialised(vec![Ok(2)]),
190        ]);
191
192        let events = market_stream_events(
193            attempts.boxed(),
194            POLICY,
195            StreamKey::new("market_stream", exchange, Some("l2")),
196            exchange,
197        )
198        .collect::<Vec<_>>()
199        .await;
200
201        let reinit_failed = |attempt| {
202            Event::Item(Err(DataError::ReinitFailed {
203                attempt,
204                error: Box::new(unacknowledged()),
205            }))
206        };
207        assert_eq!(
208            events,
209            [
210                Event::Item(Ok(1)),
211                Event::Reconnecting(exchange),
212                reinit_failed(1),
213                reinit_failed(2),
214                Event::Item(Ok(2)),
215                Event::Reconnecting(exchange),
216            ]
217        );
218    }
219
220    #[tokio::test(start_paused = true)]
221    async fn a_terminal_error_ends_the_stream_before_its_remaining_items() {
222        let exchange = ExchangeId::BinanceSpot;
223        let terminal = DataError::InvalidSequence {
224            prev_last_update_id: 1,
225            first_update_id: 3,
226        };
227        let attempts = stream::iter([initialised(vec![Ok(1), Err(terminal), Ok(2)])]);
228
229        let events = market_stream_events(
230            attempts.boxed(),
231            POLICY,
232            StreamKey::new("market_stream", exchange, Some("l2")),
233            exchange,
234        )
235        .collect::<Vec<_>>()
236        .await;
237
238        assert_eq!(events, [Event::Item(Ok(1)), Event::Reconnecting(exchange)]);
239    }
240}