Rust Waker Example: Understanding and Implementing Wakers in Async Programming

In Rust, asynchronous programming is built around futures and the async/await syntax. Central to this is the concept of a Waker, which allows the executor to wake up tasks when they’re ready to proceed. This article explains what a Waker is, how it works, and provides a step-by-step example of creating and using a custom Waker.


What is a Waker?

A Waker is a core component in Rust's async runtime. It is used to:

  1. Wake Tasks: Notify the executor that a task is ready to make progress.
  2. Control Execution: Ensure tasks are polled only when necessary, improving efficiency.

Wakers are part of the std::task module and are tightly integrated with the Future trait.


Waker in the Async Ecosystem

  1. Future: Represents an asynchronous computation that may not have completed yet.
  2. Context: Provides access to the Waker associated with the current task.
  3. Waker: Triggers the task's executor to poll the task again.

Implementing a Custom Waker

To demonstrate how a Waker works, we’ll create a simple example where a Waker notifies the executor when a future is ready.


Step 1: Create the Custom Future

Create a struct that implements the Future trait:

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

struct MyFuture {
    is_ready: bool,
}

impl MyFuture {
    fn new() -> Self {
        Self { is_ready: false }
    }

    fn make_ready(&mut self) {
        self.is_ready = true;
    }
}

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

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.is_ready {
            Poll::Ready("Future is now ready!")
        } else {
            // Register the waker for the current task
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

Step 2: Implement a Custom Waker

Create a Waker that uses a simple callback to wake the task:

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

fn create_waker(callback: Arc<Mutex<bool>>) -> Waker {
    // Create a RawWaker from a pointer and a vtable
    let raw_waker = RawWaker::new(
        Arc::into_raw(callback) as *const (),
        &RawWakerVTable::new(
            clone_callback,
            wake_callback,
            wake_by_ref_callback,
            drop_callback,
        ),
    );

    // Convert RawWaker to a Waker
    unsafe { Waker::from_raw(raw_waker) }
}

// Callback functions for the Waker's vtable
unsafe fn clone_callback(ptr: *const ()) -> RawWaker {
    let arc = Arc::from_raw(ptr as *const Mutex<bool>);
    let clone = Arc::clone(&arc);
    std::mem::forget(arc);
    RawWaker::new(
        Arc::into_raw(clone) as *const (),
        &RawWakerVTable::new(
            clone_callback,
            wake_callback,
            wake_by_ref_callback,
            drop_callback,
        ),
    )
}

unsafe fn wake_callback(ptr: *const ()) {
    let arc = Arc::from_raw(ptr as *const Mutex<bool>);
    *arc.lock().unwrap() = true;
    std::mem::forget(arc);
}

unsafe fn wake_by_ref_callback(ptr: *const ()) {
    let arc = Arc::from_raw(ptr as *const Mutex<bool>);
    *arc.lock().unwrap() = true;
    std::mem::forget(arc);
}

unsafe fn drop_callback(ptr: *const ()) {
    drop(Arc::from_raw(ptr as *const Mutex<bool>));
}

Step 3: Use the Waker with the Future

Tie everything together by creating a task that uses the Waker:

use std::task::{Context, Poll};

fn main() {
    // Shared state for the custom Waker
    let ready_state = Arc::new(Mutex::new(false));
    let waker = create_waker(ready_state.clone());

    let mut my_future = MyFuture::new();
    let mut cx = Context::from_waker(&waker);

    // Poll the future
    match Pin::new(&mut my_future).poll(&mut cx) {
        Poll::Ready(result) => println!("{}", result),
        Poll::Pending => {
            println!("Future is not ready. Waking the task...");
        }
    }

    // Simulate making the future ready
    my_future.make_ready();

    // Poll the future again
    match Pin::new(&mut my_future).poll(&mut cx) {
        Poll::Ready(result) => println!("{}", result),
        Poll::Pending => println!("Future is still not ready."),
    }
}

How the Example Works

  1. Future Implementation:
    • The MyFuture struct represents an asynchronous task.
    • The poll method checks whether the task is ready and registers the Waker.
  2. Custom Waker:
    • The Waker uses an Arc<Mutex<bool>> to simulate waking the task.
  3. Executor Simulation:
    • By manually polling the future, the example shows how Wakers notify tasks.

Output

When running the example, you’ll see:

Future is not ready. Waking the task...
Future is now ready!

Best Practices

  1. Use Existing Executors:
    • Most async runtimes (e.g., Tokio, async-std) handle Waker implementation internally.
  2. Avoid Blocking:
    • Wakers should not perform blocking operations to maintain efficiency.
  3. Understand Memory Safety:
    • Custom Waker implementations must handle raw pointers carefully to avoid memory leaks or unsafe behavior.

Conclusion

Wakers are a crucial part of Rust’s async ecosystem, enabling efficient task scheduling. While most developers rely on built-in executors, understanding how Wakers work helps demystify Rust’s async model and allows for advanced customizations when necessary. With this guide, you’re ready to explore deeper into Rust’s async capabilities!