Rust 异步编程：Waker与唤醒机制的剖析

2301_81549453

468人浏览 · 2025-10-30 18:40:00

2301_81549453 · 2025-10-30 18:40:00 发布

Rust 异步编程：Waker与唤醒机制的剖析

引言

Waker是Rust异步运行时的神经系统，它解决了一个核心问题：当异步任务未就绪时，如何高效地等待，并在条件满足时精确唤醒？ 理解Waker机制，是掌握Rust异步编程精髓的关键一步。🔥

Waker的设计哲学

在传统的阻塞模型中，线程会被操作系统挂起，直到I/O完成。但Rust的异步模型采用了不同的策略：

轮询模型（Polling）：Future不会主动通知完成，而是等待执行器轮询
唤醒机制（Waking）：当资源就绪时，通过Waker通知执行器重新轮询
零成本抽象：Waker的设计让运行时可以高效管理数百万并发任务

pub struct Context<'a> {
    waker: &'a Waker,
    // 其他字段...
}

impl Waker {
    pub fn wake(self);
    pub fn wake_by_ref(&self);
}

Context携带Waker传递给poll方法，Future可以克隆Waker并在稍后某个时刻调用wake()来通知执行器重新轮询。

实践深度一：手动实现基于Waker的Timer

让我们从零实现一个Timer Future，深入理解Waker的工作原理：

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 TimerState {
    completed: bool,
    waker: Option<Waker>,
}

pub struct Timer {
    state: Arc<Mutex<TimerState>>,
}

impl Timer {
    pub fn new(duration: Duration) -> Self {
        let state = Arc::new(Mutex::new(TimerState {
            completed: false,
            waker: None,
        }));

        let state_clone = Arc::clone(&state);
        
        // 启动后台线程，模拟定时器
        thread::spawn(move || {
            thread::sleep(duration);
            
            let mut state = state_clone.lock().unwrap();
            state.completed = true;
            
            // 关键：唤醒等待的Future
            if let Some(waker) = state.waker.take() {
                waker.wake();
            }
        });

        Timer { state }
    }
}

impl Future for Timer {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut state = self.state.lock().unwrap();
        
        if state.completed {
            return Poll::Ready(());
        }
        
        // 保存Waker，以便稍后唤醒
        state.waker = Some(cx.waker().clone());
        Poll::Pending
    }
}

// 使用示例
async fn test_timer() {
    println!("Timer started at {:?}", Instant::now());
    Timer::new(Duration::from_secs(2)).await;
    println!("Timer finished at {:?}", Instant::now());
}

核心机制解析：

第一次poll时，Timer未完成，保存Waker并返回Pending
后台线程睡眠结束后，调用waker.wake()
执行器收到唤醒信号，重新poll这个Future
第二次poll时，completed为true，返回Ready

实践深度二：实现支持多个Waker的Channel

在生产环境中，一个资源可能需要唤醒多个等待者。让我们实现一个简化的channel：

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::task::Waker;

struct ChannelState<T> {
    queue: VecDeque<T>,
    wakers: Vec<Waker>,
    closed: bool,
}

pub struct Sender<T> {
    state: Arc<Mutex<ChannelState<T>>>,
}

pub struct Receiver<T> {
    state: Arc<Mutex<ChannelState<T>>>,
}

pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let state = Arc::new(Mutex::new(ChannelState {
        queue: VecDeque::new(),
        wakers: Vec::new(),
        closed: false,
    }));

    (
        Sender {
            state: Arc::clone(&state),
        },
        Receiver { state },
    )
}

impl<T> Sender<T> {
    pub fn send(&self, value: T) -> Result<(), &'static str> {
        let mut state = self.state.lock().unwrap();
        
        if state.closed {
            return Err("Channel closed");
        }
        
        state.queue.push_back(value);
        
        // 唤醒所有等待的接收者（广播模式）
        for waker in state.wakers.drain(..) {
            waker.wake();
        }
        
        Ok(())
    }
}

impl<T> Receiver<T> {
    pub async fn recv(&self) -> Option<T> {
        RecvFuture {
            state: Arc::clone(&self.state),
        }
        .await
    }
}

struct RecvFuture<T> {
    state: Arc<Mutex<ChannelState<T>>>,
}

impl<T> Future for RecvFuture<T> {
    type Output = Option<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut state = self.state.lock().unwrap();
        
        // 尝试从队列取出数据
        if let Some(value) = state.queue.pop_front() {
            return Poll::Ready(Some(value));
        }
        
        // 如果channel已关闭且队列为空
        if state.closed {
            return Poll::Ready(None);
        }
        
        // 注册Waker并返回Pending
        let waker = cx.waker().clone();
        
        // 去重：避免重复注册同一个Waker
        if !state.wakers.iter().any(|w| w.will_wake(&waker)) {
            state.wakers.push(waker);
        }
        
        Poll::Pending
    }
}

关键设计点：

使用Vec<Waker>支持多个等待者
will_wake()方法避免重复注册
wake()消费Waker，wake_by_ref()允许多次唤醒

实践深度三：自定义执行器与Waker实现

为了完全理解Waker，让我们实现一个简化的单线程执行器：

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::task::{RawWaker, RawWakerVTable, Waker};

type TaskId = usize;

struct Task {
    id: TaskId,
    future: Pin<Box<dyn Future<Output = ()>>>,
}

pub struct Executor {
    ready_queue: Arc<Mutex<VecDeque<TaskId>>>,
    tasks: Mutex<HashMap<TaskId, Task>>,
    next_id: AtomicUsize,
}

impl Executor {
    pub fn new() -> Self {
        Executor {
            ready_queue: Arc::new(Mutex::new(VecDeque::new())),
            tasks: Mutex::new(HashMap::new()),
            next_id: AtomicUsize::new(0),
        }
    }

    pub fn spawn(&self, future: impl Future<Output = ()> + 'static) {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let task = Task {
            id,
            future: Box::pin(future),
        };
        
        self.tasks.lock().unwrap().insert(id, task);
        self.ready_queue.lock().unwrap().push_back(id);
    }

    fn create_waker(&self, task_id: TaskId) -> Waker {
        let ready_queue = Arc::clone(&self.ready_queue);
        let raw_waker = Self::create_raw_waker(task_id, ready_queue);
        unsafe { Waker::from_raw(raw_waker) }
    }

    fn create_raw_waker(
        task_id: TaskId,
        ready_queue: Arc<Mutex<VecDeque<TaskId>>>,
    ) -> RawWaker {
        let data = Box::into_raw(Box::new((task_id, ready_queue)));
        
        RawWaker::new(
            data as *const (),
            &RawWakerVTable::new(
                Self::clone_waker,
                Self::wake,
                Self::wake_by_ref,
                Self::drop_waker,
            ),
        )
    }

    unsafe fn clone_waker(data: *const ()) -> RawWaker {
        let (task_id, ready_queue) = &*(data as *const (TaskId, Arc<Mutex<VecDeque<TaskId>>>));
        Self::create_raw_waker(*task_id, Arc::clone(ready_queue))
    }

    unsafe fn wake(data: *const ()) {
        let (task_id, ready_queue) = Box::from_raw(data as *mut (TaskId, Arc<Mutex<VecDeque<TaskId>>>));
        ready_queue.lock().unwrap().push_back(task_id);
    }

    unsafe fn wake_by_ref(data: *const ()) {
        let (task_id, ready_queue) = &*(data as *const (TaskId, Arc<Mutex<VecDeque<TaskId>>>));
        ready_queue.lock().unwrap().push_back(*task_id);
    }

    unsafe fn drop_waker(data: *const ()) {
        let _ = Box::from_raw(data as *mut (TaskId, Arc<Mutex<VecDeque<TaskId>>>));
    }

    pub fn run(&self) {
        loop {
            let task_id = {
                let mut queue = self.ready_queue.lock().unwrap();
                if let Some(id) = queue.pop_front() {
                    id
                } else {
                    break; // 无任务，退出
                }
            };

            let waker = self.create_waker(task_id);
            let mut context = Context::from_waker(&waker);

            let mut tasks = self.tasks.lock().unwrap();
            if let Some(task) = tasks.get_mut(&task_id) {
                match task.future.as_mut().poll(&mut context) {
                    Poll::Ready(()) => {
                        tasks.remove(&task_id);
                    }
                    Poll::Pending => {
                        // 任务仍在等待，Waker已注册
                    }
                }
            }
        }
    }
}

RawWaker机制详解：