Tokio运行时是rust生态异步应用程序的核心,有三个关键的服务

  • I/O事件循环:用于驱动I/O资源并分发事件
  • 调度器:执行异步任务
  • 定时器:定时调度任务

一个简单的示例

下面是一个简单的使用tokio编写的控制台输入输出的echo程序

use std::io;
 
use tokio::io::{AsyncBufReadExt, BufReader};
 
#[tokio::main]
async fn main() -> io::Result<()> {
    let stdin = tokio::io::stdin();
    let reader = BufReader::new(stdin);
    let mut lines = reader.lines();
 
    while let Some(line) = lines.next_line().await? {
        if line.is_empty() {
            continue;
        }
        match line.trim() {
            "exit" | "quit" => {
                break;
            }
            v => {
                println!("hello, {}", v)
            }
        }
    }
    Ok(())
}

使用#[tokio::main]这个宏来自动生成tokio运行时的main方法了,将上面的宏展开对应的源码如下

fn main() -> io::Result<()> {
    let body = async {
        let stdin = tokio::io::stdin();
        let reader = BufReader::new(stdin);
        let mut lines = reader.lines();
        while let Some(line) = lines.next_line().await? {
            if line.is_empty() {
                continue;
            }
            match line.trim() {
                "exit" | "quit" => {
                    break;
                }
                v => {
                    std::io::_print(std::format_args_nl!("hello, {}", v));
                }
            }
        }
        Ok(())
    };
    #[cfg(all())]
    #[allow(
        clippy::expect_used,
        clippy::diverging_sub_expression,
        clippy::needless_return,
        clippy::unwrap_in_result
    )]
    {
        return tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .expect("Failed building the Runtime")
            .block_on(body);
    }
    #[cfg(not(all()))]
    {
        panic!("fell through checks")
    }
}

展开之后就可以看到就是将原来的异步的main方法作为一个异步方法的方法体,然后创建tokio运行时,将这个异步方法交给异步运行时执行。

异步运行时创建

上面的展开代码创建了一个多线程的调用器,这也是tokio的默认行为。tokio还是支持创建单线程的,也介绍了如何进行选择

是否需要工作窃取或多线程调度器?
    │ Yes               │ No
    │                   │
    ▼                   ▼
多线程运行时       是否需要执行 !Send Future?
                        │ Yes               │ No
                        │                   │
                        ▼                   ▼
                本地运行时(不稳定)     当前线程运行时

这里就只以多线程运行时进行了解。

#[derive(Debug)]
pub struct Runtime {
    /// Task scheduler
    scheduler: Scheduler,
 
    /// Handle to runtime, also contains driver handles
    handle: Handle,
 
    /// Blocking pool handle, used to signal shutdown
    blocking_pool: BlockingPool,
}

运行时的结构非常简洁,就只有三个字段,所有的逻辑都被封装在底层了。

  • scheduler:就是调度器。负责任务的调度。多线程版本还涉及工作窃取
  • handle:运行时句柄。提供对运行时的访问入口,核心逻辑都在这个句柄当中
  • blocking_pool:阻塞线程池。用来执行阻塞的、cpu密集型或者同步的代码,防止阻塞异步任务调度。这些阻塞的任务会在独立的线程当中运行的。

tokio使用Builder来负责构建出这个运行时

pub struct Builder {
    /// Runtime type
    kind: Kind,
 
    /// Whether or not to enable the I/O driver
    enable_io: bool,
    nevents: usize,
 
    /// Whether or not to enable the time driver
    enable_time: bool,
 
    /// Whether or not the clock should start paused.
    start_paused: bool,
 
    /// 多线程调度器使用的工作线程数
    worker_threads: Option<usize>,
 
    /// Cap on thread usage.
    max_blocking_threads: usize,
 
	/// 各种钩子函数
    /// Name fn used for threads spawned by the runtime.
    pub(super) thread_name: ThreadNameFn,
 
    /// Stack size used for threads spawned by the runtime.
    pub(super) thread_stack_size: Option<usize>,
 
    /// Callback to run after each thread starts.
    pub(super) after_start: Option<Callback>,
    
	...
 
    /// To run after each task is terminated.
    pub(super) after_termination: Option<TaskCallback>,
 
    /// Customizable keep alive timeout for `BlockingPool`
    pub(super) keep_alive: Option<Duration>,
	...
}

创建多线程运行时的流程如下

impl Builder {
	fn build_threaded_runtime(&mut self) -> io::Result<Runtime> {
		// 1. 创建IO事件驱动
		let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?;
		// 2. 创建阻塞线程池
		let blocking_pool =
			blocking::create_blocking_pool(self, self.max_blocking_threads + worker_threads);
		
		// 3. 创建调度器和运行时句柄
		let (scheduler, handle, launch) = MultiThread::new(
                worker_threads,
                driver,
                driver_handle,
                blocking_spawner,
				......
            );
        let handle = Handle { inner: scheduler::Handle::MultiThread(handle) };
        // 4. 启动工作线程
		let _enter = handle.enter();
		launch.launch();
		
		Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool))
	}
}

这样后台执行异步线程启动完成了,用于上层业务使用的异步运行时操作的入口对象Runtime也出创建完成了。

block_ont提交主任务

通过block_on方法提交了一个异步任务给异步运行时并阻塞当前线程直到任务完成,便从此作为入口来看一个任务从提交到完成的整个流程是怎样的。

impl MultiThread {
    pub(crate) fn block_on<F>(&self, handle: &scheduler::Handle, future: F) -> F::Output
    where
        F: Future,
    {
        crate::runtime::context::enter_runtime(handle, true, |blocking| {
            blocking.block_on(future).expect("failed to park thread")
        })
    }
}

上面的代码本身的就已经说明了逻辑,那就是进入异步运行时,拿到一个BlockingRegionGuard,在这个阻塞区域内执行异步任务。

pub(crate) fn enter_runtime<F, R>(handle: &scheduler::Handle, allow_block_in_place: bool, f: F) -> R
where
    F: FnOnce(&mut BlockingRegionGuard) -> R,
{
	// CONTEXT是一个线程局部变量,每个线程都有自己的实例
    let maybe_guard = CONTEXT.with(|c| {
        if c.runtime.get().is_entered() {
			// tokio不允许嵌套运行时
            None
        } else {
            // 设置当前上下文进入运行时
            c.runtime.set(EnterRuntime::Entered {
                allow_block_in_place,
            });
 
            // 生成新的随机数种子,用来保证每个线程的任务ID的随机和唯一性
            let rng_seed = handle.seed_generator().next_seed();
 
            // Swap the RNG seed
            let mut rng = c.rng.get().unwrap_or_else(FastRand::new);
            let old_seed = rng.replace_seed(rng_seed);
            c.rng.set(Some(rng));
			// 设置当前CONTEXT的随机数种子为新生成的,运行时句柄为外部传入的
			// 并保存旧的到Guard中,方便退出之后还原。
            Some(EnterRuntimeGuard {
                blocking: BlockingRegionGuard::new(),
                handle: c.set_current(handle),
                old_seed,
            })
        }
    });
 
	// 执行任务
    if let Some(mut guard) = maybe_guard {
        return f(&mut guard.blocking);
    }
 
    panic!(
        "Cannot start a runtime from within a runtime. This happens \
            because a function (like `block_on`) attempted to block the \
            current thread while the thread is being used to drive \
            asynchronous tasks."
    );
}

上面的逻辑就是将当前线程的tokio上下文中的随机数种子以及运行时句柄设置为新的,并通过EnterRuntimeGuard保留旧的和当前上下文执行所需的结构。这个结构体一定是实现了Drop的,在Drop负责还原Context

impl Drop for EnterRuntimeGuard {
    fn drop(&mut self) {
        CONTEXT.with(|c| {
            assert!(c.runtime.get().is_entered());
            c.runtime.set(EnterRuntime::NotEntered);
            // Replace the previous RNG seed
            let mut rng = c.rng.get().unwrap_or_else(FastRand::new);
            rng.replace_seed(self.old_seed.clone());
            c.rng.set(Some(rng));
        });
    }
}

继续回到EnterRuntimeGuard,现在异步任务由这个对象提交,本身的内容非常简单,就是创建一个CachedParkThread对象通过这个对象执行异步任务

pub(crate) struct BlockingRegionGuard {
    _p: PhantomData<NotSendOrSync>,
}
 
impl BlockingRegionGuard {
    /// Blocks the thread on the specified future, returning the value with
    /// which that future completes.
    pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError>
    where
        F: std::future::Future,
    {
        use crate::runtime::park::CachedParkThread;
 
        let mut park = CachedParkThread::new();
        park.block_on(f)
    }
}

由于我们需要阻塞任务直到执行,那么就需要有一个机制用来控制线程的阻塞和唤醒,这个机制就是条件变量。每个线程都有一个线程局部的ParkThread用来控制阻塞,并可以转为UnparkThread进行唤醒,这两者内部共享相同的数据,CachedParkThread则是两者的统一门面。

#[derive(Debug)]
struct Inner {
    state: AtomicUsize,
    mutex: Mutex<()>,
    condvar: Condvar,
}
 
#[derive(Debug)]
pub(crate) struct ParkThread {
    inner: Arc<Inner>,
}
 
/// Unblocks a thread that was blocked by `ParkThread`.
#[derive(Clone, Debug)]
pub(crate) struct UnparkThread {
    inner: Arc<Inner>,
}
 
 
tokio_thread_local! {
    static CURRENT_PARKER: ParkThread = ParkThread::new();
}
 
/// Blocks the current thread using a condition variable.
#[derive(Debug)]
pub(crate) struct CachedParkThread {
    _anchor: PhantomData<Rc<()>>,
}
 
impl CachedParkThread {
	// 获取waker
    pub(crate) fn waker(&self) -> Result<Waker, AccessError> {
        self.unpark().map(UnparkThread::into_waker)
    }
	// 内部方法,用来获取当前线程的UnparkThread
    fn unpark(&self) -> Result<UnparkThread, AccessError> {
        self.with_current(ParkThread::unpark)
    }
	// 阻塞当前线程
    pub(crate) fn park(&mut self) {
        self.with_current(|park_thread| park_thread.inner.park())
            .unwrap();
    }
	// 带有超时时间地阻塞当前线程
    pub(crate) fn park_timeout(&mut self, duration: Duration) {
        self.with_current(|park_thread| park_thread.inner.park_timeout(duration))
            .unwrap();
    }
 
    /// 对当前ParkThread执行指定地方法
    fn with_current<F, R>(&self, f: F) -> Result<R, AccessError>
    where
        F: FnOnce(&ParkThread) -> R,
    {
        CURRENT_PARKER.try_with(|inner| f(inner))
    }
}
 

最核心地方法是这里的block_on方法

impl CachedParkThread {
    pub(crate) fn block_on<F: Future>(&mut self, f: F) -> Result<F::Output, AccessError> {
        use std::task::Context;
        use std::task::Poll::Ready;
 
        let waker = self.waker()?;
        let mut cx = Context::from_waker(&waker);
 
        pin!(f);
 
        loop {
            if let Ready(v) = crate::task::coop::budget(|| f.as_mut().poll(&mut cx)) {
                return Ok(v);
            }
 
            self.park();
        }
    }
}

可以看到这里就是不断循环阻塞直到异步任务执行结束。

这样来看,通过Runtimeblock_on方法提交的任务就是阻塞在提交任务的线程,直到任务结束。

前面在tokio运行时初始化的时候看到启动了工作线程,但是这里通过block_on提交的任务却是在当前线程阻塞运行完成。这是因为我们可以将tokio的任务分为两种,主任务和子任务。

  • 主任务:通过block_on提交,进入运行时上下文,阻塞直到主任务执行结束避免应用程序直接结束
  • 子任务:通过spawn生成。这里会走任务队列调度执行的逻辑。

下面先绘制一下主任务的执行流程

graph TD
    A[开始] --> B[Runtime::block_on]
    B --> C[进入运行时上下文]
    C --> D[创建CachedParkThread]
    D --> E[获取Waker]
    E --> F[循环poll任务]
    F --> G{任务是否完成?}
    G -->|否| H[阻塞当前线程]
    H --> F
    G -->|是| I[返回结果]
    I --> J[退出运行时上下文]
    J --> K[结束]

spawn生成子任务

在主任务中可以不断生成子任务提交给运行时执行,tokio源码的echo示例程序如下

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Allow passing an address to listen on as the first argument of this
    // program, but otherwise we'll just set up our TCP listener on
    // 127.0.0.1:8080 for connections.
    let addr = env::args()
        .nth(1)
        .unwrap_or_else(|| DEFAULT_ADDR.to_string());
 
    // Next up we create a TCP listener which will listen for incoming
    // connections. This TCP listener is bound to the address we determined
    // above and must be associated with an event loop.
    let listener = TcpListener::bind(&addr).await?;
    println!("Listening on: {addr}");
 
    loop {
        // Asynchronously wait for an inbound socket.
        let (mut socket, addr) = listener.accept().await?;
 
        // And this is where much of the magic of this server happens. We
        // crucially want all clients to make progress concurrently, rather than
        // blocking one on completion of another. To achieve this we use the
        // `tokio::spawn` function to execute the work in the background.
        //
        // Essentially here we're executing a new task to run concurrently,
        // which will allow all of our clients to be processed concurrently.
 
        tokio::spawn(async move {
            let mut buf = vec![0; BUFFER_SIZE];
 
            // In a loop, read data from the socket and write the data back.
            loop {
                match socket.read(&mut buf).await {
                    Ok(0) => {
                        // Connection closed by peer
                        return;
                    }
                    Ok(n) => {
                        // Write the data back. If writing fails, log the error and exit.
                        if let Err(e) = socket.write_all(&buf[0..n]).await {
                            eprintln!("Failed to write to socket {}: {}", addr, e);
                            return;
                        }
                    }
                    Err(e) => {
                        eprintln!("Failed to read from socket {}: {}", addr, e);
                        return;
                    }
                }
            }
        });
    }
}
 

可以看到主任务就是就是一个循环,不断获取监听到客户端连接,然后通过spawn提交子任务,单独处理这个客户端连接的请求

#[track_caller]
pub fn spawn<F>(future: F) -> JoinHandle<F::Output>
where
	F: Future + Send + 'static,
	F::Output: Send + 'static,
{
	let fut_size = std::mem::size_of::<F>();
	if fut_size > BOX_FUTURE_THRESHOLD {
		spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size))
	} else {
		spawn_inner(future, SpawnMeta::new_unnamed(fut_size))
	}
}

spawn的逻辑比较简单,就是检测一下协程的大小,然后决定将协程保存在栈上还是堆上。因为Rust使用的是协作式的无栈协程,所有的协程本身上都是一个实现了Future trait的状态机,当协程复杂的时候通常对应的状态机大小也比较大。

  • 小Future栈分配,来避免堆分配的开销。小Future更好地利用CPU缓存局部性
  • 大Future堆分配,来避免栈溢出。堆分配的Future更容易在线程间移动

spawn_inner就是单纯地先包装一下用来生成id,追踪关键信息,核心的还是调用handlespawn方法

    #[track_caller]
    pub(super) fn spawn_inner<T>(future: T, meta: SpawnMeta<'_>) -> JoinHandle<T::Output>
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static,
    {
        use crate::runtime::{context, task};
 
        #[cfg(all(
            tokio_unstable,
            feature = "taskdump",
            feature = "rt",
            target_os = "linux",
            any(
                target_arch = "aarch64",
                target_arch = "x86",
                target_arch = "x86_64"
            )
        ))]
        let future = task::trace::Trace::root(future);
        let id = task::Id::next();
        let task = crate::util::trace::task(future, "task", meta, id.as_u64());
 
        match context::with_current(|handle| handle.spawn(task, id, meta.spawned_at)) {
            Ok(join_handle) => join_handle,
            Err(e) => panic!("{}", e),
        }
    }

看多线程的Handle实现

impl Handle {
    /// Spawns a future onto the thread pool
    pub(crate) fn spawn<F>(
        me: &Arc<Self>,
        future: F,
        id: task::Id,
        spawned_at: SpawnLocation,
    ) -> JoinHandle<F::Output>
    where
        F: crate::future::Future + Send + 'static,
        F::Output: Send + 'static,
    {
        Self::bind_new_task(me, future, id, spawned_at)
    }
}
 
    #[track_caller]
    pub(super) fn bind_new_task<T>(
        me: &Arc<Self>,
        future: T,
        id: task::Id,
        spawned_at: SpawnLocation,
    ) -> JoinHandle<T::Output>
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static,
    {
	    // 提交Future到任务集合当中
        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
		// 执行spawn的hook函数
        me.task_hooks.spawn(&TaskMeta {
            id,
            spawned_at,
            _phantom: Default::default(),
        });
		// 调度任务
        me.schedule_option_task_without_yield(notified);
		// 返回handle
        handle
    }

逻辑就是将Future提交到任务队列当中并启动调用之后就返回。

工作线程

在异步运行时的最后启动了后台的工作线程,执行的方法如下

fn run(worker: Arc<Worker>) {
	// 捕获panic退出进程
    #[allow(dead_code)]
    struct AbortOnPanic;
 
    impl Drop for AbortOnPanic {
        fn drop(&mut self) {
            if std::thread::panicking() {
                eprintln!("worker thread panicking; aborting process");
                std::process::abort();
            }
        }
    }
 
    // Catching panics on worker threads in tests is quite tricky. Instead, when
    // debug assertions are enabled, we just abort the process.
    #[cfg(debug_assertions)]
    let _abort_on_panic = AbortOnPanic;
 
    // Acquire a core. If this fails, then another thread is running this
    // worker and there is nothing further to do.
    // 获取工作核心(博涵任务队列等核心状态)
    let core = match worker.core.take() {
        Some(core) => core,
        None => return,
    };
 
    // 记录当前线程ID到监控指标中,用于调试和性能监控
    worker.handle.shared.worker_metrics[worker.index].set_thread_id(thread::current().id());
 
    let handle = scheduler::Handle::MultiThread(worker.handle.clone());
	// 进入运行时上下文
    crate::runtime::context::enter_runtime(&handle, true, |_| {
        // Set the worker context.
        // 创建调度上下文
        let cx = scheduler::Context::MultiThread(Context {
            worker,
            core: RefCell::new(None),
            defer: Defer::new(),
        });
		// 设置当前线程的调度器上下文
        context::set_scheduler(&cx, || {
            let cx = cx.expect_multi_thread();
 
            // This should always be an error. It only returns a `Result` to support
            // using `?` to short circuit.
            // 主循环,不断从任务队列获取并执行任务
            assert!(cx.run(core).is_err());
 
            // Check if there are any deferred tasks to notify. This can happen when
            // the worker core is lost due to `block_in_place()` being called from
            // within the task.
            // 检查并唤醒任何延迟的任务
            cx.defer.wake();
        });
    });
}

一个Worker只能被一个线程拿到执行,进入异步运行时,由调度器运行异步任务。

impl Context {
    fn run(&self, mut core: Box<Core>) -> RunResult {
        // Reset `lifo_enabled` here in case the core was previously stolen from
        // a task that had the LIFO slot disabled.
        self.reset_lifo_enabled(&mut core);
 
        // Start as "processing" tasks as polling tasks from the local queue
        // will be one of the first things we do.
        core.stats.start_processing_scheduled_tasks();
 
        while !core.is_shutdown {
            self.assert_lifo_enabled_is_correct(&core);
 
            if core.is_traced {
                core = self.worker.handle.trace_core(core);
            }
 
            // Increment the tick
            core.tick();
 
            // Run maintenance, if needed
            core = self.maintenance(core);
 
            // First, check work available to the current worker.
            if let Some(task) = core.next_task(&self.worker) {
                core = self.run_task(task, core)?;
                continue;
            }
 
            // We consumed all work in the queues and will start searching for work.
            core.stats.end_processing_scheduled_tasks();
 
            // There is no more **local** work to process, try to steal work
            // from other workers.
            if let Some(task) = core.steal_work(&self.worker) {
                // Found work, switch back to processing
                core.stats.start_processing_scheduled_tasks();
                core = self.run_task(task, core)?;
            } else {
                // Wait for work
                core = if !self.defer.is_empty() {
                    self.park_yield(core)
                } else {
                    self.park(core)
                };
                core.stats.start_processing_scheduled_tasks();
            }
        }
 
        #[cfg(all(tokio_unstable, feature = "time"))]
        {
            match self.worker.handle.timer_flavor {
                TimerFlavor::Traditional => {}
                TimerFlavor::Alternative => {
                    util::time_alt::shutdown_local_timers(
                        &mut core.time_context.wheel,
                        &mut core.time_context.canc_rx,
                        self.worker.handle.take_remote_timers(),
                        &self.worker.handle.driver,
                    );
                }
            }
        }
 
        core.pre_shutdown(&self.worker);
        // Signal shutdown
        self.worker.handle.shutdown_core(core);
        Err(())
    }
}

在调度器的run方法里不断循环来执行异步任务,直到调度器关闭。

  • tick递增,判断是否需要运行维护任务。处理一下状态检测
  • 有限从线程的本地队列当中获取任务并执行后开启下一轮循环
  • 本地队列为空,则尝试从其他工作线程窃取任务
  • 本地和远端都没有任务,那么当前线程就进入等待状态了。等待还分为两种
    • park_yield:有延迟任务时就走这个,带超时时间的Poll
    • park:没有延迟任务就走这个暂停,阻塞直到有事件触发