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
23pub const STREAM_RECONNECTION_POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
25 backoff_ms_initial: 125,
26 backoff_multiplier: 2,
27 backoff_ms_max: 60000,
28};
29
30pub type MarketStreamResult<InstrumentKey, Kind> =
33 reconnect::Event<ExchangeId, Result<MarketEvent<InstrumentKey, Kind>, DataError>>;
34
35pub type MarketStreamEvent<InstrumentKey, Kind> =
38 reconnect::Event<ExchangeId, MarketEvent<InstrumentKey, Kind>>;
39
40pub 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 let exchange = Exchange::ID;
73
74 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 = 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
109fn 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)] mod 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}