Rust Future 实现原理

/*
 * [dependencies]
 * chrono = "0.4.41"
 * futures = "0.3.24"
 * async-std = "1.13.1"
 */
use std::{thread, time::Duration};

use chrono::Local;

async fn async_hello_11() {
    println!(
        "[{}] [{:?}] hello async 11!",
        Local::now().timestamp_millis(),
        thread::current().id()
    );
}
async fn async_hello_12() {
    println!(
        "[{}][{:?}] hello async 12!",
        Local::now().timestamp_millis(),
        thread::current().id()
    );
}
async fn async_hello_1() {
    async_hello_11().await;
    tokio::time::sleep(Duration::from_secs(2)).await;
    async_hello_12().await;
}
async fn async_hello_2() {
    println!(
        "[{}][{:?}] hello async 2 !",
        Local::now().timestamp_millis(),
        thread::current().id()
    );
}
#[tokio::main]
async fn main() {
    let f1 = async_hello_1();
    let f2 = async_hello_2();
    tokio::join!(f1, f2);
}

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

use tokio::time::Sleep;
// 异步对应的状态机
enum FooState {
    // begin
    Started,
    // await Future1,ie Sleep
    Await1,
    // The future has completed.
    Done,
}
struct FooFuture {
    state: FooState,
    // 这里的 pin 后面会展开介绍，先可以看做为 Sleep(impl Future)类型
    sleep: Pin<Box<Sleep>>,
}

impl Future for FooFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        loop {
            match self.state {
                FooState::Started => {
                    println!("start async");
                    self.state = FooState::Await1;
                }
                FooState::Await1 => match self.sleep.as_mut().poll(cx) {
                    Poll::Ready(_) => {
                        self.state = FooState::Done;
                    }
                    Poll::Pending => {
                        // 重新加入调度，后面会详细介绍
                        cx.waker().wake_by_ref();               
                        return Poll::Pending;
                    }
                },
                FooState::Done => {
                    println!("end async");
                    return Poll::Ready(());
                }
            }
        }
    }
}
fn foo_sync(n: u64) -> FooFuture {
    let duration = Duration::from_secs(n);
    let sleep = Box::pin(tokio::time::sleep(duration));
    FooFuture {
        state: FooState::Started,
        sleep,
    }
}

// This function:
async fn foo() {
    step_one().await;
    step_two().await;
}
// 生成一个这样的类型:
enum Foo {
    First(StepOne),
    Second(StepTwo),
}

// 所以这个函数:
async fn recursive() {
    recursive().await;
    recursive().await;
}

// 生成一个这样的类型:
enum Recursive {
    First(Recursive),
    Second(Recursive),
}

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

struct Delay {
    when: Instant,
}

impl Future for Delay {
    type Output = &'static str;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<&'static str> {
        if Instant::now() >= self.when {
            println!("Delay complete");
            Poll::Ready("done")
        } else {
            // 此处触发一次唤醒
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

#[tokio::main]
async fn main() {
    let when = Instant::now() + Duration::from_secs(2);
    let future = Delay { when };

    let out = future.await;
    assert_eq!(out, "done");
}

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use std::thread;

struct Delay {
    when: Instant,
}

impl Future for Delay {
    type Output = &'static str;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>)
        -> Poll<&'static str>
    {
        if Instant::now() >= self.when {
            println!("Hello world");
            Poll::Ready("done")
        } else {
            // Get a handle to the waker for the current task
            let waker = cx.waker().clone();
            let when = self.when;

            // Spawn a timer thread.
            thread::spawn(move || {
                let now = Instant::now();

                if now < when {
                    thread::sleep(when - now);
                }
                waker.wake();
            });
            Poll::Pending
        }
    }
}

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::thread;
use std::time::{Duration, Instant};

struct Delay {
    when: Instant,
    // This is Some when we have spawned a thread, and None otherwise.
    waker: Option<Arc<Mutex<Waker>>>,
}

impl Future for Delay {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        // 判断是否 ready
        if Instant::now() >= self.when {
            return Poll::Ready(());
        }

        // 如果是初次 poll，创建计时器线程。
        // 否则确保 waker 是正确的
        if let Some(waker) = &self.waker {
            let mut waker = waker.lock().unwrap();

            // 检查存储的 waker 是否满足当前的 waker(是否会 wake 相同的 task)
            if !waker.will_wake(cx.waker()) {
                // 更新 waker
                *waker = cx.waker().clone();
            }
        } else {
            let when = self.when;
            let waker = Arc::new(Mutex::new(cx.waker().clone()));
            self.waker = Some(waker.clone());
            thread::spawn(move || {
                let now = Instant::now();
                if now < when {
                    thread::sleep(when - now);
                }
                let waker = waker.lock().unwrap();
                waker.wake_by_ref();
            });
        }
        Poll::Pending
    }
}

struct MiniTokio {
    scheduled: mpsc::Receiver<Arc<Task>>,
    sender: mpsc::Sender<Arc<Task>>,
}
impl MiniTokio {
    fn run(&self) {
        while let Ok(task) = self.scheduled.recv() {
            task.poll();
        }
    }
    fn new() -> MiniTokio {
        let (sender, scheduled) = mpsc::channel();
        MiniTokio { scheduled, sender }
    }
    fn spawn<F>(&self, future: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        Task::spawn(future, &self.sender);
    }
}

struct Task {
    task_future: Mutex<TaskFuture>,
    executor: mpsc::Sender<Arc<Task>>,
}
/// 将我们的 schedule 函数与 std::task::Waker 连接起来。
/// 标准库提供了一个低级 API，可以使用 手动 vtable 构建 来做到这一点。
/// 此策略为实现者提供了最大的灵活性，但需要大量不安全的样板代码。
/// 我们将使用 futures crate 提供的 ArcWake 实用程序，而不是直接使用 RawWakerVTable
impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.schedule();
    }
}
impl Task {
    /// 当在 waker 上调用 wake() 时，任务被推送到通道的发送端
    fn schedule(self: &Arc<Self>) {
        self.executor.send(self.clone());
    }
    /// 轮询
    fn poll(self: Arc<Self>) {
        // 准备 future 上下文
        let waker = task::waker(self.clone());
        let mut cx = Context::from_waker(&waker);
        
        let mut task_future = self.task_future.try_lock().unwrap();
        // Poll the inner future
        task_future.poll(&mut cx);
    }
    fn spawn<F>(future: F, sender: &mpsc::Sender<Arc<Task>>)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        // 创建一个 task
        let task = Arc::new(Task {
            task_future: Mutex::new(TaskFuture::new(future)),
            executor: sender.clone(),
        });
        // 开始调度
        let _ = sender.send(task);
    }
}
/// A structure holding a future and the result of
/// the latest call to its `poll` method.
struct TaskFuture {
    future: Pin<Box<dyn Future<Output = ()> + Send>>,
    poll: Poll<()>,
}
impl TaskFuture {
    fn new(future: impl Future<Output = ()> + Send + 'static) -> TaskFuture {
        TaskFuture {
            future: Box::pin(future),
            poll: Poll::Pending,
        }
    }
    fn poll(&mut self, cx: &mut Context<'_>) {
        // 防止重复唤醒内部的 future
        if self.poll.is_pending() {
            self.poll = self.future.as_mut().poll(cx);
        }
    }
}