rust异步编程:futures-rs之futures-executor

概述

futures-executor它引用了futures-core、futures-task、futures-util,用于task的执行,主要提供以下功能:

• 线程池 (线程和task的关系M:N )
  • ThreadPool::spawn_ok
• 生成其他任务(spawn task)
  • Spawn::spawn_obj
  • LocalSpawn::spawn_local_obj (用于!Send future)
• 单线程执行
  • LocalSpawn::spawn_local_obj (spawn non-Send tasks,对应LocalPool::spawner,可执行多个任务)
  • executor::block_on (执行单个任务)

主要提供以下api

mod local_pool;
#[cfg(feature = "std")]
pub use crate::local_pool::{block_on, block_on_stream, BlockingStream, LocalPool, LocalSpawner};

#[cfg(feature = "thread-pool")]
#[cfg(feature = "std")]
mod unpark_mutex;
#[cfg(feature = "thread-pool")]
#[cfg(feature = "std")]
mod thread_pool;
#[cfg(feature = "thread-pool")]
#[cfg(feature = "std")]
pub use crate::thread_pool::{ThreadPool, ThreadPoolBuilder};

#[cfg(feature = "std")]
mod enter;
#[cfg(feature = "std")]
pub use crate::enter::{enter, Enter, EnterError};

block_on

下面看看block_on它的实现细节:

pub fn block_on<F: Future>(f: F) -> F::Output {
    pin_mut!(f);
    run_executor(|cx| f.as_mut().poll(cx))
}

先pin住future,然后调用run_executor,里面提供一个Fn,在Fn里执行future的调用

fn run_executor<T, F: FnMut(&mut Context<'_>) -> Poll<T>>(mut f: F) -> T {
    let _enter = enter().expect(
        "cannot execute `LocalPool` executor from within \
         another executor",
    );

    CURRENT_THREAD_NOTIFY.with(|thread_notify| {
        let waker = waker_ref(thread_notify);
        let mut cx = Context::from_waker(&waker);
        loop {
            if let Poll::Ready(t) = f(&mut cx) {
                return t;
           }
            let unparked = thread_notify.unparked.swap(false, Ordering::Acquire);
            if !unparked {
                thread::park();
                thread_notify.unparked.store(false, Ordering::Release);
            }
        }
    })
}

里面有个loop循环,通过不断轮循future::poll,直到返回Ready,如果返回Pending,则让出cpu,让executor执行其他的任务,这里使用Thread的park()、unpark()方法实现,如果是单线程的话 ,线程pack后如果没有unpack,会一直挂在那.

pub struct Enter {
    _priv: (),
}
pub struct EnterError {
    _priv: (),
}
/// let enter = enter().expect("...");
/// /* run task */
/// drop(enter);
thread_local!(static ENTERED: Cell<bool> = Cell::new(false));
pub fn enter() -> Result<Enter, EnterError> {
    ENTERED.with(|c| {
        if c.get() {
            Err(EnterError { _priv: () })
        } else {
            c.set(true);

            Ok(Enter { _priv: () })
        }
    })
}
impl Drop for Enter {
    fn drop(&mut self) {
        ENTERED.with(|c| {
            assert!(c.get());
            c.set(false);
        });
    }
}

enter方法实现很简单,只是给当前线程打了个标签,标识是不是已经进来过了,在executor执行任务之前调用,执行完任务后drop掉(将标识复位),因此可以调用多次,但不能嵌套调用

pub(crate) struct ThreadNotify {
    /// The (single) executor thread.
    thread: Thread,
    unparked: AtomicBool,
}

thread_local! {
    static CURRENT_THREAD_NOTIFY: Arc<ThreadNotify> = Arc::new(ThreadNotify {
        thread: thread::current(),
        unparked: AtomicBool::new(false),
    });
}

impl ArcWake for ThreadNotify {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        // Make sure the wakeup is remembered until the next `park()`.
        let unparked = arc_self.unparked.swap(true, Ordering::Relaxed);
        if !unparked {
            arc_self.thread.unpark();
        }
    }
}

引入了thread_local!来存放ThreadNotify,ThreadNotify里面包含了当前线程和一个Atomic标识(是否unparked),当调用wake_by_ref时,将标识unparked设为true,如果原来是false,则需要调用Thread的unpark方法.
但是整个block_on流程走下来,发现并没有调用wake,如果线程yield了需要自己调用wake(),最后面会有一个例子说明如何调用.

ThreadPool

接下来再看看ThreadPool的实现

    /// ```
    /// use futures::executor::ThreadPool;
    ///
    /// let pool = ThreadPool::new().unwrap();
    ///
    /// let future = async { /* ... */ };
    /// pool.spawn_ok(future);
    /// ```
    pub fn spawn_ok<Fut>(&self, future: Fut)
    where
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.spawn_obj_ok(FutureObj::new(Box::new(future)))
    }

这里将Future包装在FutureObj里,再调用spawn_obj_ok

  pub fn spawn_obj_ok(&self, future: FutureObj<'static, ()>) {
        let task = Task {
            future,
            wake_handle: Arc::new(WakeHandle {
                exec: self.clone(),
                mutex: UnparkMutex::new(),
            }),
            exec: self.clone(),
        };
        self.state.send(Message::Run(task));
    }

这里只是构造一个Task,然后通过channel将任务发送出,看看它们的结构定义

pub struct ThreadPool {
    state: Arc<PoolState>,
}

struct PoolState {
    tx: Mutex<mpsc::Sender<Message>>,
    rx: Mutex<mpsc::Receiver<Message>>,
    cnt: AtomicUsize,
    size: usize,
}

pub struct ThreadPoolBuilder {
    pool_size: usize,
    stack_size: usize,
    name_prefix: Option<String>,
    after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
    before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>,
}

struct Task {
    future: FutureObj<'static, ()>,
    exec: ThreadPool,
    wake_handle: Arc<WakeHandle>,
}

struct WakeHandle {
    mutex: UnparkMutex<Task>,
    exec: ThreadPool,
}

pub(crate) struct UnparkMutex<D> {
    // The state of task execution (state machine described below)
    status: AtomicUsize,

    // The actual task data, accessible only in the POLLING state
    inner: UnsafeCell<Option<D>>,
}

UnparkMutex里的主status要包含以下几种状态

// 任务被block,等待一个事件
const WAITING: usize = 0;       // --> POLLING
//线程正在主动轮循任务,其他感兴趣的事件到来 应该将其设置为REPOLL
const POLLING: usize = 1;       // --> WAITING, REPOLL, or COMPLETE
// 正在轮循该任务,但是完成后需要重新轮循该任务以确保能监测到所有的事件
const REPOLL: usize = 2;        // --> POLLING
// 任务已经执行完成(成功或者error/panic)
const COMPLETE: usize = 3;      // No transitions out

上面只看到了将任务发送到channel中去,那么在哪消费的?,接下来看看ThreadPool的创建

impl ThreadPool {
    pub fn new() -> Result<ThreadPool, io::Error> {
        ThreadPoolBuilder::new().create()
    }
 }

impl ThreadPoolBuilder {
    /// Create a default thread pool configuration.
    ///
    /// See the other methods on this type for details on the defaults.
    pub fn new() -> Self {
        Self {
            pool_size: cmp::max(1, num_cpus::get()),
            stack_size: 0,
            name_prefix: None,
            after_start: None,
            before_stop: None,
        }
    }
     pub fn create(&mut self) -> Result<ThreadPool, io::Error> {
        let (tx, rx) = mpsc::channel();
        let pool = ThreadPool {
            state: Arc::new(PoolState {
                tx: Mutex::new(tx),
                rx: Mutex::new(rx),
                cnt: AtomicUsize::new(1),
                size: self.pool_size,
            }),
        };

        for counter in 0..self.pool_size {
            let state = pool.state.clone();
            let after_start = self.after_start.clone();
            let before_stop = self.before_stop.clone();
            let mut thread_builder = thread::Builder::new();
            if let Some(ref name_prefix) = self.name_prefix {
                thread_builder = thread_builder.name(format!("{}{}", name_prefix, counter));
            }
            if self.stack_size > 0 {
                thread_builder = thread_builder.stack_size(self.stack_size);
            }
            thread_builder.spawn(move || state.work(counter, after_start, before_stop))?;
        }
        Ok(pool)
    }

在ThreadPoolBuilder中可以看到,它最终会调用ThreadPool里的state的work方法

impl PoolState {
    fn send(&self, msg: Message) {
        self.tx.lock().unwrap().send(msg).unwrap();
    }

    fn work(&self,
            idx: usize,
            after_start: Option<Arc<dyn Fn(usize) + Send + Sync>>,
            before_stop: Option<Arc<dyn Fn(usize) + Send + Sync>>) {
        let _scope = enter().unwrap();
        if let Some(after_start) = after_start {
            after_start(idx);
        }
        loop {
            let msg = self.rx.lock().unwrap().recv().unwrap();
            match msg {
                Message::Run(task) => task.run(),
                Message::Close => break,
            }
        }
        if let Some(before_stop) = before_stop {
            before_stop(idx);
        }
    }
}

在这个loop循坏里取得task,并执行task,再看看run方法

 fn run(self) {
        let Task { mut future, wake_handle, mut exec } = self;
        let waker = waker_ref(&wake_handle);
        let mut cx = Context::from_waker(&waker);

        // Safety: The ownership of this `Task` object is evidence that
        // we are in the `POLLING`/`REPOLL` state for the mutex.
        unsafe {
            wake_handle.mutex.start_poll();

            loop {
                let res = future.poll_unpin(&mut cx);
                match res {
                    Poll::Pending => {}
                    Poll::Ready(()) => return wake_handle.mutex.complete(),
                }
                let task = Task {
                    future,
                    wake_handle: wake_handle.clone(),
                    exec,
                };
                match wake_handle.mutex.wait(task) {
                    Ok(()) => return, // we've waited
                    Err(task) => { // someone's notified us
                        future = task.future;
                        exec = task.exec;
                    }
                }
            }
        }
    }

调用start_poll将状态改为POLLING,紧接着调用future.poll_unpin获取执行结果:

• 如果是Ready则将状态改为COMPLETE并return出loop,紧接着继续收取其他的Message::run(task) 处理;
• 如果是Pending,则重新构建一个新的Task(状态未变),将Task的状态设置为WAITING,设置成功直接退出loop,如果设置失败,则直接将当前任务的future和exec修改成新的Task,在run的loop循坏里会再次处理

fn poll_unpin(&mut self, cx: &mut Context<'_>) -> Poll<Self::Output>
    where
        Self: Unpin,
    {
        Pin::new(self).poll(cx)
    }

poll_unpin只是将Unpin的future包装成Pin再polll

pub(crate) unsafe fn wait(&self, data: D) -> Result<(), D> {
        *self.inner.get() = Some(data);

        match self.status.compare_exchange(POLLING, WAITING, SeqCst, SeqCst) {
            // no unparks came in while we were running
            Ok(_) => Ok(()),

            // guaranteed to be in REPOLL state; just clobber the
            // state and run again.
            Err(status) => {
                assert_eq!(status, REPOLL);
                self.status.store(POLLING, SeqCst);
                Err((*self.inner.get()).take().unwrap())
            }
        }
    }

把data也就是task先保存起来，然后把POLLING状态修改成WAITING状态返回OK。
compare_exchange可能Error，那么此时status一定是REPOLL，同时将新的Task取出并返回Err(task)

接下来看看Thread是如何被唤醒重新执行任务的

impl ArcWake for WakeHandle {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        match arc_self.mutex.notify() {
            Ok(task) => arc_self.exec.state.send(Message::Run(task)),
            Err(()) => {}
        }
    }
}

这里调用了notify,唤醒后重新将任务发送到channel,等待PoolState中loop循坏取出

pub(crate) fn notify(&self) -> Result<D, ()> {
        let mut status = self.status.load(SeqCst);
        loop {
            match status {
                WAITING => {
                    match self.status.compare_exchange(WAITING, POLLING,
                                                       SeqCst, SeqCst) {
                        Ok(_) => {
                            let data = unsafe {
                                (*self.inner.get()).take().unwrap()
                            };
                            return Ok(data);
                        }
                        Err(cur) => status = cur,
                    }
                }
               POLLING => {
                    match self.status.compare_exchange(POLLING, REPOLL,
                                                       SeqCst, SeqCst) {
                        Ok(_) => return Err(()),
                        Err(cur) => status = cur,
                    }
                }
                _ => return Err(()),
            }
        }
    }

• 如果是WAITING状态 则修改成POLLING
• 如果是POLLING状态,则重新标记为REPOLL
• 如果是其他状态,则说明任务已经被调度执行或者完成

代码中并没有调用wake_by_ref或wake方法,和block_on类似,也没有自动通知机制,需要我们自己调用,否则任务将会阻塞在那

wake使用示例

use futures::executor::{ThreadPool,block_on,ThreadPoolBuilder};
use futures::future::{ Future};
use futures::task::{Context, Poll,};
use std::thread;

use std::pin::Pin;

fn main(){
    const NUM: usize = 10;

    struct Yield {
        rem: usize,
    }

    impl Future for Yield {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.rem == 0 {
                println!("{}:done",thread::current().name().unwrap());
                Poll::Ready(())
            } else {
                println!("self.rem={}",self.rem);
                self.rem -= 1;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }
   
    for _ in 0..NUM {
        let y = Yield { rem: NUM };
        block_on(y) 
    }

    for _ in 0..NUM {
        let y = Yield { rem: NUM };
        let pool = ThreadPoolBuilder::new().name_prefix("pool-").create().unwrap();
        pool.spawn_ok(y);   
    }
}

• 从控制台打印中可以看出block_on里的future运行在main线程中,spawn_ok里的future运行在pool-x(pool-0,pool-1…)线程中
• 如果将cx.waker().wake_by_ref()注释,future永远无法Ready,打印self.rem=10后,这时候block_on的话将会阻塞main线程,导致程序卡住;注释掉block_on使用pool.spawn_ok的话,这时候每个子线程打印self.rem=10,不管子线程执行的任务有没有Ready,不会阻塞主线程,随着spawn_ok执行完成,程序退出.