path: root/tokio/src/task/harness.rs

use crate::loom::alloc::Track;
use crate::loom::cell::CausalCheck;
use crate::task::core::{Cell, Core, Header, Trailer};
use crate::task::state::Snapshot;
use crate::task::{JoinError, Schedule, Task};

use std::future::Future;
use std::marker::PhantomData;
use std::mem::{ManuallyDrop, MaybeUninit};
use std::ptr::NonNull;
use std::task::{Poll, Waker};

/// Typed raw task handle
pub(super) struct Harness<T: Future, S: 'static> {
    cell: NonNull<Cell<T>>,
    _p: PhantomData<S>,
}

impl<T, S> Harness<T, S>
where
    T: Future,
    S: 'static,
{
    pub(super) unsafe fn from_raw(ptr: *mut ()) -> Harness<T, S> {
        debug_assert!(!ptr.is_null());

        Harness {
            cell: NonNull::new_unchecked(ptr as *mut Cell<T>),
            _p: PhantomData,
        }
    }

    fn header(&self) -> &Header {
        unsafe { &self.cell.as_ref().header }
    }

    fn trailer(&self) -> &Trailer {
        unsafe { &self.cell.as_ref().trailer }
    }

    fn core(&mut self) -> &mut Core<T> {
        unsafe { &mut self.cell.as_mut().core }
    }
}

impl<T, S> Harness<T, S>
where
    T: Future,
    S: Schedule,
{
    /// Poll the inner future.
    ///
    /// All necessary state checks and transitions are performed.
    ///
    /// Panics raised while polling the future are handled.
    ///
    /// Returns `true` if the task needs to be scheduled again
    ///
    /// # Safety
    ///
    /// The pointer returned by the `executor` fn must be castable to `*mut S`
    pub(super) unsafe fn poll(mut self, executor: &mut dyn FnMut() -> Option<NonNull<()>>) -> bool {
        use std::panic;

        // Transition the task to the running state.
        let res = self.header().state.transition_to_running();

        if res.is_canceled() {
            // The task was concurrently canceled.
            self.do_cancel(res);
            return false;
        }

        let join_interest = res.is_join_interested();
        debug_assert!(join_interest || !res.has_join_waker());

        // Get the cell components
        let cell = &mut self.cell.as_mut();
        let header = &cell.header;
        let core = &mut cell.core;

        // If the task's executor pointer is not yet set, then set it here. This
        // is safe because a) this is the only time the value is set. b) at this
        // point, there are no outstanding wakers which might access the
        // field concurrently.
        if header.executor().is_none() {
            // We don't want the destructor to run because we don't really
            // own the task here.
            let task = ManuallyDrop::new(Task::from_raw(header.into()));
            // Call the scheduler's bind callback
            let executor = executor().expect("first poll must happen from an executor");
            executor.cast::<S>().as_ref().bind(&task);
            header.executor.with_mut(|ptr| *ptr = Some(executor.cast()));
        }

        // The transition to `Running` done above ensures that a lock on the
        // future has been obtained. This also ensures the `*mut T` pointer
        // contains the future (as opposed to the output) and is initialized.

        let res = header.future_causality.with_mut(|_| {
            panic::catch_unwind(panic::AssertUnwindSafe(|| {
                struct Guard<'a, T: Future> {
                    core: &'a mut Core<T>,
                    polled: bool,
                }

                impl<T: Future> Drop for Guard<'_, T> {
                    fn drop(&mut self) {
                        if !self.polled {
                            self.core.transition_to_consumed();
                        }
                    }
                }

                let mut guard = Guard {
                    core,
                    polled: false,
                };

                let res = guard.core.poll::<S>(header);

                // prevent the guard from dropping the future
                guard.polled = true;

                res
            }))
        });

        match res {
            Ok(Poll::Ready(out)) => {
                self.complete(executor, join_interest, Ok(out));
                false
            }
            Ok(Poll::Pending) => {
                let res = self.header().state.transition_to_idle();

                if res.is_canceled() {
                    self.do_cancel(res);
                    false
                } else {
                    res.is_notified()
                }
            }
            Err(err) => {
                self.complete(executor, join_interest, Err(JoinError::panic2(err)));
                false
            }
        }
    }

    pub(super) unsafe fn drop_task(mut self) {
        let might_drop_join_waker_on_release = self.might_drop_join_waker_on_release();

        // Read the join waker cell just to have it
        let (join_waker, check) = self.read_join_waker();

        // transition the task to released
        let res = self.header().state.release_task();

        assert!(res.is_terminal(), "state = {:?}", res);

        if might_drop_join_waker_on_release && !res.is_join_interested() {
            debug_assert!(res.has_join_waker());

            // Its our responsibility to drop the waker
            check.check();
            let _ = join_waker.assume_init();
        }

        if res.is_final_ref() {
            self.dealloc();
        }
    }

    unsafe fn dealloc(self) {
        // Check causality
        self.header().executor.with_mut(|_| {});
        self.header().future_causality.with_mut(|_| {});
        self.trailer().waker.with_mut(|_| {
            // we can't check the contents of this cell as it is considered
            // "uninitialized" data at this point.
        });

        drop(Box::from_raw(self.cell.as_ptr()));
    }

    // ===== join handle =====

    pub(super) unsafe fn read_output(
        mut self,
        dst: *mut Track<super::Result<T::Output>>,
        state: Snapshot,
    ) {
        if state.is_canceled() {
            dst.write(Track::new(Err(JoinError::cancelled2())));
        } else {
            self.core().read_output(dst);
        }

        // Before transitioning the state, the waker must be read. It is
        // possible that, after the transition, we are responsible for dropping
        // the waker but before the waker can be read from the struct, the
        // struct is deallocated.
        let (waker, check) = self.read_join_waker();

        // The operation counts as dropping the join handle
        let res = self.header().state.complete_join_handle();

        if res.is_released() {
            // We are responsible for freeing the waker handle
            check.check();
            drop(waker.assume_init());
        }

        if res.is_final_ref() {
            self.dealloc();
        }
    }

    pub(super) fn store_join_waker(&self, waker: &Waker) -> Snapshot {
        unsafe {
            self.trailer().waker.with_mut(|ptr| {
                (*ptr).as_mut_ptr().replace(Some