diff --git a/src/lib.rs b/src/lib.rs index 79271a1..41146ac 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ pub mod tracker; #[cfg(feature = "either")] pub mod either; pub mod rc; +pub mod scoped; use std::fmt::{Debug, Formatter}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -20,6 +21,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +pub use crate::scoped::{JoinError as ScopedJoinError, Scope, ScopeExecutor, ScopedJoinHandle}; + #[cfg_attr(feature = "tokio", allow(dead_code))] pub(crate) struct CompletionGuard { finished: Arc, @@ -285,6 +288,16 @@ impl Debug for CommunicationTask { } impl CommunicationTask { + pub(crate) fn new( + task_handle: AbortableJoinHandle<()>, + channel_tx: futures::channel::mpsc::Sender, + ) -> Self { + Self { + _task_handle: task_handle, + _channel_tx: channel_tx, + } + } + /// Send a message to the task pub async fn send(&mut self, data: T) -> std::io::Result<()> { self._channel_tx @@ -335,6 +348,16 @@ impl Debug for UnboundedCommunicationTask { } impl UnboundedCommunicationTask { + pub(crate) fn new( + task_handle: AbortableJoinHandle<()>, + channel_tx: futures::channel::mpsc::UnboundedSender, + ) -> Self { + Self { + _task_handle: task_handle, + _channel_tx: channel_tx, + } + } + /// Send a message to task pub fn send(&mut self, data: T) -> std::io::Result<()> { self._channel_tx @@ -624,6 +647,83 @@ pub trait Executor { _channel_tx: tx, } } + + /// Create a structured-concurrency scope in which tasks may be spawned + /// that borrow from the enclosing stack frame. + /// + /// Unlike [`Executor::spawn`], tasks spawned on the [`Scope`] are driven + /// cooperatively by the returned future so they may borrow any data that outlives + /// the `'env` lifetime. Every task is either completed or canceled before + /// `scope` returns, so borrows never outlive the stack frame. + /// + /// This is the async analogue of [`std::thread::scope`]. + /// + /// # Example + /// + /// ```no_run + /// # async fn run() { + /// use async_rt::Executor; + /// use async_rt::rt::tokio::TokioExecutor; + /// + /// let executor = TokioExecutor; + /// let data = vec![1, 2, 3, 4]; + /// let sum = executor + /// .scope(async |s| { + /// let a = s.spawn(async { data[0] + data[1] }); + /// let b = s.spawn(async { data[2] + data[3] }); + /// a.await.unwrap() + b.await.unwrap() + /// }) + /// .await; + /// assert_eq!(sum, 10); + /// # } + /// ``` + fn scope<'env, F, T>(&self, f: F) -> impl Future + where + F: for<'scope> AsyncFnOnce(&'scope Scope<'scope, 'env>) -> T, + { + scoped::scope(f) + } + + /// Run an async closure with a scoped [`Executor`] wrapper that + /// forwards spawns to this executor, waits for all spawned tasks + /// to finish when the closure returns, and aborts any outstanding + /// tasks if the scope future itself is cancelled. + /// + /// Unlike [`Executor::scope`], tasks run on the real executor (so + /// they get real parallelism) but must be `Send + 'static`. + /// + /// # Panics in spawned tasks + /// + /// Panics inside tasks spawned via this scope are **not** propagated + /// to the caller. See [`scoped::executor_scope`] for details on + /// per-backend behaviour. If you need to react to a task panic, + /// await its [`JoinHandle`] directly. + /// + /// # Example + /// + /// ```no_run + /// # async fn run() { + /// use async_rt::Executor; + /// use async_rt::rt::tokio::TokioExecutor; + /// + /// let executor = TokioExecutor; + /// let total = executor + /// .executor_scope(async |s| { + /// let a = s.spawn(async { 1 + 2 }); + /// let b = s.spawn(async { 3 + 4 }); + /// a.await.unwrap() + b.await.unwrap() + /// }) + /// .await; + /// assert_eq!(total, 10); + /// # } + /// ``` + fn executor_scope<'scope, F, T>(&'scope self, f: F) -> impl Future + where + Self: Sized, + F: AsyncFnOnce(&ScopeExecutor<'scope, Self>) -> T, + { + scoped::executor_scope(self, f) + } } pub trait ExecutorBlocking: Executor { diff --git a/src/scoped.rs b/src/scoped.rs new file mode 100644 index 0000000..65d3935 --- /dev/null +++ b/src/scoped.rs @@ -0,0 +1,1117 @@ +//! Structured concurrency scopes for async tasks. +//! +//! A [`Scope`] lets you spawn futures that borrow from the enclosing stack +//! frame, analogous to [`std::thread::scope`]. Unlike thread scope, the +//! tasks are driven cooperatively by the scope's own future, they make +//! progress whenever the scope is polled so no runtime-level `spawn` is +//! involved and no `unsafe` is needed to extend lifetimes. +//! +//! The [`scope`] function is the entry point. Inside the async closure, +//! call [`Scope::spawn`] to enqueue a task; the returned [`ScopedJoinHandle`] +//! resolves with the task's output. Any tasks still running when the user +//! closure's future completes are drained before `scope` returns, so every +//! borrow is released before the stack frame goes away. + +use crate::{ + AbortableJoinHandle, CommunicationTask, CompletionGuard, Executor, InnerJoinHandle, JoinHandle, + UnboundedCommunicationTask, +}; +use core::future::{Future, poll_fn}; +use core::marker::PhantomData; +use core::pin::Pin; +use core::task::{Context, Poll}; +use futures::channel::mpsc::{Receiver, UnboundedReceiver}; +use futures::channel::oneshot; +use futures::future::{AbortHandle, Abortable, BoxFuture}; +use futures::stream::FuturesUnordered; +use futures::task::AtomicWaker; +use futures::{FutureExt, StreamExt}; +use parking_lot::Mutex; +use pollable_map::optional::Optional; +use std::sync::atomic::AtomicBool; +use std::sync::{Arc, Weak}; + +struct ScopeState<'scope> { + inbox: Mutex>>, + waker: AtomicWaker, +} + +/// A scope within which tasks can be spawned that borrow from the enclosing +/// stack frame. +pub struct Scope<'scope, 'env: 'scope> { + state: Weak>, + _scope: PhantomData<&'scope mut &'scope ()>, + _env: PhantomData<&'env mut &'env ()>, +} + +impl<'scope, 'env> Scope<'scope, 'env> { + fn new() -> Self { + Self { + state: Weak::new(), + _scope: PhantomData, + _env: PhantomData, + } + } + + fn push(&self, task: BoxFuture<'scope, ()>) { + if let Some(state) = self.state.upgrade() { + state.inbox.lock().push(task); + state.waker.wake(); + } + } + + /// Spawn a task into this scope. + /// + /// The future may borrow any data that outlives its lifetime `'env`. The task will + /// be polled cooperatively alongside the scope's user closure and any + /// other spawned tasks. + pub fn spawn(&'scope self, fut: Fut) -> ScopedJoinHandle + where + Fut: Future + Send + 'scope, + Fut::Output: Send + 'scope, + { + let (tx, rx) = oneshot::channel(); + let wrapped: BoxFuture<'scope, ()> = async move { + let output = fut.await; + // If the receiver was dropped, the caller doesn't care about + // the output so we will discard it. + let _ = tx.send(output); + } + .boxed(); + + self.push(wrapped); + + ScopedJoinHandle { rx } + } + + /// Spawn a task into this scope and return an [`AbortableJoinHandle`]. + /// + /// This mirrors [`Executor::spawn_abortable`] for cooperatively + /// scheduled tasks. The returned handle aborts the task when all + /// references to it have been dropped. + pub fn spawn_abortable(&'scope self, fut: Fut) -> AbortableJoinHandle + where + Fut: Future + Send + 'scope, + Fut::Output: Send + 'scope, + { + let (abort_handle, abort_reg) = AbortHandle::new_pair(); + let abortable = Abortable::new(fut, abort_reg); + let (tx, rx) = oneshot::channel(); + let finished = Arc::new(AtomicBool::new(false)); + let completion = CompletionGuard::new(finished.clone()); + + let wrapped: BoxFuture<'scope, ()> = async move { + let _completion = completion; + let val = abortable.await; + let _ = tx.send(val); + } + .boxed(); + self.push(wrapped); + + let join = JoinHandle { + inner: InnerJoinHandle::CustomHandle { + inner: Optional::new(rx), + handle: abort_handle, + finished, + }, + }; + AbortableJoinHandle::from(join) + } + + /// Spawn a task into this scope without keeping a handle to it. + /// + /// Equivalent to [`Executor::dispatch`] for scoped tasks. + pub fn dispatch(&'scope self, fut: Fut) + where + Fut: Future + Send + 'scope, + Fut::Output: Send + 'scope, + { + drop(self.spawn(fut)); + } + + /// Spawn a message-driven coroutine into this scope. + /// + /// Equivalent to [`Executor::spawn_coroutine`] for scoped tasks. + pub fn spawn_coroutine(&'scope self, f: F) -> CommunicationTask + where + F: FnMut(T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + T: Send + 'scope, + { + self.spawn_coroutine_with_buffer(1, f) + } + + /// Like [`Scope::spawn_coroutine`] but with a configurable channel + /// buffer. + pub fn spawn_coroutine_with_buffer( + &'scope self, + buffer: usize, + mut f: F, + ) -> CommunicationTask + where + F: FnMut(T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + T: Send + 'scope, + { + let (tx, mut rx) = futures::channel::mpsc::channel(buffer); + let task_handle = self.spawn_abortable(async move { + while let Some(message) = rx.next().await { + f(message).await; + } + }); + CommunicationTask::new(task_handle, tx) + } + + /// Spawn an unbounded message-driven coroutine into this scope. + /// + /// Equivalent to [`Executor::spawn_unbounded_coroutine`] for scoped + /// tasks. + pub fn spawn_unbounded_coroutine( + &'scope self, + mut f: F, + ) -> UnboundedCommunicationTask + where + F: FnMut(T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + T: Send + 'scope, + { + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + let task_handle = self.spawn_abortable(async move { + while let Some(message) = rx.next().await { + f(message).await; + } + }); + UnboundedCommunicationTask::new(task_handle, tx) + } + + /// Spawn a message-driven coroutine with caller-provided context. + /// + /// If the context must be borrowed across awaits, use + /// [`Scope::spawn_coroutine_with_receiver_and_context`]. + pub fn spawn_coroutine_with_context( + &'scope self, + context: C, + f: F, + ) -> CommunicationTask + where + F: FnMut(&mut C, T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + C: Send + 'scope, + T: Send + 'scope, + { + self.spawn_coroutine_with_buffer_and_context(context, 1, f) + } + + /// Like [`Scope::spawn_coroutine_with_context`] but with a configurable + /// channel buffer. + pub fn spawn_coroutine_with_buffer_and_context( + &'scope self, + context: C, + buffer: usize, + mut f: F, + ) -> CommunicationTask + where + F: FnMut(&mut C, T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + C: Send + 'scope, + T: Send + 'scope, + { + let (tx, mut rx) = futures::channel::mpsc::channel(buffer); + let task_handle = self.spawn_abortable(async move { + let mut context = context; + while let Some(message) = rx.next().await { + f(&mut context, message).await; + } + }); + CommunicationTask::new(task_handle, tx) + } + + /// Spawn an unbounded message-driven coroutine with caller-provided + /// context. + pub fn spawn_unbounded_coroutine_with_context( + &'scope self, + context: C, + mut f: F, + ) -> UnboundedCommunicationTask + where + F: FnMut(&mut C, T) -> Fut + Send + 'scope, + Fut: Future + Send + 'scope, + C: Send + 'scope, + T: Send + 'scope, + { + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + let task_handle = self.spawn_abortable(async move { + let mut context = context; + while let Some(message) = rx.next().await { + f(&mut context, message).await; + } + }); + UnboundedCommunicationTask::new(task_handle, tx) + } + + /// Spawn a coroutine that receives the bounded channel directly. + pub fn spawn_coroutine_with_receiver(&'scope self, f: F) -> CommunicationTask + where + F: FnMut(Receiver) -> Fut, + Fut: Future + Send + 'scope, + { + self.spawn_coroutine_with_receiver_and_buffer(1, f) + } + + /// Like [`Scope::spawn_coroutine_with_receiver`] but with a configurable + /// channel buffer. + pub fn spawn_coroutine_with_receiver_and_buffer( + &'scope self, + buffer: usize, + mut f: F, + ) -> CommunicationTask + where + F: FnMut(Receiver) -> Fut, + Fut: Future + Send + 'scope, + { + let (tx, rx) = futures::channel::mpsc::channel(buffer); + let task_handle = self.spawn_abortable(f(rx)); + CommunicationTask::new(task_handle, tx) + } + + /// Spawn a coroutine that receives caller-provided context and the + /// bounded channel directly. + pub fn spawn_coroutine_with_receiver_and_context( + &'scope self, + context: C, + f: F, + ) -> CommunicationTask + where + F: FnMut(C, Receiver) -> Fut, + Fut: Future + Send + 'scope, + { + self.spawn_coroutine_with_receiver_buffer_and_context(context, 1, f) + } + + /// Like [`Scope::spawn_coroutine_with_receiver_and_context`] but with a + /// configurable channel buffer. + pub fn spawn_coroutine_with_receiver_buffer_and_context( + &'scope self, + context: C, + buffer: usize, + mut f: F, + ) -> CommunicationTask + where + F: FnMut(C, Receiver) -> Fut, + Fut: Future + Send + 'scope, + { + let (tx, rx) = futures::channel::mpsc::channel(buffer); + let task_handle = self.spawn_abortable(f(context, rx)); + CommunicationTask::new(task_handle, tx) + } + + /// Spawn a coroutine that receives the unbounded channel directly. + pub fn spawn_unbounded_coroutine_with_receiver( + &'scope self, + mut f: F, + ) -> UnboundedCommunicationTask + where + F: FnMut(UnboundedReceiver) -> Fut, + Fut: Future + Send + 'scope, + { + let (tx, rx) = futures::channel::mpsc::unbounded(); + let task_handle = self.spawn_abortable(f(rx)); + UnboundedCommunicationTask::new(task_handle, tx) + } + + /// Spawn a coroutine that receives caller-provided context and the + /// unbounded channel directly. + pub fn spawn_unbounded_coroutine_with_receiver_and_context( + &'scope self, + context: C, + mut f: F, + ) -> UnboundedCommunicationTask + where + F: FnMut(C, UnboundedReceiver) -> Fut, + Fut: Future + Send + 'scope, + { + let (tx, rx) = futures::channel::mpsc::unbounded(); + let task_handle = self.spawn_abortable(f(context, rx)); + UnboundedCommunicationTask::new(task_handle, tx) + } +} + +/// Drive a scope once: absorb any inbox entries, then poll the active +/// set until no more immediate progress can be made. +fn drive_scope<'scope>( + active: &mut FuturesUnordered>, + state: &ScopeState<'scope>, + cx: &mut Context<'_>, +) -> (bool, bool) { + let mut made_progress = false; + loop { + state.waker.register(cx.waker()); + + // Take the queued tasks while holding the lock only long enough to + // swap the inbox, never while polling user code. + let incoming = std::mem::take(&mut *state.inbox.lock()); + active.extend(incoming); + match active.poll_next_unpin(cx) { + Poll::Ready(Some(())) => made_progress = true, + Poll::Ready(None) => { + state.waker.register(cx.waker()); + if state.inbox.lock().is_empty() { + return (made_progress, true); + } + } + Poll::Pending => { + state.waker.register(cx.waker()); + if state.inbox.lock().is_empty() { + return (made_progress, false); + } + } + } + } +} + +/// A handle to a task spawned on a [`Scope`]. +/// +/// Awaiting the handle yields the task's output. If the scope is dropped +/// before the task finishes (for example, because the scope future was +/// cancelled), awaiting yields [`JoinError::Cancelled`]. +/// +/// The handle does not borrow the spawned future itself. It can therefore +/// outlive the scope when its output type is also `'static`; borrowed output +/// types retain their normal lifetime restrictions. +pub struct ScopedJoinHandle { + rx: oneshot::Receiver, +} + +/// Error returned when a scoped task was cancelled before it produced an output +#[derive(Debug)] +pub enum JoinError { + /// The task was cancelled before it produced an output. + Cancelled, +} + +impl core::fmt::Display for JoinError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + JoinError::Cancelled => f.write_str("scoped task was cancelled"), + } + } +} + +impl core::error::Error for JoinError {} + +impl Future for ScopedJoinHandle { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match Pin::new(&mut self.rx).poll(cx) { + Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), + Poll::Ready(Err(_)) => Poll::Ready(Err(JoinError::Cancelled)), + Poll::Pending => Poll::Pending, + } + } +} + +/// Create a scope for spawning tasks that borrow from the enclosing stack +/// frame. +/// +/// The provided async closure receives a reference to the [`Scope`] and +/// returns a future. That future is driven alongside any tasks spawned on +/// the scope; when it completes, any still-running spawned tasks are +/// drained before `scope` returns its value. +/// +/// # Example +/// +/// ```no_run +/// # async fn run() { +/// use async_rt::scoped::scope; +/// +/// let data = vec![1, 2, 3, 4]; +/// let data = &data; +/// let sum = scope(async |s| { +/// let a = s.spawn(async { data[0] + data[1] }); +/// let b = s.spawn(async { data[2] + data[3] }); +/// a.await.unwrap() + b.await.unwrap() +/// }) +/// .await; +/// assert_eq!(sum, 10); +/// # } +/// ``` +pub async fn scope<'env, F, T>(f: F) -> T +where + F: for<'scope> AsyncFnOnce(&'scope Scope<'scope, 'env>) -> T, +{ + // Declaration order is part of the safety invariant: `active` is dropped + // first, then `state` (and its queued tasks), and finally `scope`. + let mut scope = Scope::new(); + let state = Arc::new(ScopeState { + inbox: Mutex::new(Vec::new()), + waker: AtomicWaker::new(), + }); + scope.state = Arc::downgrade(&state); + // The active task set lives on the driver, not inside `Scope`. Only + // `drive_scope` touches it, so no lock guards it, and the inbox mutex + // is enough for the spawn side. + let mut active = FuturesUnordered::new(); + + let result = { + let user_fut = f(&scope); + let mut user_fut = std::pin::pin!(user_fut); + + poll_fn(|cx| { + loop { + if let Poll::Ready(r) = user_fut.as_mut().poll(cx) { + return Poll::Ready(r); + } + let (made_progress, _empty) = drive_scope(&mut active, &state, cx); + if !made_progress { + return Poll::Pending; + } + } + }) + .await + }; + + // Drain any remaining spawned tasks before the scope goes away, so + // every borrow of the data is released before this frame unwinds. + poll_fn(|cx| { + let (_made_progress, empty) = drive_scope(&mut active, &state, cx); + if empty { + Poll::Ready(()) + } else { + Poll::Pending + } + }) + .await; + + result +} + +/// An [`Executor`] wrapper that tracks spawned tasks so they can be +/// canceled when a scope ends. +/// +/// `ScopeExecutor` itself implements [`Executor`], so it can be passed to +/// any code that takes `Executor` every spawn performed through +/// that code will be tracked by the scope. +/// +/// Because submitted futures must still be `Send + 'static` (enforced by +/// [`Executor::spawn`]), this flavour of scope does *not* allow borrowing +/// from the enclosing stack frame. Use [`scope`] for that. +pub struct ScopeExecutor<'scope, E> { + inner: &'scope E, + task_handles: Mutex>>, + _scope: PhantomData<&'scope mut &'scope ()>, +} + +impl<'scope, E> ScopeExecutor<'scope, E> { + fn new(inner: &'scope E) -> Self { + Self { + inner, + task_handles: Mutex::new(Vec::new()), + _scope: PhantomData, + } + } + + /// Abort every task currently tracked by this scope. + fn abort_all(&self) { + for handle in self.task_handles.lock().iter() { + handle.abort(); + } + } +} + +impl Drop for ScopeExecutor<'_, E> { + fn drop(&mut self) { + self.abort_all(); + } +} + +impl<'scope, E> Executor for ScopeExecutor<'scope, E> +where + E: Executor, +{ + fn spawn(&self, future: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let (abort_handle, abort_registration) = AbortHandle::new_pair(); + let abortable = Abortable::new(future, abort_registration); + let (tx, rx) = oneshot::channel(); + let finished = Arc::new(AtomicBool::new(false)); + let completion = CompletionGuard::new(finished.clone()); + let wrapped = async move { + let _completion = completion; + let val = abortable.await; + let _ = tx.send(val); + }; + + // Track an abort-on-drop handle so cancellation remains effective even + // after the handles are drained from `ScopeExecutor` for joining. + let task_handle: AbortableJoinHandle<()> = self.inner.spawn(wrapped).into(); + self.task_handles.lock().push(task_handle); + + JoinHandle { + inner: InnerJoinHandle::CustomHandle { + inner: Optional::new(rx), + handle: abort_handle, + finished, + }, + } + } +} + +/// Run an async closure with a scoped [`Executor`] wrapper. +/// +/// If the future containing `executor_scope` itself is dropped +/// (canceled externally), it will abort every outstanding task. +/// +/// Unlike [`scope`], futures spawned here must be `Send + 'static` +/// they're going onto the real executor, so they can't borrow from the +/// enclosing stack frame. +/// +/// # Panics in spawned tasks +/// +/// Panics inside spawned tasks are **not** propagated to the scope. +/// Unlike [`std::thread::scope`], `executor_scope` does not re-raise +/// panics collected from its tasks: the panic is caught at the task +/// boundary by the underlying runtime, and the scope simply treats the +/// task as finished. +/// +/// If you need to react to a task panic, poll or `.await` the returned +/// [`JoinHandle`] yourself and check its result and don't rely on the +/// scope to surface it. +/// +/// # Example +/// +/// ```no_run +/// # async fn run() { +/// use async_rt::Executor; +/// use async_rt::rt::tokio::TokioExecutor; +/// +/// let executor = TokioExecutor; +/// let total = executor +/// .executor_scope(async |s| { +/// let a = s.spawn(async { 1 + 2 }); +/// let b = s.spawn(async { 3 + 4 }); +/// a.await.unwrap() + b.await.unwrap() +/// }) +/// .await; +/// assert_eq!(total, 10); +/// # } +/// ``` +pub async fn executor_scope<'scope, E, F, T>(executor: &'scope E, f: F) -> T +where + E: Executor, + F: AsyncFnOnce(&ScopeExecutor<'scope, E>) -> T, +{ + let scope_exec = ScopeExecutor::new(executor); + let result = f(&scope_exec).await; + + // Wait for every spawned task to complete + let handles: Vec<_> = scope_exec.task_handles.lock().drain(..).collect(); + for handle in handles { + let _ = handle.await; + } + + result +} + +#[cfg(test)] +mod tests { + use super::*; + #[cfg(feature = "tokio")] + use futures_timer::Delay; + #[cfg(feature = "tokio")] + use std::time::Duration; + + #[tokio::test] + async fn borrows_stack_data() { + let data = vec![1, 2, 3, 4]; + let data = &data; + let sum = scope(async |s: &Scope<'_, '_>| { + let a = s.spawn(async move { data[0] + data[1] }); + let b = s.spawn(async move { data[2] + data[3] }); + a.await.unwrap() + b.await.unwrap() + }) + .await; + assert_eq!(sum, 10); + } + + #[tokio::test] + async fn drains_unawaited_tasks() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let counter = AtomicUsize::new(0); + let counter_ref = &counter; + scope(async |s: &Scope<'_, '_>| { + for _ in 0..8 { + s.spawn(async move { + counter_ref.fetch_add(1, Ordering::SeqCst); + }); + } + }) + .await; + assert_eq!(counter.load(Ordering::SeqCst), 8); + } + + #[tokio::test] + async fn returns_closure_value() { + let v: i32 = scope(async |_s: &Scope<'_, '_>| 42).await; + assert_eq!(v, 42); + } + + #[tokio::test] + async fn join_handle_yields_output() { + let out = scope(async |s: &Scope<'_, '_>| { + let h = s.spawn(async { "hello" }); + h.await.unwrap() + }) + .await; + assert_eq!(out, "hello"); + } + + #[tokio::test] + async fn abortable_handle_reports_completion_without_being_polled() { + scope(async |s: &Scope<'_, '_>| { + let handle = s.spawn_abortable(async {}); + + // Yield the user future so the cooperative driver can complete + // the child without polling its join handle. + crate::task::yield_now().await; + + assert!(handle.is_finished()); + }) + .await; + } + + #[test] + fn pushing_task_wakes_scope_driver() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::{Wake, Waker}; + + struct WakeCounter(AtomicUsize); + + impl Wake for WakeCounter { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let mut scope = Scope::new(); + let state = Arc::new(ScopeState { + inbox: Mutex::new(Vec::new()), + waker: AtomicWaker::new(), + }); + scope.state = Arc::downgrade(&state); + + let wake_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(wake_counter.clone()); + state.waker.register(&waker); + + let scope_ref = &scope; + scope.push( + async move { + // Consume the driver's registered waker while tasks are + // being polled, forcing drive_scope to register it again. + scope_ref.push(async {}.boxed()); + } + .boxed(), + ); + + assert_eq!(state.inbox.lock().len(), 1); + assert_eq!(wake_counter.0.load(Ordering::SeqCst), 1); + + let mut active = FuturesUnordered::new(); + let mut context = Context::from_waker(&waker); + let (_, empty) = drive_scope(&mut active, &state, &mut context); + assert!(empty); + assert_eq!(wake_counter.0.load(Ordering::SeqCst), 2); + + // The nested push consumed the previous registration. This push + // only wakes the driver if drive_scope registered again afterward. + scope.push(async {}.boxed()); + assert_eq!(wake_counter.0.load(Ordering::SeqCst), 3); + + let (_, empty) = drive_scope(&mut active, &state, &mut context); + assert!(empty); + + let state_ref = Arc::downgrade(&state); + scope.push( + poll_fn(move |_cx| { + // Model a push that occurs after registration but before + // the inbox is drained: its wake is consumed, then the + // newly active task returns Pending. + state_ref.upgrade().unwrap().waker.wake(); + Poll::<()>::Pending + }) + .boxed(), + ); + assert_eq!(wake_counter.0.load(Ordering::SeqCst), 4); + + let (_, empty) = drive_scope(&mut active, &state, &mut context); + assert!(!empty); + let wakes_after_pending = wake_counter.0.load(Ordering::SeqCst); + assert!(wakes_after_pending > 4); + + // The Pending path must have registered once more before returning. + scope.push(async {}.boxed()); + assert_eq!( + wake_counter.0.load(Ordering::SeqCst), + wakes_after_pending + 1 + ); + } + + #[tokio::test] + async fn many_concurrent_tasks_complete() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let counter = AtomicUsize::new(0); + let counter_ref = &counter; + let total: usize = scope(async |s: &Scope<'_, '_>| { + let handles: Vec<_> = (0..32) + .map(|i| { + s.spawn(async move { + counter_ref.fetch_add(1, Ordering::SeqCst); + i + }) + }) + .collect(); + let mut sum = 0usize; + for h in handles { + sum += h.await.unwrap(); + } + sum + }) + .await; + assert_eq!(total, (0..32).sum()); + assert_eq!(counter.load(Ordering::SeqCst), 32); + } + + #[tokio::test] + async fn child_task_can_spawn_nested_task() { + let result = scope(async |s: &Scope<'_, '_>| { + let outer = s.spawn(async move { + let inner = s.spawn(async { 41usize }); + inner.await.unwrap() + 1 + }); + outer.await.unwrap() + }) + .await; + + assert_eq!(result, 42); + } + + #[tokio::test] + async fn drains_unawaited_nested_task() { + use std::sync::atomic::{AtomicBool, Ordering}; + + let nested_ran = AtomicBool::new(false); + let nested_ran_ref = &nested_ran; + + scope(async |s: &Scope<'_, '_>| { + s.dispatch(async move { + s.dispatch(async move { + nested_ran_ref.store(true, Ordering::SeqCst); + }); + }); + }) + .await; + + assert!(nested_ran.load(Ordering::SeqCst)); + } + + #[cfg(feature = "tokio")] + #[tokio::test] + async fn executor_scope_runs_tasks() { + use crate::rt::tokio::TokioExecutor; + let executor = TokioExecutor; + let total = executor + .executor_scope(async |s| { + let a = s.spawn(async { 1 + 2 }); + let b = s.spawn(async { 3 + 4 }); + a.await.unwrap() + b.await.unwrap() + }) + .await; + assert_eq!(total, 10); + } + + #[cfg(feature = "tokio")] + #[tokio::test(flavor = "current_thread")] + async fn executor_scope_handle_reports_completion_without_being_polled() { + use crate::rt::tokio::TokioExecutor; + + let executor = TokioExecutor; + executor + .executor_scope(async |s| { + let (completed_tx, completed_rx) = oneshot::channel(); + let handle = s.spawn(async move { + let _ = completed_tx.send(()); + }); + + completed_rx.await.unwrap(); + + // On the current-thread runtime, the spawned wrapper finishes + // before the task it woke can be polled again. + assert!(handle.is_finished()); + }) + .await; + } + + #[tokio::test] + async fn scope_spawn_coroutine_receives_messages() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let total = AtomicUsize::new(0); + let total_ref = &total; + + scope(async |s: &Scope<'_, '_>| { + let mut task = s.spawn_coroutine(|value| async move { + total_ref.fetch_add(value, Ordering::SeqCst); + }); + for v in [1usize, 2, 3, 4] { + task.send(v).await.unwrap(); + } + drop(task); // closes channel → coroutine exits cleanly + }) + .await; + + assert_eq!(total.load(Ordering::SeqCst), 10); + } + + #[tokio::test] + async fn scope_receiver_coroutine_receives_messages() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let total = AtomicUsize::new(0); + let total_ref = &total; + + scope(async |s: &Scope<'_, '_>| { + let mut task = s.spawn_coroutine_with_receiver(|mut rx| async move { + while let Some(value) = rx.next().await { + total_ref.fetch_add(value, Ordering::SeqCst); + } + }); + for value in [1usize, 2, 3, 4] { + task.send(value).await.unwrap(); + } + drop(task); + }) + .await; + + assert_eq!(total.load(Ordering::SeqCst), 10); + } + + #[tokio::test] + async fn scope_coroutine_api_matches_executor() { + use futures::future::ready; + + scope(async |s: &Scope<'_, '_>| { + let task = s.spawn_coroutine_with_buffer(2, |_value: usize| ready(())); + drop(task); + + let task = s.spawn_unbounded_coroutine(|_value: usize| ready(())); + drop(task); + + let task = + s.spawn_coroutine_with_context(0usize, |context: &mut usize, value: usize| { + *context += value; + ready(()) + }); + drop(task); + + let task = s.spawn_coroutine_with_buffer_and_context( + 0usize, + 2, + |context: &mut usize, value: usize| { + *context += value; + ready(()) + }, + ); + drop(task); + + let task = s.spawn_unbounded_coroutine_with_context( + 0usize, + |context: &mut usize, value: usize| { + *context += value; + ready(()) + }, + ); + drop(task); + + let task = + s.spawn_coroutine_with_receiver_and_buffer(2, |_rx: Receiver| async {}); + drop(task); + + let task = s.spawn_coroutine_with_receiver_and_context( + 0usize, + |_context, _rx: Receiver| async {}, + ); + drop(task); + + let task = s.spawn_coroutine_with_receiver_buffer_and_context( + 0usize, + 2, + |_context, _rx: Receiver| async {}, + ); + drop(task); + + let task = + s.spawn_unbounded_coroutine_with_receiver(|_rx: UnboundedReceiver| async {}); + drop(task); + + let task = s.spawn_unbounded_coroutine_with_receiver_and_context( + 0usize, + |_context, _rx: UnboundedReceiver| async {}, + ); + drop(task); + }) + .await; + } + + #[tokio::test] + async fn scope_dispatch_runs_fire_and_forget() { + use std::sync::atomic::{AtomicBool, Ordering}; + let flag = AtomicBool::new(false); + let flag_ref = &flag; + + scope(async |s: &Scope<'_, '_>| { + s.dispatch(async move { + flag_ref.store(true, Ordering::SeqCst); + }); + }) + .await; + + assert!(flag.load(Ordering::SeqCst)); + } + + #[cfg(feature = "tokio")] + #[tokio::test] + async fn executor_scope_drains_unawaited_tasks() { + use crate::rt::tokio::TokioExecutor; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + let executor = TokioExecutor; + let flag = Arc::new(AtomicBool::new(false)); + let flag_clone = flag.clone(); + + executor + .executor_scope(async move |s| { + // Spawn a task but do NOT await its handle. + // executor_scope should wait for it to complete before + // returning. + let _h = s.spawn(async move { + Delay::new(Duration::from_millis(50)).await; + flag_clone.store(true, Ordering::SeqCst); + }); + }) + .await; + + assert!( + flag.load(Ordering::SeqCst), + "unawaited task should have completed before executor_scope returned" + ); + } + + #[cfg(feature = "tokio")] + #[tokio::test] + async fn executor_scope_swallows_task_panic() { + use crate::rt::tokio::TokioExecutor; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let executor = TokioExecutor; + let sibling_done = Arc::new(AtomicUsize::new(0)); + let sibling_done_clone = sibling_done.clone(); + + let result = executor + .executor_scope(async move |s| { + // Task A: panics. Its JoinHandle is dropped without + // being awaited. the panic goes to the runtime. + let _panicker = s.spawn(async { + panic!("deliberate test panic"); + }); + // Task B: completes normally. The scope should still + // drain it before returning. + let _sibling = s.spawn(async move { + sibling_done_clone.fetch_add(1, Ordering::SeqCst); + }); + 42usize + }) + .await; + + assert_eq!(result, 42); + assert_eq!(sibling_done.load(Ordering::SeqCst), 1); + } + + #[cfg(feature = "tokio")] + #[tokio::test] + async fn executor_scope_aborts_on_external_cancel() { + use crate::rt::tokio::TokioExecutor; + use futures::future::Either; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + let executor = TokioExecutor; + let flag = Arc::new(AtomicBool::new(false)); + let flag_clone = flag.clone(); + + { + let scope_fut = executor.executor_scope(async move |s| { + let _h = s.spawn(async move { + Delay::new(Duration::from_millis(200)).await; + flag_clone.store(true, Ordering::SeqCst); + }); + futures::future::pending::<()>().await; + }); + + let scope_fut = std::pin::pin!(scope_fut); + let timer = std::pin::pin!(Delay::new(Duration::from_millis(30))); + let result = futures::future::select(scope_fut, timer).await; + assert!( + matches!(result, Either::Right(_)), + "timer should have won the race" + ); + } + + // Give the (supposedly aborted) task plenty of time. + Delay::new(Duration::from_millis(300)).await; + assert!( + !flag.load(Ordering::SeqCst), + "task should have been aborted by scope drop" + ); + } + + #[cfg(feature = "tokio")] + #[tokio::test] + async fn executor_scope_aborts_when_cancelled_during_drain() { + use crate::rt::tokio::TokioExecutor; + use futures::future::{Either, pending, select}; + + let executor = TokioExecutor; + let (started_tx, started_rx) = oneshot::channel(); + let (held_tx, held_rx) = oneshot::channel::<()>(); + + let scope_fut = Box::pin(executor.executor_scope(async move |s| { + let _handle = s.spawn(async move { + let _held_until_task_drop = held_tx; + let _ = started_tx.send(()); + pending::<()>().await; + }); + })); + + let scope_fut = match select(scope_fut, started_rx).await { + Either::Right((Ok(()), scope_fut)) => scope_fut, + Either::Left(_) => panic!("scope unexpectedly completed"), + Either::Right((Err(_), _)) => panic!("child task never started"), + }; + + drop(scope_fut); + + match select(held_rx, Delay::new(Duration::from_secs(1))).await { + Either::Left((Err(_), _)) => {} + Either::Left((Ok(_), _)) => unreachable!("child never sends a value"), + Either::Right(_) => panic!("child remained detached after scope cancellation"), + } + } +}