kestrel_protocol_timer/
task.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::Arc;
5use tokio::sync::oneshot;
6
7/// 全局唯一的任务 ID 生成器
8static NEXT_TASK_ID: AtomicU64 = AtomicU64::new(1);
9
10/// 定时器任务的唯一标识符
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub struct TaskId(u64);
13
14impl TaskId {
15    /// 生成一个新的唯一任务 ID（内部使用）
16    pub(crate) fn new() -> Self {
17        TaskId(NEXT_TASK_ID.fetch_add(1, Ordering::Relaxed))
18    }
19
20    /// 获取任务 ID 的数值
21    pub fn as_u64(&self) -> u64 {
22        self.0
23    }
24}
25
26impl Default for TaskId {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32/// 定时器回调 trait
33/// 
34/// 实现此 trait 的类型可以作为定时器的回调函数使用。
35/// 
36/// # 示例
37/// 
38/// ```
39/// use kestrel_protocol_timer::TimerCallback;
40/// use std::future::Future;
41/// use std::pin::Pin;
42/// 
43/// struct MyCallback;
44/// 
45/// impl TimerCallback for MyCallback {
46///     fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
47///         Box::pin(async {
48///             println!("Timer callback executed!");
49///         })
50///     }
51/// }
52/// ```
53pub trait TimerCallback: Send + Sync + 'static {
54    /// 执行回调，返回一个 Future
55    fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>>;
56}
57
58/// 为闭包实现 TimerCallback trait
59/// 支持 Fn() -> Future 类型的闭包（可以多次调用，适合周期性任务）
60impl<F, Fut> TimerCallback for F
61where
62    F: Fn() -> Fut + Send + Sync + 'static,
63    Fut: Future<Output = ()> + Send + 'static,
64{
65    fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
66        Box::pin(self())
67    }
68}
69
70/// 回调包装器类型
71pub type CallbackWrapper = Arc<dyn TimerCallback>;
72
73/// 完成通知器，用于在任务完成时发送通知
74pub struct CompletionNotifier(pub oneshot::Sender<()>);
75
76/// 定时器任务
77/// 
78/// 用户通过两步式 API 使用：
79/// 1. 使用 `TimerTask::new()` 创建任务
80/// 2. 使用 `TimerWheel::register()` 或 `TimerService::register()` 注册任务
81pub struct TimerTask {
82    /// 任务唯一标识符
83    pub(crate) id: TaskId,
84    
85    /// 用户指定的延迟时间
86    pub(crate) delay: std::time::Duration,
87    
88    /// 到期时间（相对于时间轮的 tick 数）
89    pub(crate) deadline_tick: u64,
90    
91    /// 轮次计数（用于超出时间轮范围的任务）
92    pub(crate) rounds: u32,
93    
94    /// 异步回调函数（可选）
95    pub(crate) callback: Option<CallbackWrapper>,
96    
97    /// 完成通知器（用于在任务完成时发送通知，注册时创建）
98    pub(crate) completion_notifier: Option<CompletionNotifier>,
99}
100
101impl TimerTask {
102    /// 创建新的定时器任务
103    /// 
104    /// # 参数
105    /// - `delay`: 延迟时间
106    /// - `callback`: 回调函数（可选）
107    /// 
108    /// # 示例
109    /// ```no_run
110    /// use kestrel_protocol_timer::TimerTask;
111    /// use std::time::Duration;
112    /// use std::sync::Arc;
113    /// 
114    /// // 创建带回调的任务
115    /// let callback = Arc::new(|| async {
116    ///     println!("Timer fired!");
117    /// });
118    /// let task = TimerTask::new(Duration::from_secs(1), Some(callback));
119    /// 
120    /// // 创建仅通知的任务
121    /// let task = TimerTask::new(Duration::from_secs(1), None);
122    /// ```
123    pub fn new(delay: std::time::Duration, callback: Option<CallbackWrapper>) -> Self {
124        Self {
125            id: TaskId::new(),
126            delay,
127            deadline_tick: 0,
128            rounds: 0,
129            callback,
130            completion_notifier: None,
131        }
132    }
133
134    /// 获取任务 ID
135    /// 
136    /// # 示例
137    /// ```no_run
138    /// use kestrel_protocol_timer::TimerTask;
139    /// use std::time::Duration;
140    /// 
141    /// let task = TimerTask::new(Duration::from_secs(1), None);
142    /// let task_id = task.get_id();
143    /// println!("Task ID: {:?}", task_id);
144    /// ```
145    pub fn get_id(&self) -> TaskId {
146        self.id
147    }
148
149    /// 内部方法：准备注册（在注册时由时间轮调用）
150    pub(crate) fn prepare_for_registration(
151        &mut self,
152        completion_notifier: CompletionNotifier,
153        deadline_tick: u64,
154        rounds: u32,
155    ) {
156        self.completion_notifier = Some(completion_notifier);
157        self.deadline_tick = deadline_tick;
158        self.rounds = rounds;
159    }
160
161    /// 获取回调函数的克隆（如果存在）
162    pub(crate) fn get_callback(&self) -> Option<CallbackWrapper> {
163        self.callback.as_ref().map(Arc::clone)
164    }
165}
166
167/// 任务位置信息，用于取消操作
168#[derive(Debug, Clone)]
169pub struct TaskLocation {
170    pub slot_index: usize,
171    /// 任务在槽位 Vec 中的索引位置（用于 O(1) 取消）
172    pub vec_index: usize,
173    #[allow(dead_code)]
174    pub task_id: TaskId,
175}
176
177impl TaskLocation {
178    pub fn new(slot_index: usize, vec_index: usize, task_id: TaskId) -> Self {
179        Self {
180            slot_index,
181            vec_index,
182            task_id,
183        }
184    }
185}
186
kestrel_protocol_timer/task.rs

kestrel_protocol_timer/
task.rs
