use std::error::Error;
use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;
use pin_project_lite::pin_project;
use recycle_box::{coerce_box, RecycleBox};
use crate::channel::{ChannelId, Sender};
use crate::executor::Executor;
use crate::model::{InputFn, Model};
use crate::time::{MonotonicTime, TearableAtomicTime};
use crate::util::priority_queue::PriorityQueue;
use crate::util::sync_cell::SyncCellReader;
pub(crate) type SchedulerQueue = PriorityQueue<(MonotonicTime, ChannelId), Box<dyn ScheduledEvent>>;
pub trait Deadline {
fn into_time(self, now: MonotonicTime) -> MonotonicTime;
}
impl Deadline for Duration {
#[inline(always)]
fn into_time(self, now: MonotonicTime) -> MonotonicTime {
now + self
}
}
impl Deadline for MonotonicTime {
#[inline(always)]
fn into_time(self, _: MonotonicTime) -> MonotonicTime {
self
}
}
pub struct Scheduler<M: Model> {
sender: Sender<M>,
scheduler_queue: Arc<Mutex<SchedulerQueue>>,
time: SyncCellReader<TearableAtomicTime>,
}
impl<M: Model> Scheduler<M> {
pub(crate) fn new(
sender: Sender<M>,
scheduler_queue: Arc<Mutex<SchedulerQueue>>,
time: SyncCellReader<TearableAtomicTime>,
) -> Self {
Self {
sender,
scheduler_queue,
time,
}
}
pub fn time(&self) -> MonotonicTime {
self.time.try_read().expect("internal simulation error: could not perform a synchronized read of the simulation time")
}
pub fn schedule_event<F, T, S>(
&self,
deadline: impl Deadline,
func: F,
arg: T,
) -> Result<(), SchedulingError>
where
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
S: Send + 'static,
{
let now = self.time();
let time = deadline.into_time(now);
if now >= time {
return Err(SchedulingError::InvalidScheduledTime);
}
let sender = self.sender.clone();
schedule_event_at_unchecked(time, func, arg, sender, &self.scheduler_queue);
Ok(())
}
pub fn schedule_keyed_event<F, T, S>(
&self,
deadline: impl Deadline,
func: F,
arg: T,
) -> Result<EventKey, SchedulingError>
where
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
S: Send + 'static,
{
let now = self.time();
let time = deadline.into_time(now);
if now >= time {
return Err(SchedulingError::InvalidScheduledTime);
}
let sender = self.sender.clone();
let event_key =
schedule_keyed_event_at_unchecked(time, func, arg, sender, &self.scheduler_queue);
Ok(event_key)
}
pub fn schedule_periodic_event<F, T, S>(
&self,
deadline: impl Deadline,
period: Duration,
func: F,
arg: T,
) -> Result<(), SchedulingError>
where
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
let now = self.time();
let time = deadline.into_time(now);
if now >= time {
return Err(SchedulingError::InvalidScheduledTime);
}
if period.is_zero() {
return Err(SchedulingError::NullRepetitionPeriod);
}
let sender = self.sender.clone();
schedule_periodic_event_at_unchecked(
time,
period,
func,
arg,
sender,
&self.scheduler_queue,
);
Ok(())
}
pub fn schedule_keyed_periodic_event<F, T, S>(
&self,
deadline: impl Deadline,
period: Duration,
func: F,
arg: T,
) -> Result<EventKey, SchedulingError>
where
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
let now = self.time();
let time = deadline.into_time(now);
if now >= time {
return Err(SchedulingError::InvalidScheduledTime);
}
if period.is_zero() {
return Err(SchedulingError::NullRepetitionPeriod);
}
let sender = self.sender.clone();
let event_key = schedule_periodic_keyed_event_at_unchecked(
time,
period,
func,
arg,
sender,
&self.scheduler_queue,
);
Ok(event_key)
}
}
impl<M: Model> fmt::Debug for Scheduler<M> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Scheduler").finish_non_exhaustive()
}
}
#[derive(Clone, Debug)]
#[must_use = "prefer unkeyed scheduling methods if the event is never cancelled"]
pub struct EventKey {
is_cancelled: Arc<AtomicBool>,
}
impl EventKey {
pub(crate) fn new() -> Self {
Self {
is_cancelled: Arc::new(AtomicBool::new(false)),
}
}
pub(crate) fn is_cancelled(&self) -> bool {
self.is_cancelled.load(Ordering::Relaxed)
}
pub fn cancel(self) {
self.is_cancelled.store(true, Ordering::Relaxed);
}
}
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum SchedulingError {
InvalidScheduledTime,
NullRepetitionPeriod,
}
impl fmt::Display for SchedulingError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidScheduledTime => write!(
fmt,
"the scheduled time should be in the future of the current simulation time"
),
Self::NullRepetitionPeriod => write!(fmt, "the repetition period cannot be zero"),
}
}
}
impl Error for SchedulingError {}
pub(crate) fn schedule_event_at_unchecked<M, F, T, S>(
time: MonotonicTime,
func: F,
arg: T,
sender: Sender<M>,
scheduler_queue: &Mutex<SchedulerQueue>,
) where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
S: Send + 'static,
{
let channel_id = sender.channel_id();
let event_dispatcher = Box::new(new_event_dispatcher(func, arg, sender));
let mut scheduler_queue = scheduler_queue.lock().unwrap();
scheduler_queue.insert((time, channel_id), event_dispatcher);
}
pub(crate) fn schedule_keyed_event_at_unchecked<M, F, T, S>(
time: MonotonicTime,
func: F,
arg: T,
sender: Sender<M>,
scheduler_queue: &Mutex<SchedulerQueue>,
) -> EventKey
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
S: Send + 'static,
{
let event_key = EventKey::new();
let channel_id = sender.channel_id();
let event_dispatcher = Box::new(KeyedEventDispatcher::new(
event_key.clone(),
func,
arg,
sender,
));
let mut scheduler_queue = scheduler_queue.lock().unwrap();
scheduler_queue.insert((time, channel_id), event_dispatcher);
event_key
}
pub(crate) fn schedule_periodic_event_at_unchecked<M, F, T, S>(
time: MonotonicTime,
period: Duration,
func: F,
arg: T,
sender: Sender<M>,
scheduler_queue: &Mutex<SchedulerQueue>,
) where
M: Model,
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
let channel_id = sender.channel_id();
let event_dispatcher = Box::new(PeriodicEventDispatcher::new(func, arg, sender, period));
let mut scheduler_queue = scheduler_queue.lock().unwrap();
scheduler_queue.insert((time, channel_id), event_dispatcher);
}
pub(crate) fn schedule_periodic_keyed_event_at_unchecked<M, F, T, S>(
time: MonotonicTime,
period: Duration,
func: F,
arg: T,
sender: Sender<M>,
scheduler_queue: &Mutex<SchedulerQueue>,
) -> EventKey
where
M: Model,
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
let event_key = EventKey::new();
let channel_id = sender.channel_id();
let event_dispatcher = Box::new(PeriodicKeyedEventDispatcher::new(
event_key.clone(),
func,
arg,
sender,
period,
));
let mut scheduler_queue = scheduler_queue.lock().unwrap();
scheduler_queue.insert((time, channel_id), event_dispatcher);
event_key
}
pub(crate) trait ScheduledEvent: Send {
fn is_cancelled(&self) -> bool;
fn next(&self) -> Option<(Box<dyn ScheduledEvent>, Duration)>;
fn into_future(self: Box<Self>) -> Pin<Box<dyn Future<Output = ()> + Send>>;
fn spawn_and_forget(self: Box<Self>, executor: &Executor);
}
pin_project! {
pub(crate) struct EventDispatcher<F> {
#[pin]
fut: F,
}
}
fn new_event_dispatcher<M, F, T, S>(
func: F,
arg: T,
sender: Sender<M>,
) -> EventDispatcher<impl Future<Output = ()>>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
let fut = dispatch_event(func, arg, sender);
EventDispatcher { fut }
}
impl<F> Future for EventDispatcher<F>
where
F: Future,
{
type Output = F::Output;
#[inline(always)]
fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.project().fut.poll(cx)
}
}
impl<F> ScheduledEvent for EventDispatcher<F>
where
F: Future<Output = ()> + Send + 'static,
{
fn is_cancelled(&self) -> bool {
false
}
fn next(&self) -> Option<(Box<dyn ScheduledEvent>, Duration)> {
None
}
fn into_future(self: Box<Self>) -> Pin<Box<dyn Future<Output = ()> + Send>> {
Box::into_pin(self)
}
fn spawn_and_forget(self: Box<Self>, executor: &Executor) {
executor.spawn_and_forget(*self);
}
}
pub(crate) struct PeriodicEventDispatcher<M, F, T, S>
where
M: Model,
{
func: F,
arg: T,
sender: Sender<M>,
period: Duration,
_input_kind: PhantomData<S>,
}
impl<M, F, T, S> PeriodicEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
fn new(func: F, arg: T, sender: Sender<M>, period: Duration) -> Self {
Self {
func,
arg,
sender,
period,
_input_kind: PhantomData,
}
}
}
impl<M, F, T, S> ScheduledEvent for PeriodicEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
fn is_cancelled(&self) -> bool {
false
}
fn next(&self) -> Option<(Box<dyn ScheduledEvent>, Duration)> {
let event = Box::new(Self::new(
self.func.clone(),
self.arg.clone(),
self.sender.clone(),
self.period,
));
Some((event, self.period))
}
fn into_future(self: Box<Self>) -> Pin<Box<dyn Future<Output = ()> + Send>> {
let Self {
func, arg, sender, ..
} = *self;
Box::pin(dispatch_event(func, arg, sender))
}
fn spawn_and_forget(self: Box<Self>, executor: &Executor) {
let Self {
func, arg, sender, ..
} = *self;
let fut = dispatch_event(func, arg, sender);
executor.spawn_and_forget(fut);
}
}
pub(crate) struct KeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
event_key: EventKey,
func: F,
arg: T,
sender: Sender<M>,
_input_kind: PhantomData<S>,
}
impl<M, F, T, S> KeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
fn new(event_key: EventKey, func: F, arg: T, sender: Sender<M>) -> Self {
Self {
event_key,
func,
arg,
sender,
_input_kind: PhantomData,
}
}
}
impl<M, F, T, S> ScheduledEvent for KeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
S: Send + 'static,
{
fn is_cancelled(&self) -> bool {
self.event_key.is_cancelled()
}
fn next(&self) -> Option<(Box<dyn ScheduledEvent>, Duration)> {
None
}
fn into_future(self: Box<Self>) -> Pin<Box<dyn Future<Output = ()> + Send>> {
let Self {
event_key,
func,
arg,
sender,
..
} = *self;
Box::pin(dispatch_keyed_event(event_key, func, arg, sender))
}
fn spawn_and_forget(self: Box<Self>, executor: &Executor) {
let Self {
event_key,
func,
arg,
sender,
..
} = *self;
let fut = dispatch_keyed_event(event_key, func, arg, sender);
executor.spawn_and_forget(fut);
}
}
pub(crate) struct PeriodicKeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
event_key: EventKey,
func: F,
arg: T,
sender: Sender<M>,
period: Duration,
_input_kind: PhantomData<S>,
}
impl<M, F, T, S> PeriodicKeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
fn new(event_key: EventKey, func: F, arg: T, sender: Sender<M>, period: Duration) -> Self {
Self {
event_key,
func,
arg,
sender,
period,
_input_kind: PhantomData,
}
}
}
impl<M, F, T, S> ScheduledEvent for PeriodicKeyedEventDispatcher<M, F, T, S>
where
M: Model,
F: for<'a> InputFn<'a, M, T, S> + Clone,
T: Send + Clone + 'static,
S: Send + 'static,
{
fn is_cancelled(&self) -> bool {
self.event_key.is_cancelled()
}
fn next(&self) -> Option<(Box<dyn ScheduledEvent>, Duration)> {
let event = Box::new(Self::new(
self.event_key.clone(),
self.func.clone(),
self.arg.clone(),
self.sender.clone(),
self.period,
));
Some((event, self.period))
}
fn into_future(self: Box<Self>) -> Pin<Box<dyn Future<Output = ()> + Send>> {
let Self {
event_key,
func,
arg,
sender,
..
} = *self;
Box::pin(dispatch_keyed_event(event_key, func, arg, sender))
}
fn spawn_and_forget(self: Box<Self>, executor: &Executor) {
let Self {
event_key,
func,
arg,
sender,
..
} = *self;
let fut = dispatch_keyed_event(event_key, func, arg, sender);
executor.spawn_and_forget(fut);
}
}
async fn dispatch_event<M, F, T, S>(func: F, arg: T, sender: Sender<M>)
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
let _ = sender
.send(
move |model: &mut M,
scheduler,
recycle_box: RecycleBox<()>|
-> RecycleBox<dyn Future<Output = ()> + Send + '_> {
let fut = func.call(model, arg, scheduler);
coerce_box!(RecycleBox::recycle(recycle_box, fut))
},
)
.await;
}
async fn dispatch_keyed_event<M, F, T, S>(event_key: EventKey, func: F, arg: T, sender: Sender<M>)
where
M: Model,
F: for<'a> InputFn<'a, M, T, S>,
T: Send + Clone + 'static,
{
let _ = sender
.send(
move |model: &mut M,
scheduler,
recycle_box: RecycleBox<()>|
-> RecycleBox<dyn Future<Output = ()> + Send + '_> {
let fut = async move {
if !event_key.is_cancelled() {
func.call(model, arg, scheduler).await;
}
};
coerce_box!(RecycleBox::recycle(recycle_box, fut))
},
)
.await;
}