mod config;
mod error;
mod task;
mod wheel;
mod timer;
mod service;
pub use config::{
BatchConfig,
HierarchicalWheelConfig,
ServiceConfig, ServiceConfigBuilder,
TimerConfig, TimerConfigBuilder,
WheelConfig, WheelConfigBuilder,
};
pub use error::TimerError;
pub use task::{CallbackWrapper, CompletionNotifier, TaskCompletionReason, TaskId, TimerCallback, TimerTask};
pub use timer::{BatchHandle, BatchHandleIter, CompletionReceiver, TimerHandle, TimerWheel};
pub use service::TimerService;
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[tokio::test]
async fn test_basic_timer() {
let timer = TimerWheel::with_defaults();
let counter = Arc::new(AtomicU32::new(0));
let counter_clone = Arc::clone(&counter);
let task = TimerWheel::create_task(
Duration::from_millis(50),
Some(CallbackWrapper::new(move || {
let counter = Arc::clone(&counter_clone);
async move {
counter.fetch_add(1, Ordering::SeqCst);
}
})),
);
timer.register(task);
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_multiple_timers() {
let timer = TimerWheel::with_defaults();
let counter = Arc::new(AtomicU32::new(0));
for i in 0..10 {
let counter_clone = Arc::clone(&counter);
let task = TimerWheel::create_task(
Duration::from_millis(10 * (i + 1)),
Some(CallbackWrapper::new(move || {
let counter = Arc::clone(&counter_clone);
async move {
counter.fetch_add(1, Ordering::SeqCst);
}
})),
);
timer.register(task);
}
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(counter.load(Ordering::SeqCst), 10);
}
#[tokio::test]
async fn test_timer_cancellation() {
let timer = TimerWheel::with_defaults();
let counter = Arc::new(AtomicU32::new(0));
let mut handles = Vec::new();
for _ in 0..5 {
let counter_clone = Arc::clone(&counter);
let task = TimerWheel::create_task(
Duration::from_millis(100),
Some(CallbackWrapper::new(move || {
let counter = Arc::clone(&counter_clone);
async move {
counter.fetch_add(1, Ordering::SeqCst);
}
})),
);
let handle = timer.register(task);
handles.push(handle);
}
for i in 0..3 {
let cancel_result = handles[i].cancel();
assert!(cancel_result);
}
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(counter.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_completion_notification_once() {
let timer = TimerWheel::with_defaults();
let counter = Arc::new(AtomicU32::new(0));
let counter_clone = Arc::clone(&counter);
let task = TimerWheel::create_task(
Duration::from_millis(50),
Some(CallbackWrapper::new(move || {
let counter = Arc::clone(&counter_clone);
async move {
counter.fetch_add(1, Ordering::SeqCst);
}
})),
);
let handle = timer.register(task);
handle.into_completion_receiver().0.await.expect("Should receive completion notification");
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_notify_only_timer_once() {
let timer = TimerWheel::with_defaults();
let task = TimerTask::new(Duration::from_millis(50), None);
let handle = timer.register(task);
handle.into_completion_receiver().0.await.expect("Should receive completion notification");
}
#[tokio::test]
async fn test_batch_completion_notifications() {
let timer = TimerWheel::with_defaults();
let counter = Arc::new(AtomicU32::new(0));
let callbacks: Vec<(Duration, Option<CallbackWrapper>)> = (0..5)
.map(|i| {
let counter = Arc::clone(&counter);
let delay = Duration::from_millis(50 + i * 10);
let callback = CallbackWrapper::new(move || {
let counter = Arc::clone(&counter);
async move {
counter.fetch_add(1, Ordering::SeqCst);
}
});
(delay, Some(callback))
})
.collect();
let tasks = TimerWheel::create_batch_with_callbacks(callbacks);
let batch = timer.register_batch(tasks);
let receivers = batch.into_completion_receivers();
for rx in receivers {
rx.await.expect("Should receive completion notification");
}
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(counter.load(Ordering::SeqCst), 5);
}
#[tokio::test]
async fn test_completion_reason_expired() {
let timer = TimerWheel::with_defaults();
let task = TimerTask::new(Duration::from_millis(50), None);
let handle = timer.register(task);
let result = handle.into_completion_receiver().0.await.expect("Should receive completion notification");
assert_eq!(result, TaskCompletionReason::Expired);
}
#[tokio::test]
async fn test_completion_reason_cancelled() {
let timer = TimerWheel::with_defaults();
let task = TimerTask::new(Duration::from_secs(10), None);
let handle = timer.register(task);
let cancelled = handle.cancel();
assert!(cancelled);
let result = handle.into_completion_receiver().0.await.expect("Should receive completion notification");
assert_eq!(result, TaskCompletionReason::Cancelled);
}
#[tokio::test]
async fn test_batch_completion_reasons() {
let timer = TimerWheel::with_defaults();
let tasks: Vec<_> = (0..5)
.map(|_| TimerTask::new(Duration::from_secs(10), None))
.collect();
let batch = timer.register_batch(tasks);
let task_ids: Vec<_> = batch.task_ids.clone();
let mut receivers = batch.into_completion_receivers();
timer.cancel_batch(&task_ids[0..3]);
for rx in receivers.drain(0..3) {
let result = rx.await.expect("Should receive completion notification");
assert_eq!(result, TaskCompletionReason::Cancelled);
}
timer.cancel_batch(&task_ids[3..5]);
for rx in receivers {
let result = rx.await.expect("Should receive completion notification");
assert_eq!(result, TaskCompletionReason::Cancelled);
}
}
}