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 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/// 定时器任务
77pub struct TimerTask {
78    /// 任务唯一标识符
79    pub id: TaskId,
80    
81    /// 到期时间（相对于时间轮的 tick 数）
82    pub deadline_tick: u64,
83    
84    /// 轮次计数（用于超出时间轮范围的任务）
85    pub rounds: u32,
86    
87    /// 异步回调函数（可选）
88    pub callback: Option<CallbackWrapper>,
89    
90    /// 完成通知器（用于在任务完成时发送通知）
91    pub completion_notifier: CompletionNotifier,
92}
93
94impl TimerTask {
95    /// 创建一次性定时器任务
96    pub fn once(
97        deadline_tick: u64,
98        rounds: u32,
99        callback: Option<CallbackWrapper>,
100        completion_notifier: CompletionNotifier,
101    ) -> Self {
102        Self {
103            id: TaskId::new(),
104            deadline_tick,
105            rounds,
106            callback,
107            completion_notifier,
108        }
109    }
110
111    /// 获取回调函数的克隆（如果存在）
112    pub fn get_callback(&self) -> Option<CallbackWrapper> {
113        self.callback.as_ref().map(Arc::clone)
114    }
115}
116
117/// 任务位置信息，用于取消操作
118#[derive(Debug, Clone)]
119pub struct TaskLocation {
120    pub slot_index: usize,
121    /// 任务在槽位 Vec 中的索引位置（用于 O(1) 取消）
122    pub vec_index: usize,
123    #[allow(dead_code)]
124    pub task_id: TaskId,
125}
126
127impl TaskLocation {
128    pub fn new(slot_index: usize, vec_index: usize, task_id: TaskId) -> Self {
129        Self {
130            slot_index,
131            vec_index,
132            task_id,
133        }
134    }
135}
136
kestrel_protocol_timer/task.rs

kestrel_protocol_timer/
task.rs
