polyc_runtime/
retry.rs

1//! Rate-limit-aware retry for edge platform REST clients (#795).
2//!
3//! Every edge's `*_api.rs` dials a platform REST API (Slack, Telegram, …)
4//! that can answer a write with a rate-limit signal — an HTTP
5//! `429`, usually paired with a `Retry-After` header or an equivalent
6//! body-level field. Before this module, no client recognized that signal:
7//! a rate-limited `chat.postMessage` or an approval-card `chat.update` just
8//! surfaced as an ordinary API error and the reply (or the approval-card
9//! edit) was silently dropped.
10//!
11//! [`retry_rate_limited`] is the one shared loop every `*_api.rs` wraps its
12//! attempt in. The caller decides, per attempt, whether the platform's
13//! response means "done" or "rate limited, wait this long" — this module
14//! only owns the wait-and-retry mechanics, since each platform signals a
15//! rate limit differently (status code, header, or JSON body field).
16
17use std::{future::Future, time::Duration};
18
19/// One attempt's outcome for [`retry_rate_limited`], returned by the
20/// caller's per-attempt closure.
21#[derive(Debug)]
22pub enum RetryableError<E> {
23    /// The platform signaled a rate limit. Wait `after`, then retry; `error`
24    /// is what [`retry_rate_limited`] returns if this was the last allowed
25    /// attempt.
26    RateLimited {
27        /// How long to wait before the next attempt (from the platform's
28        /// `Retry-After` or equivalent, or a caller-chosen default when the
29        /// platform didn't say).
30        after: Duration,
31        /// The error to surface if retries are exhausted while still rate
32        /// limited.
33        error: E,
34    },
35    /// Not retryable — surface immediately without waiting.
36    Fatal(E),
37}
38
39/// Runs `attempt` until it returns `Ok`, returns a
40/// [`RetryableError::Fatal`], or exhausts `max_retries` retries after
41/// repeated [`RetryableError::RateLimited`] outcomes.
42///
43/// Sleeps for the signaled `after` duration between a `RateLimited` outcome
44/// and the next attempt — honoring the platform's own backoff request
45/// (e.g. Slack's `Retry-After` header, Telegram's `retry_after`
46/// body field) rather than a fixed schedule.
47///
48/// # Errors
49///
50/// Returns the error from the attempt that ended the loop: the `Fatal`
51/// error immediately, or the last `RateLimited` error once `max_retries` is
52/// exhausted.
53pub async fn retry_rate_limited<T, E, F, Fut>(mut max_retries: u32, mut attempt: F) -> Result<T, E>
54where
55    F: FnMut() -> Fut,
56    Fut: Future<Output = Result<T, RetryableError<E>>>,
57{
58    loop {
59        match attempt().await {
60            Ok(value) => return Ok(value),
61            Err(RetryableError::Fatal(error)) => return Err(error),
62            Err(RetryableError::RateLimited { after, error }) => {
63                if max_retries == 0 {
64                    return Err(error);
65                }
66                max_retries -= 1;
67                tokio::time::sleep(after).await;
68            }
69        }
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use std::{
76        sync::atomic::{AtomicU32, Ordering},
77        time::Instant,
78    };
79
80    use super::{RetryableError, retry_rate_limited};
81
82    #[tokio::test]
83    async fn retries_once_after_rate_limit_then_succeeds() {
84        let attempts = AtomicU32::new(0);
85        let started = Instant::now();
86        let result: Result<&str, &str> = retry_rate_limited(3, || {
87            let n = attempts.fetch_add(1, Ordering::SeqCst);
88            async move {
89                if n == 0 {
90                    Err(RetryableError::RateLimited {
91                        after: std::time::Duration::from_millis(30),
92                        error: "rate limited",
93                    })
94                } else {
95                    Ok("ok")
96                }
97            }
98        })
99        .await;
100        assert_eq!(result, Ok("ok"));
101        assert_eq!(attempts.load(Ordering::SeqCst), 2);
102        assert!(
103            started.elapsed() >= std::time::Duration::from_millis(30),
104            "must actually wait out the signaled backoff before retrying"
105        );
106    }
107
108    #[tokio::test]
109    async fn fatal_error_is_not_retried() {
110        let attempts = AtomicU32::new(0);
111        let result: Result<&str, &str> = retry_rate_limited(3, || {
112            attempts.fetch_add(1, Ordering::SeqCst);
113            async { Err(RetryableError::Fatal("boom")) }
114        })
115        .await;
116        assert_eq!(result, Err("boom"));
117        assert_eq!(attempts.load(Ordering::SeqCst), 1);
118    }
119
120    #[tokio::test]
121    async fn exhausting_retries_surfaces_the_last_rate_limit_error() {
122        let attempts = AtomicU32::new(0);
123        let result: Result<&str, &str> = retry_rate_limited(2, || {
124            attempts.fetch_add(1, Ordering::SeqCst);
125            async {
126                Err(RetryableError::RateLimited {
127                    after: std::time::Duration::from_millis(1),
128                    error: "still limited",
129                })
130            }
131        })
132        .await;
133        assert_eq!(result, Err("still limited"));
134        // Initial attempt + 2 retries.
135        assert_eq!(attempts.load(Ordering::SeqCst), 3);
136    }
137}
polyc_runtime/retry.rs

polyc_runtime/
retry.rs
