use crate::async_engine::AsyncEngine;
use crate::async_engine::SpawnedFuture;
use crate::io_uring::IoUring;
use crate::io_uring::IoUringError;
use crate::time::Time;
use crate::utils::buf::TypedBuf;
use crate::utils::copyhashmap::CopyHashMap;
use crate::utils::errorfmt::ErrorFmt;
use crate::utils::hash_map_ext::HashMapExt;
use crate::utils::numcell::NumCell;
use crate::utils::stack::Stack;
use jay_algorithms::oserror::OsError;
use jay_algorithms::oserror::OsErrorExt;
use jay_algorithms::oserror::OsErrorExt2;
use std::cell::Cell;
use std::cell::RefCell;
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::pin::Pin;
use std::rc::Rc;
use std::task::Context;
use std::task::Poll;
use std::task::Waker;
use std::time::Duration;
use thiserror::Error;
use uapi::OwnedFd;
use uapi::c;
#[derive(Debug, Error)]
pub enum WheelError {
#[error("Could not create the timerfd")]
CreateFailed(#[source] OsError),
#[error("Could not set the timerfd")]
SetFailed(#[source] OsError),
#[error("The timer wheel is already destroyed")]
Destroyed,
#[error("Could not read from the timerfd")]
Read(#[source] IoUringError),
}
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
struct WheelEntry {
expiration: Time,
id: u64,
}
pub struct Wheel {
data: Rc<WheelData>,
}
impl Drop for Wheel {
fn drop(&mut self) {
self.data.kill();
}
}
struct WheelTimeoutData {
id: Cell<u64>,
expired: Cell<Option<Result<(), WheelError>>>,
wheel: Rc<WheelData>,
waker: Cell<Option<Waker>>,
}
impl WheelTimeoutData {
fn complete(&self, res: Result<(), WheelError>) {
self.expired.set(Some(res));
if let Some(waker) = self.waker.take() {
waker.wake();
}
}
}
pub struct WheelTimeoutFuture {
data: Rc<WheelTimeoutData>,
}
impl Drop for WheelTimeoutFuture {
fn drop(&mut self) {
self.data.wheel.dispatchers.remove(&self.data.id.get());
self.data.waker.set(None);
if !self.data.wheel.destroyed.get() {
self.data.expired.take();
self.data.wheel.cached_futures.push(self.data.clone());
}
}
}
impl Future for WheelTimeoutFuture {
type Output = Result<(), WheelError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if let Some(res) = self.data.expired.take() {
Poll::Ready(res)
} else {
self.data.waker.set(Some(cx.waker().clone()));
Poll::Pending
}
}
}
pub struct WheelData {
destroyed: Cell<bool>,
ring: Rc<IoUring>,
eng: Rc<AsyncEngine>,
fd: Rc<OwnedFd>,
next_id: NumCell<u64>,
start: Time,
current_expiration: Cell<Option<Time>>,
dispatchers: CopyHashMap<u64, Rc<WheelTimeoutData>>,
expirations: RefCell<BinaryHeap<Reverse<WheelEntry>>>,
dispatcher: Cell<Option<SpawnedFuture<()>>>,
cached_futures: Stack<Rc<WheelTimeoutData>>,
}
impl Wheel {
pub fn new(eng: &Rc<AsyncEngine>, ring: &Rc<IoUring>) -> Result<Rc<Self>, WheelError> {
let fd = uapi::timerfd_create(c::CLOCK_MONOTONIC, c::TFD_CLOEXEC)
.map(Rc::new)
.map_os_err(WheelError::CreateFailed)?;
let data = Rc::new(WheelData {
destroyed: Cell::new(false),
ring: ring.clone(),
eng: eng.clone(),
fd,
next_id: NumCell::new(1),
start: eng.now(),
current_expiration: Default::default(),
dispatchers: Default::default(),
expirations: Default::default(),
dispatcher: Default::default(),
cached_futures: Default::default(),
});
data.dispatcher
.set(Some(eng.spawn("wheel", data.clone().dispatch())));
Ok(Rc::new(Wheel { data }))
}
pub fn clear(&self) {
self.data.kill();
}
fn future(&self) -> WheelTimeoutFuture {
let data = self.data.cached_futures.pop().unwrap_or_else(|| {
Rc::new(WheelTimeoutData {
id: Cell::new(0),
expired: Cell::new(None),
wheel: self.data.clone(),
waker: Cell::new(None),
})
});
data.id.set(self.data.next_id.fetch_add(1));
WheelTimeoutFuture { data }
}
pub fn timeout(&self, ms: u64) -> WheelTimeoutFuture {
if self.data.destroyed.get() {
return WheelTimeoutFuture {
data: Rc::new(WheelTimeoutData {
id: Cell::new(0),
expired: Cell::new(Some(Err(WheelError::Destroyed))),
wheel: self.data.clone(),
waker: Default::default(),
}),
};
}
let future = self.future();
let now = self.data.eng.now();
let expiration = (now + Duration::from_millis(ms)).round_to_ms();
let current = self.data.current_expiration.get();
if current.is_none() || expiration - self.data.start < current.unwrap() - self.data.start {
let res = uapi::timerfd_settime(
self.data.fd.raw(),
c::TFD_TIMER_ABSTIME,
&c::itimerspec {
it_interval: uapi::pod_zeroed(),
it_value: expiration.0,
},
);
if let Err(e) = res.to_os_error() {
future.data.expired.set(Some(Err(WheelError::SetFailed(e))));
return future;
}
self.data.current_expiration.set(Some(expiration));
}
self.data.expirations.borrow_mut().push(Reverse(WheelEntry {
expiration,
id: future.data.id.get(),
}));
self.data
.dispatchers
.set(future.data.id.get(), future.data.clone());
future
}
}
impl WheelData {
fn kill(&self) {
self.destroyed.set(true);
self.dispatcher.set(None);
self.cached_futures.take();
for dispatcher in self.dispatchers.lock().drain_values() {
dispatcher.complete(Err(WheelError::Destroyed));
}
}
async fn dispatch(self: Rc<Self>) {
let mut n = TypedBuf::new();
loop {
if let Err(e) = self.dispatch_once(&mut n).await {
log::error!("Could not dispatch wheel expirations: {}", ErrorFmt(e));
self.kill();
return;
}
}
}
async fn dispatch_once(&self, n: &mut TypedBuf<u64>) -> Result<(), WheelError> {
if let Err(e) = self.ring.read(&self.fd, n.buf()).await {
return Err(WheelError::Read(e));
}
let now = self.eng.now();
let dist = now - self.start;
{
let mut expirations = self.expirations.borrow_mut();
while let Some(Reverse(entry)) = expirations.peek() {
if entry.expiration - self.start > dist {
break;
}
if let Some(dispatcher) = self.dispatchers.remove(&entry.id) {
dispatcher.complete(Ok(()));
}
expirations.pop();
}
self.current_expiration.set(None);
while let Some(Reverse(entry)) = expirations.peek() {
if self.dispatchers.get(&entry.id).is_some() {
uapi::timerfd_settime(
self.fd.raw(),
c::TFD_TIMER_ABSTIME,
&c::itimerspec {
it_interval: uapi::pod_zeroed(),
it_value: entry.expiration.0,
},
)
.map_os_err(WheelError::SetFailed)?;
self.current_expiration.set(Some(entry.expiration));
break;
}
expirations.pop();
}
}
Ok(())
}
}