smol 为我们提供了一个异步运行时的实现,用来执行 async_task 封装的异步任务,当然如果你觉得它的实现不满足你的要求,你也可以根据自己需求自己实现。以下解析基于 async-executor 1.13.2
目标
通过阅读 async-executor 的源码后,我们应该理解 async-executor 如何进行任务调度以及通过多线程并行处理多个异步任务的。
async-executor 主要提供了两个异步运行时的实现,Executor 和 LocalExecutor,Executor 可以实现多线程调度异步任务,要求任务是 Send 的,LocalExecutor 为单线程异步执行器,不要求异步任务为 Send,LocalExecutor 基本上是 Executor 的包装,只是多了一个 _marker: PhantomData<Rc<()>>,让运行时 !Send。
入口
smol::block_on:阻塞当前线程,直到内部的 future 完成,这个的内部实现与我们现在探讨的内容关系不大,里面做的更多的是线程的阻塞和唤醒工作。
smol::spawn:这个是产生异步任务的最主要方式,我们主要查看其内部实现。
pub fn spawn<T: Send + 'static>(future: impl Future<Output = T> + Send + 'static) -> Task<T> {
static GLOBAL: OnceCell<Executor<'_>> = OnceCell::new();
fn global() -> &'static Executor<'static> {
GLOBAL.get_or_init_blocking(|| {
let num_threads = {
// Parse SMOL_THREADS or default to 1.
std::env::var("SMOL_THREADS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(1)
};
for n in 1..=num_threads {
thread::Builder::new()
.name(format!("smol-{}", n))
.spawn(|| loop {
catch_unwind(|| block_on(global().run(future::pending::<()>()))).ok();
})
.expect("cannot spawn executor thread");
}
// Prevent spawning another thread by running the process driver on this thread.
let ex = Executor::new();
#[cfg(not(target_os = "espidf"))]
ex.spawn(async_process::driver()).detach();
ex
})
}
global().spawn(future)
}
可以看到代码里面用了一个静态变量来存储异步运行时,通过 OnceCell 来确保异步运行时只会被初始化一次。在初始化代码中可以看到默认是启动一个线程来执行异步任务的,如果需要更多的线程同时来执行异步任务,则需要设置环境变量 SMOL_THREADS。现在主要的关注点就来到这段代码上了:
catch_unwind(|| block_on(global().run(future::pending::<()>()))).ok();
可以看到在每个线程中,都执行了一个 Executor::run 函数,并传入了一个永远不会完成的 future,通过这个操作,spawn 产生的任务就可以在多个线程中调度了。对于此处的疑问我们先放一放,要等后面看了运行时的实现之后就会豁然开朗了。
State
- queue 全局任务队列,一个异步运行时只会存在一个
- local_queues 本地任务队列,这个就是每个线程自己的任务队列,取决于线程数,当我们调用
Executor::run时,就会基于当前线程创建一个本地队列,并把它放入到这里 - notified 记录 sleepers 中的 Waker 有没有被通知过,本质上是一种优化手段,避免频繁的获取 sleepers 的锁
- sleepers 记录着当前哪些调度器正在休眠中,存储外部 poll 调度器时的 Waker,当有新的任务需要调度时,唤醒外部重新 poll 调度器
- active 存储所有的异步任务 Waker,当异步运行时 Drop 时唤醒所有异步任务,避免任务被遗忘卡住
/// The state of a executor.
struct State {
/// The global queue.
queue: ConcurrentQueue<Runnable>,
/// Local queues created by runners.
local_queues: RwLock<Vec<Arc<ConcurrentQueue<Runnable>>>>,
/// Set to `true` when a sleeping ticker is notified or no tickers are sleeping.
notified: AtomicBool,
/// A list of sleeping tickers.
sleepers: Mutex<Sleepers>,
/// Currently active tasks.
active: Mutex<Slab<Waker>>,
}
run
这个函数是我们在最上面看到的那个 run 函数,其内部也非常简单,构造了一个 Runner ,然后用一个永不完成的 future 不断的从 Runner 中取 Runnable 进行执行。在之前的篇章中,我们已经了解了执行 Runnable 其实就是执行某个异步任务中按 await 区间划分后的某一段区间的代码块,用于推进异步任务继续向下执行。这里还使用到了一个随机数生成器,这个主要是和任务偷取相关,用来产生偷取的随机区间。
pub async fn run<T>(&self, future: impl Future<Output = T>) -> T {
let mut runner = Runner::new(self);
let mut rng = fastrand::Rng::new();
// A future that runs tasks forever.
let run_forever = async {
loop {
for _ in 0..200 {
let runnable = runner.runnable(&mut rng).await;
runnable.run();
}
future::yield_now().await;
}
};
// Run `future` and `run_forever` concurrently until `future` completes.
future.or(run_forever).await
}
notify
前面提到,当调度器取不到任何的 Runnable 时(所有异步任务都处于挂起状态或者没有任何的异步任务),会将自己挂起,那么自然要有唤醒机制将调度器重新唤醒。
/// Notifies a sleeping ticker.
#[inline]
fn notify(&self) {
if self
.notified
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
let waker = self.sleepers.lock().unwrap().notify();
if let Some(w) = waker {
w.wake();
}
}
}
Runner
/// A worker in a work-stealing executor.
///
/// This is just a ticker that also has an associated local queue for improved cache locality.
struct Runner<'a> {
/// The executor state.
state: &'a State,
/// Inner ticker.
ticker: Ticker<'a>,
/// The local queue.
local: Arc<ConcurrentQueue<Runnable>>,
/// Bumped every time a runnable task is found.
ticks: usize,
}
new
从这里可以看出,我们创建 Runner 时,会将自己的任务队列放到 State 中的 local_queues 中,这样就为 Runner 做任务窃取提供了方便, Runner 可以通过 local_queues 访问到其它线程中创建的任务队列,当某个 Runner 本地队列中没有任务时,就可以尝试去偷其它 Runner 本地任务队列中的任务进行执行。这个是非常典型的任务窃取思想。
/// Creates a runner and registers it in the executor state.
fn new(state: &State) -> Runner<'_> {
let runner = Runner {
state,
ticker: Ticker::new(state),
local: Arc::new(ConcurrentQueue::bounded(512)),
ticks: 0,
};
state
.local_queues
.write()
.unwrap()
.push(runner.local.clone());
runner
}
runnable
这里 Ticker::runnable_with 接收一个搜索函数,其目的就是从队列中取出一个 Runnable 用于执行,在这个函数里面就看到了具体的任务窃取实现逻辑,窃取逻辑还是比较通俗易懂的。
- 先从本地队列中获取数据
- 本地队列没有从全局队列获取
- 全局队列没有从其它
Runner中的本地队列中偷一定数量的任务到自己的队列中 - ticks 到达一定计数的时候,从全局队列偷取一部分任务到自己的队列中,这个情况是防止本地队列一直有任务的情况下,全局队列中的任务被饿死
/// Waits for the next runnable task to run.
async fn runnable(&mut self, rng: &mut fastrand::Rng) -> Runnable {
let runnable = self
.ticker
.runnable_with(|| {
// Try the local queue.
if let Ok(r) = self.local.pop() {
return Some(r);
}
// Try stealing from the global queue.
if let Ok(r) = self.state.queue.pop() {
steal(&self.state.queue, &self.local);
return Some(r);
}
// Try stealing from other runners.
let local_queues = self.state.local_queues.read().unwrap();
// Pick a random starting point in the iterator list and rotate the list.
let n = local_queues.len();
let start = rng.usize(..n);
let iter = local_queues
.iter()
.chain(local_queues.iter())
.skip(start)
.take(n);
// Remove this runner's local queue.
let iter = iter.filter(|local| !Arc::ptr_eq(local, &self.local));
// Try stealing from each local queue in the list.
for local in iter {
steal(local, &self.local);
if let Ok(r) = self.local.pop() {
return Some(r);
}
}
None
})
.await;
// Bump the tick counter.
self.ticks = self.ticks.wrapping_add(1);
if self.ticks % 64 == 0 {
// Steal tasks from the global queue to ensure fair task scheduling.
steal(&self.state.queue, &self.local);
}
runnable
}
Drop
Runner 发生 Drop 时,会先从 local_queues 中移除自己的任务队列,然后将本地队列中的任务重新调度一次,这样这个 Runner 中的任务就会放到其它的任务队列中。
impl Drop for Runner<'_> {
fn drop(&mut self) {
// Remove the local queue.
self.state
.local_queues
.write()
.unwrap()
.retain(|local| !Arc::ptr_eq(local, &self.local));
// Re-schedule remaining tasks in the local queue.
while let Ok(r) = self.local.pop() {
r.schedule();
}
}
}
Ticker
Ticker 的作用提供一个异步函数用于提取一个可用的 Runnable 用于执行,当提取不到 Runnable 时,将次异步函数挂起,等待适当的时机恢复。
/// Runs task one by one.
struct Ticker<'a> {
/// The executor state.
state: &'a State,
/// Set to a non-zero sleeper ID when in sleeping state.
///
/// States a ticker can be in:
/// 1) Woken.
/// 2a) Sleeping and unnotified.
/// 2b) Sleeping and notified.
sleeping: usize,
}
sleep
挂起此 Ticker,这里的逻辑就有一些绕了,sleeping 记录的是当 Ticker 挂起时,其在 sleepers 中的 id,我们可以根据此 id 取得用于唤醒操作的 Waker。在已经有 id 的情况下,表示当前 Ticker 已经有过挂起操作,我们需要根据 id 更新 sleepers 中的存储的 Waker,根据 update 的注释,当我们根据 id 找到了对应的 Waker,那么会将之前的 Waker 替换为当前的,并且说明在挂起之后没有其它函数做过唤醒操作,这个时候 update 返回 false,如果根据 id 找不到对应的 Waker,那么说明这个 Ticker 在挂起后又被唤醒了,返回 true,这个时候根据 sleep 的定义,就不能返回 false 了。
结合 sleep 的逻辑和调用 sleep 处的逻辑,我们可以知道,Ticker 要经过两次 sleep 调用才会真正的进行挂起。为什么要这样做呢?其实就是为了避免 Ticker 挂起后又立马唤醒这种情况,要两次拿不到 Runnable 数据,才真正的进行挂起。
/// Moves the ticker into sleeping and unnotified state.
///
/// Returns `false` if the ticker was already sleeping and unnotified.
fn sleep(&mut self, waker: &Waker) -> bool {
let mut sleepers = self.state.sleepers.lock().unwrap();
match self.sleeping {
// Move to sleeping state.
0 => {
self.sleeping = sleepers.insert(waker);
}
// Already sleeping, check if notified.
id => {
if !sleepers.update(id, waker) {
return false;
}
}
}
self.state
.notified
.store(sleepers.is_notified(), Ordering::Release);
true
}
wake
将 Ticker 编程唤醒状态,逻辑很好理解,这里就不多解释了。
/// Moves the ticker into woken state.
fn wake(&mut self) {
if self.sleeping != 0 {
let mut sleepers = self.state.sleepers.lock().unwrap();
sleepers.remove(self.sleeping);
self.state
.notified
.store(sleepers.is_notified(), Ordering::Release);
}
self.sleeping = 0;
}
runnable_with
runnable_with 被调用的地方之前在 Runner 中已经看过了,包括 search 函数中提取 Runnable 的逻辑,这里主要的逻辑就是在取不到 Runnable 的时候,将 Ticker 进行挂起,恢复之后调用 self.wake() 将自身的挂起状态清除。同时唤醒其它可能正在挂起的 Ticker,这里主要是避免当前 Ticker 取到的 Runnable 执行耗时太长,而其它 Ticker 又在挂起的状态,导致后续的 Runnable 无法得到及时的执行这种情况。同时 self.state.notify() 对 Ticker 的唤醒不是一次性唤醒所有的 Ticker,而是一个一个进行唤醒的,这样避免了惊群效应。
/// Waits for the next runnable task to run, given a function that searches for a task.
async fn runnable_with(&mut self, mut search: impl FnMut() -> Option<Runnable>) -> Runnable {
future::poll_fn(|cx| {
loop {
match search() {
None => {
// Move to sleeping and unnotified state.
if !self.sleep(cx.waker()) {
// If already sleeping and unnotified, return.
return Poll::Pending;
}
}
Some(r) => {
// Wake up.
self.wake();
// Notify another ticker now to pick up where this ticker left off, just in
// case running the task takes a long time.
self.state.notify();
return Poll::Ready(r);
}
}
}
})
.await
}
}
Drop
当 Ticker Drop 时,如果这个 Ticker 正在挂起中,并且这个 Ticker 通过 State::notify 操作进行了唤醒操作,那么此时就必须要在进行一次 notify 操作了,因为此 Ticker 已经不会响应 notify 操作了,前面我们也说过了, Ticker 是一个一个进行唤醒的,如果不这样操作,那么就会造成漏唤醒。严重的情况下任务就卡住了。
impl Drop for Ticker<'_> {
fn drop(&mut self) {
// If this ticker is in sleeping state, it must be removed from the sleepers list.
if self.sleeping != 0 {
let mut sleepers = self.state.sleepers.lock().unwrap();
let notified = sleepers.remove(self.sleeping);
self.state
.notified
.store(sleepers.is_notified(), Ordering::Release);
// If this ticker was notified, then notify another ticker.
if notified {
drop(sleepers);
self.state.notify();
}
}
}
}
Sleepers
Sleepers 存储的就是处于挂起状态的 Ticker 关联的 Waker。
/// A list of sleeping tickers.
struct Sleepers {
/// Number of sleeping tickers (both notified and unnotified).
count: usize,
/// IDs and wakers of sleeping unnotified tickers.
///
/// A sleeping ticker is notified when its waker is missing from this list.
wakers: Vec<(usize, Waker)>,
/// Reclaimed IDs.
free_ids: Vec<usize>,
}
insert
向 Sleepers 中插入一个 Waker,并返回与之关联的 id,这里采用了 id 复用的机制,当从 Sleepers 中移除 Waker 时,会将这个 id 放入 free_ids 池子中,以供下次使用。
/// Inserts a new sleeping ticker.
fn insert(&mut self, waker: &Waker) -> usize {
let id = match self.free_ids.pop() {
Some(id) => id,
None => self.count + 1,
};
self.count += 1;
self.wakers.push((id, waker.clone()));
id
}
update
这个函数的功能主要是根据已有的 id 更新与之相关联的 Waker,如果找不到,说明这个 Ticker 被唤醒了,Waker 被移除了。
/// Re-inserts a sleeping ticker's waker if it was notified.
///
/// Returns `true` if the ticker was notified.
fn update(&mut self, id: usize, waker: &Waker) -> bool {
for item in &mut self.wakers {
if item.0 == id {
item.1.clone_from(waker);
return false;
}
}
self.wakers.push((id, waker.clone()));
true
}
remove
根据 id 移除 Waker,同样的,如果根据 id 没有找到对应的 Waker,那么它就是通过 notify 的方式移除掉了。
/// Removes a previously inserted sleeping ticker.
///
/// Returns `true` if the ticker was notified.
fn remove(&mut self, id: usize) -> bool {
self.count -= 1;
self.free_ids.push(id);
for i in (0..self.wakers.len()).rev() {
if self.wakers[i].0 == id {
self.wakers.remove(i);
return false;
}
}
true
}
is_notified
这个函数比较有意思,这个函数的值是存储到 State 的 notified 中的,需要结合 Sleepers::notify 和 Ticker::wake 来一起看,主要关注 self.count 的值变化。当我们调用 Sleepers::notify 之后,当 wakers 的数量和 count 相等时,会从 wakers 中 pop 一个 Waker,但是不会改变 count。当调用这个 Waker 的 wake 方法时,就会再次唤醒 Ticker::runnable_with 这个异步函数,这个函数里面就会调用 Ticker:wake 这个函数移除自身的挂起状态,里面会调用 Sleepers::remove,这个函数里面会修改 count 的数量,此时 wakers 和 count 的数量就又相等了。
/// Returns `true` if a sleeping ticker is notified or no tickers are sleeping.
fn is_notified(&self) -> bool {
self.count == 0 || self.count > self.wakers.len()
}
notify
结合 is_notified 的解释,这里只有 wakers 的长度和 count 相等时,才弹出 Waker,这样做的目的就是一个一个的唤醒 Ticker,避免惊群效应,唤醒的 Waker 会再次通过 State::notify 唤醒可能正在挂起的另外一个 Waker。
/// Returns notification waker for a sleeping ticker.
///
/// If a ticker was notified already or there are no tickers, `None` will be returned.
fn notify(&mut self) -> Option<Waker> {
if self.wakers.len() == self.count {
self.wakers.pop().map(|item| item.1)
} else {
None
}
}
Executor
spawn_inner
这个函数还是比较好理解的,在 Slab<Waker> 中预留一个 Waker 空位,将传入的 future 用 AsyncCallOnDrop 包裹,以便在 future Drop 的时候从 Slab<Waker> 中移除这个 Waker,后面就是将 future 用 async-task 里面的结构包裹,传入一个 schedule 函数。
/// Spawn a future while holding the inner lock.
///
/// # Safety
///
/// If this is an `Executor`, `F` and `T` must be `Send`.
unsafe fn spawn_inner<T: 'a>(
&self,
future: impl Future<Output = T> + 'a,
active: &mut Slab<Waker>,
) -> Task<T> {
// Remove the task from the set of active tasks when the future finishes.
let entry = active.vacant_entry();
let index = entry.key();
let state = self.state_as_arc();
let future = AsyncCallOnDrop::new(future, move || drop(state.active().try_remove(index)));
// Create the task and register it in the set of active tasks.
//
// SAFETY:
//
// If `future` is not `Send`, this must be a `LocalExecutor` as per this
// function's unsafe precondition. Since `LocalExecutor` is `!Sync`,
// `try_tick`, `tick` and `run` can only be called from the origin
// thread of the `LocalExecutor`. Similarly, `spawn` can only be called
// from the origin thread, ensuring that `future` and the executor share
// the same origin thread. The `Runnable` can be scheduled from other
// threads, but because of the above `Runnable` can only be called or
// dropped on the origin thread.
//
// `future` is not `'static`, but we make sure that the `Runnable` does
// not outlive `'a`. When the executor is dropped, the `active` field is
// drained and all of the `Waker`s are woken. Then, the queue inside of
// the `Executor` is drained of all of its runnables. This ensures that
// runnables are dropped and this precondition is satisfied.
//
// `self.schedule()` is `Send`, `Sync` and `'static`, as checked below.
// Therefore we do not need to worry about what is done with the
// `Waker`.
let (runnable, task) = Builder::new()
.propagate_panic(true)
.spawn_unchecked(|()| future, self.schedule());
entry.insert(runnable.waker());
runnable.schedule();
task
}
schedule
从这里看到,所有的 Runnable 目前都是插入到全局队列当中的,再靠 work-stealing 机制偷取任务到自己本地的队列。这里最好的方式还是看当前任务在哪个线程中调度,就直接把这个 Runnable 放到对应线程的本地队列最好。我也查看了对应的提交记录,确实是有相关的实现,不过因为有 BUG ,最后回滚了。
/// Returns a function that schedules a runnable task when it gets woken up.
fn schedule(&self) -> impl Fn(Runnable) + Send + Sync + 'static {
let state = self.state_as_arc();
// TODO: If possible, push into the current local queue and notify the ticker.
move |runnable| {
state.queue.push(runnable).unwrap();
state.notify();
}
}
结束
smol 核心的东西基本上都说完了,还有一些函数因为不是那么重要,就省略了,建议最好还是自己完整的把源码看一遍,文章无法做到面面俱到。