aria2_core/engine/
timer.rs1use 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 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}