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

kestrel_protocol_timer/
task.rs
