poll 与唤醒

pub(super) fn poll(self) {
    // We pass our ref-count to `poll_inner`.
    match self.poll_inner() {
        PollFuture::Notified => {
            // The `poll_inner` call has given us two ref-counts back.
            // We give one of them to a new task and call `yield_now`.
            self.core()
                .scheduler
                .yield_now(Notified(self.get_new_task()));

            // The remaining ref-count is now dropped. We kept the extra
            // ref-count until now to ensure that even if the `yield_now`
            // call drops the provided task, the task isn't deallocated
            // before after `yield_now` returns.
            self.drop_reference();
        }
        PollFuture::Complete => {
            self.complete();
        }
        PollFuture::Dealloc => {
            self.dealloc();
        }
        PollFuture::Done => (),
    }
}

fn poll_inner(&self) -> PollFuture {
    use super::state::{TransitionToIdle, TransitionToRunning};

    match self.state().transition_to_running() {
        TransitionToRunning::Success => {
            // ...
            let header_ptr = self.header_ptr();
            let waker_ref = waker_ref::<S>(&header_ptr);
            let cx = Context::from_waker(&waker_ref);
            let res = poll_future(self.core(), cx);

            if res == Poll::Ready(()) {
                // The future completed. Move on to complete the task.
                return PollFuture::Complete;
            }

            let transition_res = self.state().transition_to_idle();
            if let TransitionToIdle::Cancelled = transition_res {
                // The transition to idle failed because the task was
                // cancelled during the poll.
                cancel_task(self.core());
            }
            transition_result_to_poll_future(transition_res)
        }
        TransitionToRunning::Cancelled => {
            cancel_task(self.core());
            PollFuture::Complete
        }
        TransitionToRunning::Failed => PollFuture::Done,
        TransitionToRunning::Dealloc => PollFuture::Dealloc,
    }
}

pub(super) fn waker_ref<S>(header: &NonNull<Header>) -> WakerRef<'_, S>
where
    S: Schedule,
{
    // `Waker::will_wake` uses the VTABLE pointer as part of the check. This
    // means that `will_wake` will always return false when using the current
    // task's waker. (discussion at rust-lang/rust#66281).
    //
    // To fix this, we use a single vtable. Since we pass in a reference at this
    // point and not an *owned* waker, we must ensure that `drop` is never
    // called on this waker instance. This is done by wrapping it with
    // `ManuallyDrop` and then never calling drop.
    let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker(*header))) };

    WakerRef {
        waker,
        _p: PhantomData,
    }
}

unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
    // Safety: `ptr` was created from a `Header` pointer in function `waker_ref`.
    let header = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
    // ...（tracing）
    unsafe { header.as_ref() }.state.ref_inc();
    raw_waker(header)
}

static WAKER_VTABLE: RawWakerVTable =
    RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);

pub(super) fn wake_by_val(&self) {
    use super::state::TransitionToNotifiedByVal;

    match self.state().transition_to_notified_by_val() {
        TransitionToNotifiedByVal::Submit => {
            // The caller has given us a ref-count, and the transition has
            // created a new ref-count, so we now hold two. We turn the new
            // ref-count Notified and pass it to the call to `schedule`.
            //
            // The old ref-count is retained for now to ensure that the task
            // is not dropped during the call to `schedule` if the call
            // drops the task it was given.
            self.schedule();

            // Now that we have completed the call to schedule, we can
            // release our ref-count.
            self.drop_reference();
        }
        TransitionToNotifiedByVal::Dealloc => {
            self.dealloc();
        }
        TransitionToNotifiedByVal::DoNothing => {}
    }
}

with_current(|maybe_cx| {
    if let Some(cx) = maybe_cx {
        // Make sure the task is part of the **current** scheduler.
        if self.ptr_eq(&cx.worker.handle) {
            // And the current thread still holds a core
            if let Some(core) = cx.core.borrow_mut().as_mut() {
                self.schedule_local(core, task, is_yield);
                return;
            }
        }
    }

    // Otherwise, use the inject queue.
    self.push_remote_task(task);
    self.notify_parked_remote();
});

fn poll_future<T: Future, S: Schedule>(core: &Core<T, S>, cx: Context<'_>) -> Poll<()> {
    // Poll the future.
    let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        struct Guard<'a, T: Future, S: Schedule> {
            core: &'a Core<T, S>,
        }
        impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> {
            fn drop(&mut self) {
                // If the future panics on poll, we drop it inside the panic
                // guard.
                self.core.drop_future_or_output();
            }
        }
        let guard = Guard { core };
        let res = guard.core.poll(cx);
        mem::forget(guard);
        res
    }));

    // Prepare output for being placed in the core stage.
    let output = match output {
        Ok(Poll::Pending) => return Poll::Pending,
        Ok(Poll::Ready(output)) => Ok(output),
        Err(panic) => Err(panic_to_error(&core.scheduler, core.task_id, panic)),
    };

    // Catch and ignore panics if the future panics on drop.
    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        core.store_output(output);
    }));

    if res.is_err() {
        core.scheduler.unhandled_panic();
    }

    Poll::Ready(())
}

#[tokio::main]
async fn main() {
  let err = tokio::spawn(async { panic!("boom") }).await.unwrap_err();
  println!("{err}");   // task 13 panicked with message "boom"
  println!("{err:?}"); // JoinError::Panic(Id(13), "boom", ...)
}

fn complete(self) {
    // The future has completed and its output has been written to the task
    // stage. We transition from running to complete.
    let snapshot = self.state().transition_to_complete();

    // We catch panics here in case dropping the future or waking the
    // JoinHandle panics.
    let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        if !snapshot.is_join_interested() {
            // The `JoinHandle` is not interested in the output of
            // this task. It is our responsibility to drop the
            // output. The join waker was already dropped by the
            // `JoinHandle` before.
            self.core().drop_future_or_output();
        } else if snapshot.is_join_waker_set() {
            // Notify the waker. ...
            self.trailer().wake_join();

            // ...
            if !self
                .state()
                .unset_waker_after_complete()
                .is_join_interested()
            {
                // SAFETY: We have COMPLETE=1 and JOIN_INTEREST=0, so
                // we have exclusive access to the waker.
                unsafe { self.trailer().set_waker(None) };
            }
        }
    }));

    // ...（task_terminate_callback）

    // The task has completed execution and will no longer be scheduled.
    let num_release = self.release();

    if self.state().transition_to_terminal(num_release) {
        self.dealloc();
    }
}

fn release(&self) -> usize {
    // We don't actually increment the ref-count here, but the new task is
    // never destroyed, so that's ok.
    let me = ManuallyDrop::new(self.get_new_task());

    if let Some(task) = self.core().scheduler.release(&me) {
        mem::forget(task);
        2
    } else {
        1
    }
}

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

struct WakeSelf {
  polls: Arc<AtomicUsize>,
}

impl Future for WakeSelf {
  type Output = ();
  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
    let n = self.polls.fetch_add(1, Ordering::SeqCst) + 1;
    if n == 1 {
      // 第一次 poll：叫醒自己三次，然后返回 Pending
      cx.waker().wake_by_ref();
      cx.waker().wake_by_ref();
      cx.waker().wake_by_ref();
    }
    // 之后的 poll 不再叫醒任何人
    Poll::Pending
  }
}

#[tokio::main]
async fn main() {
  let polls = Arc::new(AtomicUsize::new(0));
  tokio::spawn(WakeSelf { polls: polls.clone() });
  tokio::time::sleep(Duration::from_millis(50)).await;
  println!("polls = {}", polls.load(Ordering::SeqCst));
}

polls = 2