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, time::Duration};
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    ///
42    /// To also honour a wait the source of a failure asked for, such as a venue's rate-limit
43    /// hint, use [`with_reconnect_backoff_reporting_hinted`](Self::with_reconnect_backoff_reporting_hinted).
44    fn with_reconnect_backoff_reporting<St, InitError>(
45        self,
46        policy: ReconnectionBackoffPolicy,
47        stream_key: StreamKey,
48    ) -> impl Stream<Item = Result<St, ReinitFailure<InitError>>>
49    where
50        Self: Stream<Item = Result<St, InitError>>,
51        St: Stream,
52        InitError: Debug,
53    {
54        self.with_reconnect_backoff_reporting_hinted(policy, stream_key, |_| None)
55    }
56
57    /// As [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting), but the
58    /// next attempt after a failure is held back for at least the wait `retry_after` reads from
59    /// it: the longer of that wait and the backoff. A failure whose wait is `None` waits the
60    /// backoff alone.
61    ///
62    /// The hint does not change the backoff, which grows with each consecutive failure as usual.
63    /// It is a floor for one wait, so an attempt is never made before the source said it could
64    /// succeed, nor sooner than the policy allows. It is not capped: a caller that wants a limit
65    /// clamps the wait in `retry_after`.
66    fn with_reconnect_backoff_reporting_hinted<St, InitError, FnRetryAfter>(
67        self,
68        policy: ReconnectionBackoffPolicy,
69        stream_key: StreamKey,
70        retry_after: FnRetryAfter,
71    ) -> impl Stream<Item = Result<St, ReinitFailure<InitError>>>
72    where
73        Self: Stream<Item = Result<St, InitError>>,
74        St: Stream,
75        InitError: Debug,
76        FnRetryAfter: Fn(&InitError) -> Option<Duration>,
77    {
78        self.enumerate()
79            .scan(
80                ReconnectionState::from(policy),
81                move |state, (attempt, result)| {
82                    let next =
83                        match result {
84                            Ok(stream) => {
85                                info!(attempt, ?stream_key, "successfully initialised Stream");
86                                state.reset_backoff();
87                                Either::Left(stream::once(future::ready(Ok(stream))))
88                            }
89                            Err(error) => {
90                                let consecutive = state.record_failure();
91                                let hint = retry_after(&error);
92                                warn!(
93                                    attempt,
94                                    consecutive,
95                                    ?stream_key,
96                                    ?error,
97                                    ?hint,
98                                    "failed to re-initialise Stream"
99                                );
100                                let backoff = state.generate_sleep_future(hint);
101                                state.multiply_backoff();
102                                let failure = ReinitFailure {
103                                    attempt: consecutive,
104                                    error,
105                                };
106                                // Yield the failure, then hold the next attempt back until the
107                                // backoff has elapsed: `flatten` drains this before pulling it
108                                Either::Right(stream::once(future::ready(Err(failure))).chain(
109                                    stream::once(backoff).filter_map(|()| future::ready(None)),
110                                ))
111                            }
112                        };
113                    future::ready(Some(next))
114                },
115            )
116            .flatten()
117    }
118
119    /// Terminates the inner [`Stream`] if the encountered error is determined to be unrecoverable
120    /// by the provided closure. This will cause the [`ReconnectingStream`] to re-initialise the
121    /// inner [`Stream`].
122    fn with_termination_on_error<St, T, E, FnIsTerminal>(
123        self,
124        is_terminal: FnIsTerminal,
125        stream_key: StreamKey,
126    ) -> impl Stream<Item = impl Stream<Item = Result<T, E>>>
127    where
128        Self: Stream<Item = St>,
129        St: Stream<Item = Result<T, E>>,
130        FnIsTerminal: Fn(&E) -> bool + Copy,
131    {
132        self.map(move |stream| terminate_on_error(stream, is_terminal, stream_key))
133    }
134
135    /// Maps every [`ReconnectingStream`] `Stream::Item` into an [`reconnect::Event::Item`](Event),
136    /// and chain a [`reconnect::Event::Reconnecting`](Event)
137    ///
138    /// Pairs with [`with_reconnect_backoff`](Self::with_reconnect_backoff), which discards failed
139    /// re-initialisations. To report them as events, use
140    /// [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting) with
141    /// [`with_reconnection_events_reporting`](Self::with_reconnection_events_reporting).
142    fn with_reconnection_events<St, Origin>(
143        self,
144        origin: Origin,
145    ) -> impl Stream<Item = Event<Origin, St::Item>>
146    where
147        Self: Stream<Item = St>,
148        St: Stream,
149        Origin: Clone + 'static,
150    {
151        self.map(move |stream| with_trailing_reconnecting(stream, origin.clone()))
152            .flatten()
153    }
154
155    /// Maps the initialisation attempts yielded by
156    /// [`with_reconnect_backoff_reporting`](Self::with_reconnect_backoff_reporting) into
157    /// [`reconnect::Event`](Event)s.
158    ///
159    /// Each initialised stream's items become [`Event::Item`]s, followed by one
160    /// [`Event::Reconnecting`] when it ends. Each failed attempt becomes the single
161    /// [`Event::Item`] that `on_failure` makes of it. A failed attempt ends no stream, so it adds
162    /// no further [`Event::Reconnecting`].
163    fn with_reconnection_events_reporting<St, InitError, Origin, FnOnFailure>(
164        self,
165        origin: Origin,
166        on_failure: FnOnFailure,
167    ) -> impl Stream<Item = Event<Origin, St::Item>>
168    where
169        Self: Stream<Item = Result<St, ReinitFailure<InitError>>>,
170        St: Stream,
171        Origin: Clone,
172        FnOnFailure: Fn(ReinitFailure<InitError>) -> St::Item,
173    {
174        self.map(move |initialised| match initialised {
175            Ok(stream) => Either::Left(with_trailing_reconnecting(stream, origin.clone())),
176            Err(failure) => Either::Right(stream::once(future::ready(Event::Item(on_failure(
177                failure,
178            ))))),
179        })
180        .flatten()
181    }
182
183    /// Handles all encountered errors with the provided closure before filtering them out,
184    /// returning a [`Stream`] of the Ok values. Useful for logging recoverable errors before
185    /// continuing.
186    fn with_error_handler<FnOnErr, Origin, T, E>(
187        self,
188        op: FnOnErr,
189    ) -> impl Stream<Item = Event<Origin, T>>
190    where
191        Self: Stream<Item = Event<Origin, Result<T, E>>>,
192        FnOnErr: Fn(E) + 'static,
193    {
194        self.filter_map(move |event| {
195            std::future::ready(match event {
196                Event::Reconnecting(origin) => Some(Event::Reconnecting(origin)),
197                Event::Item(Ok(item)) => Some(Event::Item(item)),
198                Event::Item(Err(error)) => {
199                    op(error);
200                    None
201                }
202            })
203        })
204    }
205
206    /// Future for forwarding items in [`Self`] to the provided channel [`Tx`].
207    fn forward_to<Transmitter>(self, tx: Transmitter) -> impl Future<Output = ()> + Send
208    where
209        Self: Stream + Sized + Send,
210        Self::Item: Into<Transmitter::Item>,
211        Transmitter: Tx + Send + 'static,
212    {
213        tokio_stream::StreamExt::map_while(self, move |event| tx.send(event.into()).ok()).collect()
214    }
215}
216
217impl<T> ReconnectingStream for T where T: Stream {}
218
219/// A failed re-initialisation of a [`ReconnectingStream`], yielded by
220/// [`ReconnectingStream::with_reconnect_backoff_reporting`].
221#[derive(Debug, Clone, PartialEq, Eq)]
222pub struct ReinitFailure<InitError> {
223    /// Consecutive failed attempts since the stream last initialised, starting at 1.
224    pub attempt: u32,
225    /// Why this attempt failed.
226    pub error: InitError,
227}
228
229/// Ends `stream` at the first error `is_terminal` classifies as terminal, so the
230/// [`ReconnectingStream`] re-initialises it.
231pub(crate) fn terminate_on_error<St, T, E, FnIsTerminal>(
232    stream: St,
233    is_terminal: FnIsTerminal,
234    stream_key: StreamKey,
235) -> impl Stream<Item = Result<T, E>>
236where
237    St: Stream<Item = Result<T, E>>,
238    FnIsTerminal: Fn(&E) -> bool,
239{
240    tokio_stream::StreamExt::map_while(stream, move |result| match result {
241        Ok(item) => Some(Ok(item)),
242        Err(error) if is_terminal(&error) => {
243            error!(
244                ?stream_key,
245                "MarketStream encountered terminal error that requires reconnecting"
246            );
247            None
248        }
249        Err(error) => Some(Err(error)),
250    })
251}
252
253/// Maps every item of one initialised `stream` into an [`Event::Item`], followed by an
254/// [`Event::Reconnecting`] when it ends.
255pub(crate) fn with_trailing_reconnecting<St, Origin>(
256    stream: St,
257    origin: Origin,
258) -> impl Stream<Item = Event<Origin, St::Item>>
259where
260    St: Stream,
261{
262    stream
263        .map(Event::Item)
264        .chain(stream::once(future::ready(Event::Reconnecting(origin))))
265}
266
267/// Initialise a [`ReconnectingStream`] using the provided initialisation closure.
268pub async fn init_reconnecting_stream<FnInit, St, FnInitError, FnInitFut>(
269    init_stream: FnInit,
270) -> Result<impl Stream<Item = Result<St, FnInitError>>, FnInitError>
271where
272    FnInit: Fn() -> FnInitFut,
273    FnInitFut: Future<Output = Result<St, FnInitError>>,
274{
275    let initial = init_stream().await?;
276    let reconnections = futures::stream::repeat_with(init_stream).then(convert::identity);
277
278    Ok(futures::stream::once(future::ready(Ok(initial))).chain(reconnections))
279}
280
281/// Reconnection backoff policy for a [`ReconnectingStream::with_reconnect_backoff`].
282#[derive(
283    Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize, Constructor,
284)]
285pub struct ReconnectionBackoffPolicy {
286    /// Initial backoff millisecond duration after the first `Stream` disconnection.
287    ///
288    /// This value then scales with the `backoff_multiplier` in the case of repeated failed
289    /// `Stream` reconnection attempts.
290    pub backoff_ms_initial: u64,
291
292    /// Scaling factor for the backoff duration in the case of repeated `Stream` reconnection
293    /// attempts.
294    pub backoff_multiplier: u8,
295
296    /// Maximum possible backoff duration between reconnection attempts.
297    pub backoff_ms_max: u64,
298}
299
300#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
301struct ReconnectionState {
302    policy: ReconnectionBackoffPolicy,
303    backoff_ms_current: u64,
304    consecutive_failures: u32,
305}
306
307impl From<ReconnectionBackoffPolicy> for ReconnectionState {
308    fn from(policy: ReconnectionBackoffPolicy) -> Self {
309        Self {
310            backoff_ms_current: policy.backoff_ms_initial,
311            policy,
312            consecutive_failures: 0,
313        }
314    }
315}
316
317impl ReconnectionState {
318    fn reset_backoff(&mut self) {
319        self.backoff_ms_current = self.policy.backoff_ms_initial;
320        self.consecutive_failures = 0;
321    }
322
323    /// Count one more failed attempt, returning the consecutive total.
324    fn record_failure(&mut self) -> u32 {
325        self.consecutive_failures = self.consecutive_failures.saturating_add(1);
326        self.consecutive_failures
327    }
328
329    fn multiply_backoff(&mut self) {
330        let next = self.backoff_ms_current * self.policy.backoff_multiplier as u64;
331        let next_capped = std::cmp::min(next, self.policy.backoff_ms_max);
332        self.backoff_ms_current = next_capped;
333    }
334
335    /// The wait before the next attempt: the current backoff, or `at_least` when that is longer.
336    fn generate_sleep_future(&self, at_least: Option<Duration>) -> tokio::time::Sleep {
337        let backoff = Duration::from_millis(self.backoff_ms_current);
338        tokio::time::sleep(at_least.map_or(backoff, |hint| hint.max(backoff)))
339    }
340}
341
342#[cfg(test)]
343#[allow(clippy::unwrap_used)] // Test code: panics on bad input are acceptable
344mod tests {
345    use super::*;
346    use rustrade_instrument::exchange::ExchangeId;
347    use std::time::Duration;
348    use tokio::time::Instant;
349
350    const POLICY: ReconnectionBackoffPolicy = ReconnectionBackoffPolicy {
351        backoff_ms_initial: 100,
352        backoff_multiplier: 2,
353        backoff_ms_max: 1_000,
354    };
355
356    fn key() -> StreamKey {
357        StreamKey::new_general("test", ExchangeId::Simulated)
358    }
359
360    type Attempt = Result<stream::Iter<std::vec::IntoIter<u8>>, &'static str>;
361
362    fn ok(items: &[u8]) -> Attempt {
363        Ok(stream::iter(items.to_vec()))
364    }
365
366    /// Each yielded attempt, as the failure's consecutive count or `0` for a success.
367    fn attempts<St>(results: &[Result<St, ReinitFailure<&'static str>>]) -> Vec<u32> {
368        results
369            .iter()
370            .map(|result| {
371                result
372                    .as_ref()
373                    .map_or_else(|failure| failure.attempt, |_| 0)
374            })
375            .collect()
376    }
377
378    #[tokio::test(start_paused = true)]
379    async fn failures_count_consecutively_and_reset_on_success() {
380        let results = stream::iter([ok(&[1]), Err("a"), Err("b"), ok(&[2]), Err("c")])
381            .with_reconnect_backoff_reporting(POLICY, key())
382            .collect::<Vec<_>>()
383            .await;
384
385        assert_eq!(attempts(&results), [0, 1, 2, 0, 1]);
386        let Err(failure) = &results[2] else {
387            panic!("expected a failure")
388        };
389        assert_eq!(failure.error, "b");
390    }
391
392    #[tokio::test(start_paused = true)]
393    async fn a_failure_is_yielded_at_once_and_the_backoff_delays_the_next_attempt() {
394        let started = Instant::now();
395        let mut results = Box::pin(
396            stream::iter([Err("a"), Err("b"), ok(&[1])])
397                .with_reconnect_backoff_reporting(POLICY, key()),
398        );
399
400        assert!(results.next().await.unwrap().is_err());
401        assert_eq!(started.elapsed(), Duration::ZERO);
402        assert!(results.next().await.unwrap().is_err());
403        assert_eq!(started.elapsed(), Duration::from_millis(100));
404        assert!(results.next().await.unwrap().is_ok());
405        assert_eq!(started.elapsed(), Duration::from_millis(300));
406    }
407
408    #[tokio::test(start_paused = true)]
409    async fn a_hinted_wait_longer_than_the_backoff_holds_the_next_attempt_back() {
410        // "wait" asks for 500 ms; "short" for 50 ms, under the backoff; "none" for nothing.
411        let hint = |error: &&str| match *error {
412            "wait" => Some(Duration::from_millis(500)),
413            "short" => Some(Duration::from_millis(50)),
414            _ => None,
415        };
416        let started = Instant::now();
417        let mut results = Box::pin(
418            stream::iter([Err("wait"), Err("none"), Err("short"), ok(&[1])])
419                .with_reconnect_backoff_reporting_hinted(POLICY, key(), hint),
420        );
421
422        assert!(results.next().await.unwrap().is_err());
423        assert_eq!(started.elapsed(), Duration::ZERO);
424        // The hint outlasts the 100 ms backoff, so it sets the wait.
425        assert!(results.next().await.unwrap().is_err());
426        assert_eq!(started.elapsed(), Duration::from_millis(500));
427        // The hint did not stop the backoff growing: 200 ms now.
428        assert!(results.next().await.unwrap().is_err());
429        assert_eq!(started.elapsed(), Duration::from_millis(700));
430        // A hint shorter than the 400 ms backoff does not shorten it.
431        assert!(results.next().await.unwrap().is_ok());
432        assert_eq!(started.elapsed(), Duration::from_millis(1_100));
433    }
434
435    #[tokio::test(start_paused = true)]
436    async fn the_discarding_backoff_yields_only_initialised_streams() {
437        let streams = stream::iter([ok(&[1]), Err("a"), ok(&[2])])
438            .with_reconnect_backoff(POLICY, key())
439            .flatten()
440            .collect::<Vec<_>>()
441            .await;
442
443        assert_eq!(streams, [1, 2]);
444    }
445
446    #[tokio::test]
447    async fn a_failure_becomes_one_item_and_only_an_ended_stream_adds_reconnecting() {
448        let origin = ExchangeId::Simulated;
449        let failure = |attempt| ReinitFailure {
450            attempt,
451            error: "refused",
452        };
453        let events = stream::iter([
454            Ok(stream::iter(vec![1_u32])),
455            Err(failure(1)),
456            Err(failure(2)),
457            Ok(stream::iter(vec![2])),
458        ])
459        .with_reconnection_events_reporting(origin, |failure| 100 + failure.attempt)
460        .collect::<Vec<_>>()
461        .await;
462
463        assert_eq!(
464            events,
465            [
466                Event::Item(1),
467                Event::Reconnecting(origin),
468                Event::Item(101),
469                Event::Item(102),
470                Event::Item(2),
471                Event::Reconnecting(origin),
472            ]
473        );
474    }
475}
rustrade_data/streams/reconnect/stream.rs

rustrade_data/streams/reconnect/
stream.rs
