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
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>,
66 Instrument: InstrumentData + Display,
67 Kind: SubscriptionKind + Display,
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 =
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
101fn 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)] mod 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}