use super::{TaskID, instrumentation::TaskInstrumentation, task_local::InFlightTaskCancellation};
use crate::hint::Hint;
use crate::observer::{ObserverNotified, ObserverSender};
use crate::{Priority, SomeExecutor, SomeLocalExecutor, SomeStaticExecutor};
use std::convert::Infallible;
use std::fmt::{Debug, Formatter};
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
use super::task_local::{
IS_CANCELLED, TASK_EXECUTOR, TASK_ID, TASK_LABEL, TASK_PRIORITY, TASK_STATIC_EXECUTOR,
};
pub(super) struct TaskMetadata {
pub(super) priority: Priority,
pub(super) task_id: TaskID,
pub(super) poll_after: crate::sys::Instant,
}
pub(super) struct TaskState<'a, F, N> {
pub(super) sender: &'a mut ObserverSender<F, N>,
pub(super) label: &'a mut Option<String>,
pub(super) cancellation: &'a Option<InFlightTaskCancellation>,
pub(super) executor: &'a mut Option<Box<dyn SomeExecutor<ExecutorNotifier = Infallible>>>,
pub(super) instrumentation: &'a mut TaskInstrumentation,
}
pub struct SpawnedTask<F, ONotifier, Executor>
where
F: Future,
{
pub(super) task: F,
pub(super) sender: ObserverSender<F::Output, ONotifier>,
pub(super) phantom: PhantomData<Executor>,
pub(super) poll_after: crate::sys::Instant,
pub(super) hint: Hint,
pub(super) label: Option<String>,
pub(super) cancellation: Option<InFlightTaskCancellation>,
pub(super) executor: Option<Box<dyn SomeExecutor<ExecutorNotifier = Infallible>>>,
pub(super) priority: Priority, pub(super) task_id: TaskID, pub(super) instrumentation: TaskInstrumentation,
}
pub struct SpawnedLocalTask<F, ONotifier, Executor>
where
F: Future,
{
pub(super) task: F,
pub(super) sender: ObserverSender<F::Output, ONotifier>,
pub(super) poll_after: crate::sys::Instant,
pub(super) executor: PhantomData<Executor>,
pub(super) label: Option<String>,
pub(super) cancellation: Option<InFlightTaskCancellation>,
pub(super) hint: Hint,
pub(super) priority: Priority,
pub(super) task_id: TaskID,
pub(super) instrumentation: TaskInstrumentation,
}
pub struct SpawnedStaticTask<F, ONotifier, Executor>
where
F: Future,
{
pub(super) task: F,
pub(super) sender: ObserverSender<F::Output, ONotifier>,
pub(super) poll_after: crate::sys::Instant,
pub(super) executor: PhantomData<Executor>,
pub(super) label: Option<String>,
pub(super) cancellation: Option<InFlightTaskCancellation>,
pub(super) hint: Hint,
pub(super) priority: Priority,
pub(super) task_id: TaskID,
pub(super) instrumentation: TaskInstrumentation,
}
impl<F: Future, ONotifier, ENotifier> SpawnedTask<F, ONotifier, ENotifier> {
pub fn hint(&self) -> Hint {
self.hint
}
pub fn label(&self) -> &str {
self.label.as_ref().expect("Future is polling")
}
pub fn priority(&self) -> Priority {
self.priority
}
pub fn poll_after(&self) -> crate::sys::Instant {
self.poll_after
}
pub fn task_id(&self) -> TaskID {
self.task_id
}
pub fn into_future(self) -> F {
self.task
}
}
impl<F: Future, ONotifier, Executor> SpawnedLocalTask<F, ONotifier, Executor> {
pub fn hint(&self) -> Hint {
self.hint
}
pub fn label(&self) -> &str {
self.label.as_ref().expect("Future is polling")
}
pub fn priority(&self) -> Priority {
self.priority
}
pub fn poll_after(&self) -> crate::sys::Instant {
self.poll_after
}
pub fn task_id(&self) -> TaskID {
self.task_id
}
pub fn into_future(self) -> F {
self.task
}
}
impl<F: Future, ONotifier, Executor> SpawnedStaticTask<F, ONotifier, Executor> {
pub fn hint(&self) -> Hint {
self.hint
}
pub fn label(&self) -> &str {
self.label.as_ref().expect("Future is polling")
}
pub fn priority(&self) -> Priority {
self.priority
}
pub fn poll_after(&self) -> crate::sys::Instant {
self.poll_after
}
pub fn task_id(&self) -> TaskID {
self.task_id
}
pub fn into_future(self) -> F {
self.task
}
pub fn poll<S>(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
static_executor: Option<&mut S>,
mut some_executor: Option<Box<dyn SomeExecutor<ExecutorNotifier = Infallible> + 'static>>,
) -> std::task::Poll<()>
where
ONotifier: ObserverNotified<F::Output>,
S: SomeStaticExecutor,
{
let poll_after = self.poll_after();
let (future, sender, label, priority, cancellation, task_id, instrumentation) = unsafe {
let unchecked = self.get_unchecked_mut();
let future = Pin::new_unchecked(&mut unchecked.task);
let sender = &mut unchecked.sender;
let label = &mut unchecked.label;
let priority = unchecked.priority;
let cancellation = &mut unchecked.cancellation;
let task_id = unchecked.task_id;
let instrumentation = &mut unchecked.instrumentation;
(
future,
sender,
label,
priority,
cancellation,
task_id,
instrumentation,
)
};
let metadata = TaskMetadata {
priority,
task_id,
poll_after,
};
let state = TaskState {
sender,
label,
cancellation,
executor: &mut some_executor,
instrumentation,
};
common_poll(
future,
state,
None::<&mut Infallible>,
static_executor,
metadata,
cx,
)
}
}
impl<F, ONotifier, E> Future for SpawnedStaticTask<F, ONotifier, E>
where
F: Future,
ONotifier: ObserverNotified<F::Output>,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
SpawnedStaticTask::poll::<Infallible>(self, cx, None, None)
}
}
pub(super) fn common_poll<'l, F, N, L, S>(
future: Pin<&mut F>,
state: TaskState<'_, F::Output, N>,
_local_executor: Option<&mut L>,
static_executor: Option<&mut S>,
metadata: TaskMetadata,
cx: &mut Context,
) -> std::task::Poll<()>
where
F: Future,
N: ObserverNotified<F::Output>,
L: SomeLocalExecutor<'l>,
S: SomeStaticExecutor,
{
assert!(
metadata.poll_after <= crate::sys::Instant::now(),
"Conforming executors should not poll tasks before the poll_after time."
);
if state.sender.observer_cancelled() {
state.instrumentation.cancelled();
return Poll::Ready(());
}
let new_executor = state.executor.as_ref().map(|e| e.clone_box());
let new_static_executor = static_executor.map(|e| e.clone_box());
let (old_label, old_cancel, old_priority, old_id, old_executor, old_static_executor) = unsafe {
(
TASK_LABEL.with_mut(|l| std::mem::replace(l, state.label.clone())),
IS_CANCELLED.with_mut(|c| std::mem::replace(c, state.cancellation.clone())),
TASK_PRIORITY.with_mut(|p| p.replace(metadata.priority)),
TASK_ID.with_mut(|i| i.replace(metadata.task_id)),
TASK_EXECUTOR.with_mut(|e| e.replace(new_executor)),
TASK_STATIC_EXECUTOR.with_mut(|e| e.replace(new_static_executor)),
)
};
struct RestoreGuard<F: FnOnce()>(Option<F>);
impl<F: FnOnce()> Drop for RestoreGuard<F> {
fn drop(&mut self) {
(self.0.take().expect("Guard dropped twice"))()
}
}
let restore = RestoreGuard(Some(move || unsafe {
TASK_LABEL.with_mut(|l| *l = old_label);
IS_CANCELLED.with_mut(|c| *c = old_cancel);
TASK_PRIORITY.with_mut(|p| *p = old_priority);
TASK_ID.with_mut(|i| *i = old_id);
TASK_EXECUTOR.with_mut(|e| *e = old_executor);
TASK_STATIC_EXECUTOR.with_mut(|e| *e = old_static_executor);
}));
let active_poll = state.instrumentation.begin_poll();
#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
let instrumented_waker = state.instrumentation.instrument_waker(cx.waker());
#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
let mut instrumented_context = Context::from_waker(&instrumented_waker);
let logwise_context = logwise::context::enter(state.instrumentation.context());
#[cfg(any(feature = "logwise-forensic", feature = "logwise-performance"))]
let r = future.poll(&mut instrumented_context);
#[cfg(not(any(feature = "logwise-forensic", feature = "logwise-performance")))]
let r = future.poll(cx);
drop(logwise_context);
drop(active_poll);
drop(restore);
match r {
Poll::Ready(r) => {
state.instrumentation.completed();
state.sender.send(r);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
impl<F: Future, ONotifier, E> SpawnedTask<F, ONotifier, E>
where
ONotifier: ObserverNotified<F::Output>,
{
pub fn poll<'l, L>(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
local_executor: Option<&mut L>,
) -> std::task::Poll<()>
where
L: SomeLocalExecutor<'l>,
{
let poll_after = self.poll_after();
let (future, sender, label, priority, cancellation, task_id, executor, instrumentation) = unsafe {
let unchecked = self.get_unchecked_mut();
let future = Pin::new_unchecked(&mut unchecked.task);
let sender = Pin::new_unchecked(&mut unchecked.sender);
let label = Pin::new_unchecked(&mut unchecked.label);
let priority = Pin::new_unchecked(&mut unchecked.priority);
let cancellation = Pin::new_unchecked(&mut unchecked.cancellation);
let task_id = unchecked.task_id;
let executor = Pin::new_unchecked(&mut unchecked.executor);
let instrumentation = &mut unchecked.instrumentation;
(
future,
sender,
label,
priority,
cancellation,
task_id,
executor,
instrumentation,
)
};
let metadata = TaskMetadata {
priority: *priority.get_mut(),
task_id,
poll_after,
};
let state = TaskState {
sender: sender.get_mut(),
label: label.get_mut(),
cancellation: cancellation.get_mut(),
executor: executor.get_mut(),
instrumentation,
};
common_poll(
future,
state,
local_executor,
None::<&mut Infallible>,
metadata,
cx,
)
}
}
impl<F, ONotifier, E> Future for SpawnedTask<F, ONotifier, E>
where
F: Future,
ONotifier: ObserverNotified<F::Output>,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
SpawnedTask::poll::<Infallible>(self, cx, None)
}
}
impl<'executor, F, ONotifier, Executor: SomeLocalExecutor<'executor>>
SpawnedLocalTask<F, ONotifier, Executor>
where
F: Future,
ONotifier: ObserverNotified<F::Output>,
{
pub fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
executor: &mut Executor,
mut some_executor: Option<Box<dyn SomeExecutor<ExecutorNotifier = Infallible> + 'static>>,
) -> std::task::Poll<()> {
let poll_after = self.poll_after();
let (future, sender, label, priority, cancellation, task_id, instrumentation) = unsafe {
let unchecked = self.get_unchecked_mut();
let future = Pin::new_unchecked(&mut unchecked.task);
let sender = Pin::new_unchecked(&mut unchecked.sender);
let label = Pin::new_unchecked(&mut unchecked.label);
let priority = Pin::new_unchecked(&mut unchecked.priority);
let cancellation = Pin::new_unchecked(&mut unchecked.cancellation);
let task_id = unchecked.task_id;
let instrumentation = &mut unchecked.instrumentation;
(
future,
sender,
label,
priority,
cancellation,
task_id,
instrumentation,
)
};
let metadata = TaskMetadata {
priority: *priority.get_mut(),
task_id,
poll_after,
};
let state = TaskState {
sender: sender.get_mut(),
label: label.get_mut(),
cancellation: cancellation.get_mut(),
executor: &mut some_executor,
instrumentation,
};
common_poll(
future,
state,
Some(executor),
None::<&mut Infallible>,
metadata,
cx,
)
}
}
impl<F: Future, N, E> Debug for SpawnedTask<F, N, E> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SpawnedTask")
.field("poll_after", &self.poll_after)
.field("hint", &self.hint)
.field("label", &self.label)
.field("priority", &self.priority)
.field("task_id", &self.task_id)
.finish()
}
}
impl<F: Future, N, E> Debug for SpawnedLocalTask<F, N, E> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SpawnedLocalTask")
.field("poll_after", &self.poll_after)
.field("hint", &self.hint)
.field("label", &self.label)
.field("priority", &self.priority)
.field("task_id", &self.task_id)
.finish()
}
}
impl<F: Future, N, E> Debug for SpawnedStaticTask<F, N, E> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SpawnedStaticTask")
.field("poll_after", &self.poll_after)
.field("hint", &self.hint)
.field("label", &self.label)
.field("priority", &self.priority)
.field("task_id", &self.task_id)
.finish()
}
}
impl<F: Future, N, E> AsRef<F> for SpawnedTask<F, N, E> {
fn as_ref(&self) -> &F {
&self.task
}
}
impl<F: Future, N, E> AsMut<F> for SpawnedTask<F, N, E> {
fn as_mut(&mut self) -> &mut F {
&mut self.task
}
}
impl<F: Future, N, E> AsRef<F> for SpawnedLocalTask<F, N, E> {
fn as_ref(&self) -> &F {
&self.task
}
}
impl<F: Future, N, E> AsMut<F> for SpawnedLocalTask<F, N, E> {
fn as_mut(&mut self) -> &mut F {
&mut self.task
}
}
impl<F: Future, N, E> AsRef<F> for SpawnedStaticTask<F, N, E> {
fn as_ref(&self) -> &F {
&self.task
}
}
impl<F: Future, N, E> AsMut<F> for SpawnedStaticTask<F, N, E> {
fn as_mut(&mut self) -> &mut F {
&mut self.task
}
}
use crate::task::dyn_spawned::DynSpawnedTask;
impl<'a, F, N, E> AsRef<dyn DynSpawnedTask<Infallible> + 'a> for SpawnedTask<F, N, E>
where
N: ObserverNotified<F::Output>,
F: Future + 'a,
F: Send,
N: Send,
E: Send + 'a,
F::Output: Send,
{
fn as_ref(&self) -> &(dyn DynSpawnedTask<Infallible> + 'a) {
self
}
}
impl<'a, F, N, E> AsMut<dyn DynSpawnedTask<Infallible> + 'a> for SpawnedTask<F, N, E>
where
N: ObserverNotified<F::Output>,
F: Future + 'a,
F: Send,
N: Send,
E: Send + 'a,
F::Output: Send,
{
fn as_mut(&mut self) -> &mut (dyn DynSpawnedTask<Infallible> + 'a) {
self
}
}
#[cfg(test)]
mod tests {
#[cfg(not(target_arch = "wasm32"))]
#[test]
fn panic_in_task_restores_task_locals() {
use crate::task::{Configuration, Task};
use crate::task::{TASK_ID, TASK_LABEL};
let task: Task<_, std::convert::Infallible> = Task::with_notifications(
"panicky".to_string(),
Configuration::default(),
None,
async {
panic!("boom");
},
);
let mut executor = crate::current_executor::current_executor();
let (spawned, observer) = task.spawn(&mut executor);
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
wasm_lite_std::block_on(spawned);
}));
assert!(result.is_err());
TASK_ID.with(|id| assert!(id.is_none()));
TASK_LABEL.with(|l| assert!(l.is_none()));
drop(observer);
}
#[cfg(not(target_arch = "wasm32"))]
#[test]
fn child_without_executor_does_not_inherit_parent_task_executor() {
use crate::SomeStaticExecutor;
use crate::observer::Observer;
use crate::task::{Configuration, TASK_EXECUTOR, Task};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
let child_saw_parent_executor = Arc::new(AtomicBool::new(false));
let child_saw_parent_executor_clone = child_saw_parent_executor.clone();
let parent = Task::without_notifications(
"parent".to_string(),
Configuration::default(),
async move {
let child = Task::without_notifications(
"child without executor".to_string(),
Configuration::default(),
async move {
let inherited = TASK_EXECUTOR.with(|current| {
current.and_then(|executor| executor.as_ref()).is_some()
});
child_saw_parent_executor_clone.store(inherited, Ordering::Relaxed);
},
);
let mut static_executor =
crate::static_last_resort::StaticLastResortExecutor::new();
static_executor.spawn_static(child).detach();
},
);
let mut parent_executor = crate::last_resort::LastResortExecutor::new();
let (spawned, observer) = parent.spawn(&mut parent_executor);
wasm_lite_std::block_on(spawned);
observer.detach();
assert!(
!child_saw_parent_executor.load(Ordering::Relaxed),
"a child polled without an executor inherited TASK_EXECUTOR from its parent"
);
}
}