1use crate::error::{Aria2Error, RecoverableError};
2use std::sync::atomic::{AtomicU32, Ordering};
3use tracing::warn;
4
5pub 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}