kcode_telegram_request_policy/
lib.rs

1#![doc = include_str!("../Documentation.md")]
2
3use std::{future::Future, time::Duration};
4
5use teloxide::{DownloadError, RequestError};
6
7const MAX_ATTEMPTS: usize = 5;
8const INITIAL_NETWORK_BACKOFF_MILLIS: u64 = 250;
9const MAX_NETWORK_BACKOFF_MILLIS: u64 = 2_000;
10
11pub fn request_error_class(error: &RequestError) -> &'static str {
12    match error {
13        RequestError::Api(_) => "telegram_api",
14        RequestError::MigrateToChatId(_) => "telegram_migrate",
15        RequestError::RetryAfter(_) => "telegram_rate_limit",
16        RequestError::Network(network) if network.is_timeout() => "telegram_network_timeout",
17        RequestError::Network(network) if network.is_connect() => "telegram_network_connect",
18        RequestError::Network(_) => "telegram_network",
19        RequestError::InvalidJson { .. } => "telegram_invalid_json",
20        RequestError::Io(error) => io_error_class(error.kind()),
21    }
22}
23
24fn download_error_class(error: &DownloadError) -> &'static str {
25    match error {
26        DownloadError::Network(network) if network.is_timeout() => {
27            "telegram_download_network_timeout"
28        }
29        DownloadError::Network(network) if network.is_connect() => {
30            "telegram_download_network_connect"
31        }
32        DownloadError::Network(_) => "telegram_download_network",
33        DownloadError::Io(error) => io_error_class(error.kind()),
34    }
35}
36
37fn transient_io(kind: std::io::ErrorKind) -> bool {
38    matches!(
39        kind,
40        std::io::ErrorKind::TimedOut
41            | std::io::ErrorKind::ConnectionRefused
42            | std::io::ErrorKind::ConnectionReset
43            | std::io::ErrorKind::ConnectionAborted
44            | std::io::ErrorKind::NotConnected
45            | std::io::ErrorKind::BrokenPipe
46            | std::io::ErrorKind::UnexpectedEof
47            | std::io::ErrorKind::Interrupted
48            | std::io::ErrorKind::WouldBlock
49    )
50}
51
52fn network_backoff(attempt: usize) -> Duration {
53    let exponent = u32::try_from(attempt.saturating_sub(1)).unwrap_or(u32::MAX);
54    let multiplier = 2_u64.checked_pow(exponent).unwrap_or(u64::MAX);
55    Duration::from_millis(
56        INITIAL_NETWORK_BACKOFF_MILLIS
57            .saturating_mul(multiplier)
58            .min(MAX_NETWORK_BACKOFF_MILLIS),
59    )
60}
61
62fn request_retry_delay(error: &RequestError, attempt: usize) -> Option<Duration> {
63    match error {
64        RequestError::RetryAfter(delay) => Some(delay.duration()),
65        RequestError::Network(_) => Some(network_backoff(attempt)),
66        RequestError::Io(error) if transient_io(error.kind()) => Some(network_backoff(attempt)),
67        RequestError::Api(_)
68        | RequestError::MigrateToChatId(_)
69        | RequestError::InvalidJson { .. }
70        | RequestError::Io(_) => None,
71    }
72}
73
74fn download_retry_delay(error: &DownloadError, attempt: usize) -> Option<Duration> {
75    match error {
76        DownloadError::Network(_) => Some(network_backoff(attempt)),
77        DownloadError::Io(error) if transient_io(error.kind()) => Some(network_backoff(attempt)),
78        DownloadError::Io(_) => None,
79    }
80}
81
82async fn retry_operation<T, E, Attempt, AttemptFuture, Classify, Delay, Sleep, SleepFuture>(
83    operation: &'static str,
84    mut attempt_operation: Attempt,
85    classify: Classify,
86    delay: Delay,
87    sleep: Sleep,
88) -> Result<T, E>
89where
90    Attempt: FnMut() -> AttemptFuture,
91    AttemptFuture: Future<Output = Result<T, E>>,
92    Classify: Fn(&E) -> &'static str,
93    Delay: Fn(&E, usize) -> Option<Duration>,
94    Sleep: Fn(Duration) -> SleepFuture,
95    SleepFuture: Future<Output = ()>,
96{
97    for attempt in 1..=MAX_ATTEMPTS {
98        match attempt_operation().await {
99            Ok(value) => return Ok(value),
100            Err(error) => {
101                let Some(wait) = delay(&error, attempt) else {
102                    return Err(error);
103                };
104                if attempt == MAX_ATTEMPTS {
105                    return Err(error);
106                }
107                tracing::debug!(
108                    operation,
109                    attempt,
110                    error_class = classify(&error),
111                    "Transient Telegram operation failed; retrying"
112                );
113                sleep(wait).await;
114            }
115        }
116    }
117    unreachable!("the bounded Telegram retry loop always returns")
118}
119
120pub async fn retry_request<T, Attempt, AttemptFuture>(
121    operation: &'static str,
122    attempt: Attempt,
123) -> Result<T, RequestError>
124where
125    Attempt: FnMut() -> AttemptFuture,
126    AttemptFuture: Future<Output = Result<T, RequestError>>,
127{
128    retry_operation(
129        operation,
130        attempt,
131        request_error_class,
132        request_retry_delay,
133        tokio::time::sleep,
134    )
135    .await
136}
137
138pub async fn retry_download<T, Attempt, AttemptFuture>(
139    operation: &'static str,
140    attempt: Attempt,
141) -> Result<T, DownloadError>
142where
143    Attempt: FnMut() -> AttemptFuture,
144    AttemptFuture: Future<Output = Result<T, DownloadError>>,
145{
146    retry_operation(
147        operation,
148        attempt,
149        download_error_class,
150        download_retry_delay,
151        tokio::time::sleep,
152    )
153    .await
154}
155
156fn io_error_class(kind: std::io::ErrorKind) -> &'static str {
157    match kind {
158        std::io::ErrorKind::TimedOut => "io_timeout",
159        std::io::ErrorKind::ConnectionRefused => "io_connection_refused",
160        std::io::ErrorKind::ConnectionReset => "io_connection_reset",
161        std::io::ErrorKind::ConnectionAborted => "io_connection_aborted",
162        std::io::ErrorKind::NotConnected => "io_not_connected",
163        std::io::ErrorKind::BrokenPipe => "io_broken_pipe",
164        std::io::ErrorKind::UnexpectedEof => "io_unexpected_eof",
165        std::io::ErrorKind::PermissionDenied => "io_permission_denied",
166        std::io::ErrorKind::NotFound => "io_not_found",
167        _ => "io_other",
168    }
169}
170
171pub fn anyhow_error_class(error: &anyhow::Error) -> &'static str {
172    for cause in error.chain() {
173        if let Some(request_error) = cause.downcast_ref::<RequestError>() {
174            return request_error_class(request_error);
175        }
176        if let Some(download_error) = cause.downcast_ref::<DownloadError>() {
177            return download_error_class(download_error);
178        }
179        if let Some(io_error) = cause.downcast_ref::<std::io::Error>() {
180            return io_error_class(io_error.kind());
181        }
182    }
183    "local_processing"
184}
185
186#[cfg(test)]
187mod tests {
188    use std::sync::{
189        Arc,
190        atomic::{AtomicUsize, Ordering},
191    };
192
193    use teloxide::types::{ChatId, Seconds};
194
195    use super::*;
196
197    fn transient_error() -> RequestError {
198        RequestError::Io(Arc::new(std::io::Error::new(
199            std::io::ErrorKind::ConnectionReset,
200            "test-only transient failure",
201        )))
202    }
203
204    #[tokio::test]
205    async fn transient_operation_succeeds_within_five_total_attempts() {
206        let attempts = AtomicUsize::new(0);
207        let value = retry_operation(
208            "test",
209            || {
210                let attempt = attempts.fetch_add(1, Ordering::SeqCst) + 1;
211                async move {
212                    if attempt < 3 {
213                        Err(transient_error())
214                    } else {
215                        Ok("ok")
216                    }
217                }
218            },
219            request_error_class,
220            request_retry_delay,
221            |_| async {},
222        )
223        .await
224        .unwrap();
225        assert_eq!(value, "ok");
226        assert_eq!(attempts.load(Ordering::SeqCst), 3);
227    }
228
229    #[tokio::test]
230    async fn transient_operation_stops_after_five_total_attempts() {
231        let attempts = AtomicUsize::new(0);
232        let result = retry_operation(
233            "test",
234            || {
235                attempts.fetch_add(1, Ordering::SeqCst);
236                async { Err::<(), _>(transient_error()) }
237            },
238            request_error_class,
239            request_retry_delay,
240            |_| async {},
241        )
242        .await;
243        assert!(result.is_err());
244        assert_eq!(attempts.load(Ordering::SeqCst), MAX_ATTEMPTS);
245    }
246
247    #[tokio::test]
248    async fn permanent_operation_is_attempted_once() {
249        let attempts = AtomicUsize::new(0);
250        let result = retry_operation(
251            "test",
252            || {
253                attempts.fetch_add(1, Ordering::SeqCst);
254                async { Err::<(), _>(RequestError::MigrateToChatId(ChatId(-100))) }
255            },
256            request_error_class,
257            request_retry_delay,
258            |_| async {},
259        )
260        .await;
261        assert!(result.is_err());
262        assert_eq!(attempts.load(Ordering::SeqCst), 1);
263    }
264
265    #[test]
266    fn retry_after_is_respected_exactly() {
267        let delay = request_retry_delay(&RequestError::RetryAfter(Seconds::from_seconds(17)), 1);
268        assert_eq!(delay, Some(Duration::from_secs(17)));
269    }
270}
kcode_telegram_request_policy/lib.rs

kcode_telegram_request_policy/
lib.rs
