先看下源码中的文档，如何使用这个 crate：
Task abstraction for building executors.
To spawn a future onto an executor, we first need to allocate it on the heap and keep some state attached to it. The state indicates whether the future is ready for polling, waiting to be woken up, or completed. Such a stateful future is called a task.
All executors have a queue that holds scheduled tasks:
let (sender, receiver) = flume::unbounded();
A task is created using either spawn(), spawn_local(), or spawn_unchecked() which return a Runnable and a Task:
// A future that will be spawned.
let future = async { 1 + 2 };
// A function that schedules the task when it gets woken up.
let schedule = move |runnable| sender.send(runnable).unwrap();
// Construct a task.
let (runnable, task) = async_task::spawn(future, schedule);
// Push the task into the queue by invoking its schedule function.
runnable.schedule();
The Runnable is used to poll the task's future, and the Task is used to await its output.
Finally, we need a loop that takes scheduled tasks from the queue and runs them:
for runnable in receiver {
runnable.run();
}
Method run() polls the task's future once. Then, the Runnable vanishes and only reappears when its Waker wakes the task, thus scheduling it to be run again.
可以看到，我们定义的协程在 spawn 后会产生两个结构，一个是 Runnable，另外一个是 Task，Runnable 用于对协程进行轮询， Task 则用于取得协程的结果。当协程需要执行的时候，就会产生一个 Runnable，当调用 Runnable的 run 函数时，就会执行一次协程的 poll，所以 Runnable 在哪执行，协程就在哪里执行，并且每次执行的线程也可能不一样（取决于调度器实现），例如上面的代码的队列类型是 SPMC（Single Producer Multi Consumer） 类型的，那么由于 Consumer 在不同的线程中接收 Runnable 并执行，当协程从挂起状态恢复时，会产生一个 Runnable 用于 poll 协程，这个 Runnable 就可能会在其它线程中执行了。
写在之前的话
看源码一定要带着问题去看，多想想为什么，这样才能更好的帮助你理解代码，实在看不懂的地方可以先放一放，先把握整体，后面再死抠细节。关于 async-task 建议读者多关注任务的生命周期以及任务的状态变化。
Waker
Waker 是 Rust 的 Future 为我们提供的协程唤醒机制，它被封装到了 poll 参数中的 Context 中，我们可以调用 Waker 中的 wake 进行任务唤醒。 Waker 里面封装的是 RawWaker，它允许我们定义自己的唤醒函数以及存储自定义的数据。在 async-task 中，我们 poll 内部任务的 Waker 就是自己定义的，这样才能实现当内部任务里面调用 wake 时，我们可以通过 schedule 这个函数进行任务的再次调度。
RawTask
看源码的一种方式是从我们平常调用的函数入口处一步一步向内分析，这里我们也采用这种方式，当我们调用 async_task::spawn 生成异步任务时，内部会调用到 Buidler 里面的 spawn_unchecked 函数：
pub fn spawn<F, S>(future: F, schedule: S) -> (Runnable, Task<F::Output>)
where
F: Future + Send + 'static,
F::Output: Send + 'static,
S: Schedule + Send + Sync + 'static,
{
unsafe { spawn_unchecked(future, schedule) }
}
pub unsafe fn spawn_unchecked<F, S>(future: F, schedule: S) -> (Runnable, Task<F::Output>)
where
F: Future,
S: Schedule,
{
Builder::new().spawn_unchecked(move |()| future, schedule)
}
pub unsafe fn spawn_unchecked<'a, F, Fut, S>(
self,
future: F,
schedule: S,
) -> (Runnable<M>, Task<Fut::Output, M>)
where
F: FnOnce(&'a M) -> Fut,
Fut: Future + 'a,
S: Schedule<M>,
M: 'a,
{
// Allocate large futures on the heap.
let ptr = if mem::size_of::<Fut>() >= 2048 {
let future = |meta| {
let future = future(meta);
Box::pin(future)
};
RawTask::<_, Fut::Output, S, M>::allocate(future, schedule, self)
} else {
RawTask::<Fut, Fut::Output, S, M>::allocate(future, schedule, self)
};
let runnable = Runnable::from_raw(ptr);
let task = Task {
ptr,
_marker: PhantomData,
};
(runnable, task)
}
}
这里我们可以看到，外部的 future 和 schedule 都传入了 RawTask::allocate 函数，返回了一个 RawTask 的指针，这个 RawTask 就是 async-task 里面比较核心的结构了。
定义
/// Raw pointers to the fields inside a task.
pub(crate) struct RawTask<F, T, S, M> {
/// The task header.
pub(crate) header: *const Header<M>,
/// The schedule function.
pub(crate) schedule: *const S,
/// The future.
pub(crate) future: *mut F,
/// The output of the future.
pub(crate) output: *mut Result<T, Panic>,
}
RawTask 里的字段包含四个部分，并且都是指针，使用指针的原因是因为 RawTask 需要多处共享，且生命周期不稳定，这里是手动控制的里面数据的生命周期，并且可以达到更加紧凑的内存布局的目的。实际上 future 和 output 是使用的同一块内存，类似于 C 中的 union，当我们看过 allocate 就清楚了。 future 和 schedule 是我们外部传入的数据，而 output 是 future 完成后写入的结果，最后只有一个 header 不知道是什么。
/// The header of a task.
///
/// This header is stored in memory at the beginning of the heap-allocated task.
pub(crate) struct Header<M> {
/// Current state of the task.
///
/// Contains flags representing the current state and the reference count.
pub(crate) state: AtomicUsize,
/// The task that is blocked on the `Task` handle.
///
/// This waker needs to be woken up once the task completes or is closed.
pub(crate) awaiter: UnsafeCell<Option<Waker>>,
/// The virtual table.
///
/// In addition to the actual waker virtual table, it also contains pointers to several other
/// methods necessary for bookkeeping the heap-allocated task.
pub(crate) vtable: &'static TaskVTable,
/// Metadata associated with the task.
///
/// This metadata may be provided to the user.
pub(crate) metadata: M,
/// Whether or not a panic that occurs in the task should be propagated.
#[cfg(feature = "std")]
pub(crate) propagate_panic: bool,
}
- state 使用一系列的比特位来存储任务的状态
/// Set if the task is scheduled for running.
///
/// A task is considered to be scheduled whenever its `Runnable` exists.
///
/// This flag can't be set when the task is completed. However, it can be set while the task is
/// running, in which case it will be rescheduled as soon as polling finishes.
pub(crate) const SCHEDULED: usize = 1 << 0;
/// Set if the task is running.
///
/// A task is in running state while its future is being polled.
///
/// This flag can't be set when the task is completed. However, it can be in scheduled state while
/// it is running, in which case it will be rescheduled as soon as polling finishes.
pub(crate) const RUNNING: usize = 1 << 1;
/// Set if the task has been completed.
///
/// This flag is set when polling returns `Poll::Ready`. The output of the future is then stored
/// inside the task until it becomes closed. In fact, `Task` picks up the output by marking
/// the task as closed.
///
/// This flag can't be set when the task is scheduled or running.
pub(crate) const COMPLETED: usize = 1 << 2;
/// Set if the task is closed.
///
/// If a task is closed, that means it's either canceled or its output has been consumed by the
/// `Task`. A task becomes closed in the following cases:
///
/// 1. It gets canceled by `Runnable::drop()`, `Task::drop()`, or `Task::cancel()`.
/// 2. Its output gets awaited by the `Task`.
/// 3. It panics while polling the future.
/// 4. It is completed and the `Task` gets dropped.
pub(crate) const CLOSED: usize = 1 << 3;
/// Set if the `Task` still exists.
///
/// The `Task` is a special case in that it is only tracked by this flag, while all other
/// task references (`Runnable` and `Waker`s) are tracked by the reference count.
pub(crate) const TASK: usize = 1 << 4;
/// Set if the `Task` is awaiting the output.
///
/// This flag is set while there is a registered awaiter of type `Waker` inside the task. When the
/// task gets closed or completed, we need to wake the awaiter. This flag can be used as a fast
/// check that tells us if we need to wake anyone.
pub(crate) const AWAITER: usize = 1 << 5;
/// Set if an awaiter is being registered.
///
/// This flag is set when `Task` is polled and we are registering a new awaiter.
pub(crate) const REGISTERING: usize = 1 << 6;
/// Set if the awaiter is being notified.
///
/// This flag is set when notifying the awaiter. If an awaiter is concurrently registered and
/// notified, whichever side came first will take over the reposibility of resolving the race.
pub(crate) const NOTIFYING: usize = 1 << 7;
/// A single reference.
///
/// The lower bits in the state contain various flags representing the task state, while the upper
/// bits contain the reference count. The value of `REFERENCE` represents a single reference in the
/// total reference count.
///
/// Note that the reference counter only tracks the `Runnable` and `Waker`s. The `Task` is
/// tracked separately by the `TASK` flag.
pub(crate) const REFERENCE: usize = 1 << 8;
从代码中状态来看，这些状态分别占用一个 usize 的不同比特位。SCHEDULED RUNNING COMPLETED CLOSED 这四个状态是用来表示任务的状态的，TASK 是用于表示 Task 这个句柄是否存在（async_task::spawn返回的这个结构）， AWAITER REGISTERING NOTIFYING 是和 Future 中 Waker 相关的状态，REFERENCE 是一个引用计数，用于追踪产生的 Runnable 和 Waker 结构（任务运行时，可能会产生多个这样的结构），这两个结构里面都有 RawTask 的指针，因为我们的 RawTask 的生命周期是手动管理的，所以需要用引用计数来进行 RawTask 的生命周期管理，当引用计数降为0的时候，就需要考虑（注意是考虑，不是绝对的，还有其它条件）对 RawTask 进行析构了。
- awaiter
当我们
pollasync_task::spawn产生的这个task时，会将外部的Waker存储起来，当任务结束时，通知外部再次poll获取结果。 - vtable
RawTask的虚表，里面有些是构造RawWaker要用的，有些是RawTask内部使用。 - metadata 构造任务时可以传入一些额外数据
- propagate_panic
Future在 panic 时，决定 panic 的传播形式，在poll时就在内部 panic 或者在 task await 处 panic
allocate
/// Allocates a task with the given `future` and `schedule` function.
///
/// It is assumed that initially only the `Runnable` and the `Task` exist.
pub(crate) fn allocate<'a, Gen: FnOnce(&'a M) -> F>(
future: Gen,
schedule: S,
builder: crate::Builder<M>,
) -> NonNull<()>
where
F: 'a,
M: 'a,
{
// Compute the layout of the task for allocation. Abort if the computation fails.
//
// n.b. notgull: task_layout now automatically aborts instead of panicking
// 这里会根据 RawTask 的结构计算出一个布局数据，指示如何进行内存分配和析构以及如何通过
// 便宜访问 RawTask 中不同的字段
let task_layout = Self::task_layout();
unsafe {
// Allocate enough space for the entire task.
let ptr = match NonNull::new(alloc::alloc::alloc(task_layout.layout) as *mut ()) {
// 这里的 abort 是一个双重 panic，内存分配失败时会直接结束程序
None => abort(),
Some(p) => p,
};
// 通过裸指针构造 RawTask 结构，这里会使用 layout 的数据获得字段数据起始位置的指针
let raw = Self::from_ptr(ptr.as_ptr());
let crate::Builder {
metadata,
#[cfg(feature = "std")]
propagate_panic,
} = builder;
// 分批好各个字段的内存后，向字段中写入数据
// Write the header as the first field of the task.
(raw.header as *mut Header<M>).write(Header {
// 初始状态 通过函数的注释了解到，假设分配的时候 Runnable 和 Task
// 都存在，所以会设置下面的状态，为什么有 REFERENCE？因为已经有一个
// Runnable 了，请牢记这一点
state: AtomicUsize::new(SCHEDULED | TASK | REFERENCE),
awaiter: UnsafeCell::new(None),
vtable: &TaskVTable {
schedule: Self::schedule,
drop_future: Self::drop_future,
get_output: Self::get_output,
drop_ref: Self::drop_ref,
destroy: Self::destroy,
run: Self::run,
clone_waker: Self::clone_waker,
layout_info: &Self::TASK_LAYOUT,
},
metadata,
#[cfg(feature = "std")]
propagate_panic,
});
// Write the schedule function as the third field of the task.
(raw.schedule as *mut S).write(schedule);
// Generate the future, now that the metadata has been pinned in place.
let future = abort_on_panic(|| future(&(*raw.header).metadata));
// Write the future as the fourth field of the task.
raw.future.write(future);
ptr
}
}
这里顺带提一下 abort 和 abort_on_panic 这两个函数：
/// Aborts the process.
///
/// To abort, this function simply panics while panicking.
pub(crate) fn abort() -> ! {
struct Panic;
impl Drop for Panic {
fn drop(&mut self) {
panic!("aborting the process");
}
}
let _panic = Panic;
panic!("aborting the process");
}
/// Calls a function and aborts if it panics.
///
/// This is useful in unsafe code where we can't recover from panics.
#[inline]
pub(crate) fn abort_on_panic<T>(f: impl FnOnce() -> T) -> T {
struct Bomb;
impl Drop for Bomb {
fn drop(&mut self) {
abort();
}
}
let bomb = Bomb;
let t = f();
mem::forget(bomb);
t
}
abort 使用的是双重 panic，通过在析构的时候再次 panic 结束程序（Rust 默认行为）。abort_on_panic 通过构造一个炸弹（Bomb）在函数调用之前，在炸弹的析构函数中 panic，如果函数调用成功，则使用 mem::forget 阻止 Bomb 的析构函数执行。如果调用 f 过程中发生 panic，那么 mem::forget 就得不到执行，在执行清理的时候， Bomb 的析构函数就会执行，从而终止程序。
wake
调用 wake 会对协程进行调度，如果当前协程正在轮询中，则设置 SCHEDULED 标志位，等待轮询后再次轮询，反之生成一个 Runnable 等待后续执行。
在这个函数中，可能第一段代码会让你感到有点疑惑，如果 schedule 函数是一个捕获了变量的闭包时，为什么直接调用的是 wake_by_ref 然后再调用 drop_waker？注释中提到和引用计数有关，那么通过比对 wake 和 wake_by_ref 发现调用 schedule 的方式不同，wake 中在 schedule 是闭包的时候，会执行一次 clone_waker 操作。我感觉是这里的问题，如果你有其它的理解也欢迎留言讨论。
/// Wakes a waker.
unsafe fn wake(ptr: *const ()) {
// This is just an optimization. If the schedule function has captured variables, then
// we'll do less reference counting if we wake the waker by reference and then drop it.
if mem::size_of::<S>() > 0 {
Self::wake_by_ref(ptr);
Self::drop_waker(ptr);
return;
}
let raw = Self::from_ptr(ptr);
// 取得任务当前的状态
let mut state = (*raw.header).state.load(Ordering::Acquire);
// 下面对 state 的修改都是基于 weak 版本的 compare_exchange 进行的，
// 所以需要套一层循环，在失败的时候继续重试
loop {
// If the task is completed or closed, it can't be woken up.
// 任务已经完成或者被关闭，不能再进行轮询了，直接走 drop_waker
if state & (COMPLETED | CLOSED) != 0 {
// Drop the waker.
Self::drop_waker(ptr);
break;
}
// If the task is already scheduled, we just need to synchronize with the thread that
// will run the task by "publishing" our current view of the memory.
// 任务已经在调度中，重新同步一次当前状态，因为 state 也可能在我们代码执行这个过程中改变的，
// 所以这个操作是必要的，虽然看起来有点奇怪
if state & SCHEDULED != 0 {
// Update the state without actually modifying it.
match (*raw.header).state.compare_exchange_weak(
state,
state,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// Drop the waker.
Self::drop_waker(ptr);
break;
}
Err(s) => state = s,
}
} else {
// Mark the task as scheduled.
// 当前任务不在调度中，更新状态，注意这里并没有增加引用计数，因为 wake 的语义
// 是执行唤醒并消耗掉这个 Waker，所以 wake 之后引用计数会-1，如果state & RUNNING == 0
// 时，会在 schedule 中生成一个 Runnable，引用计数会+1，抵消掉了。
match (*raw.header).state.compare_exchange_weak(
state,
state | SCHEDULED,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// If the task is not yet scheduled and isn't currently running, now is the
// time to schedule it.
if state & RUNNING == 0 {
// Schedule the task.
Self::schedule(ptr, ScheduleInfo::new(false));
} else {
// Drop the waker.
// 任务在 poll 中的情况下我们仅仅是标记了 SCHEDULED 这个状态，poll 完成后
// 会检查是否有 SCHEDULED 这个标记，如果有，会重新进行调度，这样就避免了
// poll 并发执行，并且可以将 poll 期间的 wake 进行合并，只需要一次 wake 就
// 可以了
Self::drop_waker(ptr);
}
break;
}
Err(s) => state = s,
}
}
}
}
请牢记 wake 的语义，执行 wake 之后需要将 RawWaker 中的资源进行释放。
wake_by_ref
wake_by_ref 于 wake 类似，只是调用后 waker 还在，注意其中的引用计数的处理区别就行了。在 state & RUNNING == 0 时，会生成一个 Runnable，所以需要将引用计数+1。
/// Wakes a waker by reference.
unsafe fn wake_by_ref(ptr: *const ()) {
let raw = Self::from_ptr(ptr);
let mut state = (*raw.header).state.load(Ordering::Acquire);
loop {
// If the task is completed or closed, it can't be woken up.
if state & (COMPLETED | CLOSED) != 0 {
break;
}
// If the task is already scheduled, we just need to synchronize with the thread that
// will run the task by "publishing" our current view of the memory.
if state & SCHEDULED != 0 {
// Update the state without actually modifying it.
match (*raw.header).state.compare_exchange_weak(
state,
state,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(s) => state = s,
}
} else {
// If the task is not running, we can schedule right away.
let new = if state & RUNNING == 0 {
(state | SCHEDULED) + REFERENCE
} else {
state | SCHEDULED
};
// Mark the task as scheduled.
match (*raw.header).state.compare_exchange_weak(
state,
new,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// If the task is not running, now is the time to schedule.
if state & RUNNING == 0 {
// If the reference count overflowed, abort.
if state > isize::MAX as usize {
abort();
}
// Schedule the task. There is no need to call `Self::schedule(ptr)`
// because the schedule function cannot be destroyed while the waker is
// still alive.
let task = Runnable::from_raw(NonNull::new_unchecked(ptr as *mut ()));
(*raw.schedule).schedule(task, ScheduleInfo::new(false));
}
break;
}
Err(s) => state = s,
}
}
}
}
clone_waker
这个比较简单，只是处理了以下引用计数。
/// Clones a waker.
unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
let raw = Self::from_ptr(ptr);
// Increment the reference count. With any kind of reference-counted data structure,
// relaxed ordering is appropriate when incrementing the counter.
let state = (*raw.header).state.fetch_add(REFERENCE, Ordering::Relaxed);
// If the reference count overflowed, abort.
if state > isize::MAX as usize {
abort();
}
RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)
}
drop_waker
销毁一个 Waker ，从状态中扣减引用计数，当引用计数变为0时，并且 TASK 句柄不存在，则需要走资源释放的流程。这里有两个分支，如果任务既没有 COMPLETED 也没有 CLOSED 标志位，那么重新调度一次，资源交由执行器释放（此时任务可能还在运行，所以不能在这里释放资源），另外一个分支则是有 COMPLETED 或者 CLOSED 标志位，此时任务是不会再次运行的，这里就可以放心的走资源释放的逻辑了。
/// Drops a waker.
///
/// This function will decrement the reference count. If it drops down to zero, the associated
/// `Task` has been dropped too, and the task has not been completed, then it will get
/// scheduled one more time so that its future gets dropped by the executor.
#[inline]
unsafe fn drop_waker(ptr: *const ()) {
let raw = Self::from_ptr(ptr);
// Decrement the reference count.
let new = (*raw.header).state.fetch_sub(REFERENCE, Ordering::AcqRel) - REFERENCE;
// If this was the last reference to the task and the `Task` has been dropped too,
// then we need to decide how to destroy the task.
if new & !(REFERENCE - 1) == 0 && new & TASK == 0 {
if new & (COMPLETED | CLOSED) == 0 {
// If the task was not completed nor closed, close it and schedule one more time so
// that its future gets dropped by the executor.
(*raw.header)
.state
.store(SCHEDULED | CLOSED | REFERENCE, Ordering::Release);
Self::schedule(ptr, ScheduleInfo::new(false));
} else {
// Otherwise, destroy the task right away.
Self::destroy(ptr);
}
}
}
drop_ref
减少一个任务的引用，这个函数和 drop_waker 的逻辑很相似，但是这里并没有判断任务的状态是否为 COMPLETED 或者 CLOSED，其实我查看了这个函数的调用处，都是在 RawTask::run 中调用的，其实就是在 Runnable 执行后减少 Runnable 的引用，所以引用计数减少到0了可以直接走资源销毁的逻辑，不需要再调度一次，因为本身就在 run 函数中了。
/// Drops a task reference (`Runnable` or `Waker`).
///
/// This function will decrement the reference count. If it drops down to zero and the
/// associated `Task` handle has been dropped too, then the task gets destroyed.
#[inline]
unsafe fn drop_ref(ptr: *const ()) {
let raw = Self::from_ptr(ptr);
// Decrement the reference count.
let new = (*raw.header).state.fetch_sub(REFERENCE, Ordering::AcqRel) - REFERENCE;
// If this was the last reference to the task and the `Task` has been dropped too,
// then destroy the task.
if new & !(REFERENCE - 1) == 0 && new & TASK == 0 {
Self::destroy(ptr);
}
}
scheudle
调用这个函数本质上就是调用我们外部传入的 scheudle 函数，这个函数的唯一作用就是将 Runnable 放到队列中等待执行器执行。
这段代码中可能有一个比较让人迷惑的地方，那就是在于 mem::size_of::<S>() > 0 这个判断已经里面的代码是做什么用的，其实这里是在判断我们外部传入的 schedule 函数捕获了变量，如果捕获了变量，那么这个函数的生命周期会随着我们任务的销毁而销毁，它本质上就成了一个匿名类型而不是函数。如果我们在函数内部手动 Drop 了这个 Runnable，又恰好这是最后一个引用计数，那么我们的 (*raw.schedule).schedule 调用就很可能变得非法。所以这里通过克隆 Waker 的方式增加了一个引用计数阻止在 schedule 过程中可能出现资源释放的情况。
/// Schedules a task for running.
///
/// This function doesn't modify the state of the task. It only passes the task reference to
/// its schedule function.
unsafe fn schedule(ptr: *const (), info: ScheduleInfo) {
let raw = Self::from_ptr(ptr);
// If the schedule function has captured variables, create a temporary waker that prevents
// the task from getting deallocated while the function is being invoked.
let _waker;
if mem::size_of::<S>() > 0 {
_waker = Waker::from_raw(Self::clone_waker(ptr));
}
let task = Runnable::from_raw(NonNull::new_unchecked(ptr as *mut ()));
(*raw.schedule).schedule(task, info);
}
drop_future
这个没什么好说的，直接调用 future 的析构函数进行资源释放。
/// Drops the future inside a task.
#[inline]
unsafe fn drop_future(ptr: *const ()) {
let raw = Self::from_ptr(ptr);
// We need a safeguard against panics because the destructor can panic.
abort_on_panic(|| {
raw.future.drop_in_place();
})
}
get_output
获取 future 执行完成后的结果的指针。
/// Returns a pointer to the output inside a task.
unsafe fn get_output(ptr: *const ()) -> *const () {
let raw = Self::from_ptr(ptr);
raw.output as *const ()
}
destroy
析构 RawTask 并归还申请的内存，这里我们需要注意的是这里没有调用 future 和 output 这两个字段的析构函数，这两个字段是在其它地方处理的。
/// Cleans up task's resources and deallocates it.
///
/// The schedule function will be dropped, and the task will then get deallocated.
/// The task must be closed before this function is called.
#[inline]
unsafe fn destroy(ptr: *const ()) {
let raw = Self::from_ptr(ptr);
let task_layout = Self::task_layout();
// We need a safeguard against panics because destructors can panic.
abort_on_panic(|| {
// Drop the header along with the metadata.
(raw.header as *mut Header<M>).drop_in_place();
// Drop the schedule function.
(raw.schedule as *mut S).drop_in_place();
});
// Finally, deallocate the memory reserved by the task.
alloc::alloc::dealloc(ptr as *mut u8, task_layout.layout);
}
run
run 是 RawTask 里面最长的函数了，大部分分析会写到代码的注释里面，少部分写在代码外面作为补充。run 函数就是实际推动我们协程进行状态流转的核心函数了，直到这个协程完成或者被取消为止。
/// Runs a task.
///
/// If polling its future panics, the task will be closed and the panic will be propagated into
/// the caller.
unsafe fn run(ptr: *const ()) -> bool {
let raw = Self::from_ptr(ptr);
// 构造一个 Waker，里面存的数据就是 RawTask 的指针，注意这里的 RAW_WAKER_VTABLE 的定义我们
// 并没有在文章中贴出来，自己去看一下就明白了，调用这里 waker 的 wake 或者 wake_by_ref
// 就会调用到上面我们介绍的那些同名函数
// 注意这里使用了 ManuallyDrop 来阻止 Waker 执行析构函数，因为我们的 Waker 是通过
// Waker::from_raw 创建的，此 Waker 并没有数据的所有权，此 Waker 析构会错误的改变
// 引用计数器的状态
// Create a context from the raw task pointer and the vtable inside the its header.
let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)));
let cx = &mut Context::from_waker(&waker);
let mut state = (*raw.header).state.load(Ordering::Acquire);
// Update the task's state before polling its future.
loop {
// If the task has already been closed, drop the task reference and return.
if state & CLOSED != 0 {
// Drop the future.
Self::drop_future(ptr);
// Mark the task as unscheduled.
let state = (*raw.header).state.fetch_and(!SCHEDULED, Ordering::AcqRel);
// Take the awaiter out.
let mut awaiter = None;
if state & AWAITER != 0 {
awaiter = (*raw.header).take(None);
}
// Drop the task reference.
// 因为我们是从 Runnable 里面执行 RawTask 的 run 函数的，所以我们要-1，
// 因为 Runnable 已经没了
Self::drop_ref(ptr);
// Notify the awaiter that the future has been dropped.
// 这里是在通知外部的等待者，也就是我们 task await 的地方，
// 让外部 poll 可以得知任务已经关闭
if let Some(w) = awaiter {
abort_on_panic(|| w.wake());
}
return false;
}
// Mark the task as unscheduled and running.
// 下面就是要真正的 poll 内部的 future 了，先把状态改以下
match (*raw.header).state.compare_exchange_weak(
state,
(state & !SCHEDULED) | RUNNING,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// Update the state because we're continuing with polling the future.
state = (state & !SCHEDULED) | RUNNING;
break;
}
Err(s) => state = s,
}
}
// Poll the inner future, but surround it with a guard that closes the task in case polling
// panics.
// If available, we should also try to catch the panic so that it is propagated correctly.
// 这里还是采用之前 abort 函数类似的方式，避免我们在执行 poll 的时候 panic，导致资源得不到
// 释放，这种情况由 Guard 负责释放资源
let guard = Guard(raw);
// Panic propagation is not available for no_std.
#[cfg(not(feature = "std"))]
let poll = <F as Future>::poll(Pin::new_unchecked(&mut *raw.future), cx).map(Ok);
#[cfg(feature = "std")]
let poll = {
// Check if we should propagate panics.
if (*raw.header).propagate_panic {
// Use catch_unwind to catch the panic.
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
<F as Future>::poll(Pin::new_unchecked(&mut *raw.future), cx)
})) {
Ok(Poll::Ready(v)) => Poll::Ready(Ok(v)),
Ok(Poll::Pending) => Poll::Pending,
Err(e) => Poll::Ready(Err(e)),
}
} else {
<F as Future>::poll(Pin::new_unchecked(&mut *raw.future), cx).map(Ok)
}
};
// poll 成功，阻止 Guard 析构函数执行
mem::forget(guard);
match poll {
Poll::Ready(out) => {
// Replace the future with its output.
// 拿到结果了，这里首先是执行了 future 的销毁函数，然后才写入结果，前面我们说过
// future 和 output 是共用的一块内存，因此顺序很重要
Self::drop_future(ptr);
raw.output.write(out);
// The task is now completed.
loop {
// If the `Task` is dropped, we'll need to close it and drop the output.
// 这里主要是看 `Task` 句柄还在不在，如果外部的 `Task` 句柄已经不在了，
// 那么这个结果是返回不到外部的，所以直接需要再加一个 CLOSED 表示任务关闭
let new = if state & TASK == 0 {
(state & !RUNNING & !SCHEDULED) | COMPLETED | CLOSED
} else {
(state & !RUNNING & !SCHEDULED) | COMPLETED
};
// Mark the task as not running and completed.
match (*raw.header).state.compare_exchange_weak(
state,
new,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// If the `Task` is dropped or if the task was closed while running,
// now it's time to drop the output.
// TASK 句柄不在了或者任务已经关闭的情况直接析构 future 的输出
// 哪些情况会导致 CLOSED 在状态定义的地方说得很清楚
if state & TASK == 0 || state & CLOSED != 0 {
// Drop the output.
abort_on_panic(|| raw.output.drop_in_place());
}
// Take the awaiter out.
let mut awaiter = None;
if state & AWAITER != 0 {
awaiter = (*raw.header).take(None);
}
// Drop the task reference.
Self::drop_ref(ptr);
// Notify the awaiter that the future has been dropped.
// 唤醒外部的等待着拿结果了
if let Some(w) = awaiter {
abort_on_panic(|| w.wake());
}
break;
}
Err(s) => state = s,
}
}
}
Poll::Pending => {
let mut future_dropped = false;
// The task is still not completed.
loop {
// If the task was closed while running, we'll need to unschedule in case it
// was woken up and then destroy it.
let new = if state & CLOSED != 0 {
state & !RUNNING & !SCHEDULED
} else {
state & !RUNNING
};
if state & CLOSED != 0 && !future_dropped {
// The thread that closed the task didn't drop the future because it was
// running so now it's our responsibility to do so.
Self::drop_future(ptr);
future_dropped = true;
}
// Mark the task as not running.
match (*raw.header).state.compare_exchange_weak(
state,
new,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(state) => {
// If the task was closed while running, we need to notify the awaiter.
// If the task was woken up while running, we need to schedule it.
// Otherwise, we just drop the task reference.
if state & CLOSED != 0 {
// Take the awaiter out.
let mut awaiter = None;
if state & AWAITER != 0 {
awaiter = (*raw.header).take(None);
}
// Drop the task reference.
Self::drop_ref(ptr);
// Notify the awaiter that the future has been dropped.
if let Some(w) = awaiter {
abort_on_panic(|| w.wake());
}
} else if state & SCHEDULED != 0 {
// The thread that woke the task up didn't reschedule it because
// it was running so now it's our responsibility to do so.
// 这里在前面 wake 函数说到过，如果 task 正在 RUNNING 的情况下，
// 只会加一个 SCHEDULED 状态，在这里 poll 结束后会再调度一次
// 同时注意这里没有调用 Self::drop_ref(ptr)，因为 schedule 又会
// 产生一个 Runnable，正好和本次的 Runnable 的计数抵消了
Self::schedule(ptr, ScheduleInfo::new(true));
return true;
} else {
// Drop the task reference.
Self::drop_ref(ptr);
}
break;
}
Err(s) => state = s,
}
}
}
}
return false;
// 下面是 Guard 的定义，就是在 Guard 的 Drop trait 中修改任务的状态，并释放资源，
// 因为走到 Guard 的 Drop 了，那就一定是 poll future panic 了，如果看懂了上面的代码，
// 下面的也不难理解
/// A guard that closes the task if polling its future panics.
struct Guard<F, T, S, M>(RawTask<F, T, S, M>)
where
F: Future<Output = T>,
S: Schedule<M>;
impl<F, T, S, M> Drop for Guard<F, T, S, M>
where
F: Future<Output = T>,
S: Schedule<M>,
{
fn drop(&mut self) {
let raw = self.0;
let ptr = raw.header as *const ();
unsafe {
let mut state = (*raw.header).state.load(Ordering::Acquire);
loop {
// If the task was closed while running, then unschedule it, drop its
// future, and drop the task reference.
if state & CLOSED != 0 {
// The thread that closed the task didn't drop the future because it
// was running so now it's our responsibility to do so.
RawTask::<F, T, S, M>::drop_future(ptr);
// Mark the task as not running and not scheduled.
(*raw.header)
.state
.fetch_and(!RUNNING & !SCHEDULED, Ordering::AcqRel);
// Take the awaiter out.
let mut awaiter = None;
if state & AWAITER != 0 {
awaiter = (*raw.header).take(None);
}
// Drop the task reference.
RawTask::<F, T, S, M>::drop_ref(ptr);
// Notify the awaiter that the future has been dropped.
if let Some(w) = awaiter {
abort_on_panic(|| w.wake());
}
break;
}
// Mark the task as not running, not scheduled, and closed.
match (*raw.header).state.compare_exchange_weak(
state,
(state & !RUNNING & !SCHEDULED) | CLOSED,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(state) => {
// Drop the future because the task is now closed.
RawTask::<F, T, S, M>::drop_future(ptr);
// Take the awaiter out.
let mut awaiter = None;
if state & AWAITER != 0 {
awaiter = (*raw.header).take(None);
}
// Drop the task reference.
RawTask::<F, T, S, M>::drop_ref(ptr);
// Notify the awaiter that the future has been dropped.
if let Some(w) = awaiter {
abort_on_panic(|| w.wake());
}
break;
}
Err(s) => state = s,
}
}
}
}
}
}
到这里我们 RawTask 就看完了，其实看到这里，我们还没看到这个 RawTask 是如何和外部进行交互的，如何将 future 的结果交给外部等待的 task ，这里我们看到 RawTask 始终都是由 run 函数驱动的，也就是我们通过不断的调用 run 就可以一直推动我们传入的 future 直到结束。
Runnable
Runnalbe 在文章最初我们就接触到了，我们在 async_task::spawn 的调用中，除了需要传入 future 之外，还有一个 schedule 函数，这个函数负责将 Runnable 发送到 Executor 所在的线程进行执行。在 Executor 线程中执行 Runnable::run 就会执行到我们 RawTask::run 从而推进协程的进行。
定义
pub struct Runnable<M = ()> {
/// A pointer to the heap-allocated task.
pub(crate) ptr: NonNull<()>,
/// A marker capturing generic type `M`.
pub(crate) _marker: PhantomData<M>,
}
这个结构本身很简单，其实就是持有一个 RawTask 的指针而已。
schedule
pub fn schedule(self) {
let ptr = self.ptr.as_ptr();
let header = ptr as *const Header<M>;
mem::forget(self);
unsafe {
((*header).vtable.schedule)(ptr, ScheduleInfo::new(false));
}
}
这里就直接调用到 RawTask 中的 schedule 去了，需要注意的就是 mem::forget 这个点，如果我们 Runnable 构造出来了，但是不调用 schedule，而是直接丢弃，那么就会造成潜在的资源泄露，因此在 Runnable 中的 Drop 中还有东西，当 Runnable 被丢弃的情况，也会清理资源。所以这里我们把 Runnable 转交走了就不需要执行 Runnable 的析构函数了。
run
这个函数就是简单的传递 Runnable 给 RawTask 执行。
pub fn run(self) -> bool {
let ptr = self.ptr.as_ptr();
let header = ptr as *const Header<M>;
mem::forget(self);
unsafe { ((*header).vtable.run)(ptr) }
}
drop
整个代码还是比较好理解的，代码上已经自带了注释了，解释的比较清楚。不过在这里我们假设一种场景，会不会出现 future 在 poll 的时候，这里恰好调用的 drop_future 呢？其实分析下来是不会的，一个任务的Runnable 至多会存在一个，也就是说不会在执行 Runnable 的同时还有额外的 Runnable 产生，因为在 wake 函数中已经判断了当前任务的状态，只有非 RUNNING 的时候才会调度，如果任务已经在 RUNNING 中，那么 wake 函数只会给任务加上 SCHEDULED 的状态。
impl<M> Drop for Runnable<M> {
fn drop(&mut self) {
let ptr = self.ptr.as_ptr();
let header = self.header();
unsafe {
let mut state = header.state.load(Ordering::Acquire);
loop {
// If the task has been completed or closed, it can't be canceled.
if state & (COMPLETED | CLOSED) != 0 {
break;
}
// Mark the task as closed.
match header.state.compare_exchange_weak(
state,
state | CLOSED,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(s) => state = s,
}
}
// Drop the future.
(header.vtable.drop_future)(ptr);
// Mark the task as unscheduled.
let state = header.state.fetch_and(!SCHEDULED, Ordering::AcqRel);
// Notify the awaiter that the future has been dropped.
if state & AWAITER != 0 {
(*header).notify(None);
}
// Drop the task reference.
(header.vtable.drop_ref)(ptr);
}
}
}
Task
当我们调用 async_task::spawn 后拿到的就是这个结构，我们可以用来将我们生成的异步任务分离或者取得结果以及取消。它本身也是一个 Future，只不过是我们生成的异步任务的一层壳，我们的异步任务本质上还是靠的 Runable::schedule 进行调度的。这层壳相当于只起到一个管理的作用。
定义
可以看到它的定义也是十分的简单，就仅仅只是持有了一个 RawTask 的指针而已，因此在上面的代码中，我们发现，当 Runnable 和 Waker 的引用计数归零后，还需要判断 TASK 这个句柄是否存在，才能决定是否要进行资源的释放。
#[must_use = "tasks get canceled when dropped, use `.detach()` to run them in the background"]
pub struct Task<T, M = ()> {
/// A raw task pointer.
pub(crate) ptr: NonNull<()>,
/// A marker capturing generic types `T` and `M`.
pub(crate) _marker: PhantomData<(T, M)>,
}
为 Task 实现的 Future 则是委派给了 poll_task。
impl<T, M> Future for Task<T, M> {
type Output = T;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.poll_task(cx) {
Poll::Ready(t) => Poll::Ready(t.expect("Task polled after completion")),
Poll::Pending => Poll::Pending,
}
}
}
poll_task
poll_task 就是和 RawTask 打交道的地方了，当外部执行一次 poll 时，我们就会将它提供的 Waker 保存起来，然后查看我们任务的状态，如果任务还在执行中，就返回 Poll::Pending，后续只有任务完成或者取消的时候才会 poll 了，因为我们内部的任务其实是靠 Runnable 来调度的，当我们内部的任务完成后就会调用这个存储的 Waker 让外部再次 poll 拿取结果。
/// Polls the task to retrieve its output.
///
/// Returns `Some` if the task has completed or `None` if it was closed.
///
/// A task becomes closed in the following cases:
///
/// 1. It gets canceled by `Runnable::drop()`, `Task::drop()`, or `Task::cancel()`.
/// 2. Its output gets awaited by the `Task`.
/// 3. It panics while polling the future.
/// 4. It is completed and the `Task` gets dropped.
fn poll_task(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
let ptr = self.ptr.as_ptr();
let header = ptr as *const Header<M>;
unsafe {
let mut state = (*header).state.load(Ordering::Acquire);
loop {
// If the task has been closed, notify the awaiter and return `None`.
if state & CLOSED != 0 {
// If the task is scheduled or running, we need to wait until its future is
// dropped.
if state & (SCHEDULED | RUNNING) != 0 {
// Replace the waker with one associated with the current task.
(*header).register(cx.waker());
// Reload the state after registering. It is possible changes occurred just
// before registration so we need to check for that.
state = (*header).state.load(Ordering::Acquire);
// If the task is still scheduled or running, we need to wait because its
// future is not dropped yet.
if state & (SCHEDULED | RUNNING) != 0 {
return Poll::Pending;
}
}
// Even though the awaiter is most likely the current task, it could also be
// another task.
(*header).notify(Some(cx.waker()));
return Poll::Ready(None);
}
// If the task is not completed, register the current task.
if state & COMPLETED == 0 {
// Replace the waker with one associated with the current task.
(*header).register(cx.waker());
// Reload the state after registering. It is possible that the task became
// completed or closed just before registration so we need to check for that.
state = (*header).state.load(Ordering::Acquire);
// If the task has been closed, restart.
if state & CLOSED != 0 {
continue;
}
// If the task is still not completed, we're blocked on it.
if state & COMPLETED == 0 {
return Poll::Pending;
}
}
// Since the task is now completed, mark it as closed in order to grab its output.
match (*header).state.compare_exchange(
state,
state | CLOSED,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// Notify the awaiter. Even though the awaiter is most likely the current
// task, it could also be another task.
if state & AWAITER != 0 {
(*header).notify(Some(cx.waker()));
}
// Take the output from the task.
let output = ((*header).vtable.get_output)(ptr) as *mut Result<T, Panic>;
let output = output.read();
// Propagate the panic if the task panicked.
let output = match output {
Ok(output) => output,
#[allow(unreachable_patterns)]
Err(panic) => {
#[cfg(feature = "std")]
std::panic::resume_unwind(panic);
#[cfg(not(feature = "std"))]
match panic {}
}
};
return Poll::Ready(Some(output));
}
Err(s) => state = s,
}
}
}
}
稍后我们会分析 Header 中的 register 和 notify 的代码。
set_detached
smol 的 spawn 行为和 Tokio 的不太像， Tokio 是默认在后台执行，除非你手动取消，而 smol 中你想要达到将一个任务持续在后台执行的话，需要调用 detach 进行分离，否则当 Task 句柄被丢弃后，整个任务就会被取消掉。分离任务的方式也很简单，我们调用 detach 的时候，会阻止 Task 的析构函数执行，然后将 TASK 状态删掉就行了。
pub fn detach(self) {
let mut this = self;
let _out = this.set_detached();
mem::forget(this);
}
/// Puts the task in detached state.
fn set_detached(&mut self) -> Option<Result<T, Panic>> {
let ptr = self.ptr.as_ptr();
let header = ptr as *const Header<M>;
unsafe {
// A place where the output will be stored in case it needs to be dropped.
let mut output = None;
// Optimistically assume the `Task` is being detached just after creating the task.
// This is a common case so if the `Task` is datached, the overhead of it is only one
// compare-exchange operation.
// 分离成功这种情况是不需要处理的，只有分离失败的情况才会处理
if let Err(mut state) = (*header).state.compare_exchange_weak(
SCHEDULED | TASK | REFERENCE,
SCHEDULED | REFERENCE,
Ordering::AcqRel,
Ordering::Acquire,
) {
loop {
// If the task has been completed but not yet closed, that means its output
// must be dropped.
if state & COMPLETED != 0 && state & CLOSED == 0 {
// Mark the task as closed in order to grab its output.
match (*header).state.compare_exchange_weak(
state,
state | CLOSED,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// Read the output.
// 这里将 future 的结果创建了一个副本，后续对 future 的析构不会影响
// 到 output
output = Some(
(((*header).vtable.get_output)(ptr) as *mut Result<T, Panic>)
.read(),
);
// Update the state variable because we're continuing the loop.
state |= CLOSED;
}
Err(s) => state = s,
}
} else {
// If this is the last reference to the task and it's not closed, then
// close it and schedule one more time so that its future gets dropped by
// the executor.
// 注意这里虽然 REFERENCE 为0了，但是 task 还在，所以 task 里面持有的 RawTask
// 指针就是最后一个引用了，这里是采用的再调度一次，由 executor 负责执行清理任务
// 注意这种情况下添加了一个引用计数，这个引用计数是下面 Runnable 的
let new = if state & (!(REFERENCE - 1) | CLOSED) == 0 {
SCHEDULED | CLOSED | REFERENCE
} else {
state & !TASK
};
// Unset the `TASK` flag.
match (*header).state.compare_exchange_weak(
state,
new,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// If this is the last reference to the task, we need to either
// schedule dropping its future or destroy it.
if state & !(REFERENCE - 1) == 0 {
if state & CLOSED == 0 {
((*header).vtable.schedule)(ptr, ScheduleInfo::new(false));
} else {
((*header).vtable.destroy)(ptr);
}
}
break;
}
Err(s) => state = s,
}
}
}
}
output
}
}
set_canceled
取消任务，会将原来的 Task 的包一层，将 原来 Task 的 poll 替换掉，因为取消后可能会拿不到任务的结果，因为任务可能还没有完成，原来那种情况会 panic。
impl<T, M> Future for FallibleTask<T, M> {
type Output = Option<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.task.poll_task(cx)
}
}
pub async fn cancel(self) -> Option<T> {
let mut this = self;
this.set_canceled();
this.fallible().await
}
取消的本质还是为任务挂一个 CLOSED 的状态，资源释放都是在 RawTask::run 中进行的。
/// Puts the task in canceled state.
fn set_canceled(&mut self) {
let ptr = self.ptr.as_ptr();
let header = ptr as *const Header<M>;
unsafe {
let mut state = (*header).state.load(Ordering::Acquire);
loop {
// If the task has been completed or closed, it can't be canceled.
if state & (COMPLETED | CLOSED) != 0 {
break;
}
// If the task is not scheduled nor running, we'll need to schedule it.
// 同样注意这里的引用计数变化
let new = if state & (SCHEDULED | RUNNING) == 0 {
(state | SCHEDULED | CLOSED) + REFERENCE
} else {
state | CLOSED
};
// Mark the task as closed.
match (*header).state.compare_exchange_weak(
state,
new,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
// If the task is not scheduled nor running, schedule it one more time so
// that its future gets dropped by the executor.
if state & (SCHEDULED | RUNNING) == 0 {
((*header).vtable.schedule)(ptr, ScheduleInfo::new(false));
}
// Notify the awaiter that the task has been closed.
if state & AWAITER != 0 {
(*header).notify(None);
}
break;
}
Err(s) => state = s,
}
}
}
}
Header
关于 Header 这里我们主要看和 awaiter 相关的内容。awaiter 也是一个 Waker，实际存储的是外部在 poll 我们这个 task 时，从 Context 参数中克隆出来的 Waker。主要是在任务 COMPLETED 或者 CLOSED 的时候通知外部再次 poll 从而推进外部的协程继续往下执行。
register
这个函数单独看可能会觉得有些地方写得莫名其妙，特别是第二段 loop，这个要结合 notify 函数一起看。原因是我们在 register 的途中，可能会发生 notify 调用，进行 awaiter 的唤醒，为了避免并发修改 awaiter 字段，notify 那边是检测到有 REGISTERING 标志时，跳过后续的操作，只是增加一个 NOTIFYING 标记。所以当我们这里把 Waker 写入到 awaiter 字段后，检测到了 NOTIFYING 标志，就重新把 Waker 从 awaiter 中取出来，然后才取消 REGISTERING 的状态。
/// Registers a new awaiter blocked on this task.
///
/// This method is called when `Task` is polled and it has not yet completed.
#[inline]
pub(crate) fn register(&self, waker: &Waker) {
// Load the state and synchronize with it.
let mut state = self.state.fetch_or(0, Ordering::Acquire);
loop {
// There can't be two concurrent registrations because `Task` can only be polled
// by a unique pinned reference.
debug_assert!(state & REGISTERING == 0);
// If we're in the notifying state at this moment, just wake and return without
// registering.
if state & NOTIFYING != 0 {
abort_on_panic(|| waker.wake_by_ref());
return;
}
// Mark the state to let other threads know we're registering a new awaiter.
match self.state.compare_exchange_weak(
state,
state | REGISTERING,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
state |= REGISTERING;
break;
}
Err(s) => state = s,
}
}
// Put the waker into the awaiter field.
unsafe {
abort_on_panic(|| (*self.awaiter.get()) = Some(waker.clone()));
}
// This variable will contain the newly registered waker if a notification comes in before
// we complete registration.
let mut waker = None;
loop {
// If there was a notification, take the waker out of the awaiter field.
if state & NOTIFYING != 0 {
if let Some(w) = unsafe { (*self.awaiter.get()).take() } {
abort_on_panic(|| waker = Some(w));
}
}
// The new state is not being notified nor registered, but there might or might not be
// an awaiter depending on whether there was a concurrent notification.
// 这里需要解释以下，有点绕，如果我们没有从 awaiter 里面取到 Waker，那就表示没有
// NOTIFYING 状态，所以要设置 AWAITER，因为 awaiter 字段里面还有值
// 如果取到了，那么就是有 NOTIFYING 状态，我们就直接清空 AWAITER 的状态了，在下面直接
// 对取出的 Waker 进行通知
let new = if waker.is_none() {
(state & !NOTIFYING & !REGISTERING) | AWAITER
} else {
state & !NOTIFYING & !REGISTERING & !AWAITER
};
match self
.state
.compare_exchange_weak(state, new, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => break,
Err(s) => state = s,
}
}
// If there was a notification during registration, wake the awaiter now.
if let Some(w) = waker {
abort_on_panic(|| w.wake());
}
}
take
take 函数结合 register 函数一起看的话，还是比较容易理解的，这个函数的目的就是取出 awaiter 中的 Waker，返回出去，如果 awaiter 中的 Waker 和传入的 current 相同，那么就返回 None。这是一种不需要通知的情况，我们可以关注一下什么时候传入了 current，都是在 poll_task 中传入的，并且 notify 调用的时机都是任务状态发生改变的时候，其目的是唤醒不在当前上下文中的 awaiter，避免其一直阻塞下去。我们需要知道的是每次 poll_task 调用，Context 中的 Waker 不一定是相同的。
/// Takes the awaiter blocked on this task.
///
/// If there is no awaiter or if it is the same as the current waker, returns `None`.
#[inline]
pub(crate) fn take(&self, current: Option<&Waker>) -> Option<Waker> {
// Set the bit indicating that the task is notifying its awaiter.
let state = self.state.fetch_or(NOTIFYING, Ordering::AcqRel);
// If the task was not notifying or registering an awaiter...
if state & (NOTIFYING | REGISTERING) == 0 {
// Take the waker out.
let waker = unsafe { (*self.awaiter.get()).take() };
// Unset the bit indicating that the task is notifying its awaiter.
self.state
.fetch_and(!NOTIFYING & !AWAITER, Ordering::Release);
// Finally, notify the waker if it's different from the current waker.
if let Some(w) = waker {
match current {
None => return Some(w),
Some(c) if !w.will_wake(c) => return Some(w),
Some(_) => abort_on_panic(|| drop(w)),
}
}
}
None
}
notify
notify 函数比较简单，就是调用 take 取出 Waker 进行唤醒。
/// Notifies the awaiter blocked on this task.
///
/// If the awaiter is the same as the current waker, it will not be notified.
#[inline]
pub(crate) fn notify(&self, current: Option<&Waker>) {
if let Some(w) = self.take(current) {
abort_on_panic(|| w.wake());
}
}
尾
如果你仔细读完整篇文章，并自己通读一遍源码，我想你脑中已经对 async_task 的架构有一个全貌了。同时我们也可以试着向自己提出一些问题并试着做出解答，例如，为什么要单独 drop_future ？而不是等整个 RawTask 生命周期结束的时候再进行析构呢？通过这些提问同样会加深你对源码的理解。
