brokapi_tbank/
stream.rs

1//! Stream supervision: open, read, reconnect with backoff. No protocol logic here.
2//!
3//! A server-side stream can end at any time: the server closes it, the network breaks, or it goes
4//! silent. [`supervise`] reopens it with a growing delay ([`backoff`]) and tells the consumer on every
5//! successful open, so the consumer can resend its state. It keeps trying until it is cancelled (or the
6//! receiving end of the channel is dropped), except when the broker refuses the stream in a way a repeat
7//! cannot cure: a rejected token, no permission, the limit of open streams (80001, see [`fatal_open`]).
8//! Then it reports the reason once and stops.
9//!
10//! # Examples
11//!
12//! ```
13//! use std::time::Duration;
14//! use brokapi_tbank::stream::backoff;
15//!
16//! // 500 ms doubling up to 30 s.
17//! assert_eq!(backoff(0), Duration::from_millis(500));
18//! assert_eq!(backoff(2), Duration::from_secs(2));
19//! assert_eq!(backoff(40), Duration::from_secs(30));
20//! ```
21
22use std::future::Future;
23use std::sync::Mutex;
24use std::time::Duration;
25
26use tokio::sync::mpsc;
27use tokio_util::sync::CancellationToken;
28use tonic::{Status, Streaming};
29use tracing::{debug, warn};
30
31use crate::error::map_status;
32
33/// Delay before reconnect attempt `n` (0-based): 500 ms doubling to 30 s. The same range at BCS and
34/// Finam: none of the three brokers documents a reason for a different one.
35pub fn backoff(attempt: u32) -> Duration {
36    let ms = 500u64.saturating_mul(1u64 << attempt.min(6));
37    Duration::from_millis(ms.min(30_000))
38}
39
40/// Whether the refusal to open a stream is final: a repeat will not help, so the supervisor stops and
41/// reports it. `Some(reason)` for a rejected token (`Unauthenticated`), no permission
42/// (`PermissionDenied`) and the limit of open streams (80001, `streams_limit`). Everything else (the
43/// network, `Unavailable`, a rate limit, the lifetime limit 80005) is worth another try.
44///
45/// The status is read with [`map_status`], as every other answer of the broker.
46///
47/// # Examples
48///
49/// ```
50/// use brokapi_tbank::stream::fatal_open;
51/// use tonic::Status;
52///
53/// assert!(fatal_open(&Status::unauthenticated("40003")).is_some());
54/// assert!(fatal_open(&Status::resource_exhausted("80001")).is_some());
55/// assert!(fatal_open(&Status::unavailable("down")).is_none());
56/// assert!(fatal_open(&Status::resource_exhausted("80002")).is_none());
57/// ```
58pub fn fatal_open(status: &Status) -> Option<String> {
59    use brokapi_core::Error;
60    match map_status(status) {
61        Error::Config(e) => Some(format!("unauthenticated: {e}")),
62        Error::Rejected { code, message }
63            if code == "streams_limit" || code == "PermissionDenied" =>
64        {
65            Some(format!("{code}: {message}"))
66        }
67        _ => None,
68    }
69}
70
71/// The last reason a stream closed, for error messages of calls that wait on a stream.
72#[derive(Default)]
73pub struct LastError(Mutex<Option<String>>);
74
75impl LastError {
76    /// Remembers the reason the stream closed, replacing an earlier one.
77    pub fn set(&self, msg: impl Into<String>) {
78        *self.0.lock().expect("poisoned") = Some(msg.into());
79    }
80
81    /// Forgets the reason, after the stream reopened.
82    pub fn clear(&self) {
83        *self.0.lock().expect("poisoned") = None;
84    }
85
86    /// The last reason, if the stream is down now.
87    ///
88    /// # Examples
89    ///
90    /// ```
91    /// use brokapi_tbank::stream::LastError;
92    ///
93    /// let last = LastError::default();
94    /// assert_eq!(last.get(), None);
95    /// last.set("order_state: closed by the server");
96    /// assert_eq!(last.get().as_deref(), Some("order_state: closed by the server"));
97    /// last.clear();
98    /// assert_eq!(last.get(), None);
99    /// ```
100    pub fn get(&self) -> Option<String> {
101        self.0.lock().expect("poisoned").clone()
102    }
103}
104
105/// Sleeps `d` unless cancelled. `false` when cancelled.
106///
107/// # Examples
108///
109/// ```
110/// use std::time::Duration;
111/// use brokapi_tbank::stream::sleep_or_cancel;
112/// use tokio_util::sync::CancellationToken;
113///
114/// # #[tokio::main(flavor = "current_thread")]
115/// # async fn main() {
116/// let cancel = CancellationToken::new();
117/// assert!(sleep_or_cancel(Duration::from_millis(1), &cancel).await);
118/// cancel.cancel();
119/// assert!(!sleep_or_cancel(Duration::from_secs(60), &cancel).await);
120/// # }
121/// ```
122pub async fn sleep_or_cancel(d: Duration, cancel: &CancellationToken) -> bool {
123    tokio::select! {
124        _ = cancel.cancelled() => false,
125        _ = tokio::time::sleep(d) => true,
126    }
127}
128
129/// Keeps one server-side stream open. `connected` goes out after every successful open (then the
130/// consumer knows a break happened if it is not the first), `wrap` turns each message into a frame.
131/// A stream silent for `stale_after` is closed and reopened (pings arrive on a healthy stream).
132///
133/// The pause before a reopen grows with each failure ([`backoff`]) and goes back to the start only after
134/// the stream has delivered a message, not after the open alone: a stream that opens and breaks at once
135/// is not tried again at full speed. When the open is refused for good ([`fatal_open`]), `failed(reason)`
136/// is sent once and the supervisor ends.
137#[allow(clippy::too_many_arguments)]
138pub async fn supervise<M, F, Fut, Frame>(
139    name: &'static str,
140    open: F,
141    connected: impl Fn() -> Frame,
142    wrap: impl Fn(M) -> Frame,
143    failed: impl Fn(String) -> Frame,
144    tx: mpsc::Sender<Frame>,
145    stale_after: Duration,
146    last_error: std::sync::Arc<LastError>,
147    cancel: CancellationToken,
148) where
149    F: Fn() -> Fut,
150    Fut: Future<Output = Result<Streaming<M>, Status>>,
151{
152    let mut attempt = 0u32;
153    loop {
154        if cancel.is_cancelled() {
155            return;
156        }
157        let opened = tokio::select! {
158            _ = cancel.cancelled() => return,
159            r = open() => r,
160        };
161        let reason = match opened {
162            Ok(mut stream) => {
163                last_error.clear();
164                debug!(stream = name, "stream open");
165                if tx.send(connected()).await.is_err() {
166                    return;
167                }
168
169                loop {
170                    let msg = tokio::select! {
171                        _ = cancel.cancelled() => return,
172                        m = tokio::time::timeout(stale_after, stream.message()) => m,
173                    };
174                    match msg {
175                        Ok(Ok(Some(m))) => {
176                            // The stream works: the next break starts from the short pause again.
177                            attempt = 0;
178                            if tx.send(wrap(m)).await.is_err() {
179                                return;
180                            }
181                        }
182                        Ok(Ok(None)) => break "closed by the server".to_string(),
183                        Ok(Err(status)) => {
184                            break format!("{}: {}", status.code(), status.message());
185                        }
186                        Err(_) => break format!("silent for {stale_after:?}"),
187                    }
188                }
189            }
190            Err(status) => {
191                if let Some(reason) = fatal_open(&status) {
192                    warn!(stream = name, %reason, "stream refused for good, not reopened");
193                    last_error.set(format!("{name}: {reason}"));
194                    let _ = tx.send(failed(reason)).await;
195                    return;
196                }
197                format!("open failed: {}: {}", status.code(), status.message())
198            }
199        };
200        warn!(stream = name, %reason, "stream down, reconnecting");
201        last_error.set(format!("{name}: {reason}"));
202        if !sleep_or_cancel(backoff(attempt), &cancel).await {
203            return;
204        }
205        attempt = attempt.saturating_add(1);
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn backoff_grows_and_caps() {
215        assert_eq!(backoff(0), Duration::from_millis(500));
216        assert_eq!(backoff(1), Duration::from_secs(1));
217        assert_eq!(backoff(3), Duration::from_secs(4));
218        assert_eq!(backoff(30), Duration::from_secs(30));
219    }
220
221    use std::sync::Arc;
222    use std::sync::atomic::{AtomicUsize, Ordering};
223
224    #[test]
225    fn only_a_final_refusal_stops_a_stream() {
226        // A rejected token, no right, the limit of open streams: a repeat will not help.
227        assert!(
228            fatal_open(&Status::unauthenticated("40003"))
229                .unwrap()
230                .starts_with("unauthenticated")
231        );
232        assert!(
233            fatal_open(&Status::resource_exhausted("80001"))
234                .unwrap()
235                .starts_with("streams_limit")
236        );
237        assert!(
238            fatal_open(&Status::permission_denied("no"))
239                .unwrap()
240                .starts_with("PermissionDenied")
241        );
242        // The network, a rate limit, the lifetime of a stream, a server fault: try again.
243        for retry in [
244            Status::unavailable("down"),
245            Status::deadline_exceeded("slow"),
246            Status::internal("oops"),
247            Status::resource_exhausted("80002"),
248            Status::resource_exhausted("80005"),
249            Status::not_found("x"),
250        ] {
251            assert!(fatal_open(&retry).is_none(), "{retry:?}");
252        }
253    }
254
255    type Opened = std::future::Ready<Result<Streaming<u8>, Status>>;
256
257    /// An `open` that always fails with `status`, and the number of its calls.
258    fn failing(status: Status) -> (impl Fn() -> Opened, Arc<AtomicUsize>) {
259        let calls = Arc::new(AtomicUsize::new(0));
260        let seen = Arc::clone(&calls);
261        (
262            move || {
263                seen.fetch_add(1, Ordering::SeqCst);
264                std::future::ready(Err(status.clone()))
265            },
266            calls,
267        )
268    }
269
270    #[tokio::test(start_paused = true)]
271    async fn a_refusal_for_good_is_reported_once_and_the_supervisor_ends() {
272        let (open, calls) = failing(Status::unauthenticated("40003"));
273        let (tx, mut rx) = mpsc::channel::<String>(8);
274        let last = Arc::new(LastError::default());
275        supervise(
276            "order_state",
277            open,
278            || "connected".to_string(),
279            |_m: u8| "message".to_string(),
280            |reason| format!("failed: {reason}"),
281            tx,
282            Duration::from_secs(5),
283            Arc::clone(&last),
284            CancellationToken::new(),
285        )
286        .await;
287        assert_eq!(calls.load(Ordering::SeqCst), 1, "not tried again");
288        let told = rx.recv().await.unwrap();
289        assert!(told.starts_with("failed: unauthenticated"), "{told}");
290        assert!(rx.recv().await.is_none(), "nothing else was sent");
291        assert!(last.get().unwrap().contains("unauthenticated"));
292    }
293
294    #[tokio::test(start_paused = true)]
295    async fn a_failure_that_may_pass_is_retried_with_a_growing_pause_until_cancelled() {
296        let (open, calls) = failing(Status::unavailable("down"));
297        let (tx, mut rx) = mpsc::channel::<String>(8);
298        let cancel = CancellationToken::new();
299        let last = Arc::new(LastError::default());
300        let task = tokio::spawn(supervise(
301            "trades",
302            open,
303            || "connected".to_string(),
304            |_m: u8| "message".to_string(),
305            |reason| format!("failed: {reason}"),
306            tx,
307            Duration::from_secs(5),
308            Arc::clone(&last),
309            cancel.clone(),
310        ));
311        let settle = || async {
312            for _ in 0..8 {
313                tokio::task::yield_now().await;
314            }
315        };
316        settle().await;
317        assert_eq!(calls.load(Ordering::SeqCst), 1);
318        assert!(last.get().unwrap().contains("open failed"));
319        // The first pause is 500 ms, the second 1 s: the open is not tried at full speed.
320        tokio::time::advance(Duration::from_millis(499)).await;
321        settle().await;
322        assert_eq!(calls.load(Ordering::SeqCst), 1);
323        tokio::time::advance(Duration::from_millis(1)).await;
324        settle().await;
325        assert_eq!(calls.load(Ordering::SeqCst), 2);
326        tokio::time::advance(Duration::from_millis(999)).await;
327        settle().await;
328        assert_eq!(calls.load(Ordering::SeqCst), 2);
329        tokio::time::advance(Duration::from_millis(1)).await;
330        settle().await;
331        assert_eq!(calls.load(Ordering::SeqCst), 3);
332        // Nothing was reported as final.
333        assert!(rx.try_recv().is_err());
334        cancel.cancel();
335        task.await.unwrap();
336    }
337
338    #[tokio::test]
339    async fn a_supervisor_that_is_cancelled_before_it_starts_does_nothing() {
340        let (open, calls) = failing(Status::unavailable("down"));
341        let (tx, _rx) = mpsc::channel::<String>(1);
342        let cancel = CancellationToken::new();
343        cancel.cancel();
344        supervise(
345            "x",
346            open,
347            || "connected".to_string(),
348            |_m: u8| "message".to_string(),
349            |reason| reason,
350            tx,
351            Duration::from_secs(1),
352            Arc::new(LastError::default()),
353            cancel,
354        )
355        .await;
356        assert_eq!(calls.load(Ordering::SeqCst), 0);
357    }
358}
brokapi_tbank/stream.rs

brokapi_tbank/
stream.rs
