gpui_ce_scheduler 0.3.0

use crate::{Instant, Priority, RunnableMeta, Scheduler, SessionId, Timer};
use async_task::Runnable;
use std::{
    any::Any,
    collections::HashMap,
    future::Future,
    marker::PhantomData,
    mem::ManuallyDrop,
    panic::Location,
    pin::Pin,
    rc::Rc,
    sync::{
        Arc, Mutex, Weak,
        atomic::{AtomicBool, Ordering},
    },
    task::{Context, Poll, Waker},
    thread::{self, ThreadId},
    time::Duration,
};

/// The type-erased closure that [`Scheduler::spawn_dedicated`] runs on a new
/// session.
///
/// It receives the session's [`LocalExecutor`] and returns a future whose
/// output is boxed as `Box<dyn Any + Send + Sync>`, which keeps the trait
/// object-safe. Call [`LocalExecutor::spawn_dedicated`] or
/// [`BackgroundExecutor::spawn_dedicated`] instead of building one by hand.
pub type DedicatedFn = Box<
    dyn FnOnce(LocalExecutor) -> Pin<Box<dyn Future<Output = Box<dyn Any + Send + Sync>> + 'static>>
        + Send
        + 'static,
>;

/// Tracks local tasks that a platform must cancel on their owning thread
/// before retiring a thread-bound event loop.
#[doc(hidden)]
pub struct LocalTaskScope {
    retiring: AtomicBool,
    state: Mutex<LocalTaskState>,
}

struct LocalTaskState {
    next_task_id: u64,
    tasks: HashMap<u64, Option<Waker>>,
    sealed: bool,
}

impl LocalTaskScope {
    /// Creates an empty owner-thread task scope.
    pub fn new() -> Self {
        Self {
            retiring: AtomicBool::new(false),
            state: Mutex::new(LocalTaskState {
                next_task_id: 0,
                tasks: HashMap::new(),
                sealed: false,
            }),
        }
    }

    /// Whether retirement has started.
    pub fn is_retiring(&self) -> bool {
        self.retiring.load(Ordering::Acquire)
    }

    /// Returns the number of local futures that have not yet been destroyed.
    pub fn live_task_count(&self) -> usize {
        self.state
            .lock()
            .expect("local task scope poisoned")
            .tasks
            .len()
    }

    /// Latches retirement and wakes registered tasks without holding the
    /// registry lock. The platform must drop their runnables on the owner
    /// thread before treating this scope as quiescent.
    pub fn begin_retirement(&self) {
        if self.retiring.swap(true, Ordering::AcqRel) {
            return;
        }
        let wakers = {
            let mut state = self.state.lock().expect("local task scope poisoned");
            state
                .tasks
                .values_mut()
                .filter_map(Option::take)
                .collect::<Vec<_>>()
        };
        for waker in wakers {
            waker.wake();
        }
    }

    /// Seals task admission after retirement when every registered future is destroyed.
    ///
    /// A platform calls this only after its owner-thread queues are empty. The task registry
    /// lock makes the empty check atomic with later spawn registration.
    #[doc(hidden)]
    pub fn seal_if_quiescent(&self) -> bool {
        let mut state = self.state.lock().expect("local task scope poisoned");
        if !self.is_retiring() || !state.tasks.is_empty() {
            return false;
        }
        state.sealed = true;
        true
    }

    fn register(&self) -> u64 {
        let mut state = self.state.lock().expect("local task scope poisoned");
        assert!(
            !state.sealed,
            "cannot spawn a local task after its owner scope has been sealed"
        );
        let id = state.next_task_id;
        state.next_task_id = state
            .next_task_id
            .checked_add(1)
            .expect("local task scope task ID overflow");
        state.tasks.insert(id, None);
        id
    }

    fn attach_waker(&self, id: u64, waker: Waker) {
        let mut waker = Some(waker);
        let replaced = {
            let mut state = self.state.lock().expect("local task scope poisoned");
            state
                .tasks
                .get_mut(&id)
                .and_then(|entry| entry.replace(waker.take().unwrap()))
        };
        drop(replaced);
        drop(waker);
    }

    fn unregister(&self, id: u64) {
        let removed = self
            .state
            .lock()
            .expect("local task scope poisoned")
            .tasks
            .remove(&id);
        drop(removed);
    }
}

impl Default for LocalTaskScope {
    fn default() -> Self {
        Self::new()
    }
}

struct LocalTaskRegistration {
    scope: Weak<LocalTaskScope>,
    id: u64,
}

impl Drop for LocalTaskRegistration {
    fn drop(&mut self) {
        if let Some(scope) = self.scope.upgrade() {
            scope.unregister(self.id);
        }
    }
}

struct RegisteredLocalFuture<F> {
    future: F,
    // Rust drops fields in declaration order: the user future is destroyed
    // before the registration is removed from its scope.
    _registration: LocalTaskRegistration,
}

impl<F: Future> Future for RegisteredLocalFuture<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // The future remains pinned in its original field.
        unsafe { self.map_unchecked_mut(|this| &mut this.future) }.poll(cx)
    }
}

/// An executor for one session, for futures that are not `Send`.
///
/// Tasks spawned on a `LocalExecutor` run in order on whichever thread drains
/// the dispatch destination it was built with. That is the main thread for the
/// default session, or a dedicated thread for a session created with
/// [`LocalExecutor::spawn_dedicated`]. Futures spawned here are polled and
/// dropped on the thread that spawned them.
#[derive(Clone)]
pub struct LocalExecutor {
    session_id: SessionId,
    scheduler: Arc<dyn Scheduler>,
    // Spawned tasks' schedule callbacks each hold an `Arc` clone of this
    // closure, so the destination it captures stays alive as long as work
    // could still land on it.
    dispatch: Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync>,
    retirement: Option<Arc<LocalTaskScope>>,
    dispatch_cancellation: Option<Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync>>,
    not_send: PhantomData<Rc<()>>,
}

impl LocalExecutor {
    /// Creates a local executor for `session_id` that sends each spawned task's
    /// runnables to `dispatch`.
    ///
    /// For the default session, `dispatch` usually calls
    /// [`Scheduler::schedule_local`]. The dedicated-session path passes a
    /// closure that sends to the dedicated thread's channel. `scheduler` gives
    /// the executor access to clocks, timers and the other scheduler services.
    pub fn new(
        session_id: SessionId,
        scheduler: Arc<dyn Scheduler>,
        dispatch: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
    ) -> Self {
        Self {
            session_id,
            scheduler,
            dispatch: Arc::new(dispatch),
            retirement: None,
            dispatch_cancellation: None,
            not_send: PhantomData,
        }
    }

    /// Creates a local executor whose tasks are tracked by a platform-owned
    /// retirement scope. Cancellation dispatch must enqueue runnables for
    /// owner-thread destruction and must not invoke user code inline.
    #[doc(hidden)]
    pub fn new_with_scope(
        session_id: SessionId,
        scheduler: Arc<dyn Scheduler>,
        scope: Arc<LocalTaskScope>,
        dispatch: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
        dispatch_cancellation: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
    ) -> Self {
        Self {
            session_id,
            scheduler,
            dispatch: Arc::new(dispatch),
            retirement: Some(scope),
            dispatch_cancellation: Some(Arc::new(dispatch_cancellation)),
            not_send: PhantomData,
        }
    }

    /// Returns this executor's platform retirement scope, if present.
    #[doc(hidden)]
    pub fn retirement_scope(&self) -> Option<Arc<LocalTaskScope>> {
        self.retirement.clone()
    }

    /// Returns the session this executor runs tasks for.
    pub fn session_id(&self) -> SessionId {
        self.session_id
    }

    /// Returns the scheduler that this executor's timers and clock come from.
    pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
        &self.scheduler
    }

    /// Whether this executor runs on a TestScheduler.
    pub fn is_test(&self) -> bool {
        self.scheduler.as_test().is_some()
    }

    /// Spawns a future on this session and returns a [`Task`] for its output.
    ///
    /// The future doesn't need to be `Send`. It is polled on the thread that
    /// spawned it, so it runs on the session's thread when called from there.
    /// Dropping the returned task cancels the future; call [`Task::detach`] to
    /// let it run to completion instead.
    #[track_caller]
    pub fn spawn<F>(&self, future: F) -> Task<F::Output>
    where
        F: Future + 'static,
        F::Output: 'static,
    {
        let dispatch = self.dispatch.clone();
        self.spawn_with_dispatch(future, move |runnable| dispatch(runnable))
    }

    /// Like [`Self::spawn`], but sends the task's runnables to `dispatch`
    /// instead of this executor's own destination.
    ///
    /// `dispatch` must deliver runnables back to this executor's thread. The
    /// future is polled and dropped on the thread that spawned it.
    #[track_caller]
    pub fn spawn_with_dispatch<F>(
        &self,
        future: F,
        dispatch: impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static,
    ) -> Task<F::Output>
    where
        F: Future + 'static,
        F::Output: 'static,
    {
        let location = Location::caller();
        let metadata = RunnableMeta {
            location,
            spawned: crate::SpawnTime(Instant::now()),
        };
        if let Some(scope) = self.retirement.clone() {
            let id = scope.register();
            let registration = LocalTaskRegistration {
                scope: Arc::downgrade(&scope),
                id,
            };
            let future = RegisteredLocalFuture {
                future,
                _registration: registration,
            };
            let cancellation = self
                .dispatch_cancellation
                .as_ref()
                .expect("scoped executor must have a cancellation dispatcher")
                .clone();
            let normal_dispatch: Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync> =
                Arc::new(dispatch);
            let scope_for_schedule = scope.clone();
            let schedule: Arc<dyn Fn(Runnable<RunnableMeta>) + Send + Sync> =
                Arc::new(move |runnable| {
                    if scope_for_schedule.is_retiring() {
                        cancellation(runnable);
                    } else {
                        normal_dispatch(runnable);
                    }
                });
            let (runnable, task) = spawn_local_with_source_location(
                future,
                move |runnable| schedule(runnable),
                metadata,
            );
            // Registration precedes initial scheduling. If retirement raced
            // with construction, that initial schedule goes to owner cancellation.
            scope.attach_waker(id, runnable.waker());
            runnable.schedule();
            Task(TaskState::Spawned(task))
        } else {
            let (runnable, task) = spawn_local_with_source_location(future, dispatch, metadata);
            runnable.schedule();
            Task(TaskState::Spawned(task))
        }
    }

    /// Blocks the current thread until `future` completes, and returns its output.
    ///
    /// Use this where async code can't be used, such as a test or the entry
    /// point of a session. Not available on wasm, where the main thread can't
    /// block.
    ///
    /// # Panics
    ///
    /// Panics if the scheduler stops before `future` completes, for example
    /// because `future` waits on work that can never run.
    #[cfg(not(target_family = "wasm"))]
    pub fn block_on<Fut: Future>(&self, future: Fut) -> Fut::Output {
        use std::cell::Cell;

        let output = Cell::new(None);
        let future = async {
            output.set(Some(future.await));
        };
        let mut future = std::pin::pin!(future);

        self.scheduler
            .block(Some(self.session_id), future.as_mut(), None);

        output.take().expect("block_on future did not complete")
    }

    /// Blocks the current thread until `future` completes or `timeout` elapses.
    ///
    /// Returns `Ok` with the output if the future completed in time. On timeout
    /// it returns `Err` with the unfinished future. Await that future later to
    /// keep waiting, or drop it to cancel the work. The future may already have
    /// run part of its work.
    #[cfg(not(target_family = "wasm"))]
    pub fn block_with_timeout<Fut: Future>(
        &self,
        timeout: Duration,
        future: Fut,
    ) -> Result<Fut::Output, impl Future<Output = Fut::Output> + use<Fut>> {
        use std::cell::Cell;

        let output = Cell::new(None);
        let mut future = Box::pin(future);

        {
            let future_ref = &mut future;
            let wrapper = async {
                output.set(Some(future_ref.await));
            };
            let mut wrapper = std::pin::pin!(wrapper);

            self.scheduler
                .block(Some(self.session_id), wrapper.as_mut(), Some(timeout));
        }

        match output.take() {
            Some(value) => Ok(value),
            None => Err(future),
        }
    }

    /// Returns a [`Timer`] that completes after `duration`.
    #[track_caller]
    pub fn timer(&self, duration: Duration) -> Timer {
        self.scheduler.timer(duration)
    }

    /// Returns the current time from this scheduler's clock.
    ///
    /// Under the test scheduler this is the simulated clock, which moves only
    /// when the test advances it.
    pub fn now(&self) -> Instant {
        self.scheduler.clock().now()
    }

    /// Spawns a closure on a new session that has its own [`LocalExecutor`].
    ///
    /// The closure runs on a new OS thread under the platform scheduler, or
    /// on the test scheduler's loop under `TestScheduler`. Its future can use
    /// the executor it receives to spawn more work on that thread.
    ///
    /// The returned task is the closure's work. Dropping it cancels the work,
    /// awaiting it yields the closure's return value, and [`Task::detach`]
    /// lets the work run on after the caller is done with it.
    #[track_caller]
    pub fn spawn_dedicated<F, Fut>(&self, f: F) -> Task<Fut::Output>
    where
        F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
        Fut: Future + 'static,
        Fut::Output: Send + Sync + 'static,
    {
        self.scheduler
            .clone()
            .spawn_dedicated(box_dedicated(f))
            .downcast::<Fut::Output>()
    }
}

/// Boxes the user-supplied dedicated closure into the type-erased shape
/// expected by [`Scheduler::spawn_dedicated`]. The user's `Fut::Output` is
/// boxed as `Box<dyn Any + Send + Sync>` on the dedicated side and downcast
/// back to `Fut::Output` by [`Task::downcast`] in the wrapper.
fn box_dedicated<F, Fut>(f: F) -> DedicatedFn
where
    F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
    Fut: Future + 'static,
    Fut::Output: Send + Sync + 'static,
{
    Box::new(move |executor| {
        Box::pin(async move { Box::new(f(executor).await) as Box<dyn Any + Send + Sync> })
    })
}

/// An executor that runs `Send` futures on the background thread pool.
///
/// Applications get one from `background_executor` on their context. Cloning
/// is cheap, and every clone shares the same scheduler.
#[derive(Clone)]
pub struct BackgroundExecutor {
    scheduler: Arc<dyn Scheduler>,
}

impl BackgroundExecutor {
    /// Creates a background executor that spawns onto `scheduler`.
    pub fn new(scheduler: Arc<dyn Scheduler>) -> Self {
        Self { scheduler }
    }

    /// Spawns a future on the background thread pool with the default priority.
    ///
    /// The future may run on any background thread, so it must be [`Send`]. Dropping
    /// the returned [`Task`] cancels the future. Call [`Task::detach`] to let
    /// it finish on its own. Use [`LocalExecutor::spawn`] for futures that
    /// aren't `Send` or that need to stay on the main thread.
    #[track_caller]
    pub fn spawn<F>(&self, future: F) -> Task<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        self.spawn_with_priority(Priority::default(), future)
    }

    /// Spawns a future on the background thread pool at `priority`.
    ///
    /// Like [`Self::spawn`], but the scheduler uses `priority` to order this
    /// future against other background work waiting to run.
    #[track_caller]
    pub fn spawn_with_priority<F>(&self, priority: Priority, future: F) -> Task<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let schedule = self.schedule_with_priority(priority);
        let location = Location::caller();
        let (runnable, task) = async_task::Builder::new()
            .metadata(RunnableMeta {
                location,
                spawned: crate::SpawnTime(Instant::now()),
            })
            .spawn(move |_| future, schedule);
        runnable.schedule();
        Task(TaskState::Spawned(task))
    }

    /// Spawns a future on the platform's realtime thread, for latency-sensitive
    /// work such as audio processing.
    ///
    /// The future must be [`Send`]. Use [`Self::spawn`] for everything else.
    #[track_caller]
    pub fn spawn_realtime<F>(&self, future: F) -> Task<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let location = Location::caller();
        let (tx, rx) = flume::bounded::<async_task::Runnable<RunnableMeta>>(1);

        self.scheduler.spawn_realtime(Box::new(move || {
            while let Ok(runnable) = rx.recv() {
                runnable.run();
            }
        }));

        let (runnable, task) = async_task::Builder::new()
            .metadata(RunnableMeta {
                location,
                spawned: crate::SpawnTime(Instant::now()),
            })
            .spawn(
                move |_| future,
                move |runnable| {
                    let _ = tx.send(runnable);
                },
            );
        runnable.schedule();
        Task(TaskState::Spawned(task))
    }

    /// Returns a [`Timer`] that completes after `duration`.
    #[track_caller]
    pub fn timer(&self, duration: Duration) -> Timer {
        self.scheduler.timer(duration)
    }

    /// Returns the current time from this scheduler's clock.
    ///
    /// Under the test scheduler this is the simulated clock, which moves only
    /// when the test advances it.
    pub fn now(&self) -> Instant {
        self.scheduler.clock().now()
    }

    /// Returns the scheduler that this executor spawns onto.
    pub fn scheduler(&self) -> &Arc<dyn Scheduler> {
        &self.scheduler
    }

    /// Whether this executor runs on a TestScheduler.
    pub fn is_test(&self) -> bool {
        self.scheduler.as_test().is_some()
    }

    /// Spawns a closure on a new session that has its own [`LocalExecutor`].
    ///
    /// The closure runs on a new OS thread under the platform scheduler, or on
    /// the test scheduler's loop under `TestScheduler`. Its future can use
    /// the executor it receives to spawn more work on that thread.
    ///
    /// The returned task is the closure's work. Dropping it cancels the work,
    /// awaiting it yields the closure's return value, and [`Task::detach`]
    /// lets the work run on after the caller is done with it.
    #[track_caller]
    pub fn spawn_dedicated<F, Fut>(&self, f: F) -> Task<Fut::Output>
    where
        F: FnOnce(LocalExecutor) -> Fut + Send + 'static,
        Fut: Future + 'static,
        Fut::Output: Send + Sync + 'static,
    {
        self.scheduler
            .clone()
            .spawn_dedicated(box_dedicated(f))
            .downcast::<Fut::Output>()
    }

    fn schedule_with_priority(
        &self,
        priority: Priority,
    ) -> impl Fn(Runnable<RunnableMeta>) + Send + Sync + 'static {
        let scheduler = Arc::downgrade(&self.scheduler);
        move |runnable| {
            if let Some(scheduler) = scheduler.upgrade() {
                scheduler.schedule_background_with_priority(runnable, priority);
            }
        }
    }
}

/// A handle to one long-lived session. Every future spawned on it runs on that
/// session's single thread, never on the background pool.
///
/// Use this when tasks rely on thread-local state that must stay the same
/// across spawns, or when per-thread setup is expensive enough to pay once.
/// Creating the handle starts a thread (on the web, a worker, which needs the
/// `wasm-threads` feature and is costly). Each later spawn only sends a message
/// to that thread.
///
/// Dropping the handle cancels the session. Futures that haven't run yet are
/// cancelled: awaiting one of their [`Task`]s panics, and a
/// [`FallibleTask`] resolves to `None`.
pub struct DedicatedExecutor {
    sender: flume::Sender<Runnable<RunnableMeta>>,
    _session: Task<()>,
}

impl DedicatedExecutor {
    /// Starts a dedicated session on `executor` and returns a handle that
    /// spawns futures onto it.
    ///
    /// Under `TestScheduler` the session runs on the deterministic test loop,
    /// and no real thread is created.
    #[track_caller]
    pub fn new(executor: &BackgroundExecutor) -> Self {
        let (sender, receiver) = flume::unbounded::<Runnable<RunnableMeta>>();
        let session = executor.spawn_dedicated(move |_executor| async move {
            while let Ok(runnable) = receiver.recv_async().await {
                runnable.run();
            }
        });
        Self {
            sender,
            _session: session,
        }
    }

    /// Spawns a future onto the dedicated session's thread and returns a
    /// [`Task`] for its output.
    ///
    /// The future must be [`Send`], because it moves to the session's thread.
    /// Dropping the task cancels the future, awaiting the task yields its
    /// output, and [`Task::detach`] lets it run to completion on its own.
    #[track_caller]
    pub fn spawn<F>(&self, future: F) -> Task<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let sender = self.sender.clone();
        let (runnable, task) = async_task::Builder::new()
            .metadata(RunnableMeta::new_with_callers_location())
            .spawn(
                move |_| future,
                move |runnable| {
                    let _ = sender.send(runnable);
                },
            );
        runnable.schedule();
        Task(TaskState::Spawned(task))
    }
}

/// A handle to a future that runs in the background, and to its output.
///
/// Executors return a `Task` from their `spawn` methods. `gpui::Task` is this
/// type. The handle controls the future's lifetime in three ways:
///
/// - **Dropping** the task cancels the future. It stops at its current
///   `.await` point and is never polled again. Work it already finished stays
///   finished. Keep the handle alive while you need the result.
/// - **[`Task::detach`]** gives up the handle and lets the future run to
///   completion. Its output is discarded, so use this for side effects.
/// - **Awaiting** the task waits for the output. A `Task` implements
///   [`Future`], and awaiting it consumes it. If the future panics, the panic
///   resumes in the awaiting code.
///
/// Use [`Task::ready`] for a value you already have. Use [`Task::fallible`]
/// when a cancelled task should resolve to `None` instead of panicking.
///
/// # Panics
///
/// Awaiting a task panics if the task closed without producing an output. This
/// happens when the scheduler drops the runnable before it runs, as during
/// shutdown. Use [`Task::fallible`] to get `None` in that case.
///
/// For the bigger picture, see `gpui::_documentation::concepts::executors`.
#[must_use]
pub struct Task<T>(TaskState<T>);

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, RunnableMeta>),

    /// A typed view of a [`Task<Box<dyn Any + Send + Sync>>`] obtained via
    /// [`Task::downcast`]. The inner task drives the actual work; the
    /// downcast layer just unwraps the `Box<dyn Any + Send + Sync>` on poll.
    Downcast {
        inner: Box<Task<Box<dyn Any + Send + Sync>>>,
        marker: PhantomData<fn() -> T>,
    },

    /// A task whose real handle is delivered later by another thread (see
    /// [`Task::rendezvous`]). Once delivered, polling replaces this state
    /// with the delivered task's state.
    Rendezvous(RendezvousReceiver<T>),
}

/// State shared between the two halves of a [`Task::rendezvous`] pair.
enum RendezvousState<T> {
    /// No task delivered yet; holds the consumer's waker if it polled.
    Pending(Option<Waker>),
    /// The producer delivered before the consumer consumed the task.
    Delivered(Task<T>),
    /// The consumer was dropped before delivery; delivery cancels the task.
    Cancelled,
    /// The consumer was detached before delivery; delivery detaches the task.
    Detached,
    /// The consumer took the delivered task.
    Taken,
}

pub(crate) struct RendezvousReceiver<T> {
    shared: Arc<parking_lot::Mutex<RendezvousState<T>>>,
}

impl<T> RendezvousReceiver<T> {
    fn poll_take(&self, cx: &mut Context) -> Poll<Task<T>> {
        let mut state = self.shared.lock();
        match std::mem::replace(&mut *state, RendezvousState::Taken) {
            RendezvousState::Delivered(task) => Poll::Ready(task),
            RendezvousState::Pending(_) => {
                *state = RendezvousState::Pending(Some(cx.waker().clone()));
                Poll::Pending
            }
            RendezvousState::Cancelled | RendezvousState::Detached | RendezvousState::Taken => {
                unreachable!("a rendezvous task was polled after its receiver was consumed")
            }
        }
    }

    fn is_ready(&self) -> bool {
        match &*self.shared.lock() {
            RendezvousState::Delivered(task) => task.is_ready(),
            _ => false,
        }
    }

    fn detach(self) {
        let mut state = self.shared.lock();
        match std::mem::replace(&mut *state, RendezvousState::Detached) {
            RendezvousState::Delivered(task) => {
                *state = RendezvousState::Taken;
                drop(state);
                task.detach();
            }
            // The producer applies the detached disposition on delivery.
            RendezvousState::Pending(_) => {}
            RendezvousState::Cancelled | RendezvousState::Detached | RendezvousState::Taken => {
                unreachable!("a rendezvous task was detached after its receiver was consumed")
            }
        }
    }
}

impl<T> Drop for RendezvousReceiver<T> {
    fn drop(&mut self) {
        let mut state = self.shared.lock();
        match &*state {
            RendezvousState::Pending(_) => *state = RendezvousState::Cancelled,
            RendezvousState::Delivered(_) => {
                let delivered = std::mem::replace(&mut *state, RendezvousState::Cancelled);
                drop(state);
                drop(delivered);
            }
            // `Taken`, `Detached`, and `Cancelled` record dispositions that
            // this drop must not overwrite: the receiver is also dropped as a
            // normal side effect of taking or detaching.
            _ => {}
        }
    }
}

pub(crate) struct RendezvousSender<T> {
    shared: Arc<parking_lot::Mutex<RendezvousState<T>>>,
}

impl<T> RendezvousSender<T> {
    /// Hands the real task to the rendezvous, applying the consumer's
    /// disposition if it was dropped or detached before delivery.
    pub(crate) fn deliver(self, task: Task<T>) {
        let mut state = self.shared.lock();
        match std::mem::replace(&mut *state, RendezvousState::Delivered(task)) {
            RendezvousState::Pending(waker) => {
                drop(state);
                if let Some(waker) = waker {
                    waker.wake();
                }
            }
            RendezvousState::Cancelled => {
                let delivered = std::mem::replace(&mut *state, RendezvousState::Cancelled);
                drop(state);
                drop(delivered);
            }
            RendezvousState::Detached => {
                let RendezvousState::Delivered(task) =
                    std::mem::replace(&mut *state, RendezvousState::Detached)
                else {
                    unreachable!("the delivered task was just stored");
                };
                drop(state);
                task.detach();
            }
            RendezvousState::Delivered(_) | RendezvousState::Taken => {
                unreachable!("a rendezvous task was delivered twice")
            }
        }
    }
}

impl<T> Task<T> {
    /// Returns a task that is already complete and resolves to `val`.
    ///
    /// Use it where a function returns a `Task` but the value is known now,
    /// such as a cache hit. No work is scheduled. Dropping or detaching the
    /// task has no effect.
    ///
    /// # Examples
    ///
    /// ```
    /// # use scheduler::Task;
    /// let task = Task::ready(42);
    /// assert!(task.is_ready());
    /// ```
    pub fn ready(val: T) -> Self {
        Task(TaskState::Ready(Some(val)))
    }

    /// Creates a task whose real handle arrives later through the returned
    /// sender, typically from another thread. Until delivery the task is
    /// pending; afterwards it behaves exactly like the delivered task.
    /// Dropping or detaching the task before delivery is honored on delivery.
    pub(crate) fn rendezvous() -> (Self, RendezvousSender<T>) {
        let shared = Arc::new(parking_lot::Mutex::new(RendezvousState::Pending(None)));
        (
            Task(TaskState::Rendezvous(RendezvousReceiver {
                shared: shared.clone(),
            })),
            RendezvousSender { shared },
        )
    }

    /// Wraps an `async_task::Task` that carries [`RunnableMeta`] as its metadata.
    ///
    /// The executors use this internally. Call it only when you build the
    /// `async_task` runnable yourself.
    pub fn from_async_task(task: async_task::Task<T, RunnableMeta>) -> Self {
        Task(TaskState::Spawned(task))
    }

    /// Returns `true` if the output is ready to take without waiting.
    ///
    /// A [`Task::ready`] task is always ready. A spawned task becomes ready
    /// when its future finishes, even if the output hasn't been awaited yet.
    /// This is a snapshot: a task on another thread may finish right after the
    /// call returns `false`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use scheduler::Task;
    /// assert!(Task::ready("done").is_ready());
    /// ```
    pub fn is_ready(&self) -> bool {
        match &self.0 {
            TaskState::Ready(_) => true,
            TaskState::Spawned(task) => task.is_finished(),
            TaskState::Downcast { inner, .. } => inner.is_ready(),
            TaskState::Rendezvous(receiver) => receiver.is_ready(),
        }
    }

    /// Lets the future run to completion in the background and discards its
    /// output.
    ///
    /// Use this for work whose result you don't need, such as saving a file or
    /// sending a message. Dropping the task would cancel it instead. Detaching
    /// a [`Task::ready`] task does nothing.
    pub fn detach(self) {
        match self {
            Task(TaskState::Ready(_)) => {}
            Task(TaskState::Spawned(task)) => task.detach(),
            Task(TaskState::Downcast { inner, .. }) => inner.detach(),
            Task(TaskState::Rendezvous(receiver)) => receiver.detach(),
        }
    }

    /// Converts this task into a [`FallibleTask`], which resolves to `None`
    /// instead of panicking when the task closes without an output.
    ///
    /// A panic inside the future still propagates to the awaiter.
    pub fn fallible(self) -> FallibleTask<T> {
        FallibleTask(match self.0 {
            TaskState::Ready(val) => FallibleTaskState::Ready(val),
            TaskState::Spawned(task) => FallibleTaskState::Spawned(task.fallible()),
            TaskState::Downcast { inner, .. } => FallibleTaskState::Downcast {
                inner: Box::new(inner.fallible()),
                marker: PhantomData,
            },
            TaskState::Rendezvous(receiver) => FallibleTaskState::Rendezvous(receiver),
        })
    }
}

impl Task<Box<dyn Any + Send + Sync>> {
    /// Converts a task that yields a boxed `Any` into a task that yields `T`.
    ///
    /// [`LocalExecutor::spawn_dedicated`] and [`BackgroundExecutor::spawn_dedicated`]
    /// use this to recover the closure's output type from the type-erased
    /// [`Scheduler::spawn_dedicated`].
    ///
    /// # Panics
    ///
    /// Polling the returned task panics if the boxed value is not a `T`. The
    /// check runs on poll, not in this call. A mismatch is a bug in whatever
    /// produced the inner task, because the caller chose `T`.
    pub fn downcast<T: Send + Sync + 'static>(self) -> Task<T> {
        Task(TaskState::Downcast {
            inner: Box::new(self),
            marker: PhantomData,
        })
    }
}

impl<T> std::fmt::Debug for Task<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            TaskState::Ready(_) => f.debug_tuple("Task::Ready").finish(),
            TaskState::Spawned(task) => f.debug_tuple("Task::Spawned").field(task).finish(),
            TaskState::Downcast { inner, .. } => {
                f.debug_tuple("Task::Downcast").field(inner).finish()
            }
            TaskState::Rendezvous(_) => f.debug_tuple("Task::Rendezvous").finish(),
        }
    }
}

/// A [`Task`] that resolves to `Option<T>`, and to `None` instead of panicking
/// when the task is cancelled.
///
/// Create one with [`Task::fallible`]. Like `Task`, dropping it cancels the
/// future and detaching lets it finish. Await it to get `Some(output)`, or
/// `None` if the task closed without an output. A panic in the future still
/// propagates to the awaiter.
#[must_use]
pub struct FallibleTask<T>(FallibleTaskState<T>);

enum FallibleTaskState<T> {
    /// A task that is ready to return a value
    Ready(Option<T>),

    /// A task that is currently running (wraps async_task::FallibleTask).
    Spawned(async_task::FallibleTask<T, RunnableMeta>),

    /// Mirror of [`TaskState::Downcast`] for fallible tasks.
    Downcast {
        inner: Box<FallibleTask<Box<dyn Any + Send + Sync>>>,
        marker: PhantomData<fn() -> T>,
    },

    /// Mirror of [`TaskState::Rendezvous`] for fallible tasks.
    Rendezvous(RendezvousReceiver<T>),
}

impl<T> FallibleTask<T> {
    /// Returns a fallible task that is already complete and resolves to
    /// `Some(val)`.
    pub fn ready(val: T) -> Self {
        FallibleTask(FallibleTaskState::Ready(Some(val)))
    }

    /// Lets the future run to completion in the background and discards its
    /// output. See [`Task::detach`].
    pub fn detach(self) {
        match self.0 {
            FallibleTaskState::Ready(_) => {}
            FallibleTaskState::Spawned(task) => task.detach(),
            FallibleTaskState::Downcast { inner, .. } => inner.detach(),
            FallibleTaskState::Rendezvous(receiver) => receiver.detach(),
        }
    }
}

impl<T: 'static> Future for FallibleTask<T> {
    type Output = Option<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = unsafe { self.get_unchecked_mut() };
        loop {
            match &mut this.0 {
                FallibleTaskState::Ready(val) => return Poll::Ready(val.take()),
                FallibleTaskState::Spawned(task) => return Pin::new(task).poll(cx),
                FallibleTaskState::Downcast { inner, .. } => {
                    return match Pin::new(inner.as_mut()).poll(cx) {
                        Poll::Ready(Some(boxed_any)) => Poll::Ready(Some(
                            *boxed_any
                                .downcast::<T>()
                                .expect("FallibleTask::poll: downcast type mismatch"),
                        )),
                        Poll::Ready(None) => Poll::Ready(None),
                        Poll::Pending => Poll::Pending,
                    };
                }
                FallibleTaskState::Rendezvous(receiver) => match receiver.poll_take(cx) {
                    Poll::Ready(task) => {
                        this.0 = task.fallible().0;
                        continue;
                    }
                    Poll::Pending => return Poll::Pending,
                },
            }
        }
    }
}

impl<T> std::fmt::Debug for FallibleTask<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            FallibleTaskState::Ready(_) => f.debug_tuple("FallibleTask::Ready").finish(),
            FallibleTaskState::Spawned(task) => {
                f.debug_tuple("FallibleTask::Spawned").field(task).finish()
            }
            FallibleTaskState::Downcast { inner, .. } => f
                .debug_tuple("FallibleTask::Downcast")
                .field(inner)
                .finish(),
            FallibleTaskState::Rendezvous(_) => f.debug_tuple("FallibleTask::Rendezvous").finish(),
        }
    }
}

impl<T: 'static> Future for Task<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = unsafe { self.get_unchecked_mut() };
        loop {
            match &mut this.0 {
                TaskState::Ready(val) => return Poll::Ready(val.take().unwrap()),
                TaskState::Spawned(task) => return Pin::new(task).poll(cx),
                TaskState::Downcast { inner, .. } => {
                    return match Pin::new(inner.as_mut()).poll(cx) {
                        Poll::Ready(boxed_any) => Poll::Ready(
                            *boxed_any
                                .downcast::<T>()
                                .expect("Task::poll: downcast type mismatch"),
                        ),
                        Poll::Pending => Poll::Pending,
                    };
                }
                TaskState::Rendezvous(receiver) => match receiver.poll_take(cx) {
                    Poll::Ready(task) => {
                        this.0 = task.0;
                        continue;
                    }
                    Poll::Pending => return Poll::Pending,
                },
            }
        }
    }
}

/// Variant of `async_task::spawn_local` that includes the source location of the spawn in panics.
#[track_caller]
fn spawn_local_with_source_location<Fut, S>(
    future: Fut,
    schedule: S,
    metadata: RunnableMeta,
) -> (
    async_task::Runnable<RunnableMeta>,
    async_task::Task<Fut::Output, RunnableMeta>,
)
where
    Fut: Future + 'static,
    Fut::Output: 'static,
    S: async_task::Schedule<RunnableMeta> + 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_eq!(
                self.id,
                thread_id(),
                "local task dropped by a thread that didn't spawn it. Task spawned at {}",
                self.location
            );
            // SAFETY: `inner` is wrapped in `ManuallyDrop`, so this is the only
            // place it is dropped. The thread check above ensures local futures
            // are dropped on the thread that created them.
            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> {
            // SAFETY: We don't move any fields out of `self`; this mutable
            // reference is only used to check metadata and to project the pin to
            // `inner` below.
            let this = unsafe { self.get_unchecked_mut() };
            assert!(
                this.id == thread_id(),
                "local task polled by a thread that didn't spawn it. Task spawned at {}",
                this.location
            );
            // SAFETY: `inner` is structurally pinned by `Checked`; after
            // `Checked` is pinned, `inner` is never moved. The thread check
            // above ensures the local future is only polled by its spawning
            // thread.
            unsafe { Pin::new_unchecked(&mut *this.inner).poll(cx) }
        }
    }

    let location = metadata.location;

    let future = move |_| Checked {
        id: thread_id(),
        inner: ManuallyDrop::new(future),
        location,
    };

    let builder = async_task::Builder::new().metadata(metadata);
    // SAFETY: `Checked` enforces the invariants required by `spawn_unchecked`:
    // the non-`Send` future is only polled and dropped on the thread that
    // spawned it.
    unsafe { builder.spawn_unchecked(future, schedule) }
}

#[cfg(test)]
mod local_task_scope_tests {
    use super::*;
    use crate::TestScheduler;
    use std::{
        collections::VecDeque,
        sync::{Arc, Mutex},
        task::{Poll, Waker},
        thread::ThreadId,
    };

    type Queue = Arc<Mutex<VecDeque<Runnable<RunnableMeta>>>>;

    fn queue() -> Queue {
        Arc::new(Mutex::new(VecDeque::new()))
    }

    fn pop(queue: &Queue) -> Runnable<RunnableMeta> {
        queue.lock().unwrap().pop_front().expect("queued runnable")
    }

    fn executor(scope: Arc<LocalTaskScope>, normal: Queue, cancellation: Queue) -> LocalExecutor {
        let scheduler: Arc<dyn Scheduler> =
            Arc::new(TestScheduler::new(crate::TestSchedulerConfig::default()));
        LocalExecutor::new_with_scope(
            SessionId::new(0),
            scheduler,
            scope,
            move |runnable| normal.lock().unwrap().push_back(runnable),
            move |runnable| cancellation.lock().unwrap().push_back(runnable),
        )
    }

    struct DropOnThread {
        dropped_on: Arc<Mutex<Option<ThreadId>>>,
    }

    impl Drop for DropOnThread {
        fn drop(&mut self) {
            *self.dropped_on.lock().unwrap() = Some(std::thread::current().id());
        }
    }

    #[test]
    fn sealed_scope_late_waker_cannot_dispatch_or_drop_local_work_off_owner() {
        let owner = std::thread::current().id();
        let scope = Arc::new(LocalTaskScope::new());
        let normal = queue();
        let cancellation = queue();
        let executor = executor(scope.clone(), normal.clone(), cancellation.clone());
        let saved_waker = Arc::new(Mutex::new(None::<Waker>));
        let dropped_on = Arc::new(Mutex::new(None));

        let task = {
            let saved_waker = saved_waker.clone();
            let dropped_on = dropped_on.clone();
            executor.spawn(async move {
                let _probe = DropOnThread { dropped_on };
                futures::future::poll_fn(move |cx| {
                    *saved_waker.lock().unwrap() = Some(cx.waker().clone());
                    Poll::<()>::Pending
                })
                .await;
            })
        };
        task.detach();
        pop(&normal).run();
        assert_eq!(scope.live_task_count(), 1);

        let late_waker = saved_waker.lock().unwrap().clone().unwrap();
        scope.begin_retirement();
        drop(pop(&cancellation));

        assert_eq!(scope.live_task_count(), 0);
        assert_eq!(*dropped_on.lock().unwrap(), Some(owner));
        assert!(scope.seal_if_quiescent());
        std::thread::spawn(move || late_waker.wake())
            .join()
            .expect("late wake thread");
        assert!(cancellation.lock().unwrap().is_empty());
    }

    #[test]
    fn completion_unregisters_task_and_spawn_after_retirement_is_canceled() {
        let scope = Arc::new(LocalTaskScope::new());
        let normal = queue();
        let cancellation = queue();
        let executor = executor(scope.clone(), normal.clone(), cancellation.clone());

        executor.spawn(async { 42 }).detach();
        pop(&normal).run();
        assert_eq!(scope.live_task_count(), 0);

        scope.begin_retirement();
        executor.spawn(async { 7 }).detach();
        assert!(normal.lock().unwrap().is_empty());
        assert_eq!(scope.live_task_count(), 1);
        drop(pop(&cancellation));
        assert_eq!(scope.live_task_count(), 0);
    }

    #[test]
    fn custom_idle_dispatch_is_bypassed_during_retirement() {
        let scope = Arc::new(LocalTaskScope::new());
        let normal = queue();
        let cancellation = queue();
        let executor = executor(scope.clone(), normal, cancellation.clone());
        let idle = queue();
        let idle_dispatches = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let pending_waker = Arc::new(Mutex::new(None::<Waker>));

        let task = {
            let idle_dispatches = idle_dispatches.clone();
            let idle = idle.clone();
            let pending_waker = pending_waker.clone();
            executor.spawn_with_dispatch(
                async move {
                    futures::future::poll_fn(move |cx| {
                        *pending_waker.lock().unwrap() = Some(cx.waker().clone());
                        Poll::<()>::Pending
                    })
                    .await;
                },
                move |runnable| {
                    idle_dispatches.fetch_add(1, Ordering::SeqCst);
                    idle.lock().unwrap().push_back(runnable);
                },
            )
        };
        task.detach();
        pop(&idle).run();
        assert_eq!(idle_dispatches.load(Ordering::SeqCst), 1);

        scope.begin_retirement();
        drop(pop(&cancellation));
        assert_eq!(idle_dispatches.load(Ordering::SeqCst), 1);
        assert_eq!(scope.live_task_count(), 0);
    }
}

#[cfg(test)]
mod sealed_scope_tests {
    use super::*;
    use crate::TestScheduler;
    use std::{
        collections::VecDeque,
        sync::{Arc, Mutex},
        thread::ThreadId,
    };

    #[test]
    fn sealed_scope_rejects_new_spawns_and_drops_input_on_owner() {
        let owner = std::thread::current().id();
        let scope = Arc::new(LocalTaskScope::new());
        scope.begin_retirement();
        assert!(scope.seal_if_quiescent());

        let normal = Arc::new(Mutex::new(VecDeque::new()));
        let cancellation = Arc::new(Mutex::new(VecDeque::new()));
        let scheduler: Arc<dyn Scheduler> =
            Arc::new(TestScheduler::new(crate::TestSchedulerConfig::default()));
        let executor = LocalExecutor::new_with_scope(
            SessionId::new(1),
            scheduler,
            scope,
            {
                let normal = normal.clone();
                move |runnable| normal.lock().unwrap().push_back(runnable)
            },
            {
                let cancellation = cancellation.clone();
                move |runnable| cancellation.lock().unwrap().push_back(runnable)
            },
        );
        let dropped_on = Arc::new(Mutex::new(None::<ThreadId>));
        let probe = {
            let dropped_on = dropped_on.clone();
            struct Probe(Arc<Mutex<Option<ThreadId>>>);
            impl Drop for Probe {
                fn drop(&mut self) {
                    *self.0.lock().unwrap() = Some(std::thread::current().id());
                }
            }
            Probe(dropped_on)
        };
        let future = async move {
            let _probe = probe;
        };
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            executor.spawn(future).detach();
        }));

        assert!(result.is_err());
        assert_eq!(*dropped_on.lock().unwrap(), Some(owner));
        assert!(normal.lock().unwrap().is_empty());
        assert!(cancellation.lock().unwrap().is_empty());
    }
}