Rust 异步编程基础

最新推荐文章于 2024-12-03 17:37:29 发布

Vallelonga

最新推荐文章于 2024-12-03 17:37:29 发布

阅读量1.1k

点赞数 13

文章标签： rust 网络服务器

本文链接：https://blog.csdn.net/Vallelonga/article/details/140353664

版权

本文意在作为对 Rust 异步编程圣经前两章的导读，希望能帮助读者理解 Rust 异步编程的底层逻辑。读者应该先坚持着把 Rust 异步编程圣经的前两章读上两三遍，不必做到理解，但至少先熟悉它说了什么，再带着问题回来看这个导读。

Future、async/.await

Rust 官方的基本库和 Rust 编译器提供了 Future trait 和 async/.await 语法作为整个 Rust 异步编程的基础。想要在 Rust 上进行异步编程就必须使用 Future trait 和 async/.await 语法，因而必须理解他们到底提供了怎样的功能。

圣经中提供了一个简化版的 Future trait 如下

trait SimpleFuture {
    type Output;
    fn poll(&mut self, wake: fn()) -> Poll<Self::Output>;
}

enum Poll<T> {
    Ready(T),
    Pending,
}

pub struct SocketRead<'a> {
    socket: &'a Socket,
}

impl SimpleFuture for SocketRead<'_> {
    type Output = Vec<u8>;

    fn poll(&mut self, wake: fn()) -> Poll<Self::Output> {
        if self.socket.has_data_to_read() {
            // The socket has data -- read it into a buffer and return it.
            Poll::Ready(self.socket.read_buf())
        } else {
            // The socket does not yet have data.
            //
            // Arrange for `wake` to be called once data is available.
            // When data becomes available, `wake` will be called, and the
            // user of this `Future` will know to call `poll` again and
            // receive data.
            self.socket.set_readable_callback(wake);
            Poll::Pending
        }
    }
}

Future trait 的关键在于 poll 方法。我们称实现了 Future trait 的类的对象为一个 Future 对象，Future 对象会对应一组操作，并且还可能有一个返回值。这组操作，包括返回值，都被写在了 poll 方法里。调用 Future 对象的 poll 方法就会执行这组操作。这看起来和一般的对象，一般的方法无异。

Future 的特殊之处在于，poll 方法中的操作不一定在调用 poll 方法时就全部执行完成，当它遇到需要阻塞等待的情况时，可以选择先退出。Rust 设计者设计了 Poll::Pending 这个返回值来表示这种情况。所以，你可以写一个循环来反复调用 poll 方法直至返回 Poll::Ready 来保证 Future 对象的这组操作执行完成。你可以认为这个循环就是一个最简单且无用的执行器。

之所以说这个循环执行器无用，是因为它没有起到异步的作用，它让我们阻塞在了 Future 对象对应的这组操作中。

所以 Rust 设计者为 poll 方法设计了 wake 参数来解决这个问题。具体是如何解决的，等我们说到执行器的时候再来揭晓。现在你只需要知道 Future 对象对应一组操作，这组操作通过调用 poll 方法来执行，poll 方法不一定在一次调用中执行完成。

再来看 async {}，注意 Futrue 是一个 trait 而 async {} 是一个语法。我们需要了解编译器是如何处理这个语法的，这一点是圣经里完全没有提及的。一句话概括：async {} 会被转化成一个实现 Future trait 的状态机。解释如下：

你可以这么理解：async {} 会先被编译器转化为一段代码，这段代码的核心是一个实现了 Future trait 的类，而async {} 中编写的操作就会被转化到 poll 方法里。这个转化可以说是 Rust 异步编程的精髓。可以通过阅读下面这段由 ChatGPT 生成的例子来理解。

/// 转化前
async fn multiple_awaits() -> i32 {
    let a = async { 1 }.await;
    let b = async { 2 }.await;
    let c = async { 3 }.await;
    a + b + c
}

/// 转化后
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

struct MultipleAwaits {
    state: State,
}

enum State {
    Start,
    AwaitA(Pin<Box<dyn Future<Output = i32>>>),
    AwaitB(i32, Pin<Box<dyn Future<Output = i32>>>),
    AwaitC(i32, i32, Pin<Box<dyn Future<Output = i32>>>),
    Done,
}

impl Future for MultipleAwaits {
    type Output = i32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match &mut self.state {
                State::Start => {
                    let future_a = Box::pin(async { 1 });
                    self.state = State::AwaitA(future_a);
                }
                State::AwaitA(future_a) => {
                    match future_a.as_mut().poll(cx) {
                        Poll::Ready(a) => {
                            let future_b = Box::pin(async { 2 });
                            self.state = State::AwaitB(a, future_b);
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                State::AwaitB(a, future_b) => {
                    match future_b.as_mut().poll(cx) {
                        Poll::Ready(b) => {
                            let future_c = Box::pin(async { 3 });
                            self.state = State::AwaitC(*a, b, future_c);
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                State::AwaitC(a, b, future_c) => {
                    match future_c.as_mut().poll(cx) {
                        Poll::Ready(c) => {
                            self.state = State::Done;
                            return Poll::Ready(a + b + c);
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                State::Done => panic!("poll called after completion"),
            }
        }
    }
}

fn multiple_awaits() -> impl Future<Output = i32> {
    MultipleAwaits {
        state: State::Start,
    }
}

先解释下这里用到的 .await，圣经中提到

Rust 的 Future 是惰性的：它们不会干任何事，除非它们被驱动执行。一个驱动 future 类型的方法是在 async 函数中使用 .await 调用，但这只是将问题抛到上一层。

转化前的 async{} 块中用到了 .await，转化到 poll 方法中后，这些 .await 变成了对 poll 方法的调用。所以驱动 Future 对象执行的方式还是只有调用 poll，.await 会被转化成调用 poll 的代码，这些代码等待更上层的调用。

观察转化后的 poll 方法，可以看到里面维护了一个状态机，这个状态机记录了这段异步操作执行的进展。每次 poll 都会从当前的进展开始（可能）做一定的推进。

这个转化为我们展示了使用 Future trait 的思路：用一个状态机来维护 Future 对象对应的操作执行到了哪里，当遇到可能会阻塞的操作时，可以先记下状态并返回 pending，这个阻塞的操作完成后，再从这个状态继续执行。实际代码中，我们一般不用手写这么复杂的状态机，我们最多只用写简单的底层的 Future 对象，复杂的 Future 对象交给 async{}/.await 去生成即可。

执行器 executor

上面这些知识很琐碎，大部分是对客观事实的列举。下面就要综合出一些理解，并真正的体会到异步了。另外，上面的 Future trait、async/.await 都是 Rust 语言或基本库提供的，是 Rust 异步编程基础架构。而这部分的执行器是 Rust 用户自己编写的，或由社区提供的。你可以认为执行器也是一种架构，不过它位于语言层面的架构之上。

你已经知道一个 Future 对象是怎样的，但要体会到异步，需要有多个 Future 对象才行。很简单，当一个 Future 对象的操作阻塞时，我们可以切换执行其他 Future 对象的操作，也就是 poll 其他的 Future 对象，这样就是异步了。那么，谁来在适当的时候 poll 适当的 Future 对象呢，就是执行器。

上面说最简单而无用的执行器是一个循环，循环 poll 每个 Future 对象。还提到为了让执行器有用，Future trait 的设计者为 poll 方法设计了一个 wake 参数。这个 wake 参数是一个函数，它的作用是在 Future 对象关联的操作可以有进展时，通知执行器可以 poll 这个 Future 对象。有了 wake 这个东西，执行器就可以不再无脑的循环 poll 每个 Future 对象，而只 poll 那些可以有进展的 Future 对象。

现在说得和圣经一样太过抽象，你可能想象不出 wake 函数是怎样的，它是在什么时候被调用的，执行器又是怎样的。圣经中执行器一节给出了经典的例子，可以回答这些问题，但这个例子很难看懂，我在下面给它添加一些中文注释，你也可以试着跑一跑。也可以看别人跑一跑

// src/lib.rs
use std::{
    future::Future,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
    thread,
    time::Duration,
};

pub struct TimerFuture {
    shared_state: Arc<Mutex<SharedState>>,
}

/// Shared state between the future and the waiting thread
struct SharedState {
    /// Whether or not the sleep time has elapsed
    completed: bool,

    /// The waker for the task that `TimerFuture` is running on.
    /// The thread can use this after setting `completed = true` to tell
    /// `TimerFuture`'s task to wake up, see that `completed = true`, and
    /// move forward.
    waker: Option<Waker>,
}

impl Future for TimerFuture {
    type Output = ();
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // Look at the shared state to see if the timer has already completed.
        println!("[{:?}] TimerFuture 开始 Poll ...", std::thread::current().id());
        let mut shared_state = self.shared_state.lock().unwrap();
        if shared_state.completed {
            println!("[{:?}] TimerFuture Poll 完成 ...", std::thread::current().id());
            Poll::Ready(())
        } else {
            // Set waker so that the thread can wake up the current task
            // when the timer has completed, ensuring that the future is polled
            // again and sees that `completed = true`.
            //
            // It's tempting to do this once rather than repeatedly cloning
            // the waker each time. However, the `TimerFuture` can move between
            // tasks on the executor, which could cause a stale waker pointing
            // to the wrong task, preventing `TimerFuture` from waking up
            // correctly.
            //
            // N.B. it's possible to check for this using the `Waker::will_wake`
            // function, but we omit that here to keep things simple.
            println!("[{:?}] TimerFuture Poll 未完成，设置 Waker ...", std::thread::current().id());
            shared_state.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}

impl TimerFuture {
    /// Create a new `TimerFuture` which will complete after the provided
    /// timeout.
    pub fn new(duration: Duration) -> Self {
        let shared_state = Arc::new(Mutex::new(SharedState {
            completed: false,
            waker: None,
        }));

        // Spawn the new thread
        let thread_shared_state = shared_state.clone();
        /// 这里涉及到 wake 函数由谁来调用，在什么时候被调用的问题。在例子中，这个问题的答案是这样的
        /// wake 函数由另一个线程调用，主线程在 poll 各个 Future（虽然例子中只有一个 Future）时，
        /// 每个 Future 都关联一个副线程在倒计时，倒计时结束代表这个 Future 能继续执行，于是副线程
        /// 调用 wake 通知执行器这个 Future 可以被 poll 了。
        thread::spawn(move || {
            thread::sleep(duration);
            let mut shared_state = thread_shared_state.lock().unwrap();
            // Signal that the timer has completed and wake up the last
            // task on which the future was polled, if one exists.
            shared_state.completed = true;
            println!("[{:?}] TimerFuture 完成，唤醒 Waker ...", std::thread::current().id());
            if let Some(waker) = shared_state.waker.take() {
                waker.wake()
            }
        });

        TimerFuture { shared_state }
    }
}

// src/main.rs
use {
    futures::{
        future::{BoxFuture, FutureExt},
        task::{waker_ref, ArcWake},
    },
    std::{
        future::Future,
        sync::mpsc::{sync_channel, Receiver, SyncSender},
        sync::{Arc, Mutex},
        task::{Context, Poll},
        time::Duration,
    },
    // The timer we wrote in the previous section:
    timer_future::TimerFuture,
};

/// Task executor that receives tasks off of a channel and runs them.
/// 例子中调度的核心就是一个管道，调度器守在管道的出口，出来一个任务，就把他的
/// Future 对象取出来 poll。你看看下面为调度器实现的 run 方法就知道调度器确实就
/// 干了这么个事情。
/// 我们大概能猜出 wake 函数是怎么通知调度器一个 Future 对象可以有进展的了吧。
struct Executor {
    ready_queue: Receiver<Arc<Task>>,
}

/// `Spawner` spawns new futures onto the task channel.
#[derive(Clone)]
struct Spawner {
    task_sender: SyncSender<Arc<Task>>,
}

/// A future that can reschedule itself to be polled by an `Executor`.
/// 任务是对 Future 对象的封装，有了这层封装，除了承载 Future 对象以外，它
/// 还能有点别的功能
struct Task {
    /// In-progress future that should be pushed to completion.
    ///
    /// The `Mutex` is not necessary for correctness, since we only have
    /// one thread executing tasks at once. However, Rust isn't smart
    /// enough to know that `future` is only mutated from one thread,
    /// so we need to use the `Mutex` to prove thread-safety. A production
    /// executor would not need this, and could use `UnsafeCell` instead.
    future: Mutex<Option<BoxFuture<'static, ()>>>,

    /// Handle to place the task itself back onto the task queue.
    task_sender: SyncSender<Arc<Task>>,
}

fn new_executor_and_spawner() -> (Executor, Spawner) {
    // Maximum number of tasks to allow queueing in the channel at once.
    // This is just to make `sync_channel` happy, and wouldn't be present in
    // a real executor.
    const MAX_QUEUED_TASKS: usize = 10_000;
    let (task_sender, ready_queue) = sync_channel(MAX_QUEUED_TASKS);
    (Executor { ready_queue }, Spawner { task_sender })
}

impl Spawner {
    fn spawn(&self, future: impl Future<Output = ()> + 'static + Send) {
        let future = future.boxed();
        let task = Arc::new(Task {
            future: Mutex::new(Some(future)),
            task_sender: self.task_sender.clone(),
        });
        println!("[{:?}] 将 Future 组成 Task，放入 Channel ...", std::thread::current().id());
        self.task_sender.send(task).expect("too many tasks queued");
    }
}

impl ArcWake for Task {
    /// wake 函数的实现在这里，它没有直接定义一个函数叫 wake，而是采用了这样一个有些复杂的方式，不要在意
    /// 跟你猜的一样吗，wake 函数通知调度器的方式就是把任务重新放回管道里。就这么简单。
    fn wake_by_ref(arc_self: &Arc<Self>) {
        // Implement `wake` by sending this task back onto the task channel
        // so that it will be polled again by the executor.
        println!("[{:?}] Task 被唤醒，将 Task 放回 Channel ...", std::thread::current().id());
        let cloned = arc_self.clone();
        arc_self
            .task_sender
            .send(cloned)
            .expect("too many tasks queued");
    }
}

impl Executor {
    fn run(&self) {
        println!("[{:?}] 开始执行任务 ...", std::thread::current().id());
        /// 你可能有疑问，我这唯一一个任务 pending 后，在 wake 函数及时把它塞回管道之前，管道为空
        /// 这时不会错误的退出循环吗？你需要了解，这个管道发送端的引用不为零时，从接收端取值就会阻塞，而
        /// 不会简单的返回错误值。所以即使管道为空，只要 task 还保持着对管道发送端的引用，就不会退出循环
        while let Ok(task) = self.ready_queue.recv() {
            /// 执行器取出任务，生成 wake 并找到任务中的 Future，简单看看
            println!("[{:?}] 从 Channel 中取出 Task，准备执行 ...", std::thread::current().id());
            // Take the future, and if it has not yet completed (is still Some),
            // poll it in an attempt to complete it.
            let mut future_slot = task.future.lock().unwrap();
            if let Some(mut future) = future_slot.take() {
                // Create a `LocalWaker` from the task itself
                let waker = waker_ref(&task);
                let context = &mut Context::from_waker(&*waker);
                // `BoxFuture<T>` is a type alias for
                // `Pin<Box<dyn Future<Output = T> + Send + 'static>>`.
                // We can get a `Pin<&mut dyn Future + Send + 'static>`
                // from it by calling the `Pin::as_mut` method.
                println!("[{:?}] Polling future...", std::thread::current().id());
                /// 这里终于 poll 这个 Future 了
                /// if 分支就是遇到了返回 pending 的情况，这个 Future 还没有执行完，执行器
                /// 会先去 poll 其他 Future 而不会阻塞在这里，没问题，很异步！
                /// 那这个 Future 什么时候继续执行呢？当然是 wake 函数通知执行器继续 poll 它之后了
                /// 想象下 async {} 转化后的样子，转化后的代码在 poll TimeFuture
                /// 时把前面生成的 wake 函数传给了 TimerFuture，所以 TimerFuture 开启的另一个线程
                /// 倒计时结束后调用 wake 时，就会把这个任务重新放回管道里啦。
                /// 这个任务在重新被 poll 时，async {} 转化出的状态机记录了之前的执行进展，所以不会
                /// 再次输入 howdy 了，而是会从 poll TimerFuture 开始执行，这次 poll TimerFuture 
                /// 会返回 ready，之后输出 done，这个 async {} 转化出的 Future 就执行完毕了，走 else。
                if future.as_mut().poll(context).is_pending() {
                    // We're not done processing the future, so put it
                    // back in its task to be run again in the future.
                    *future_slot = Some(future);
                    println!("[{:?}] Future 还未完成，将 Future 放回 Task ...", std::thread::current().id());
                } else {
                    println!("[{:?}] Future 已完成，不再放回 Task ...", std::thread::current().id());
                }
            }
        }
    }
}

fn main() {
    /// 初始化执行器和 spawner，主要是生成一个管道并让他俩守在两端。
    let (executor, spawner) = new_executor_and_spawner();

    // Spawn a task to print before and after waiting on a timer.
    /// 读一读 spawn 方法的实现，他以一个 Future 对象为参数，而 async {} 会被编译器转化为一个 Future
    /// 这个 Future 会被封装到一个 task 里并被送入管道，执行器就会收到这个 Future 并 poll 它了。
    /// 回忆下前面的 async {} 转化的例子，.await 被转化成对 poll 的调用。所以，
    /// 第一次 poll 这个 Future 时，输出 howdy 后会 poll TimerFuture，去看 TimerFuture 的 poll 的实现
    /// 如果这时计时器没有完成，就会返回 pending，那么外面这层 Future 也会返回 pending（async {}会转化出这种逻辑）。
    /// 我们去看看执行器的 run 里是怎么处理这种情况的。
    spawner.spawn(async {
        println!("howdy!");
        // Wait for our timer future to complete after two seconds.
        TimerFuture::new(Duration::new(2, 0)).await;
        println!("done!");
    });

    // Drop the spawner so that our executor knows it is finished and won't
    // receive more incoming tasks to run.
    drop(spawner);

    // Run the executor until the task queue is empty.
    // This will print "howdy!", pause, and then print "done!".
    executor.run();
}

上面这个例子展示了 Future 对象可能是什么样子，async/.await 可以怎么使用，wake 函数可能是什么样子，执行器可能是什么样子，可能的调用 wake 的方式是什么等等。我相信这是个简化版的例子，不过麻雀虽小五脏俱全，它可以解决你对具体概念想象不出，没有概念的问题，也展示了 Rust 异步编程的范式。