Tuesday, 16 January 2024

Futures and Pinning

Scratching the surface of how asynchronous code is implemented in Rust

tags: rust, futures, async, pin

How I like to visualize asynchronous code in Rust

Asynchronous code is useful for handling computations that are I/O bound. In Rust, the Future trait is a key .

pub trait Future {
    type Output;

    // Required method
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

By calling an asynchronous function, you return a future. The first level of abstraction I like to think of is that a future is a state machine. This machine is lazily evaluated - it doesn't do any work until it is actively polled by an executor.

Inside an async block or function, the await keyword can be used to suspend the current function's execution and wait for another future (that which is being .awaited) to complete .

futures-state-machine

The three states:

  1. Start: just a computation that could be called
  2. Pending: started computation, might not have a result yet
  3. Ready: finished computation

Analogy

Think of it like a contract to do some work that you could hand off to someone else - contractors.

The contractors won't start working on the contract until you tell them to.

Once the contractors have started, they'll keep chipping away at the contract until they're ready to deliver results.

Whilst your contractors are chipping away, you pass off responsibility to them completely. When they think they're ready, they ask you to check their work.

When they ask you to check their work, you make sure the contract is fulfilled. If it is complete, you mark the contract as delivered, otherwise you tell them to keep chipping away.

Getting more technical

There's a few things that I think are quite interesting about how this flow is implemented.

  1. The Executor
  2. Wakers
  3. Pins

Executor

The executor is the component responsible for running and managing our asynchronous tasks to completion. An executor could use any variety of scheduling algorithms; in our case, let's just think of it as maintaining a queue of tasks that are ready to poll. In practice, the queue could be any form of collection.

Tasks are added to the executor's queue when:

  • a top level future is submitted to the executor, this could be from spawning an async thread or when the main function of an asynchronous program starts.
  • a pending task is woken by its waker - signalling it's ready to be re-polled.

The executor polls each tasks' computation and handles the result:

  • If the computation has completed execution, it just has to return any resulting value and the task is then dropped from the queue.
  • If the computation is pending completion, the task is temporarily suspended by the executor. This is where the waker comes in...

What are wakers?

A notification system for tasks to reschedule themselves for polling.

The executor could just continuously keep scheduling all of its tasks in a cycle. But it doesn't - that's just wasted computation. Going back to our analogy, you're not going to keep phoning your contractors in an endless cycle - it's better use of your time to let them phone you.

Instead, it is expected that the future will store a reference to a waker. When some progress has been made, the future will call the wake method on that waker to notify the executor to reschedule the current task. If you're coming from the web world, it's akin to the difference between a client constantly polling the server, against the server just sending a message when it needs to.

What is Pin and Unpin?

The poll method of a future takes self: Pin<&mut Self>. Wrapping a type in Pin acts as a mechanism to ensure that once the data is placed in memory, it cannot be moved.

The main problem that is being solved here is for self-referential types - what are those?

The gist of self-referential types

When the compiler transforms an async fn into a state machine, the machine states are represented as structs and these need to hold all of the local variables that live across an await point within that function as fields. Consider something like:

async fn my_async_fn() -> i32 {
    let data = [1, 2, 3];
    let reference = &data[2];
    another_async_function().await;
    *reference
}

This would get transformed to:

futures-state-machine

In the first state, mapping this to memory would look along the lines of:

futures-state-machine

Now, we reach another_async_function().await. At this point, we suspend my_async_fn. To reflect that, we transition into another state and the executor could decide to move the entire state struct to a new memory location. During this transition, the local variables get moved - in this case, data and reference.

If the entire state struct is moved, data now has a new memory address. However, there's no logic that defines how to update reference which is just a raw memory address to where data[2] used to be. Uh, oh... we've got a dangling pointer.

futures-state-machine

This is precisely the problem Pin solves. By ensuring the state struct is pinned (fixed) in memory, the executor guarantees it won't be moved.

What is Unpin?

Unpin is a marker trait that says "I'm actually safe to move around in memory". Most standard types implement Unpin by default. Pin primarily exists to deal with types that are not Unpin which could include compiler generated state machines for async functions which include self-references.