Skip to main content

aria2_core/engine/
timer.rs

1use std::collections::HashMap;
2use std::time::Duration;
3use tokio::sync::mpsc;
4use tokio::time::{Instant as TokioInstant, sleep_until};
5use tracing::{debug, info};
6
7use crate::error::Result;
8
9pub type TimerId = u64;
10
11#[derive(Clone, Debug)]
12pub enum TimerEvent {
13    OneShot(TimerId),
14    Periodic(TimerId),
15}
16
17pub struct Timer {
18    #[allow(dead_code)]
19    // Timer ID stored for identification; not yet used in current implementation
20    id: TimerId,
21    next_fire: TokioInstant,
22    interval: Option<Duration>,
23}
24
25impl Timer {
26    fn new(id: TimerId, delay: Duration, interval: Option<Duration>) -> Self {
27        let next_fire = TokioInstant::now() + delay;
28        Timer {
29            id,
30            next_fire,
31            interval,
32        }
33    }
34}
35
36pub struct TimerA2 {
37    timers: HashMap<TimerId, Timer>,
38    event_tx: mpsc::UnboundedSender<TimerEvent>,
39    next_id: TimerId,
40}
41
42impl TimerA2 {
43    pub fn new() -> (Self, mpsc::UnboundedReceiver<TimerEvent>) {
44        let (event_tx, event_rx) = mpsc::unbounded_channel();
45
46        let timer_a2 = TimerA2 {
47            timers: HashMap::new(),
48            event_tx,
49            next_id: 0,
50        };
51
52        info!("Timer system initialization complete");
53        (timer_a2, event_rx)
54    }
55
56    pub fn add_timer(&mut self, delay: Duration) -> Result<TimerId> {
57        self.add_periodic_timer(delay, None)
58    }
59
60    pub fn add_periodic_timer(
61        &mut self,
62        delay: Duration,
63        interval: Option<Duration>,
64    ) -> Result<TimerId> {
65        let id = self.next_id;
66        self.next_id += 1;
67
68        let timer = Timer::new(id, delay, interval);
69        self.timers.insert(id, timer);
70
71        debug!(
72            "Adding timer #{} (delay: {:?}, interval: {:?})",
73            id, delay, interval
74        );
75        Ok(id)
76    }
77
78    pub fn remove_timer(&mut self, id: TimerId) {
79        if self.timers.remove(&id).is_some() {
80            debug!("Removing timer #{}", id);
81        }
82    }
83
84    pub async fn run(mut self) -> Result<()> {
85        info!("Timer system started");
86
87        loop {
88            if self.timers.is_empty() {
89                debug!("No active timers, waiting for new timer");
90                tokio::task::yield_now().await;
91                continue;
92            }
93
94            let now = TokioInstant::now();
95
96            let mut fired_timers: Vec<(TimerId, bool)> = Vec::new();
97
98            for (&id, timer) in &self.timers {
99                if now >= timer.next_fire {
100                    let is_periodic = timer.interval.is_some();
101                    fired_timers.push((id, is_periodic));
102                }
103            }
104
105            for (id, is_periodic) in fired_timers {
106                if is_periodic {
107                    let _ = self.event_tx.send(TimerEvent::Periodic(id));
108
109                    if let Some(timer) = self.timers.get_mut(&id)
110                        && let Some(interval) = timer.interval
111                    {
112                        timer.next_fire = TokioInstant::now() + interval;
113                    }
114                } else {
115                    let _ = self.event_tx.send(TimerEvent::OneShot(id));
116                    self.timers.remove(&id);
117                }
118            }
119
120            if !self.timers.is_empty() {
121                let next_fire = self
122                    .timers
123                    .values()
124                    .map(|t| t.next_fire)
125                    .min()
126                    .expect("至少有一个定时器");
127
128                sleep_until(next_fire).await;
129            }
130        }
131    }
132}