use crate::{Scheduler, SessionId, Timer};use futures::FutureExt as _;use std::{ future::Future, marker::PhantomData, mem::ManuallyDrop, panic::Location, pin::Pin, rc::Rc, sync::Arc, task::{Context, Poll}, thread::{self, ThreadId}, time::Duration,}; #[derive(Clone)]pub struct ForegroundExecutor { session_id: SessionId, scheduler: Arc<dyn Scheduler>, not_send: PhantomData<Rc<()>>,} impl ForegroundExecutor { pub fn new(session_id: SessionId, scheduler: Arc<dyn Scheduler>) -> Self { Self { session_id, scheduler, not_send: PhantomData, } } #[track_caller] pub fn spawn<F>(&self, future: F) -> Task<F::Output> where F: Future + 'static, F::Output: 'static, { let session_id = self.session_id; let scheduler = Arc::clone(&self.scheduler); let (runnable, task) = spawn_local_with_source_location(future, move |runnable| { scheduler.schedule_foreground(session_id, runnable); }); runnable.schedule(); Task(TaskState::Spawned(task)) } pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output { let mut output = None; self.scheduler.block( Some(self.session_id), async { output = Some(future.await) }.boxed_local(), None, ); output.unwrap() } pub fn block_with_timeout<Fut: Unpin + Future>( &self, timeout: Duration, mut future: Fut, ) -> Result<Fut::Output, Fut> { let mut output = None; self.scheduler.block( Some(self.session_id), async { output = Some((&mut future).await) }.boxed_local(), Some(timeout), ); output.ok_or(future) } pub fn timer(&self, duration: Duration) -> Timer { self.scheduler.timer(duration) }} #[derive(Clone)]pub struct BackgroundExecutor { scheduler: Arc<dyn Scheduler>,} impl BackgroundExecutor { pub fn new(scheduler: Arc<dyn Scheduler>) -> Self { Self { scheduler } } pub fn spawn<F>(&self, future: F) -> Task<F::Output> where F: Future + Send + 'static, F::Output: Send + 'static, { let scheduler = Arc::clone(&self.scheduler); let (runnable, task) = async_task::spawn(future, move |runnable| { scheduler.schedule_background(runnable); }); runnable.schedule(); Task(TaskState::Spawned(task)) } pub fn timer(&self, duration: Duration) -> Timer { self.scheduler.timer(duration) } pub fn scheduler(&self) -> &Arc<dyn Scheduler> { &self.scheduler }} /// Task is a primitive that allows work to happen in the background.////// It implements [`Future`] so you can `.await` on it.////// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows/// the task to continue running, but with no way to return a value.#[must_use]#[derive(Debug)]pub struct Task<T>(TaskState<T>); #[derive(Debug)]enum TaskState<T> { /// A task that is ready to return a value Ready(Option<T>), /// A task that is currently running. Spawned(async_task::Task<T>),} impl<T> Task<T> { /// Creates a new task that will resolve with the value pub fn ready(val: T) -> Self { Task(TaskState::Ready(Some(val))) } pub fn is_ready(&self) -> bool { match &self.0 { TaskState::Ready(_) => true, TaskState::Spawned(task) => task.is_finished(), } } /// Detaching a task runs it to completion in the background pub fn detach(self) { match self { Task(TaskState::Ready(_)) => {} Task(TaskState::Spawned(task)) => task.detach(), } }} impl<T> Future for Task<T> { type Output = T; fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> { match unsafe { self.get_unchecked_mut() } { Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()), Task(TaskState::Spawned(task)) => Pin::new(task).poll(cx), } }} /// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.////// Copy-modified from:/// <https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L405>#[track_caller]fn spawn_local_with_source_location<Fut, S>( future: Fut, schedule: S,) -> (async_task::Runnable, async_task::Task<Fut::Output, ()>)where Fut: Future + 'static, Fut::Output: 'static, S: async_task::Schedule + Send + Sync + 'static,{ #[inline] fn thread_id() -> ThreadId { std::thread_local! { static ID: ThreadId = thread::current().id(); } ID.try_with(|id| *id) .unwrap_or_else(|_| thread::current().id()) } struct Checked<F> { id: ThreadId, inner: ManuallyDrop<F>, location: &'static Location<'static>, } impl<F> Drop for Checked<F> { fn drop(&mut self) { assert!( self.id == thread_id(), "local task dropped by a thread that didn't spawn it. Task spawned at {}", self.location ); unsafe { ManuallyDrop::drop(&mut self.inner); } } } impl<F: Future> Future for Checked<F> { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { assert!( self.id == thread_id(), "local task polled by a thread that didn't spawn it. Task spawned at {}", self.location ); unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) } } } // Wrap the future into one that checks which thread it's on. let future = Checked { id: thread_id(), inner: ManuallyDrop::new(future), location: Location::caller(), }; unsafe { async_task::spawn_unchecked(future, schedule) }}
use crate::{Scheduler, SessionId, Timer};
use futures::FutureExt as _;
use std::{
future::Future,
marker::PhantomData,
mem::ManuallyDrop,
panic::Location,
pin::Pin,
rc::Rc,
sync::Arc,
task::{Context, Poll},
thread::{self, ThreadId},
time::Duration,
};
#[derive(Clone)]
pub struct ForegroundExecutor {
session_id: SessionId,
scheduler: Arc<dyn Scheduler>,
not_send: PhantomData<Rc<()>>,
}
impl ForegroundExecutor {
pub fn new(session_id: SessionId, scheduler: Arc<dyn Scheduler>) -> Self {
Self {
session_id,
scheduler,
not_send: PhantomData,
}
}
#[track_caller]
pub fn spawn<F>(&self, future: F) -> Task<F::Output>
where
F: Future + 'static,
F::Output: 'static,
{
let session_id = self.session_id;
let scheduler = Arc::clone(&self.scheduler);
let (runnable, task) = spawn_local_with_source_location(future, move |runnable| {
scheduler.schedule_foreground(session_id, runnable);
});
runnable.schedule();
Task(TaskState::Spawned(task))
}
pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output {
let mut output = None;
self.scheduler.block(
Some(self.session_id),
async { output = Some(future.await) }.boxed_local(),
None,
);
output.unwrap()
}
pub fn block_with_timeout<Fut: Unpin + Future>(
&self,
timeout: Duration,
mut future: Fut,
) -> Result<Fut::Output, Fut> {
let mut output = None;
self.scheduler.block(
Some(self.session_id),
async { output = Some((&mut future).await) }.boxed_local(),
Some(timeout),
);
output.ok_or(future)
}
pub fn timer(&self, duration: Duration) -> Timer {
self.scheduler.timer(duration)
}
}
#[derive(Clone)]
pub struct BackgroundExecutor {
scheduler: Arc<dyn Scheduler>,
}
impl BackgroundExecutor {
pub fn new(scheduler: Arc<dyn Scheduler>) -> Self {
Self { scheduler }
}
pub fn spawn<F>(&self, future: F) -> Task<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let scheduler = Arc::clone(&self.scheduler);
let (runnable, task) = async_task::spawn(future, move |runnable| {
scheduler.schedule_background(runnable);
});
runnable.schedule();
Task(TaskState::Spawned(task))
}
pub fn timer(&self, duration: Duration) -> Timer {
self.scheduler.timer(duration)
}
pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
&self.scheduler
}
}
/// Task is a primitive that allows work to happen in the background.
///
/// It implements [`Future`] so you can `.await` on it.
///
/// If you drop a task it will be cancelled immediately. Calling [`Task::detach`] allows
/// the task to continue running, but with no way to return a value.
#[must_use]
#[derive(Debug)]
pub struct Task<T>(TaskState<T>);
#[derive(Debug)]
enum TaskState<T> {
/// A task that is ready to return a value
Ready(Option<T>),
/// A task that is currently running.
Spawned(async_task::Task<T>),
}
impl<T> Task<T> {
/// Creates a new task that will resolve with the value
pub fn ready(val: T) -> Self {
Task(TaskState::Ready(Some(val)))
}
pub fn is_ready(&self) -> bool {
match &self.0 {
TaskState::Ready(_) => true,
TaskState::Spawned(task) => task.is_finished(),
}
}
/// Detaching a task runs it to completion in the background
pub fn detach(self) {
match self {
Task(TaskState::Ready(_)) => {}
Task(TaskState::Spawned(task)) => task.detach(),
}
}
}
impl<T> Future for Task<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
match unsafe { self.get_unchecked_mut() } {
Task(TaskState::Ready(val)) => Poll::Ready(val.take().unwrap()),
Task(TaskState::Spawned(task)) => Pin::new(task).poll(cx),
}
}
}
/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
///
/// Copy-modified from:
/// <https://github.com/smol-rs/async-task/blob/ca9dbe1db9c422fd765847fa91306e30a6bb58a9/src/runnable.rs#L405>
#[track_caller]
fn spawn_local_with_source_location<Fut, S>(
future: Fut,
schedule: S,
) -> (async_task::Runnable, async_task::Task<Fut::Output, ()>)
where
Fut: Future + 'static,
Fut::Output: 'static,
S: async_task::Schedule + Send + Sync + 'static,
{
#[inline]
fn thread_id() -> ThreadId {
std::thread_local! {
static ID: ThreadId = thread::current().id();
}
ID.try_with(|id| *id)
.unwrap_or_else(|_| thread::current().id())
}
struct Checked<F> {
id: ThreadId,
inner: ManuallyDrop<F>,
location: &'static Location<'static>,
}
impl<F> Drop for Checked<F> {
fn drop(&mut self) {
assert!(
self.id == thread_id(),
"local task dropped by a thread that didn't spawn it. Task spawned at {}",
self.location
);
unsafe {
ManuallyDrop::drop(&mut self.inner);
}
}
}
impl<F: Future> Future for Checked<F> {
type Output = F::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
assert!(
self.id == thread_id(),
"local task polled by a thread that didn't spawn it. Task spawned at {}",
self.location
);
unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
}
}
// Wrap the future into one that checks which thread it's on.
let future = Checked {
id: thread_id(),
inner: ManuallyDrop::new(future),
location: Location::caller(),
};
unsafe { async_task::spawn_unchecked(future, schedule) }
}