Wakers

Wakers are the missing piece. This is the system by which a resource is able to notify the waiting task that the resource has become ready to continue some operation.

Let’s look at the Future::poll definition again:

fn poll(self: Pin<&mut Self>, cx: &mut Context)
    -> Poll<Self::Output>;

The Context argument to poll has a waker() method. This method returns a Waker bound to the current task. The Waker has a wake() method. Calling this method signals to the executor that the associated task should be scheduled for execution. Resources call wake() when they transition to a ready state to notify the executor that polling the task will be able to make progress.

Updating Delay

We can update Delay to use wakers:

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
        }
    }
}

Now, once the requested duration has elapsed, the calling task is notified and the executor can ensure the task is scheduled again. The next step is to update mini-tokio to listen for wake notifications.

There are still a few remaining issues with our Delay implementation. We will fix them later.

When a future returns Poll::Pending, it must ensure that the waker is signalled at some point. Forgetting to do this results in the task hanging indefinitely.

Forgetting to wake a task after returning Poll::Pending is a common source of bugs.

Recall the first iteration of Delay. Here was the future implementation:

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 {
            // Ignore this line for now.
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

Before returning Poll::Pending, we called cx.waker().wake_by_ref(). This is to satisfy the future contract. By returning Poll::Pending, we are responsible for signalling the waker. Because we didn’t implement the timer thread yet, we signalled the waker inline. Doing so will result in the future being immediately re-scheduled, executed again, and probably not be ready to complete.

Notice that you are allowed to signal the waker more often than necessary. In this particular case, we signal the waker even though we are not ready to continue the operation at all. There is nothing wrong with this besides some wasted CPU cycles. However, this particular implementation will result in a busy loop.

Updating Mini Tokio

The next step is updating Mini Tokio to receive waker notifications. We want the executor to only run tasks when they are woken, and to do this, Mini Tokio will provide its own waker. When the waker is invoked, its associated task is queued to be executed. Mini-Tokio passes this waker to the future when it polls the future.

The updated Mini Tokio will use a channel to store scheduled tasks. Channels allow tasks to be queued for execution from any thread. Wakers must be Send and Sync, so we use the channel from the crossbeam crate, as the standard library channel is not Sync.

The Send and Sync traits are marker traits related to concurrency provided by Rust. Types that can be sent to a different thread are Send. Most types are Send, but something like Rc is not. Types that can be concurrently accessed through immutable references are Sync. A type can be Send but not Sync — a good example is Cell, which can be modified through an immutable reference, and is thus not safe to access concurrently.

For more details, see the related chapter in the Rust book.

Add the following dependency to your Cargo.toml to pull in channels.

crossbeam = "0.7"

Then, update the MiniTokio struct.

use crossbeam::channel;
use std::sync::Arc;
struct MiniTokio {
    scheduled: channel::Receiver<Arc<Task>>,
    sender: channel::Sender<Arc<Task>>,
}
struct Task {
    // This will be filled in soon.
}

Wakers are Sync and can be cloned. When wake is called, the task must be scheduled for execution. To implement this, we have a channel. When the wake() is called on the waker, the task is pushed into the send half of the channel. Our Task structure will implement the wake logic. To do this, it needs to contain both the spawned future and the channel send half.

use std::sync::{Arc, Mutex};
struct Task {
    // The `Mutex` is to make `Task` implement `Sync`. Only
    // one thread accesses `future` at any given time. The
    // `Mutex` is not required for correctness. Real Tokio
    // does not use a mutex here, but real Tokio has
    // more lines of code than can fit in a single tutorial
    // page.
    future: Mutex<Pin<Box<dyn Future<Output = ()> + Send>>>,
    executor: channel::Sender<Arc<Task>>,
}
impl Task {
    fn schedule(self: &Arc<Self>) {
        self.executor.send(self.clone());
    }
}

To schedule the task, the Arc is cloned and sent through the channel. Now, we need to hook our schedule function with std::task::Waker. The standard library provides a low-level API to do this using manual vtable construction. This strategy provides maximum flexibility to implementors, but requires a bunch of unsafe boilerplate code. Instead of using RawWakerVTable directly, we will use the ArcWake utility provided by the futures crate. This allows us to implement a simple trait to expose our Task struct as a waker.

Add the following dependency to your Cargo.toml to pull in futures.

futures = "0.3"

Then implement futures::task::ArcWake.

use futures::task::ArcWake;
use std::sync::Arc;
impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.schedule();
    }
}

When the timer thread above calls waker.wake(), the task is pushed into the channel. Next, we implement receiving and executing the tasks in the MiniTokio::run() function.

impl MiniTokio {
    fn run(&self) {
        while let Ok(task) = self.scheduled.recv() {
            task.poll();
        }
    }
    /// Initialize a new mini-tokio instance.
    fn new() -> MiniTokio {
        let (sender, scheduled) = channel::unbounded();
        MiniTokio { scheduled, sender }
    }
    /// Spawn a future onto the mini-tokio instance.
    ///
    /// The given future is wrapped with the `Task` harness and pushed into the
    /// `scheduled` queue. The future will be executed when `run` is called.
    fn spawn<F>(&self, future: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        Task::spawn(future, &self.sender);
    }
}
impl Task {
    fn poll(self: Arc<Self>) {
        // Create a waker from the `Task` instance. This
        // uses the `ArcWake` impl from above.
        let waker = task::waker(self.clone());
        let mut cx = Context::from_waker(&waker);
        // No other thread ever tries to lock the future
        let mut future = self.future.try_lock().unwrap();
        // Poll the future
        let _ = future.as_mut().poll(&mut cx);
    }
    // Spawns a new taks with the given future.
    //
    // Initializes a new Task harness containing the given future and pushes it
    // onto `sender`. The receiver half of the channel will get the task and
    // execute it.
    fn spawn<F>(future: F, sender: &channel::Sender<Arc<Task>>)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let task = Arc::new(Task {
            future: Mutex::new(Box::pin(future)),
            executor: sender.clone(),
        });
        let _ = sender.send(task);
    }
}

Multiple things are happening here. First, MiniTokio::run() is implemented. The function runs in a loop receiving scheduled tasks from the channel. As tasks are pushed into the channel when they are woken, these tasks are able to make progress when executed.

Additionally, the MiniTokio::new() and MiniTokio::spawn() functions are adjusted to use a channel rather than a VecDeque. When new tasks are spawned, they are given a clone of the sender-part of the channel, which the task can use to schedule itself on the runtime.

The Task::poll() function creates the waker using the ArcWake utility from the futures crate. The waker is used to create a task::Context. That task::Context is passed to poll.