rustrade_data/streams/reconnect/
stream.rs

1use crate::streams::{consumer::StreamKey, reconnect::Event};
2use derive_more::Constructor;
3use futures::{Stream, future::Either, stream};
4use futures_util::StreamExt;
5use rustrade_integration::channel::Tx;
6use serde::{Deserialize, Serialize};
7use std::{convert, fmt::Debug, future, future::Future};
8use tracing::{error, info, warn};
9
10/// Utilities for handling a continually reconnecting [`Stream`] initialised via the
11/// [`init_reconnecting_stream`] function.
12pub trait ReconnectingStream
13where
14    Self: Stream + Sized,
15{
16    /// Add an exponential backoff policy to an initialised [`ReconnectingStream`] using the
17    /// provided [`ReconnectionBackoffPolicy`].
18    ///
19    /// A failed re-initialisation is logged with `warn!` and discarded. Use
20    /// [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting) to receive it.
21    fn with_reconnect_backoff<St, InitError>(
22        self,
23        policy: ReconnectionBackoffPolicy,
24        stream_key: StreamKey,
25    ) -> impl Stream<Item = St>
26    where
27        Self: Stream<Item = Result<St, InitError>>,
28        St: Stream,
29        InitError: Debug,
30    {
31        self.with_reconnect_backoff_reporting(policy, stream_key)
32            .filter_map(|result| future::ready(result.ok()))
33    }
34
35    /// Add an exponential backoff policy to an initialised [`ReconnectingStream`] using the
36    /// provided [`ReconnectionBackoffPolicy`], yielding each failed re-initialisation as a
37    /// [`ReinitFailure`].
38    ///
39    /// A failure is yielded as soon as it happens; the backoff then runs before the next attempt.
40    /// [`ReinitFailure::attempt`] counts consecutive failures, and resets when an attempt succeeds.
41    fn with_reconnect_backoff_reporting<St, InitError>(
42        self,
43        policy: ReconnectionBackoffPolicy,
44        stream_key: StreamKey,
45    ) -> impl Stream<Item = Result<St, ReinitFailure<InitError>>>
46    where
47        Self: Stream<Item = Result<St, InitError>>,
48        St: Stream,
49        InitError: Debug,
50    {
51        self.enumerate()
52            .scan(
53                ReconnectionState::from(policy),
54                move |state, (attempt, result)| {
55                    let next =
56                        match result {
57                            Ok(stream) => {
58                                info!(attempt, ?stream_key, "successfully initialised Stream");
59                                state.reset_backoff();
60                                Either::Left(stream::once(future::ready(Ok(stream))))
61                            }
62                            Err(error) => {
63                                let consecutive = state.record_failure();
64                                warn!(
65                                    attempt,
66                                    consecutive,
67                                    ?stream_key,
68                                    ?error,
69                                    "failed to re-initialise Stream"
70                                );
71                                let backoff = state.generate_sleep_future();
72                                state.multiply_backoff();
73                                let failure = ReinitFailure {
74                                    attempt: consecutive,
75                                    error,
76                                };
77                                // Yield the failure, then hold the next attempt back until the
78                                // backoff has elapsed: `flatten` drains this before pulling it
79                                Either::Right(stream::once(future::ready(Err(failure))).chain(
80                                    stream::once(backoff).filter_map(|()| future::ready(None)),
81                                ))
82                            }
83                        };
84                    future::ready(Some(next))
85                },
86            )
87            .flatten()
88    }
89
90    /// Terminates the inner [`Stream`] if the encountered error is determined to be unrecoverable
91    /// by the provided closure. This will cause the [`ReconnectingStream`] to re-initialise the
92    /// inner [`Stream`].
93    fn with_termination_on_error<St, T, E, FnIsTerminal>(
94        self,
95        is_terminal: FnIsTerminal,
96        stream_key: StreamKey,
97    ) -> impl Stream<Item = impl Stream<Item = Result<T, E>>>
98    where
99        Self: Stream<Item = St>,
100        St: Stream<Item = Result<T, E>>,
101        FnIsTerminal: Fn(&E) -> bool + Copy,
102    {
103        self.map(move |stream| terminate_on_error(stream, is_terminal, stream_key))
104    }
105
106    /// Maps every [`ReconnectingStream`] `Stream::Item` into an [`reconnect::Event::Item`](Event),
107    /// and chain a [`reconnect::Event::Reconnecting`](Event)
108    ///
109    /// Pairs with [`with_reconnect_backoff`](Self::with_reconnect_backoff), which discards failed
110    /// re-initialisations. To report them as events, use
111    /// [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting) with
112    /// [`with_reconnection_events_reporting`](Self::with_reconnection_events_reporting).
113    fn with_reconnection_events<St, Origin>(
114        self,
115        origin: Origin,
116    ) -> impl Stream<Item = Event<Origin, St::Item>>
117    where
118        Self: Stream<Item = St>,
119        St: Stream,
120        Origin: Clone + 'static,
121    {
122        self.map(move |stream| with_trailing_reconnecting(stream, origin.clone()))
123            .flatten()
124    }
125
126    /// Maps the initialisation attempts yielded by
127    /// [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting) into
128    /// [`reconnect::Event`](Event)s.
129    ///
130    /// Each initialised stream's items become [`Event::Item`]s, followed by one
131    /// [`Event::Reconnecting`] when it ends. Each failed attempt becomes the single
132    /// [`Event::Item`] that `on_failure` makes of it. A failed attempt ends no stream, so it adds
133    /// no further [`Event::Reconnecting`].
134    fn with_reconnection_events_reporting<St, InitError, Origin, FnOnFailure>(
135        self,
136        origin: Origin,
137        on_failure: FnOnFailure,
138    ) -> impl Stream<Item = Event<Origin, St::Item>>
139    where
140        Self: Stream<Item = Result<St, ReinitFailure<InitError>>>,
141        St: Stream,
142        Origin: Clone,
143        FnOnFailure: Fn(ReinitFailure<InitError>) -> St::Item,
144    {
145        self.map(move |initialised| match initialised {
146            Ok(stream) => Either::Left(with_trailing_reconnecting(stream, origin.clone())),
147            Err(failure) => Either::Right(stream::once(future::ready(Event::Item(on_failure(
148                failure,
149            ))))),
150        })
151        .flatten()
152    }
153
154    /// Handles all encountered errors with the provided closure before filtering them out,
155    /// returning a [`Stream`] of the Ok values. Useful for logging recoverable errors before
156    /// continuing.
157    fn with_error_handler<FnOnErr, Origin, T, E>(
158        self,
159        op: FnOnErr,
160    ) -> impl Stream<Item = Event<Origin, T>>
161    where
162        Self: Stream<Item = Event<Origin, Result<T, E>>>,
163        FnOnErr: Fn(E) + 'static,
164    {
165        self.filter_map(move |event| {
166            std::future::ready(match event {
167                Event::Reconnecting(origin) => Some(Event::Reconnecting(origin)),
168                Event::Item(Ok(item)) => Some(Event::Item(item)),
169                Event::Item(Err(error)) => {
170                    op(error);
171                    None
172                }
173            })
174        })
175    }
176
177    /// Future for forwarding items in [`Self`] to the provided channel [`Tx`].
178    fn forward_to<Transmitter>(self, tx: Transmitter) -> impl Future<Output = ()> + Send
179    where
180        Self: Stream + Sized + Send,
181        Self::Item: Into<Transmitter::Item>,
182        Transmitter: Tx + Send + 'static,
183    {
184        tokio_stream::StreamExt::map_while(self, move |event| tx.send(event.into()).ok()).collect()
185    }
186}
187
188impl<T> ReconnectingStream for T where T: Stream {}
189
190/// A failed re-initialisation of a [`ReconnectingStream`], yielded by
191/// [`ReconnectingStream::with_reconnect_backoff_reporting`].
192#[derive(Debug, Clone, PartialEq, Eq)]
193pub struct ReinitFailure<InitError> {
194    /// Consecutive failed attempts since the stream last initialised, starting at 1.
195    pub attempt: u32,
196    /// Why this attempt failed.
197    pub error: InitError,
198}
199
200/// Ends `stream` at the first error `is_terminal` classifies as terminal, so the
201/// [`ReconnectingStream`] re-initialises it.
202pub(crate) fn terminate_on_error<St, T, E, FnIsTerminal>(
203    stream: St,
204    is_terminal: FnIsTerminal,
205    stream_key: StreamKey,
206) -> impl Stream<Item = Result<T, E>>
207where
208    St: Stream<Item = Result<T, E>>,
209    FnIsTerminal: Fn(&E) -> bool,
210{
211    tokio_stream::StreamExt::map_while(stream, move |result| match result {
212        Ok(item) => Some(Ok(item)),
213        Err(error) if is_terminal(&error) => {
214            error!(
215                ?stream_key,
216                "MarketStream encountered terminal error that requires reconnecting"
217            );
218            None
219        }
220        Err(error) => Some(Err(error)),
221    })
222}
223
224/// Maps every item of one initialised `stream` into an [`Event::Item`], followed by an
225/// [`Event::Reconnecting`] when it ends.
226pub(crate) fn with_trailing_reconnecting<St, Origin>(
227    stream: St,
228    origin: Origin,
229) -> impl Stream<Item = Event<Origin, St::Item>>
230where
231    St: Stream,
232{
233    stream
234        .map(Event::Item)
235        .chain(stream::once(future::ready(Event::Reconnecting(origin))))
236}
237
238/// Initialise a [`ReconnectingStream`] using the provided initialisation closure.
239pub async fn init_reconnecting_stream<FnInit, St, FnInitError, FnInitFut>(
240    init_stream: FnInit,
241) -> Result<impl Stream<Item = Result<St, FnInitError>>, FnInitError>
242where
243    FnInit: Fn() -> FnInitFut,
244    FnInitFut: Future<Output = Result<St, FnInitError>>,
245{
246    let initial = init_stream().await?;
247    let reconnections = futures::stream::repeat_with(init_stream).then(convert::identity);
248
249    Ok(futures::stream::once(future::ready(Ok(initial))).chain(reconnections))
250}
251
252/// Reconnection backoff policy for a [`ReconnectingStream::with_reconnect_backoff`].
253#[derive(
254    Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize, Constructor,
255)]
256pub struct ReconnectionBackoffPolicy {
257    /// Initial backoff millisecond duration after the first `Stream` disconnection.
258    ///
259    /// This value then scales with the `backoff_multiplier` in the case of repeated failed
260    /// `Stream` reconnection attempts.
261    pub backoff_ms_initial: u64,
262
263    /// Scaling factor for the backoff duration in the case of repeated `Stream` reconnection
264    /// attempts.
265    pub backoff_multiplier: u8,
266
267    /// Maximum possible backoff duration between reconnection attempts.
268    pub backoff_ms_max: u64,
269}
270
271#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
272struct ReconnectionState {
273    policy: ReconnectionBackoffPolicy,
274    backoff_ms_current: u64,
275    consecutive_failures: u32,
276}
277
278impl From<ReconnectionBackoffPolicy> for ReconnectionState {
279    fn from(policy: ReconnectionBackoffPolicy) -> Self {
280        Self {
281            backoff_ms_current: policy.backoff_ms_initial,
282            policy,
283            consecutive_failures: 0,
284        }
285    }
286}
287
288impl ReconnectionState {
289    fn reset_backoff(&mut self) {
290        self.backoff_ms_current = self.policy.backoff_ms_initial;
291        self.consecutive_failures = 0;
292    }
293
294    /// Count one more failed attempt, returning the consecutive total.
295    fn record_failure(&mut self) -> u32 {
296        self.consecutive_failures = self.consecutive_failures.saturating_add(1);
297        self.consecutive_failures
298    }
299
300    fn multiply_backoff(&mut self) {
301        let next = self.backoff_ms_current * self.policy.backoff_multiplier as u64;
302        let next_capped = std::cmp::min(next, self.policy.backoff_ms_max);
303        self.backoff_ms_current = next_capped;
304    }
305
306    fn generate_sleep_future(&self) -> tokio::time::Sleep {
307        let sleep_duration = std::time::Duration::from_millis(self.backoff_ms_current);
308        tokio::time::sleep(sleep_duration)
309    }
310}
311
312#[cfg(test)]
313#[allow(clippy::unwrap_used)] // Test code: panics on bad input are acceptable
314mod tests {
315    use super::*;
316    use rustrade_instrument::exchange::ExchangeId;
317    use std::time::Duration;
318    use tokio::time::Instant;
319
320    const POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
321        backoff_ms_initial: 100,
322        backoff_multiplier: 2,
323        backoff_ms_max: 1_000,
324    };
325
326    fn key() -> StreamKey {
327        StreamKey::new_general("test", ExchangeId::Simulated)
328    }
329
330    type Attempt = Result<stream::Iter<std::vec::IntoIter<u8>>, &'static str>;
331
332    fn ok(items: &[u8]) -> Attempt {
333        Ok(stream::iter(items.to_vec()))
334    }
335
336    /// Each yielded attempt, as the failure's consecutive count or `0` for a success.
337    fn attempts<St>(results: &[Result<St, ReinitFailure<&'static str>>]) -> Vec<u32> {
338        results
339            .iter()
340            .map(|result| {
341                result
342                    .as_ref()
343                    .map_or_else(|failure| failure.attempt, |_| 0)
344            })
345            .collect()
346    }
347
348    #[tokio::test(start_paused = true)]
349    async fn failures_count_consecutively_and_reset_on_success() {
350        let results = stream::iter([ok(&[1]), Err("a"), Err("b"), ok(&[2]), Err("c")])
351            .with_reconnect_backoff_reporting(POLICY, key())
352            .collect::<Vec<_>>()
353            .await;
354
355        assert_eq!(attempts(&results), [0, 1, 2, 0, 1]);
356        let Err(failure) = &results[2] else {
357            panic!("expected a failure")
358        };
359        assert_eq!(failure.error, "b");
360    }
361
362    #[tokio::test(start_paused = true)]
363    async fn a_failure_is_yielded_at_once_and_the_backoff_delays_the_next_attempt() {
364        let started = Instant::now();
365        let mut results = Box::pin(
366            stream::iter([Err("a"), Err("b"), ok(&[1])])
367                .with_reconnect_backoff_reporting(POLICY, key()),
368        );
369
370        assert!(results.next().await.unwrap().is_err());
371        assert_eq!(started.elapsed(), Duration::ZERO);
372        assert!(results.next().await.unwrap().is_err());
373        assert_eq!(started.elapsed(), Duration::from_millis(100));
374        assert!(results.next().await.unwrap().is_ok());
375        assert_eq!(started.elapsed(), Duration::from_millis(300));
376    }
377
378    #[tokio::test(start_paused = true)]
379    async fn the_discarding_backoff_yields_only_initialised_streams() {
380        let streams = stream::iter([ok(&[1]), Err("a"), ok(&[2])])
381            .with_reconnect_backoff(POLICY, key())
382            .flatten()
383            .collect::<Vec<_>>()
384            .await;
385
386        assert_eq!(streams, [1, 2]);
387    }
388
389    #[tokio::test]
390    async fn a_failure_becomes_one_item_and_only_an_ended_stream_adds_reconnecting() {
391        let origin = ExchangeId::Simulated;
392        let failure = |attempt| ReinitFailure {
393            attempt,
394            error: "refused",
395        };
396        let events = stream::iter([
397            Ok(stream::iter(vec![1_u32])),
398            Err(failure(1)),
399            Err(failure(2)),
400            Ok(stream::iter(vec![2])),
401        ])
402        .with_reconnection_events_reporting(origin, |failure| 100 + failure.attempt)
403        .collect::<Vec<_>>()
404        .await;
405
406        assert_eq!(
407            events,
408            [
409                Event::Item(1),
410                Event::Reconnecting(origin),
411                Event::Item(101),
412                Event::Item(102),
413                Event::Item(2),
414                Event::Reconnecting(origin),
415            ]
416        );
417    }
418}
rustrade_data/streams/reconnect/stream.rs

rustrade_data/streams/reconnect/
stream.rs
