zerolaunch_plugin_api/services/timer/
tokio_timer_manager.rs

1use crate::host::error::HostApiError;
2use crate::services::timer::timer_manager::TimerManager;
3use crate::services::timer::types::{TimerCallback, TimerId, TimerMode};
4use async_trait::async_trait;
5use dashmap::DashMap;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::sync::Notify;
10use tokio::time::MissedTickBehavior;
11
12/// 定时器条目，用于通过条件变量机制取消定时器。
13struct TimerEntry {
14    /// 条件变量通知器，用于唤醒等待中的定时器任务以提前取消。
15    notify: Arc<Notify>,
16}
17
18/// 基于 tokio 的定时器管理器实现。
19/// 内部使用 tokio::sync::Notify 作为条件变量机制，支持高效的取消操作。
20/// 所有定时器在 tokio 运行时中异步执行，不阻塞任何线程。
21pub struct TokioTimerManager {
22    /// 自增 ID 生成器
23    next_id: AtomicU64,
24    /// 活跃定时器集合（使用 Arc 包装以便在异步任务中共享）
25    timers: Arc<DashMap<TimerId, TimerEntry>>,
26}
27
28impl TokioTimerManager {
29    /// 创建 TokioTimerManager 实例。
30    pub fn new() -> Self {
31        Self {
32            next_id: AtomicU64::new(1),
33            timers: Arc::new(DashMap::new()),
34        }
35    }
36}
37
38impl Default for TokioTimerManager {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44#[async_trait]
45impl TimerManager for TokioTimerManager {
46    async fn set_timer(
47        &self,
48        delay: Duration,
49        mode: TimerMode,
50        callback: TimerCallback,
51    ) -> Result<TimerId, HostApiError> {
52        let id = TimerId::new(self.next_id.fetch_add(1, Ordering::SeqCst));
53
54        let notify = Arc::new(Notify::new());
55
56        self.timers.insert(
57            id,
58            TimerEntry {
59                notify: notify.clone(),
60            },
61        );
62
63        let timers = self.timers.clone();
64
65        tokio::spawn(async move {
66            match mode {
67                TimerMode::OneShot => {
68                    tokio::select! {
69                        _ = tokio::time::sleep(delay) => {
70                            callback(id);
71                            timers.remove(&id);
72                        }
73                        _ = notify.notified() => {
74                        }
75                    }
76                }
77                TimerMode::Interval => {
78                    let mut interval = tokio::time::interval(delay);
79                    interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
80                    interval.tick().await;
81
82                    loop {
83                        tokio::select! {
84                            _ = interval.tick() => {
85                                callback(id);
86                            }
87                            _ = notify.notified() => {
88                                break;
89                            }
90                        }
91                    }
92                }
93            }
94        });
95
96        Ok(id)
97    }
98
99    async fn cancel_timer(&self, id: TimerId) -> Result<(), HostApiError> {
100        if let Some((_, entry)) = self.timers.remove(&id) {
101            entry.notify.notify_one();
102        }
103        Ok(())
104    }
105
106    async fn cancel_all(&self) -> Result<(), HostApiError> {
107        let entries: Vec<TimerEntry> = self
108            .timers
109            .iter()
110            .map(|e| TimerEntry {
111                notify: e.notify.clone(),
112            })
113            .collect();
114        self.timers.clear();
115        for entry in entries {
116            entry.notify.notify_one();
117        }
118        Ok(())
119    }
120}
zerolaunch_plugin_api/services/timer/tokio_timer_manager.rs

zerolaunch_plugin_api/services/timer/
tokio_timer_manager.rs
