kestrel_timer/timer.rs
1pub mod handle;
2
3use crate::config::{BatchConfig, ServiceConfig, WheelConfig};
4use crate::error::TimerError;
5use crate::task::{CallbackWrapper, TaskHandle, TaskId};
6use crate::wheel::Wheel;
7use handle::{BatchHandle, BatchHandleWithCompletion, TimerHandle, TimerHandleWithCompletion};
8use parking_lot::Mutex;
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::task::JoinHandle;
12use tokio::time::Instant;
13
14/// Timing Wheel Timer Manager
15///
16/// 时间轮定时器管理器
17pub struct TimerWheel {
18 /// Timing wheel instance, wrapped in `Arc<Mutex>` for multi-threaded access
19 ///
20 /// 时间轮实例，包装在 `Arc<Mutex>` 中以支持多线程访问
21 wheel: Arc<Mutex<Wheel>>,
22
23 /// Background tick loop task handle
24 ///
25 /// 后台 tick 循环任务句柄
26 tick_handle: Option<JoinHandle<()>>,
27}
28
29impl TimerWheel {
30 /// Create a new timer manager
31 ///
32 /// # Parameters
33 /// - `config`: Timing wheel configuration
34 /// - `batch_config`: Batch operation configuration
35 ///
36 /// 创建新的定时器管理器
37 ///
38 /// # 参数
39 /// - `config`: 时间轮配置
40 /// - `batch_config`: 批量操作配置
41 ///
42 /// # Examples (示例)
43 /// ```no_run
44 /// use kestrel_timer::{TimerWheel, config::WheelConfig, TimerTask, config::BatchConfig};
45 /// use std::time::Duration;
46 ///
47 /// #[tokio::main]
48 /// async fn main() {
49 /// let config = WheelConfig::builder()
50 /// .l0_tick_duration(Duration::from_millis(10))
51 /// .l0_slot_count(512)
52 /// .l1_tick_duration(Duration::from_secs(1))
53 /// .l1_slot_count(64)
54 /// .build()
55 /// .unwrap();
56 /// let timer = TimerWheel::new(config, BatchConfig::default()).unwrap();
57 ///
58 /// // Use two-step API: allocate handle first, then register
59 /// // 使用两步 API：先分配 handle，再注册
60 /// let handle = timer.allocate_handle();
61 /// let task = TimerTask::new_oneshot(Duration::from_secs(1), None);
62 /// let _timer_handle = timer.register(handle, task).unwrap();
63 /// }
64 /// ```
65 pub fn new(config: WheelConfig, batch_config: BatchConfig) -> Result<Self, TimerError> {
66 let tick_duration = config.l0_tick_duration();
67 let wheel = Wheel::new(config, batch_config)?;
68 let wheel = Arc::new(Mutex::new(wheel));
69 let wheel_clone = wheel.clone();
70
71 // Start background tick loop
72 // 启动后台 tick 循环
73 let tick_handle = tokio::spawn(async move {
74 Self::tick_loop(wheel_clone, tick_duration).await;
75 });
76
77 Ok(Self {
78 wheel,
79 tick_handle: Some(tick_handle),
80 })
81 }
82
83 /// Create timer manager with default configuration, hierarchical mode
84 /// - L0 layer tick duration: 10ms, slot count: 512
85 /// - L1 layer tick duration: 1s, slot count: 64
86 ///
87 /// # Returns
88 /// Timer manager instance
89 ///
90 /// 使用默认配置创建定时器管理器，分层模式
91 /// - L0 层 tick 持续时间：10ms，槽数量：512
92 /// - L1 层 tick 持续时间：1s，槽数量：64
93 ///
94 /// # 返回值
95 /// 定时器管理器实例
96 ///
97 /// # Examples (示例)
98 /// ```no_run
99 /// use kestrel_timer::TimerWheel;
100 ///
101 /// #[tokio::main]
102 /// async fn main() {
103 /// let timer = TimerWheel::with_defaults();
104 /// }
105 /// ```
106 pub fn with_defaults() -> Self {
107 Self::new(WheelConfig::default(), BatchConfig::default())
108 .expect("default timer wheel configuration must be valid")
109 }
110
111 /// Create TimerService bound to this timing wheel with default configuration
112 ///
113 /// # Parameters
114 /// - `service_config`: Service configuration
115 ///
116 /// # Returns
117 /// TimerService instance bound to this timing wheel
118 ///
119 /// 创建绑定到此时间轮的 TimerService，使用默认配置
120 ///
121 /// # 参数
122 /// - `service_config`: 服务配置
123 ///
124 /// # 返回值
125 /// 绑定到此时间轮的 TimerService 实例
126 ///
127 /// # Examples (示例)
128 /// ```no_run
129 /// use kestrel_timer::{TimerWheel, TimerService, CallbackWrapper, config::ServiceConfig};
130 /// use std::time::Duration;
131 ///
132 ///
133 /// #[tokio::main]
134 /// async fn main() {
135 /// let timer = TimerWheel::with_defaults();
136 /// let mut service = timer.create_service(ServiceConfig::default());
137 ///
138 /// // Use two-step API to batch schedule timers through service
139 /// // 使用两步 API 通过服务批量调度定时器
140 /// // Step 1: Allocate handles
141 /// let handles = service.allocate_handles(5);
142 ///
143 /// // Step 2: Create tasks
144 /// let tasks: Vec<_> = (0..5)
145 /// .map(|_| {
146 /// use kestrel_timer::TimerTask;
147 /// TimerTask::new_oneshot(Duration::from_millis(100), Some(CallbackWrapper::new(|| async {})))
148 /// })
149 /// .collect();
150 ///
151 /// // Step 3: Register batch
152 /// service.register_batch(handles, tasks).unwrap();
153 ///
154 /// // Receive timeout notifications
155 /// // 接收超时通知
156 /// let mut rx = service.take_receiver().unwrap();
157 /// while let Some(task_id) = rx.recv().await {
158 /// println!("Task {:?} completed", task_id);
159 /// }
160 /// }
161 /// ```
162 pub fn create_service(&self, service_config: ServiceConfig) -> crate::service::TimerService {
163 crate::service::TimerService::new(self.wheel.clone(), service_config)
164 }
165
166 /// Create TimerService bound to this timing wheel with custom configuration
167 ///
168 /// # Parameters
169 /// - `config`: Service configuration
170 ///
171 /// # Returns
172 /// TimerService instance bound to this timing wheel
173 ///
174 /// 创建绑定到此时间轮的 TimerService，使用自定义配置
175 ///
176 /// # 参数
177 /// - `config`: 服务配置
178 ///
179 /// # 返回值
180 /// 绑定到此时间轮的 TimerService 实例
181 ///
182 /// # Examples (示例)
183 /// ```no_run
184 /// use kestrel_timer::{TimerWheel, config::ServiceConfig, TimerTask};
185 /// use std::num::NonZeroUsize;
186 ///
187 /// #[tokio::main]
188 /// async fn main() {
189 /// let timer = TimerWheel::with_defaults();
190 /// let config = ServiceConfig::builder()
191 /// .command_channel_capacity(NonZeroUsize::new(1024).unwrap())
192 /// .timeout_channel_capacity(NonZeroUsize::new(2000).unwrap())
193 /// .build();
194 /// let service = timer.create_service_with_config(config);
195 /// }
196 /// ```
197 pub fn create_service_with_config(
198 &self,
199 config: ServiceConfig,
200 ) -> crate::service::TimerService {
201 crate::service::TimerService::new(self.wheel.clone(), config)
202 }
203
204 /// Allocate a handle from DeferredMap
205 ///
206 /// # Returns
207 /// A unique handle for later insertion
208 ///
209 /// # 返回值
210 /// 用于后续插入的唯一 handle
211 pub fn allocate_handle(&self) -> TaskHandle {
212 self.wheel.lock().allocate_handle()
213 }
214
215 /// Batch allocate handles from DeferredMap
216 ///
217 /// # Parameters
218 /// - `count`: Number of handles to allocate
219 ///
220 /// # Returns
221 /// Vector of unique handles for later batch insertion
222 ///
223 /// # 参数
224 /// - `count`: 要分配的 handle 数量
225 ///
226 /// # 返回值
227 /// 用于后续批量插入的唯一 handles 向量
228 pub fn allocate_handles(&self, count: usize) -> Vec<TaskHandle> {
229 self.wheel.lock().allocate_handles(count)
230 }
231
232 /// Register timer task to timing wheel (registration phase)
233 ///
234 /// # Parameters
235 /// - `task`: Task created via `create_task()`
236 ///
237 /// # Returns
238 /// Return `Ok` with a timer handle and completion receiver. Returns
239 /// `Err(TimerError::WrongWheel)` when the handle belongs to another wheel,
240 /// or `Err(TimerError::Shutdown)` when the wheel is closed.
241 ///
242 /// 注册定时器任务到时间轮 (注册阶段)
243 ///
244 /// # 参数
245 /// - `task`: 通过 `create_task()` 创建的任务
246 ///
247 /// # 返回值
248 /// 成功时返回包含完成通知接收器的定时器句柄；handle 属于其他时间轮时
249 /// 返回 `Err(TimerError::WrongWheel)`；时间轮关闭时返回
250 /// `Err(TimerError::Shutdown)`。
251 ///
252 /// # Examples (示例)
253 /// ```no_run
254 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
255 ///
256 /// use std::time::Duration;
257 ///
258 /// #[tokio::main]
259 /// async fn main() {
260 /// let timer = TimerWheel::with_defaults();
261 ///
262 /// // Step 1: Allocate handle
263 /// let allocated_handle = timer.allocate_handle();
264 /// let task_id = allocated_handle.task_id();
265 ///
266 /// // Step 2: Create task
267 /// let task = TimerTask::new_oneshot(Duration::from_secs(1), Some(CallbackWrapper::new(|| async {
268 /// println!("Timer fired!");
269 /// })));
270 ///
271 /// // Step 3: Register task
272 /// let handle = timer.register(allocated_handle, task).unwrap();
273 ///
274 /// // Wait for timer completion
275 /// // 等待定时器完成
276 /// use kestrel_timer::CompletionReceiver;
277 /// let (rx, _handle) = handle.into_parts();
278 /// match rx {
279 /// CompletionReceiver::OneShot(receiver) => {
280 /// receiver.recv().await.unwrap();
281 /// },
282 /// _ => {}
283 /// }
284 /// }
285 /// ```
286 #[inline]
287 pub fn register(
288 &self,
289 handle: TaskHandle,
290 task: crate::task::TimerTask,
291 ) -> Result<TimerHandleWithCompletion, TimerError> {
292 let task_id = handle.task_id();
293
294 let (task, completion_rx) =
295 crate::task::TimerTaskWithCompletionNotifier::from_timer_task(task);
296
297 // Single lock to complete all operations
298 // 单次加锁完成所有操作
299 let mut wheel_guard = self.wheel.lock();
300 wheel_guard.insert_at(handle, task)?;
301
302 Ok(TimerHandleWithCompletion::new(
303 TimerHandle::new(task_id, self.wheel.clone()),
304 completion_rx,
305 ))
306 }
307
308 /// Batch register timer tasks to timing wheel (registration phase)
309 ///
310 /// # Parameters
311 /// - `handles`: Pre-allocated handles for tasks
312 /// - `tasks`: List of timer tasks
313 ///
314 /// # Returns
315 /// - `Ok(BatchHandleWithCompletion)` if all tasks are successfully registered
316 /// - `Err(TimerError::BatchLengthMismatch)` if handles and tasks lengths don't match
317 /// - `Err(TimerError::WrongWheel)` if any handle belongs to another wheel
318 /// - `Err(TimerError::Shutdown)` if the wheel is closed
319 ///
320 /// 批量注册定时器任务到时间轮 (注册阶段)
321 ///
322 /// # 参数
323 /// - `handles`: 任务的预分配 handles
324 /// - `tasks`: 定时器任务列表
325 ///
326 /// # 返回值
327 /// - `Ok(BatchHandleWithCompletion)` 如果所有任务成功注册
328 /// - `Err(TimerError::BatchLengthMismatch)` 如果 handles 和 tasks 长度不匹配
329 /// - `Err(TimerError::WrongWheel)` 如果任一 handle 属于其他时间轮
330 /// - `Err(TimerError::Shutdown)` 如果时间轮已关闭
331 ///
332 /// # Examples (示例)
333 /// ```no_run
334 /// use kestrel_timer::{TimerWheel, TimerTask};
335 /// use std::time::Duration;
336 ///
337 /// #[tokio::main]
338 /// async fn main() {
339 /// let timer = TimerWheel::with_defaults();
340 ///
341 /// // Step 1: Allocate handles
342 /// let handles = timer.allocate_handles(3);
343 ///
344 /// // Step 2: Create tasks
345 /// let tasks: Vec<_> = (0..3)
346 /// .map(|_| TimerTask::new_oneshot(Duration::from_secs(1), None))
347 /// .collect();
348 ///
349 /// // Step 3: Batch register
350 /// let batch = timer.register_batch(handles, tasks)
351 /// .expect("register_batch should succeed");
352 /// println!("Registered {} timers", batch.len());
353 /// }
354 /// ```
355 #[inline]
356 pub fn register_batch(
357 &self,
358 handles: Vec<TaskHandle>,
359 tasks: Vec<crate::task::TimerTask>,
360 ) -> Result<BatchHandleWithCompletion, TimerError> {
361 // Validate lengths match
362 if handles.len() != tasks.len() {
363 return Err(TimerError::BatchLengthMismatch {
364 handles_len: handles.len(),
365 tasks_len: tasks.len(),
366 });
367 }
368
369 let task_count = tasks.len();
370 let mut completion_rxs = Vec::with_capacity(task_count);
371 let mut task_ids = Vec::with_capacity(task_count);
372 let mut prepared_handles = Vec::with_capacity(task_count);
373 let mut prepared_tasks = Vec::with_capacity(task_count);
374
375 // Step 1: Prepare all channels and notifiers
376 for (handle, task) in handles.into_iter().zip(tasks) {
377 let task_id = handle.task_id();
378 let (task, completion_rx) =
379 crate::task::TimerTaskWithCompletionNotifier::from_timer_task(task);
380 task_ids.push(task_id);
381 completion_rxs.push(completion_rx);
382 prepared_handles.push(handle);
383 prepared_tasks.push(task);
384 }
385
386 // Step 2: Single lock, batch insert
387 {
388 let mut wheel_guard = self.wheel.lock();
389 wheel_guard.insert_batch_at(prepared_handles, prepared_tasks)?;
390 }
391
392 Ok(BatchHandleWithCompletion::new(
393 BatchHandle::new(task_ids, self.wheel.clone()),
394 completion_rxs,
395 ))
396 }
397
398 /// Cancel timer
399 ///
400 /// # Parameters
401 /// - `task_id`: Task ID
402 ///
403 /// # Returns
404 /// Returns `Ok(true)` when cancelled, `Ok(false)` when absent, or
405 /// `Err(TimerError::WrongWheel)` for an ID from another wheel.
406 ///
407 /// 取消定时器
408 ///
409 /// # 参数
410 /// - `task_id`: 任务 ID
411 ///
412 /// # 返回值
413 /// 任务存在且取消成功时返回 `Ok(true)`，任务不存在时返回 `Ok(false)`，
414 /// ID 属于其他时间轮时返回 `Err(TimerError::WrongWheel)`。
415 ///
416 /// # Examples (示例)
417 /// ```no_run
418 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
419 ///
420 /// use std::time::Duration;
421 ///
422 /// #[tokio::main]
423 /// async fn main() {
424 /// let timer = TimerWheel::with_defaults();
425 ///
426 /// // Step 1: Allocate handle
427 /// let allocated_handle = timer.allocate_handle();
428 /// let task_id = allocated_handle.task_id();
429 ///
430 /// // Step 2: Create and register task
431 /// let task = TimerTask::new_oneshot(Duration::from_secs(10), Some(CallbackWrapper::new(|| async {
432 /// println!("Timer fired!");
433 /// })));
434 /// let _handle = timer.register(allocated_handle, task).unwrap();
435 ///
436 /// // Cancel task using task ID
437 /// // 使用任务 ID 取消任务
438 /// let cancelled = timer.cancel(task_id).unwrap();
439 /// println!("Canceled successfully: {}", cancelled);
440 /// }
441 /// ```
442 #[inline]
443 pub fn cancel(&self, task_id: TaskId) -> Result<bool, TimerError> {
444 let mut wheel = self.wheel.lock();
445 wheel.cancel(task_id)
446 }
447
448 /// Batch cancel timers
449 ///
450 /// # Parameters
451 /// - `task_ids`: List of task IDs to cancel
452 ///
453 /// # Returns
454 /// Number of successfully cancelled tasks, or
455 /// `Err(TimerError::WrongWheel)` if any ID belongs to another wheel.
456 ///
457 /// 批量取消定时器
458 ///
459 /// # 参数
460 /// - `task_ids`: 要取消的任务 ID 列表
461 ///
462 /// # 返回值
463 /// 成功取消的任务数量；如果任一 ID 属于其他时间轮则返回
464 /// `Err(TimerError::WrongWheel)`，且不修改任何任务。
465 ///
466 /// # Performance Advantages
467 /// - Batch processing reduces lock contention
468 /// - Internally optimized batch cancellation operation
469 ///
470 /// # Examples (示例)
471 /// ```no_run
472 /// use kestrel_timer::{TimerWheel, TimerTask};
473 /// use std::time::Duration;
474 ///
475 /// #[tokio::main]
476 /// async fn main() {
477 /// let timer = TimerWheel::with_defaults();
478 ///
479 /// // Create multiple timers
480 /// // 创建多个定时器
481 /// let task1 = TimerTask::new_oneshot(Duration::from_secs(10), None);
482 /// let task2 = TimerTask::new_oneshot(Duration::from_secs(10), None);
483 /// let task3 = TimerTask::new_oneshot(Duration::from_secs(10), None);
484 ///
485 /// // Allocate handles and get task IDs
486 /// let h1 = timer.allocate_handle();
487 /// let h2 = timer.allocate_handle();
488 /// let h3 = timer.allocate_handle();
489 /// let task_ids = vec![h1.task_id(), h2.task_id(), h3.task_id()];
490 ///
491 /// let _h1 = timer.register(h1, task1).unwrap();
492 /// let _h2 = timer.register(h2, task2).unwrap();
493 /// let _h3 = timer.register(h3, task3).unwrap();
494 ///
495 /// // Batch cancel
496 /// // 批量取消
497 /// let cancelled = timer.cancel_batch(&task_ids).unwrap();
498 /// println!("Canceled {} timers", cancelled);
499 /// }
500 /// ```
501 #[inline]
502 pub fn cancel_batch(&self, task_ids: &[TaskId]) -> Result<usize, TimerError> {
503 let mut wheel = self.wheel.lock();
504 wheel.cancel_batch(task_ids)
505 }
506
507 /// Postpone timer
508 ///
509 /// # Parameters
510 /// - `task_id`: Task ID to postpone
511 /// - `new_delay`: New delay duration, recalculated from current time
512 /// - `callback`: New callback function, pass `None` to keep original callback, pass `Some` to replace with new callback
513 ///
514 /// # Returns
515 /// Returns `Ok(true)` when postponed, `Ok(false)` when absent, or
516 /// `Err(TimerError::WrongWheel)` for an ID from another wheel.
517 ///
518 /// 推迟定时器
519 ///
520 /// # 参数
521 /// - `task_id`: 要推迟的任务 ID
522 /// - `new_delay`: 新的延迟时间，从当前时间重新计算
523 /// - `callback`: 新的回调函数，传递 `None` 保持原始回调，传递 `Some` 替换为新的回调
524 ///
525 /// # 返回值
526 /// 任务存在且延期成功时返回 `Ok(true)`，任务不存在时返回 `Ok(false)`，
527 /// ID 属于其他时间轮时返回 `Err(TimerError::WrongWheel)`。
528 ///
529 /// # Note
530 /// - Task ID remains unchanged after postponement
531 /// - Original completion_receiver remains valid
532 ///
533 /// # 注意
534 /// - 任务 ID 在推迟后保持不变
535 /// - 原始 completion_receiver 保持有效
536 ///
537 /// # Examples (示例)
538 ///
539 /// ## Keep original callback (保持原始回调)
540 /// ```no_run
541 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
542 /// use std::time::Duration;
543 ///
544 ///
545 /// #[tokio::main]
546 /// async fn main() {
547 /// let timer = TimerWheel::with_defaults();
548 ///
549 /// // Allocate handle first
550 /// let allocated_handle = timer.allocate_handle();
551 /// let task_id = allocated_handle.task_id();
552 ///
553 /// let task = TimerTask::new_oneshot(Duration::from_secs(5), Some(CallbackWrapper::new(|| async {
554 /// println!("Timer fired!");
555 /// })));
556 /// let _handle = timer.register(allocated_handle, task).unwrap();
557 ///
558 /// // Postpone to 10 seconds after triggering, and keep original callback
559 /// // 推迟到 10 秒后触发，并保持原始回调
560 /// let success = timer
561 /// .postpone(task_id, Duration::from_secs(10), None)
562 /// .unwrap();
563 /// println!("Postponed successfully: {}", success);
564 /// }
565 /// ```
566 ///
567 /// ## Replace with new callback (替换为新的回调)
568 /// ```no_run
569 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
570 /// use std::time::Duration;
571 ///
572 /// #[tokio::main]
573 /// async fn main() {
574 /// let timer = TimerWheel::with_defaults();
575 ///
576 /// // Allocate handle first
577 /// let allocated_handle = timer.allocate_handle();
578 /// let task_id = allocated_handle.task_id();
579 ///
580 /// let task = TimerTask::new_oneshot(Duration::from_secs(5), Some(CallbackWrapper::new(|| async {
581 /// println!("Original callback!");
582 /// })));
583 /// let _handle = timer.register(allocated_handle, task).unwrap();
584 ///
585 /// // Postpone to 10 seconds after triggering, and replace with new callback
586 /// // 推迟到 10 秒后触发，并替换为新的回调
587 /// let success = timer.postpone(task_id, Duration::from_secs(10), Some(CallbackWrapper::new(|| async {
588 /// println!("New callback!");
589 /// }))).unwrap();
590 /// println!("Postponed successfully: {}", success);
591 /// }
592 /// ```
593 #[inline]
594 pub fn postpone(
595 &self,
596 task_id: TaskId,
597 new_delay: Duration,
598 callback: Option<CallbackWrapper>,
599 ) -> Result<bool, TimerError> {
600 let mut wheel = self.wheel.lock();
601 wheel.postpone_at(task_id, new_delay, callback)
602 }
603
604 /// Batch postpone timers (keep original callbacks)
605 ///
606 /// # Parameters
607 /// - `updates`: List of tuples of (task ID, new delay)
608 ///
609 /// # Returns
610 /// Number of successfully postponed tasks, or
611 /// `Err(TimerError::WrongWheel)` if any ID belongs to another wheel.
612 ///
613 /// 批量推迟定时器 (保持原始回调)
614 ///
615 /// # 参数
616 /// - `updates`: (任务 ID, 新延迟) 元组列表
617 ///
618 /// # 返回值
619 /// 成功推迟的任务数量；如果任一 ID 属于其他时间轮则返回
620 /// `Err(TimerError::WrongWheel)`，且不修改任何任务。
621 ///
622 /// # Note
623 /// - This method keeps all tasks' original callbacks unchanged
624 /// - Use `postpone_batch_with_callbacks` if you need to replace callbacks
625 ///
626 /// # 注意
627 /// - 此方法保持所有任务的原始回调不变
628 /// - 如果需要替换回调，请使用 `postpone_batch_with_callbacks`
629 ///
630 /// # Performance Advantages
631 /// - Batch processing reduces lock contention
632 /// - Internally optimized batch postponement operation
633 ///
634 /// # Examples (示例)
635 /// ```no_run
636 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
637 /// use std::time::Duration;
638 ///
639 /// #[tokio::main]
640 /// async fn main() {
641 /// let timer = TimerWheel::with_defaults();
642 ///
643 /// // Create multiple tasks with callbacks
644 /// // 创建多个带有回调的任务
645 /// let task1 = TimerTask::new_oneshot(Duration::from_secs(5), Some(CallbackWrapper::new(|| async {
646 /// println!("Task 1 fired!");
647 /// })));
648 /// let task2 = TimerTask::new_oneshot(Duration::from_secs(5), Some(CallbackWrapper::new(|| async {
649 /// println!("Task 2 fired!");
650 /// })));
651 /// let task3 = TimerTask::new_oneshot(Duration::from_secs(5), Some(CallbackWrapper::new(|| async {
652 /// println!("Task 3 fired!");
653 /// })));
654 ///
655 /// // Allocate handles and register
656 /// let h1 = timer.allocate_handle();
657 /// let h2 = timer.allocate_handle();
658 /// let h3 = timer.allocate_handle();
659 ///
660 /// let task_ids = vec![
661 /// (h1.task_id(), Duration::from_secs(10)),
662 /// (h2.task_id(), Duration::from_secs(15)),
663 /// (h3.task_id(), Duration::from_secs(20)),
664 /// ];
665 ///
666 /// timer.register(h1, task1).unwrap();
667 /// timer.register(h2, task2).unwrap();
668 /// timer.register(h3, task3).unwrap();
669 ///
670 /// // Batch postpone (keep original callbacks)
671 /// // 批量推迟 (保持原始回调)
672 /// let postponed = timer.postpone_batch(task_ids).unwrap();
673 /// println!("Postponed {} timers", postponed);
674 /// }
675 /// ```
676 #[inline]
677 pub fn postpone_batch(&self, updates: Vec<(TaskId, Duration)>) -> Result<usize, TimerError> {
678 let mut wheel = self.wheel.lock();
679 wheel.postpone_batch_at(updates)
680 }
681
682 /// Batch postpone timers (replace callbacks)
683 ///
684 /// # Parameters
685 /// - `updates`: List of tuples of (task ID, new delay, new callback)
686 ///
687 /// # Returns
688 /// Number of successfully postponed tasks, or
689 /// `Err(TimerError::WrongWheel)` if any ID belongs to another wheel.
690 ///
691 /// 批量推迟定时器 (替换回调)
692 ///
693 /// # 参数
694 /// - `updates`: (任务 ID, 新延迟, 新回调) 元组列表
695 ///
696 /// # 返回值
697 /// 成功推迟的任务数量；如果任一 ID 属于其他时间轮则返回
698 /// `Err(TimerError::WrongWheel)`，且不修改任何任务。
699 ///
700 /// # Performance Advantages
701 /// - Batch processing reduces lock contention
702 /// - Internally optimized batch postponement operation
703 ///
704 /// # Examples (示例)
705 /// ```no_run
706 /// use kestrel_timer::{TimerWheel, TimerTask, CallbackWrapper};
707 /// use std::time::Duration;
708 /// use std::sync::Arc;
709 /// use std::sync::atomic::{AtomicU32, Ordering};
710 ///
711 /// #[tokio::main]
712 /// async fn main() {
713 /// let timer = TimerWheel::with_defaults();
714 /// let counter = Arc::new(AtomicU32::new(0));
715 ///
716 /// // Create multiple timers
717 /// // 创建多个定时器
718 /// let task1 = TimerTask::new_oneshot(Duration::from_secs(5), None);
719 /// let task2 = TimerTask::new_oneshot(Duration::from_secs(5), None);
720 ///
721 /// // Allocate handles first
722 /// let h1 = timer.allocate_handle();
723 /// let h2 = timer.allocate_handle();
724 /// let id1 = h1.task_id();
725 /// let id2 = h2.task_id();
726 ///
727 /// timer.register(h1, task1).unwrap();
728 /// timer.register(h2, task2).unwrap();
729 ///
730 /// // Batch postpone and replace callbacks
731 /// // 批量推迟并替换回调
732 /// let updates: Vec<_> = vec![id1, id2]
733 /// .into_iter()
734 /// .map(|id| {
735 /// let counter = Arc::clone(&counter);
736 /// (id, Duration::from_secs(10), Some(CallbackWrapper::new(move || {
737 /// let counter = Arc::clone(&counter);
738 /// async move { counter.fetch_add(1, Ordering::SeqCst); }
739 /// })))
740 /// })
741 /// .collect();
742 /// let postponed = timer.postpone_batch_with_callbacks(updates).unwrap();
743 /// println!("Postponed {} timers", postponed);
744 /// }
745 /// ```
746 #[inline]
747 pub fn postpone_batch_with_callbacks(
748 &self,
749 updates: Vec<(TaskId, Duration, Option<CallbackWrapper>)>,
750 ) -> Result<usize, TimerError> {
751 let mut wheel = self.wheel.lock();
752 wheel.postpone_batch_with_callbacks_at(updates)
753 }
754
755 /// Core tick loop
756 ///
757 /// Background task that advances the timing wheel at regular intervals
758 ///
759 /// # Parameters
760 /// - `wheel`: Shared timing wheel instance
761 /// - `tick_duration`: Duration between ticks
762 ///
763 /// 核心 tick 循环
764 ///
765 /// 定期推进时间轮的后台任务
766 ///
767 /// # 参数
768 /// - `wheel`: 共享的时间轮实例
769 /// - `tick_duration`: tick 之间的持续时间
770 async fn tick_loop(wheel: Arc<Mutex<Wheel>>, tick_duration: Duration) {
771 let mut interval = tokio::time::interval(tick_duration);
772 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
773
774 loop {
775 interval.tick().await;
776
777 // Advance timing wheel and get expired tasks
778 // Note: wheel.advance_at() already handles completion notifications
779 let expired_tasks = {
780 let mut wheel_guard = wheel.lock();
781 wheel_guard.advance_at(Instant::now())
782 };
783
784 // Execute callbacks for expired tasks
785 // Notifications have already been sent by wheel.advance_at()
786 for task in expired_tasks {
787 if let Some(callback) = task.callback {
788 // Spawn callback execution in a separate tokio task
789 tokio::spawn(async move {
790 let future = callback.call();
791 future.await;
792 });
793 }
794 }
795 }
796 }
797
798 /// Graceful shutdown of TimerWheel
799 ///
800 /// Cancels every task still owned by this wheel, sends `Cancelled` to
801 /// retained completion receivers, and stops the background driver. The
802 /// wheel rejects all later registrations.
803 ///
804 /// 优雅关闭 TimerWheel
805 ///
806 /// 取消此时间轮中仍存在的所有任务，向仍保留的完成接收器发送
807 /// `Cancelled`，并停止后台驱动。关闭后不再接受新的注册。
808 ///
809 /// # Examples (示例)
810 /// ```no_run
811 /// # use kestrel_timer::TimerWheel;
812 /// # #[tokio::main]
813 /// # async fn main() {
814 /// let timer = TimerWheel::with_defaults();
815 ///
816 /// // Use timer... (使用定时器...)
817 ///
818 /// timer.shutdown().await;
819 /// # }
820 /// ```
821 pub async fn shutdown(mut self) {
822 self.wheel.lock().shutdown();
823
824 if let Some(handle) = self.tick_handle.take() {
825 handle.abort();
826 let _ = handle.await;
827 }
828 }
829}
830
831/// Close the wheel and abort the background tick task when TimerWheel is dropped
832///
833/// 当 TimerWheel 被销毁时关闭时间轮并中止后台 tick 任务
834impl Drop for TimerWheel {
835 fn drop(&mut self) {
836 self.wheel.lock().shutdown();
837
838 if let Some(handle) = self.tick_handle.take() {
839 handle.abort();
840 }
841 }
842}
843
844#[cfg(test)]
845mod tests {
846 use super::*;
847 use crate::task::{CompletionReceiver, TaskCompletion, TimerTask};
848 use std::sync::atomic::{AtomicU32, Ordering};
849
850 #[tokio::test]
851 async fn test_timer_creation() {
852 let _timer = TimerWheel::with_defaults();
853 }
854
855 #[tokio::test]
856 async fn test_schedule_once() {
857 use std::sync::Arc;
858 let timer = TimerWheel::with_defaults();
859 let counter = Arc::new(AtomicU32::new(0));
860 let counter_clone = Arc::clone(&counter);
861
862 let task = TimerTask::new_oneshot(
863 Duration::from_millis(50),
864 Some(CallbackWrapper::new(move || {
865 let counter = Arc::clone(&counter_clone);
866 async move {
867 counter.fetch_add(1, Ordering::SeqCst);
868 }
869 })),
870 );
871 let allocate_handle = timer.allocate_handle();
872 let _handle = timer.register(allocate_handle, task).unwrap();
873
874 // Wait for timer to trigger
875 // 等待定时器触发
876 tokio::time::sleep(Duration::from_millis(100)).await;
877 assert_eq!(counter.load(Ordering::SeqCst), 1);
878 }
879
880 #[tokio::test]
881 async fn test_non_integral_delay_does_not_fire_early() {
882 let timer = TimerWheel::with_defaults();
883 let delay = Duration::from_millis(15);
884 let request_time = Instant::now();
885 let handle = timer
886 .register(timer.allocate_handle(), TimerTask::new_oneshot(delay, None))
887 .unwrap();
888 let (receiver, _handle) = handle.into_parts();
889
890 let completion = match receiver {
891 CompletionReceiver::OneShot(receiver) => {
892 tokio::time::timeout(Duration::from_secs(1), receiver.recv())
893 .await
894 .expect("timer should complete")
895 .expect("completion channel should remain open")
896 }
897 CompletionReceiver::Periodic(_) => {
898 panic!("expected one-shot completion receiver")
899 }
900 };
901
902 assert_eq!(completion, TaskCompletion::Called);
903 assert!(
904 request_time.elapsed() >= delay,
905 "timer fired before its requested delay"
906 );
907 }
908
909 #[tokio::test]
910 async fn test_shutdown_notifies_retained_oneshot_receiver() {
911 let timer = TimerWheel::with_defaults();
912 let task = TimerTask::new_oneshot(Duration::from_secs(10), None);
913 let (receiver, _handle) = timer
914 .register(timer.allocate_handle(), task)
915 .unwrap()
916 .into_parts();
917
918 timer.shutdown().await;
919
920 match receiver {
921 CompletionReceiver::OneShot(receiver) => {
922 assert_eq!(receiver.recv().await.unwrap(), TaskCompletion::Cancelled);
923 }
924 CompletionReceiver::Periodic(_) => panic!("expected one-shot receiver"),
925 }
926 }
927
928 #[tokio::test]
929 async fn test_drop_notifies_retained_periodic_receiver() {
930 let (receiver, _handle) = {
931 let timer = TimerWheel::with_defaults();
932 timer
933 .register(
934 timer.allocate_handle(),
935 TimerTask::new_periodic(
936 Duration::from_secs(10),
937 Duration::from_secs(10),
938 None,
939 Some(std::num::NonZeroUsize::new(1).unwrap()),
940 )
941 .unwrap(),
942 )
943 .unwrap()
944 .into_parts()
945 };
946
947 match receiver {
948 CompletionReceiver::Periodic(mut receiver) => {
949 assert_eq!(receiver.recv().await, Some(TaskCompletion::Cancelled));
950 }
951 CompletionReceiver::OneShot(_) => panic!("expected periodic receiver"),
952 }
953 }
954}