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};
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>,
66    Instrument: InstrumentData + Display,
67    Kind: SubscriptionKind + Display,
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    let attempts =
89        init_reconnecting_stream(move || {
90            let subscriber = subscriber.clone();
91            let subscriptions = subscriptions.clone();
92            async move {
93                Exchange::Stream::init::<Exchange::SnapFetcher>(&subscriber, &subscriptions).await
94            }
95        })
96        .await?;
97
98    Ok(market_stream_events(attempts, policy, stream_key, exchange))
99}
100
101/// Turns a market stream's initialisation attempts into the events its consumer receives.
102///
103/// Each initialised stream's items, ended early by a terminal error, then one `Reconnecting`. A
104/// failed attempt ends no stream, so it yields `ReinitFailed` and no further `Reconnecting`.
105fn market_stream_events<Attempts, St, T>(
106    attempts: Attempts,
107    policy: ReconnectionBackoffPolicy,
108    stream_key: StreamKey,
109    exchange: ExchangeId,
110) -> impl Stream<Item = reconnect::Event<ExchangeId, Result<T, DataError>>>
111where
112    Attempts: Stream<Item = Result<St, DataError>>,
113    St: Stream<Item = Result<T, DataError>>,
114{
115    attempts
116        .with_reconnect_backoff_reporting(policy, stream_key)
117        .map(move |initialised| {
118            initialised.map(|stream| terminate_on_error(stream, DataError::is_terminal, stream_key))
119        })
120        .with_reconnection_events_reporting(exchange, |ReinitFailure { attempt, error }| {
121            Err(DataError::ReinitFailed {
122                attempt,
123                error: Box::new(error),
124            })
125        })
126}
127
128#[derive(
129    Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize, Constructor,
130)]
131pub struct StreamKey<Kind = &'static str> {
132    pub stream: &'static str,
133    pub exchange: ExchangeId,
134    pub kind: Option<Kind>,
135}
136
137impl StreamKey {
138    pub fn new_general(stream: &'static str, exchange: ExchangeId) -> Self {
139        Self::new(stream, exchange, None)
140    }
141}
142
143impl std::fmt::Debug for StreamKey {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        match self.kind {
146            None => write!(f, "{}-{}", self.stream, self.exchange),
147            Some(kind) => write!(f, "{}-{}-{}", self.stream, self.exchange, kind),
148        }
149    }
150}
151
152#[cfg(test)]
153#[allow(clippy::unwrap_used)] // Test code: panics on bad input are acceptable
154mod tests {
155    use super::*;
156    use crate::streams::reconnect::Event;
157    use futures::stream;
158
159    const POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
160        backoff_ms_initial: 1,
161        backoff_multiplier: 2,
162        backoff_ms_max: 10,
163    };
164
165    type Item = Result<u8, DataError>;
166
167    fn initialised(items: Vec<Item>) -> Result<stream::Iter<std::vec::IntoIter<Item>>, DataError> {
168        Ok(stream::iter(items))
169    }
170
171    fn unacknowledged() -> DataError {
172        DataError::SubscriptionsUnacknowledged {
173            reason: "WebSocket stream terminated unexpectedly".to_owned(),
174            unacknowledged: vec!["l2Book|XYZ:TSLA".into()],
175        }
176    }
177
178    #[tokio::test(start_paused = true)]
179    async fn each_failed_reinit_is_yielded_and_one_reconnecting_marks_each_disconnect() {
180        let exchange = ExchangeId::HyperliquidPerp;
181        let attempts = stream::iter([
182            initialised(vec![Ok(1)]),
183            Err(unacknowledged()),
184            Err(unacknowledged()),
185            initialised(vec![Ok(2)]),
186        ]);
187
188        let events = market_stream_events(
189            attempts,
190            POLICY,
191            StreamKey::new("market_stream", exchange, Some("l2")),
192            exchange,
193        )
194        .collect::<Vec<_>>()
195        .await;
196
197        let reinit_failed = |attempt| {
198            Event::Item(Err(DataError::ReinitFailed {
199                attempt,
200                error: Box::new(unacknowledged()),
201            }))
202        };
203        assert_eq!(
204            events,
205            [
206                Event::Item(Ok(1)),
207                Event::Reconnecting(exchange),
208                reinit_failed(1),
209                reinit_failed(2),
210                Event::Item(Ok(2)),
211                Event::Reconnecting(exchange),
212            ]
213        );
214    }
215
216    #[tokio::test(start_paused = true)]
217    async fn a_terminal_error_ends_the_stream_before_its_remaining_items() {
218        let exchange = ExchangeId::BinanceSpot;
219        let terminal = DataError::InvalidSequence {
220            prev_last_update_id: 1,
221            first_update_id: 3,
222        };
223        let attempts = stream::iter([initialised(vec![Ok(1), Err(terminal), Ok(2)])]);
224
225        let events = market_stream_events(
226            attempts,
227            POLICY,
228            StreamKey::new("market_stream", exchange, Some("l2")),
229            exchange,
230        )
231        .collect::<Vec<_>>()
232        .await;
233
234        assert_eq!(events, [Event::Item(Ok(1)), Event::Reconnecting(exchange)]);
235    }
236}