Skip to main content

aria2_core/
retry.rs

1use crate::error::{Aria2Error, RecoverableError};
2use std::sync::atomic::{AtomicU32, Ordering};
3use tracing::warn;
4
5// Re-export unified RetryPolicy from engine::retry_policy
6pub use crate::engine::retry_policy::RetryPolicy;
7
8#[derive(Debug, Default)]
9pub struct RetryStats {
10    total: AtomicU32,
11    timeouts: AtomicU32,
12    server_errors: AtomicU32,
13    network_failures: AtomicU32,
14    max_retries_reached: AtomicU32,
15}
16
17impl RetryStats {
18    pub fn record_retry(&self, error: &Aria2Error) {
19        self.total.fetch_add(1, Ordering::Relaxed);
20        match error {
21            Aria2Error::Recoverable(RecoverableError::Timeout) => {
22                self.timeouts.fetch_add(1, Ordering::Relaxed);
23            }
24            Aria2Error::Recoverable(RecoverableError::ServerError { .. }) => {
25                self.server_errors.fetch_add(1, Ordering::Relaxed);
26            }
27            Aria2Error::Recoverable(RecoverableError::TemporaryNetworkFailure { .. }) => {
28                self.network_failures.fetch_add(1, Ordering::Relaxed);
29            }
30            Aria2Error::Recoverable(RecoverableError::MaxTriesReached { .. }) => {
31                self.max_retries_reached.fetch_add(1, Ordering::Relaxed);
32            }
33            _ => {}
34        }
35    }
36
37    pub fn total(&self) -> u32 {
38        self.total.load(Ordering::Relaxed)
39    }
40
41    pub fn timeouts(&self) -> u32 {
42        self.timeouts.load(Ordering::Relaxed)
43    }
44
45    pub fn server_errors(&self) -> u32 {
46        self.server_errors.load(Ordering::Relaxed)
47    }
48
49    pub fn network_failures(&self) -> u32 {
50        self.network_failures.load(Ordering::Relaxed)
51    }
52
53    pub fn reset(&self) {
54        self.total.store(0, Ordering::Relaxed);
55        self.timeouts.store(0, Ordering::Relaxed);
56        self.server_errors.store(0, Ordering::Relaxed);
57        self.network_failures.store(0, Ordering::Relaxed);
58        self.max_retries_reached.store(0, Ordering::Relaxed);
59    }
60}
61
62pub struct RetryExecutor<'a> {
63    policy: &'a RetryPolicy,
64    stats: &'a RetryStats,
65}
66
67impl<'a> RetryExecutor<'a> {
68    pub fn new(policy: &'a RetryPolicy, stats: &'a RetryStats) -> Self {
69        Self { policy, stats }
70    }
71
72    pub async fn execute<F, Fut, T>(&self, mut operation: F) -> crate::error::Result<T>
73    where
74        F: FnMut(u32) -> Fut,
75        Fut: std::future::Future<Output = crate::error::Result<T>>,
76    {
77        let mut attempt = 0u32;
78        loop {
79            let result = operation(attempt).await;
80            match result {
81                Ok(value) => return Ok(value),
82                Err(error) => {
83                    if !self.policy.should_retry(attempt, &error) {
84                        warn!(
85                            "重试失败 (尝试 {}/{}, 不再重试): {}",
86                            attempt + 1,
87                            self.policy.max_tries(),
88                            error
89                        );
90                        self.stats.record_retry(&error);
91                        return Err(error);
92                    }
93                    attempt += 1;
94                    let wait = self.policy.wait_duration(attempt);
95                    warn!(
96                        "第 {} 次重试, 等待 {:?} 后执行 (原因: {})",
97                        attempt, wait, error
98                    );
99                    self.stats.record_retry(&error);
100                    tokio::time::sleep(wait).await;
101                }
102            }
103        }
104    }
105}