Rust Async: async-task源码分析

最新推荐文章于 2024-07-18 10:32:15 发布

Rust语言中文社区

最新推荐文章于 2024-07-18 10:32:15 发布

阅读量1.7k

点赞数

本文链接：https://blog.csdn.net/u012067469/article/details/103175336

版权

本文深入分析了Rust异步生态中的async-task库，探讨了Task的结构、状态管理、JoinHandle的实现以及Waker的相关函数。async-task通过一次内存分配、隐藏RawWaker和提供JoinHandle等功能，简化了Executor的实现。文章详细解释了Task的schedule、run方法，以及JoinHandle的cancel、drop和poll方法的工作原理。

摘要由CSDN通过智能技术生成

本文转载自知乎地址：https://zhuanlan.zhihu.com/p/92679351?utm_source=wechat_session&utm_medium=social&utm_oi=51691969839104

async-std是rust异步生态中的基础运行时库之一，核心理念是合理的性能 + 用户友好的api体验。经过几个月密集的开发，前些天已经发布1.0稳定版本。因此是时候来一次深入的底层源码分析。async-std的核心是一个带工作窃取的多线程Executor，而其本身的实现又依赖于async-task这个关键库，因此本文主要对async-task的源码进行分析。

当Future提交给Executor执行时，Executor需要在堆上为这个Future分配空间，同时需要给它分配一些状态信息，比如Future是否可以执行(poll)，是否在等待被唤醒，是否已经执行完成等等。我们一般把提交给Executor执行的Future和其连带的状态称为 task。async-task这个库就是对task进行抽象封装，以便于Executor的实现，其有几个创新的特性：

整个task只需要一次内存分配；
完全隐藏了RawWaker，以避免实现Executor时处理unsafe代码的麻烦；
提供了 JoinHandle，这样spawn函数对Future没有 Output=()的限制，极大方便用户使用；

使用方式

async-task只对外暴露了一个函数接口以及对应了两个返回值类型：

pub fn spawn<F, R, S, T>(future: F, schedule: S, tag: T) -> (Task<T>, JoinHandle<R, T>)where    F: Future<Output = R> + Send + 'static,    R: Send + 'static,    S: Fn(Task<T>) + Send + Sync + 'static,    T: Send + Sync + 'static,

其中，参数future表示要执行的Future，schedule是一个闭包，当task变为可执行状态时会调用这个函数以调度该task重新执行，tag是附带在该task上的额外上下文信息，比如task的名字，id等。返回值Task就是构造好的task对象，JoinHandle实现了Future，用于接收最终执行的结果。

值得注意的是spawn这个函数并不会做类似在后台进行计算的操作，而仅仅是分配内存，创建一个task出来，因此其实叫create_task反而更为恰当且好理解。

Task提供了如下几个方法：

    // 对该task进行调度
    pub fn schedule(self);
    // poll一次内部的Future，如果Future完成了，则会通知JoinHandle取结果。否则task进
    // 入等待，直到被被下一次唤醒进行重新调度执行。
    pub fn run(self);
    // 取消task的执行
    pub fn cancel(&self);
    // 返回创建时传入的tag信息
    pub fn tag(&self) -> &T;

JoinHandle实现了Future trait，同时也提供了如下几个方法：

    // 取消task的执行
    pub fn cancel(&self);
    // 返回创建时传入的tag信息
    pub fn tag(&self) -> &T;

同时，Task和JoinHandle都实现了Send+Sync，所以他们可以出现在不同的线程，并通过tag方法可以同时持有 &T，因此spawn函数对T有Sync的约束。

借助于async_task的抽象，下面的几十行代码就实现了一个共享全局任务队列的多线程Executor：

use std::future::Future;
use std::thread;

use crossbeam::channel::{unbounded, Sender};
use futures::executor;
use once_cell::sync::Lazy;

static QUEUE: Lazy<Sender<async_task::Task<()>>> = Lazy::new(|| {
    
    let (sender, receiver) = unbounded::<async_task::Task<()>>();
    for _ in 0..4 {
    
        let recv = receiver.clone();

        thread::spawn(|| {
    
            for task in recv {
    
                task.run();
            }
        });
    }

    sender
});

fn spawn<F, R>(future: F) -> async_task::JoinHandle<R, ()>
where
    F: Future<Output = R> + Send + 'static,
    R: Send + 'static,
{
    
    let schedule = |task| QUEUE.send(task).unwrap();
    let (task, handle) = async_task::spawn(future, schedule, ());

    task.schedule();

    handle
}

fn main() {
    
    let handles: Vec<_> = (0..10).map(|i| {
    
        spawn(async move {
    
            println!("Hello from task {}", i);
        })
    }).collect();

    // Wait for the tasks to finish.
    for handle in handles {
    
        executor::block_on(handle);
    }
}

Task的结构图

通常rust里的并发数据结构会包含底层的实现，一般叫Inner或者RawXXX，包含大量裸指针等unsafe操作，然后再其基础上进行类型安全包装，提供上层语义。比如channel，上层暴露出 Sender和 Receiver,其行为不一样，但内部表示是完全一样的。async-task也类似，JoinHandle, Task以及调用Future::poll时传递的Waker类型内部都共享同一个RawTask结构。由于JoinHandle本身是一个Future，整个并发结构还有第四个角色-在JoinHandle上调用poll的task传递的Waker，为避免引起混淆就称它为Awaiter吧。整个的结构图大致如下：

整个task在堆上一次分配，内存布局按Header，Tag, Schedule,Future/Output排列。由于Future和Output不同时存在，因此他们共用同一块内存。

JoinHandle：只有一个，不访问Future，可以访问Output，一旦销毁就不再生成；
Task：主要访问Future，销毁后可以继续生成，不过同一时间最多只有一个，这样可以避免潜在的多个Task对Future进行并发访问的bug；
Waker：可以存在多份，主要访问schedule数据，由于spawn函数的参数要求schedule必须是Send+Sync，因此多个waker并发调用是安全的。
Header：本身包含三个部分，state是一个原子变量，包含引用计数，task的执行状态，awaiter锁等信息；awaiter保存的是JoinHandle所在的task执行时传递的Waker，用于当Output生成后通知JoinHandle来取；vtable是一个指向静态变量的虚表指针。