use crate::error::TimerError;
use std::collections::VecDeque;
use std::future::Future;
use std::num::{NonZeroU64, NonZeroUsize};
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use tokio::time::Instant;
use deferred_map::{DefaultKey, Key};
use lite_sync::oneshot::lite::{Receiver, Sender, State, channel};
use parking_lot::Mutex;
use tokio::sync::watch;
use tokio::sync::{Notify, mpsc};
const ONESHOT_PENDING: u8 = 0;
const ONESHOT_CALLED: u8 = 1;
const ONESHOT_CANCELLED: u8 = 2;
const ONESHOT_CLOSED: u8 = 3;
const ONESHOT_RECEIVER_CLOSED: u8 = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskCompletion {
Called,
Cancelled,
}
impl State for TaskCompletion {
#[inline]
fn to_u8(&self) -> u8 {
match self {
TaskCompletion::Called => ONESHOT_CALLED,
TaskCompletion::Cancelled => ONESHOT_CANCELLED,
}
}
#[inline]
fn from_u8(value: u8) -> Option<Self> {
match value {
ONESHOT_CALLED => Some(TaskCompletion::Called),
ONESHOT_CANCELLED => Some(TaskCompletion::Cancelled),
_ => None,
}
}
#[inline]
fn pending_value() -> u8 {
ONESHOT_PENDING
}
#[inline]
fn closed_value() -> u8 {
ONESHOT_CLOSED
}
#[inline]
fn receiver_closed_value() -> u8 {
ONESHOT_RECEIVER_CLOSED
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct WheelId(NonZeroU64);
static NEXT_WHEEL_ID: AtomicU64 = AtomicU64::new(1);
impl WheelId {
#[inline]
pub(crate) fn new() -> Self {
let value = NEXT_WHEEL_ID
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
current.checked_add(1)
})
.expect("timing wheel owner ID space exhausted");
Self(NonZeroU64::new(value).expect("timing wheel owner ID must be non-zero"))
}
}
pub(crate) struct WheelOwner {
id: WheelId,
released_handles: Mutex<Vec<deferred_map::Handle>>,
closed: AtomicBool,
close_signal: watch::Sender<bool>,
}
impl WheelOwner {
#[inline]
pub(crate) fn new() -> Arc<Self> {
let (close_signal, _) = watch::channel(false);
Arc::new(Self {
id: WheelId::new(),
released_handles: Mutex::new(Vec::new()),
closed: AtomicBool::new(false),
close_signal,
})
}
#[inline]
pub(crate) fn id(&self) -> WheelId {
self.id
}
#[inline]
fn release_handle(&self, handle: deferred_map::Handle) {
self.released_handles.lock().push(handle);
}
#[inline]
pub(crate) fn take_released_handles(&self) -> Vec<deferred_map::Handle> {
std::mem::take(&mut *self.released_handles.lock())
}
#[inline]
pub(crate) fn close(&self) {
self.closed.store(true, Ordering::Release);
let _ = self.close_signal.send(true);
}
#[inline]
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
#[inline]
pub(crate) fn subscribe_closed(&self) -> watch::Receiver<bool> {
self.close_signal.subscribe()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TaskId {
wheel_id: WheelId,
key: DefaultKey,
}
impl TaskId {
#[inline]
pub(crate) fn from_key(wheel_id: WheelId, key: DefaultKey) -> Self {
Self { wheel_id, key }
}
#[inline]
pub(crate) fn key(&self) -> DefaultKey {
self.key
}
#[inline]
pub(crate) fn belongs_to(&self, wheel_id: WheelId) -> bool {
self.wheel_id == wheel_id
}
#[cfg(test)]
#[inline]
pub(crate) fn wheel_id(&self) -> WheelId {
self.wheel_id
}
#[inline]
pub fn raw(&self) -> u64 {
self.key.raw()
}
}
pub struct TaskHandle {
owner: Arc<WheelOwner>,
handle: Option<deferred_map::Handle>,
}
impl TaskHandle {
#[inline]
pub(crate) fn new(owner: Arc<WheelOwner>, handle: deferred_map::Handle) -> Self {
Self {
owner,
handle: Some(handle),
}
}
#[inline]
pub fn task_id(&self) -> TaskId {
let handle = self
.handle
.as_ref()
.expect("task handle has already been consumed");
TaskId::from_key(self.owner.id(), handle.key())
}
#[inline]
pub(crate) fn belongs_to(&self, wheel_id: WheelId) -> bool {
self.owner.id() == wheel_id
}
#[inline]
pub(crate) fn into_handle(mut self) -> deferred_map::Handle {
self.handle
.take()
.expect("task handle has already been consumed")
}
}
impl Drop for TaskHandle {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
self.owner.release_handle(handle);
}
}
}
pub trait TimerCallback: Send + Sync + 'static {
fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>>;
}
impl<F, Fut> TimerCallback for F
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
Box::pin(self())
}
}
#[derive(Clone)]
pub struct CallbackWrapper {
callback: Arc<dyn TimerCallback>,
}
impl CallbackWrapper {
#[inline]
pub fn new(callback: impl TimerCallback) -> Self {
Self {
callback: Arc::new(callback),
}
}
#[inline]
pub(crate) fn call(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
self.callback.call()
}
}
#[derive(Clone)]
pub enum TaskType {
OneShot,
Periodic {
interval: std::time::Duration,
buffer_size: Option<NonZeroUsize>,
},
}
pub enum TaskTypeWithCompletionNotifier {
OneShot {
completion_notifier: Sender<TaskCompletion>,
},
Periodic {
interval: std::time::Duration,
completion_notifier: PeriodicCompletionNotifier,
},
}
impl TaskTypeWithCompletionNotifier {
#[inline]
pub fn get_interval(&self) -> Option<std::time::Duration> {
match self {
TaskTypeWithCompletionNotifier::Periodic { interval, .. } => Some(*interval),
TaskTypeWithCompletionNotifier::OneShot { .. } => None,
}
}
}
pub(crate) struct PeriodicCompletionQueue {
state: Mutex<PeriodicCompletionQueueState>,
notify: Notify,
}
struct PeriodicCompletionQueueState {
events: VecDeque<TaskCompletion>,
capacity: Option<NonZeroUsize>,
terminal_pending: bool,
closed: bool,
receiver_alive: bool,
dropped_called: usize,
}
enum PeriodicReceiveState {
Empty,
Closed,
}
impl PeriodicCompletionQueue {
fn new(capacity: Option<NonZeroUsize>) -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(PeriodicCompletionQueueState {
events: VecDeque::new(),
capacity,
terminal_pending: false,
closed: false,
receiver_alive: true,
dropped_called: 0,
}),
notify: Notify::new(),
})
}
fn send(&self, completion: TaskCompletion) -> bool {
let mut state = self.state.lock();
if state.closed || !state.receiver_alive {
return false;
}
if completion == TaskCompletion::Cancelled {
state.terminal_pending = true;
drop(state);
self.notify.notify_waiters();
return true;
}
if state.terminal_pending {
return false;
}
if let Some(capacity) = state.capacity
&& state.events.len() >= capacity.get()
{
state.dropped_called = state.dropped_called.saturating_add(1);
return false;
}
state.events.push_back(completion);
drop(state);
self.notify.notify_one();
true
}
fn try_receive(&self) -> Result<TaskCompletion, PeriodicReceiveState> {
let mut state = self.state.lock();
if let Some(completion) = state.events.pop_front() {
return Ok(completion);
}
if state.terminal_pending {
state.terminal_pending = false;
return Ok(TaskCompletion::Cancelled);
}
if state.closed {
Err(PeriodicReceiveState::Closed)
} else {
Err(PeriodicReceiveState::Empty)
}
}
fn close_sender(&self) {
let mut state = self.state.lock();
state.closed = true;
drop(state);
self.notify.notify_waiters();
}
fn close_receiver(&self) {
let mut state = self.state.lock();
state.receiver_alive = false;
state.closed = true;
state.events.clear();
state.terminal_pending = false;
drop(state);
self.notify.notify_waiters();
}
fn dropped_called(&self) -> usize {
self.state.lock().dropped_called
}
}
pub struct PeriodicCompletionNotifier(pub(crate) Arc<PeriodicCompletionQueue>);
impl PeriodicCompletionNotifier {
#[inline]
pub(crate) fn send(&self, completion: TaskCompletion) -> bool {
self.0.send(completion)
}
}
impl Drop for PeriodicCompletionNotifier {
fn drop(&mut self) {
self.0.close_sender();
}
}
pub struct PeriodicCompletionReceiver(pub(crate) Arc<PeriodicCompletionQueue>);
impl PeriodicCompletionReceiver {
#[inline]
pub fn try_recv(&mut self) -> Result<TaskCompletion, mpsc::error::TryRecvError> {
match self.0.try_receive() {
Ok(completion) => Ok(completion),
Err(PeriodicReceiveState::Empty) => Err(mpsc::error::TryRecvError::Empty),
Err(PeriodicReceiveState::Closed) => Err(mpsc::error::TryRecvError::Disconnected),
}
}
#[inline]
pub async fn recv(&mut self) -> Option<TaskCompletion> {
loop {
let notified = self.0.notify.notified();
match self.0.try_receive() {
Ok(completion) => return Some(completion),
Err(PeriodicReceiveState::Closed) => return None,
Err(PeriodicReceiveState::Empty) => notified.await,
}
}
}
#[inline]
pub fn dropped_notifications(&self) -> usize {
self.0.dropped_called()
}
}
impl Drop for PeriodicCompletionReceiver {
fn drop(&mut self) {
self.0.close_receiver();
}
}
pub enum CompletionNotifier {
OneShot(Sender<TaskCompletion>),
Periodic(PeriodicCompletionNotifier),
}
pub enum CompletionReceiver {
OneShot(Receiver<TaskCompletion>),
Periodic(PeriodicCompletionReceiver),
}
pub struct TimerTask {
pub(crate) task_type: TaskType,
pub(crate) delay: std::time::Duration,
pub(crate) callback: Option<CallbackWrapper>,
}
impl TimerTask {
#[inline]
pub fn new_oneshot(delay: std::time::Duration, callback: Option<CallbackWrapper>) -> Self {
Self {
task_type: TaskType::OneShot,
delay,
callback,
}
}
#[inline]
pub fn new_periodic(
initial_delay: std::time::Duration,
interval: std::time::Duration,
callback: Option<CallbackWrapper>,
buffer_size: Option<NonZeroUsize>,
) -> Result<Self, TimerError> {
if interval.is_zero() {
return Err(TimerError::InvalidConfiguration {
field: "interval".to_string(),
reason: "periodic interval must be greater than zero".to_string(),
});
}
Ok(Self {
task_type: TaskType::Periodic {
interval,
buffer_size,
},
delay: initial_delay,
callback,
})
}
#[inline]
pub fn get_task_type(&self) -> &TaskType {
&self.task_type
}
#[inline]
pub fn get_interval(&self) -> Option<std::time::Duration> {
match self.task_type {
TaskType::Periodic { interval, .. } => Some(interval),
TaskType::OneShot => None,
}
}
}
pub struct TimerTaskWithCompletionNotifier {
pub(crate) task_type: TaskTypeWithCompletionNotifier,
pub(crate) delay: std::time::Duration,
pub(crate) callback: Option<CallbackWrapper>,
}
impl TimerTaskWithCompletionNotifier {
pub fn from_timer_task(task: TimerTask) -> (Self, CompletionReceiver) {
match task.task_type {
TaskType::OneShot => {
let (notifier, receiver) = channel();
(
Self {
task_type: TaskTypeWithCompletionNotifier::OneShot {
completion_notifier: notifier,
},
delay: task.delay,
callback: task.callback,
},
CompletionReceiver::OneShot(receiver),
)
}
TaskType::Periodic {
interval,
buffer_size,
} => {
let queue = PeriodicCompletionQueue::new(buffer_size);
let notifier = PeriodicCompletionNotifier(Arc::clone(&queue));
let receiver = PeriodicCompletionReceiver(queue);
(
Self {
task_type: TaskTypeWithCompletionNotifier::Periodic {
interval,
completion_notifier: notifier,
},
delay: task.delay,
callback: task.callback,
},
CompletionReceiver::Periodic(receiver),
)
}
}
}
#[inline]
pub fn into_task_type(self) -> TaskTypeWithCompletionNotifier {
self.task_type
}
#[inline]
pub fn get_interval(&self) -> Option<std::time::Duration> {
self.task_type.get_interval()
}
}
pub(crate) struct TimerTaskForWheel {
pub(crate) task_id: TaskId,
pub(crate) task: TimerTaskWithCompletionNotifier,
pub(crate) deadline_tick: u64,
pub(crate) rounds: u32,
pub(crate) deadline_at: Instant,
}
impl TimerTaskForWheel {
#[inline]
pub(crate) fn new_with_id(
task_id: TaskId,
task: TimerTaskWithCompletionNotifier,
deadline_tick: u64,
rounds: u32,
deadline_at: Instant,
) -> Self {
Self {
task_id,
task,
deadline_tick,
rounds,
deadline_at,
}
}
#[inline]
pub fn get_id(&self) -> TaskId {
self.task_id
}
#[inline]
pub fn into_task_type(self) -> TaskTypeWithCompletionNotifier {
self.task.into_task_type()
}
#[inline]
pub fn update_delay(&mut self, delay: std::time::Duration) {
self.task.delay = delay
}
#[inline]
pub fn update_callback(&mut self, callback: CallbackWrapper) {
self.task.callback = Some(callback)
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct TaskLocation {
pub slot_index: usize,
pub vec_index: usize,
pub level: u8,
}
impl TaskLocation {
#[inline(always)]
pub fn new(level: u8, slot_index: usize, vec_index: usize) -> Self {
Self {
slot_index,
vec_index,
level,
}
}
}