use crate::lockfree::{AtomicWakerSlot, MpmcStack};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicI32, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::task::{Context, Poll};
use std::thread::Thread;
use std::time::{Duration, Instant};
const WHEEL_BITS: u32 = 8;
const WHEEL_SIZE: usize = 1 << WHEEL_BITS; const WHEEL_MASK: u64 = (WHEEL_SIZE as u64) - 1;
const TICK: Duration = Duration::from_millis(1);
const PENDING: i32 = 0;
const DONE: i32 = 1;
#[repr(align(64))]
struct SleepState {
status: AtomicI32,
waker: AtomicWakerSlot,
}
#[repr(align(64))]
struct Node {
state: Arc<SleepState>,
rounds: AtomicUsize,
}
#[repr(align(64))]
struct Wheel {
buckets: Box<[MpmcStack<Arc<Node>>]>,
current_tick: AtomicU64,
pending: AtomicUsize,
start: Instant,
worker: OnceLock<Thread>,
}
static WHEEL: OnceLock<Arc<Wheel>> = OnceLock::new();
fn wheel() -> &'static Arc<Wheel> {
WHEEL.get_or_init(|| {
let mut buckets = Vec::with_capacity(WHEEL_SIZE);
for _ in 0..WHEEL_SIZE {
buckets.push(MpmcStack::new());
}
let w = Arc::new(Wheel {
buckets: buckets.into_boxed_slice(),
current_tick: AtomicU64::new(0),
pending: AtomicUsize::new(0),
start: Instant::now(),
worker: OnceLock::new(),
});
let worker_wheel = Arc::clone(&w);
let handle = std::thread::Builder::new()
.name("dtact-timer-wheel".into())
.spawn(move || worker_loop(&worker_wheel))
.expect("failed to spawn dtact-timer-wheel worker thread");
let _ = w.worker.set(handle.thread().clone());
w
})
}
fn worker_loop(w: &Arc<Wheel>) {
loop {
if w.pending.load(Ordering::Acquire) == 0 {
std::thread::park_timeout(Duration::from_millis(50));
continue;
}
let tick = w.current_tick.load(Ordering::Relaxed);
let tick_deadline = w.start + TICK * (tick as u32 + 1);
let now = Instant::now();
if tick_deadline > now {
std::thread::park_timeout(tick_deadline - now);
if Instant::now() < tick_deadline {
continue;
}
}
let tick = w.current_tick.fetch_add(1, Ordering::Relaxed) + 1;
let slot = (tick & WHEEL_MASK) as usize;
let bucket = &w.buckets[slot];
if bucket.is_empty() {
continue;
}
let batch = bucket.drain_all();
let mut fired = 0usize;
for node in batch {
let rounds = node.rounds.load(Ordering::Relaxed);
if rounds == 0 {
node.state.status.store(DONE, Ordering::Release);
node.state.waker.take_and_wake();
fired += 1;
} else {
node.rounds.fetch_sub(1, Ordering::Relaxed);
bucket.push(node);
}
}
if fired > 0 {
w.pending.fetch_sub(fired, Ordering::AcqRel);
}
}
}
fn register(deadline: Instant) -> Arc<SleepState> {
let now = Instant::now();
if deadline <= now {
return Arc::new(SleepState {
status: AtomicI32::new(DONE),
waker: AtomicWakerSlot::new(),
});
}
let w = wheel();
let state = Arc::new(SleepState {
status: AtomicI32::new(PENDING),
waker: AtomicWakerSlot::new(),
});
let remaining = deadline - now;
let ticks_from_now = remaining.as_nanos().div_ceil(TICK.as_nanos()) as u64;
let current_tick = w.current_tick.load(Ordering::Relaxed);
let target_tick = current_tick + ticks_from_now;
let slot = (target_tick & WHEEL_MASK) as usize;
let rounds = (ticks_from_now >> WHEEL_BITS) as usize;
let node = Arc::new(Node {
state: Arc::clone(&state),
rounds: AtomicUsize::new(rounds),
});
w.buckets[slot].push(node);
w.pending.fetch_add(1, Ordering::AcqRel);
if let Some(t) = w.worker.get() {
t.unpark();
}
state
}
pub struct DtactSleep {
state: Arc<SleepState>,
}
impl DtactSleep {
#[must_use]
pub fn new(duration: Duration) -> Self {
Self::until(Instant::now() + duration)
}
#[must_use]
pub fn until(deadline: Instant) -> Self {
Self {
state: register(deadline),
}
}
}
impl Future for DtactSleep {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.state.status.load(Ordering::Acquire) == DONE {
return Poll::Ready(());
}
self.state.waker.register(cx.waker());
if self.state.status.load(Ordering::Acquire) == DONE {
return Poll::Ready(());
}
Poll::Pending
}
}
#[must_use]
pub fn sleep(duration: Duration) -> DtactSleep {
DtactSleep::new(duration)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum MissedTickBehavior {
#[default]
Burst,
Skip,
Delay,
}
pub struct DtactInterval {
period: Duration,
next: Instant,
missed_tick_behavior: MissedTickBehavior,
}
impl DtactInterval {
#[must_use]
pub fn new(period: Duration) -> Self {
assert!(
period > Duration::ZERO,
"dtact-timer: interval period must be > 0"
);
Self {
period,
next: Instant::now() + period,
missed_tick_behavior: MissedTickBehavior::default(),
}
}
#[must_use]
pub const fn missed_tick_behavior(&self) -> MissedTickBehavior {
self.missed_tick_behavior
}
pub const fn set_missed_tick_behavior(&mut self, behavior: MissedTickBehavior) {
self.missed_tick_behavior = behavior;
}
pub async fn tick(&mut self) -> Instant {
DtactSleep::until(self.next).await;
let fired_at = self.next;
match self.missed_tick_behavior {
MissedTickBehavior::Burst => {
self.next += self.period;
}
MissedTickBehavior::Skip => {
let now = Instant::now();
let mut next = self.next + self.period;
while next <= now {
next += self.period;
}
self.next = next;
}
MissedTickBehavior::Delay => {
self.next = Instant::now() + self.period;
}
}
fired_at
}
}
#[must_use]
pub fn interval(period: Duration) -> DtactInterval {
DtactInterval::new(period)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TimeoutError;
impl std::fmt::Display for TimeoutError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"dtact-timer: deadline elapsed before the future completed"
)
}
}
impl std::error::Error for TimeoutError {}
pub struct DtactTimeout<F> {
inner: Pin<Box<F>>,
sleep: DtactSleep,
}
impl<F> DtactTimeout<F> {
pub fn new(duration: Duration, inner: F) -> Self {
Self {
inner: Box::pin(inner),
sleep: DtactSleep::new(duration),
}
}
}
impl<F: Future> Future for DtactTimeout<F> {
type Output = Result<F::Output, TimeoutError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
if let Poll::Ready(v) = this.inner.as_mut().poll(cx) {
return Poll::Ready(Ok(v));
}
match Pin::new(&mut this.sleep).poll(cx) {
Poll::Ready(()) => Poll::Ready(Err(TimeoutError)),
Poll::Pending => Poll::Pending,
}
}
}
pub fn timeout<F: Future>(duration: Duration, fut: F) -> DtactTimeout<F> {
DtactTimeout::new(duration, fut)
}