use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::task::Waker;
use std::time::{Duration, Instant};
const TICK_MS: u64 = 1;
const WHEEL_SIZE: usize = 256;
type TimeoutCallback = Box<dyn FnOnce() + Send + 'static>;
struct TimerEntry {
expiration: Instant,
callback: TimeoutCallback,
waker: Option<Waker>,
}
struct TimerWheel {
slots: Vec<VecDeque<TimerEntry>>,
tick_duration: Duration,
current_tick: u64,
}
impl TimerWheel {
fn new(tick_duration: Duration) -> Self {
Self {
slots: (0..WHEEL_SIZE).map(|_| VecDeque::new()).collect(),
tick_duration,
current_tick: 0,
}
}
fn add(&mut self, entry: TimerEntry, now: Instant) -> Option<TimerEntry> {
let delay = entry.expiration.saturating_duration_since(now);
let ticks = (delay.as_millis() as u64) / self.tick_duration.as_millis() as u64;
if ticks >= WHEEL_SIZE as u64 {
return Some(entry);
}
let slot = ((self.current_tick + ticks) % WHEEL_SIZE as u64) as usize;
self.slots[slot].push_back(entry);
None
}
fn tick(&mut self, now: Instant) -> Vec<TimerEntry> {
let slot_idx = (self.current_tick % WHEEL_SIZE as u64) as usize;
self.current_tick += 1;
let mut expired = Vec::new();
while let Some(entry) = self.slots[slot_idx].pop_front() {
if entry.expiration <= now {
expired.push(entry);
} else {
self.slots[slot_idx].push_back(entry);
}
}
expired
}
}
pub struct TimerWheelScheduler {
wheel_l0: TimerWheel,
wheel_l1: TimerWheel,
wheel_l2: TimerWheel,
start_time: Instant,
pending: Vec<TimerEntry>,
}
impl TimerWheelScheduler {
pub fn new() -> Self {
Self {
wheel_l0: TimerWheel::new(Duration::from_millis(TICK_MS)),
wheel_l1: TimerWheel::new(Duration::from_millis(TICK_MS * WHEEL_SIZE as u64)),
wheel_l2: TimerWheel::new(Duration::from_millis(TICK_MS * WHEEL_SIZE as u64 * WHEEL_SIZE as u64)),
start_time: Instant::now(),
pending: Vec::new(),
}
}
pub fn schedule<F>(&mut self, delay: Duration, callback: F)
where
F: FnOnce() + Send + 'static,
{
let expiration = Instant::now() + delay;
let entry = TimerEntry {
expiration,
callback: Box::new(callback),
waker: None,
};
self.add_entry(entry);
}
pub fn schedule_with_waker<F>(&mut self, delay: Duration, waker: Waker, callback: F)
where
F: FnOnce() + Send + 'static,
{
let expiration = Instant::now() + delay;
let entry = TimerEntry {
expiration,
callback: Box::new(callback),
waker: Some(waker),
};
self.add_entry(entry);
}
fn add_entry(&mut self, mut entry: TimerEntry) {
let now = Instant::now();
entry = match self.wheel_l0.add(entry, now) {
Some(e) => e,
None => return,
};
entry = match self.wheel_l1.add(entry, now) {
Some(e) => e,
None => return,
};
entry = match self.wheel_l2.add(entry, now) {
Some(e) => e,
None => return,
};
self.pending.push(entry);
}
pub fn tick(&mut self) -> usize {
let now = Instant::now();
let mut expired_count = 0;
let mut expired = self.wheel_l0.tick(now);
expired_count += expired.len();
for entry in expired.drain(..) {
if let Some(waker) = entry.waker {
waker.wake();
}
(entry.callback)();
}
if self.wheel_l0.current_tick % WHEEL_SIZE as u64 == 0 {
let l1_expired = self.wheel_l1.tick(now);
for entry in l1_expired {
if let Some(entry) = self.wheel_l0.add(entry, now) {
let _ = self.wheel_l1.add(entry, now);
}
}
}
if self.wheel_l0.current_tick % (WHEEL_SIZE as u64 * WHEEL_SIZE as u64) == 0 {
let l2_expired = self.wheel_l2.tick(now);
for entry in l2_expired {
let entry = match self.wheel_l1.add(entry, now) {
Some(e) => e,
None => continue,
};
let entry = match self.wheel_l2.add(entry, now) {
Some(e) => e,
None => continue,
};
self.pending.push(entry);
}
}
let pending: Vec<_> = self.pending.drain(..).collect();
let mut still_pending = Vec::new();
for entry in pending {
if let Some(entry) = self.add_from_pending(entry, now) {
still_pending.push(entry);
}
}
self.pending = still_pending;
expired_count
}
fn add_from_pending(&mut self, mut entry: TimerEntry, now: Instant) -> Option<TimerEntry> {
entry = match self.wheel_l0.add(entry, now) {
Some(e) => e,
None => return None,
};
entry = match self.wheel_l1.add(entry, now) {
Some(e) => e,
None => return None,
};
entry = match self.wheel_l2.add(entry, now) {
Some(e) => e,
None => return None,
};
Some(entry)
}
pub fn time_until_next(&self) -> Option<Duration> {
Some(Duration::from_millis(TICK_MS))
}
}
impl Default for TimerWheelScheduler {
fn default() -> Self {
Self::new()
}
}
pub struct GlobalTimerWheel {
inner: Arc<Mutex<TimerWheelScheduler>>,
}
impl GlobalTimerWheel {
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(TimerWheelScheduler::new())),
}
}
pub fn schedule<F>(&self, delay: Duration, callback: F)
where
F: FnOnce() + Send + 'static,
{
let mut wheel = self.inner.lock().unwrap();
wheel.schedule(delay, callback);
}
pub fn tick(&self) -> usize {
let mut wheel = self.inner.lock().unwrap();
wheel.tick()
}
}
impl Default for GlobalTimerWheel {
fn default() -> Self {
Self::new()
}
}
pub struct Sleep {
deadline: Instant,
registered: bool,
}
impl Sleep {
pub fn new(duration: Duration) -> Self {
Self {
deadline: Instant::now() + duration,
registered: false,
}
}
}
impl std::future::Future for Sleep {
type Output = ();
fn poll(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
if Instant::now() >= self.deadline {
return std::task::Poll::Ready(());
}
if !self.registered {
self.registered = true;
}
std::task::Poll::Pending
}
}
pub fn sleep(duration: Duration) -> Sleep {
Sleep::new(duration)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_timer_wheel_creation() {
let wheel = TimerWheelScheduler::new();
assert_eq!(wheel.wheel_l0.slots.len(), WHEEL_SIZE);
}
#[test]
fn test_schedule_immediate() {
let mut wheel = TimerWheelScheduler::new();
let called = Arc::new(Mutex::new(false));
let called_clone = Arc::clone(&called);
wheel.schedule(Duration::from_millis(1), move || {
*called_clone.lock().unwrap() = true;
});
std::thread::sleep(Duration::from_millis(5));
wheel.tick();
assert!(*called.lock().unwrap());
}
#[test]
fn test_timer_wheel_tick() {
let mut wheel = TimerWheelScheduler::new();
let count = Arc::new(Mutex::new(0));
let count_clone = Arc::clone(&count);
for i in 0..10 {
let c = Arc::clone(&count);
wheel.schedule(Duration::from_millis(i * 10), move || {
*c.lock().unwrap() += 1;
});
}
for _ in 0..200 {
wheel.tick();
std::thread::sleep(Duration::from_millis(1));
}
assert!(*count.lock().unwrap() > 0);
}
#[test]
fn test_global_timer_wheel() {
let wheel = GlobalTimerWheel::new();
let called = Arc::new(Mutex::new(false));
let called_clone = Arc::clone(&called);
wheel.schedule(Duration::from_millis(5), move || {
*called_clone.lock().unwrap() = true;
});
std::thread::sleep(Duration::from_millis(10));
for _ in 0..20 {
wheel.tick();
std::thread::sleep(Duration::from_millis(1));
}
assert!(*called.lock().unwrap());
}
#[test]
fn test_sleep_future() {
let sleep = sleep(Duration::from_millis(10));
}
}