Build your own async primitive

First, our primitive:

//! A oneshot channel is a synchronization primitive used to send a single
//! message from one task to another. When the sender has sent a message and it
//! is ready to be received, the oneshot notifies the receiver. The receiver is
//! a Rust async `Future` that can be `await`ed. Our implementation does not
//! depend on std or alloc.
//!
//! See https://tweedegolf.nl/blog/50/build-your-own-async-primitive 
//! for a full description of the internals.

What do we mean by all that? First, not depending on std or alloc means we can't use many common libraries. Rust async is still a new area and the libraries so far depend on allocating memory at least. That makes their primitives easier to build but also unusable on embedded devices which can't afford an allocator. I work with such devices regularly and so I've learned to avoid allocation whenever possible. We'll see that we can write our oneshot without allocation just fine.

Next, we don't want to block the entire CPU while waiting for the message to be sent. For one thing, it won't be sent if the CPU is blocked. When an operating system is available, it can schedule the sending task while the receiving task is blocked and there's no problem. Small embedded devices usually can't afford an operating system though, so you're left with implementing task switching manually. That can involve a lot of moving state from your only call stack into a struct and back so you can use the call stack to continue another task.

Rust async can do it for you: First, you choose an async executor to run your tasks. These executors can be much simpler than even an embedded OS. Then you write your tasks as normal. When you reach a point where you need to wait for another task to finish something, you just write .await. The executor will store the state of your call stack and continue with a task that isn't blocked. When the thing you were waiting for is done, the executor will wake your task up again.

Let's begin with a Oneshot struct which will contain the data both the sender and receiver need to access in order to transfer the message. It needs to be created high on the call stack so it will outlive all tasks that might reference it. Then we create a Sender and Receiver struct which contain a reference to the Oneshot. We can then give the sender to one task and the receiver to another. Since they both have a reference to the Oneshot, they must be shared references, which means we can't safely modify just anything inside. Instead, we need to use types which are safe to modify through shared references.

The message

First, we need a way to transfer the message from the sender to the receiver. We could move it directly if we wait for both the sender and receiver to be ready for the transfer: We'd know where the data is on the sender side and where it should go on the receiver side. We call this synchronous communication, both the sender and the receiver need to be ready before the communication happens. It has its uses, but in this case, we don't want the sender to wait for the receiver since it won't get a response back anyway. Instead, we'll store the message in the Oneshot where it can wait for the receiver to be ready.

We don't want to restrict our oneshot to just one kind of message, so we'll generalize it over the type of the message. We'll call that type T. When we create a new Oneshot we won't have a message for it yet so we need to tell Rust we'll initialize it later. In Rust, you'd usually use an Option<T> but we're going to track whether we have a T in a separate field so we'll use a MaybeUninit<T> instead. Finally, we need to modify the field through a shared reference.

We'll store an UnsafeCell<MaybeUninit<T>> which allows us to modify the message through a shared reference as long as we do so in a block of code marked as unsafe. An unsafe block just means we make it our own responsibility to check the code is safe instead of the compiler's. We need to do this whenever the safety logic is too complicated for the compiler to understand.

The synchronization

Next, we need to remember whether the message is ready to be received, or in other words, whether the message field has been initialized. This is not as easy as it seems. Both compilers and processors will reorder memory accesses in order to make your code run faster. They'll make sure the code in a single task seems like it ran in the right order, but they won't give the same guarantee for code running in different tasks. For example, if we had used an Option<T> to store the message, the sender might turn it into a Some() first and store the message afterwards. If the receiver happened to check in between, it would try to read a message before there was ever stored one and it would end up with a bunch of garbage.

In order to make concurrent memory accesses safe, atomics were invented. We'll remember if we have a message with an AtomicBool. Atomics make sure no task sees a partial update of the atomic. The atomic either updates completely or not at all. On top of that, it also helps us synchronize access to other memory: The value of the AtomicBool determines whether we allow the sender or the receiver to access the message. That way, they never access it at the same time.

Finally, we want to use Rust async/await to wait until the message is ready. This requires us to store a waker. To be able to update the waker from a shared reference, we use the very nice AtomicWaker from the futures library.

use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::sync::atomic::{AtomicBool, Ordering};
use futures::task::AtomicWaker;

/// Transfer a single message between tasks using async/await.
pub struct Oneshot<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    waker: AtomicWaker,
    has_message: AtomicBool,
}

The Sender and Receiver are just thin wrappers around Oneshot references. They are generic over both the type of the underlying message (T) and the lifetime of the reference ('a). Wrapping the references in a struct allows us to name them and dictate what operations they support.

pub struct Sender<'a, T>(&'a Oneshot<T>);
pub struct Receiver<'a, T>(&'a Oneshot<T>);

It may seem like we're not going very fast but choosing the right data representation is often the hardest part of programming. From a good data representation, the implementation will follow naturally. We begin with the new function which allows the creation of a new Oneshot. We have no message yet so we don't initialize it and set has_message to false.

impl<T> Oneshot<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            waker: AtomicWaker::new(),
            has_message: AtomicBool::new(false),
        }
    }

The unsafe

Next, we implement the most basic operations, putting a message in and taking it out again. We'll get to synchronization later, so these will be unsafe functions. For performance reasons, users may want to use these functions directly so we'll make them public. This is fine as long as we document what properties must hold to use the function safely.

The put function will store a message. Since we're implementing a oneshot, we only expect one communication so we assume there was no message before put gets called. This is a property a user of this unsafe function must be told about!

As soon as we set has_message to true, the receiver may try to read the message so we must make sure to write the message first. We also need to prevent the write of the message and of has_message from being reordered. We do that by storing has_message with Ordering::Release. This is a signal to the compiler that it must make sure any memory accesses before it must be finished and visible to other cores.

Finally, we use the waker in order to schedule the receiving task to continue.

    /// NOTE(unsafe): This function must not be used when the oneshot might
    /// contain a message or a `Sender` exists referencing this oneshot. This
    /// means it can't be used concurrently with itself or the latter to run
    /// will violate that constraint.
    pub unsafe fn put(&self, message: T) {
        self.message.get().write(MaybeUninit::new(message));
        self.has_message.store(true, Ordering::Release);
        self.waker.wake();
    }

The take function will check if a message is available yet and take it out of the oneshot if it is. This time, we use Ordering::Acquire in order to ensure all memory accesses after loading has_message really happen after it. For every synchronization between tasks, a Release-Acquire pair is needed. Sometimes synchronization needs to go both ways and Ordering::AcqRel can be used to get both effects.

When we're done taking the message, we Release its memory again by setting has_message to false. If two instances of take run concurrently, they might both reach the message read before setting has_message to false and thus duplicate the message so we need to disallow that in the safety contract.

On the other hand, if it is run concurrently with put, according to the contract of put, there must be no message beforehand. If take loads has_message first, it finds no message and returns. If put writes has_message first, there is guaranteed to be a message ready so take takes it without issue. Therefor, a single take can be safely run concurrently with a single put. This is exactly what we need for a oneshot channel. We never expect either function to run more than once but a single message can be transferred between tasks safely. Now to enforce that safety in the Rust type system.

    /// NOTE(unsafe): This function must not be used concurrently with itself,
    /// but it can be used concurrently with put.
    pub unsafe fn take(&self) -> Option<T> {
        if self.has_message.load(Ordering::Acquire) {
            let message = self.message.get().read().assume_init();
            self.has_message.store(false, Ordering::Release);
            Some(message)
        } else {
            None
        }
    }

The split

The split function splits the oneshot into a sender and a receiver. The sender can be used to send one message and the receiver to receive one. The split function takes a unique reference to the oneshot and returns two shared references to it with the same lifetime 'a. This means Rust will consider the oneshot to be uniquely borrowed until it is sure both the sender's and receiver's lifetime have ended. It prevents us from making more than one pair!

However, once the lifetimes end, the oneshot will no longer be borrowed, so nothing prevents someone from calling split again. Is that a problem? Not really: It means the Oneshot can be reused to send another single message. We still call it a oneshot because the sender and receiver it creates can only be used once.

There is one thing we need to take care of however: The receiving task might stop caring and drop its Receiver before the sender sends its message. In that case, the oneshot will still contain a message after the lifetime of the references end. In order to allow the Oneshot to be reused again, we need to remove that message.

Note that simply setting has_message to false is a problem because the message is stored in a MaybeUninit which doesn't know by itself whether it has a message and so it also doesn't know when it should drop the message. Values that are forgotten without being dropped can leak resources or even cause undefined behavior! We always need to make sure the message is initialized before setting has_message to true and dropped or moved elsewhere before setting it to false.

The safe thing to do is to take any message out of the MaybeUninit into a type that drops implicitly again. However, take is unsafe so before we use it, we must check its contract: The contract states take may not be used concurrently with itself. In this case, we can be sure of that because split owns the unique reference to the Oneshot and so nobody else could have a reference through which they could call take.

    /// Split the Oneshot into a Sender and a Receiver. The Sender can send one
    /// message. The Receiver is a Future and can be used to await that message.
    /// If the Receiver is dropped before taking the message, the message is
    /// dropped as well. This function can be called again after the lifetimes
    /// of the Sender and Receiver end in order to send a new message.
    pub fn split<'a>(&'a mut self) -> (Sender<'a, T>, Receiver<'a, T>) {
        unsafe { self.take() };
        (Sender(self), Receiver(self))
    }

Even if a user is going to track when it's safe to send and receive manually, they might still want to make use of the Sender and/or the Receiver. We can allow creating them separately with functions marked unsafe. This allows a user to create a Receiver while using put directly and preventing the overhead of a Sender for example.

    /// NOTE(unsafe): There must be no more than one `Receiver` at a time.
    /// `take` should not be called while a `Receiver` exists.
    pub unsafe fn recv<'a>(&'a self) -> Receiver<'a, T> {
        Receiver(self)
    }

    /// NOTE(unsafe): There must be no more than one `Sender` at a time. `put`
    /// should not be called while a `Sender` exists. The `Oneshot` must be
    /// empty when the `Sender` is created.
    pub unsafe fn send<'a>(&'a self) -> Sender<'a, T> {
        Sender(self)
    }

We define a simple is_empty function to check if a message was sent yet. This would ordinarily not be useful since the send would know whether it has sent anything and the point of the receiver is that it finds out when it tries to receive. It is still useful for debugging and assertions however. Since we don't know what is being checked exactly and the performance of debugging is not an issue, we should use the AcqRel ordering here.

    pub fn is_empty(&self) -> bool {
        !self.has_message.load(Ordering::AcqRel)
    }
}

If we drop the Oneshot, we need to drop any message it might have as well.

impl<T> Drop for Oneshot<T> {
    fn drop(&mut self) {
        unsafe { self.take() };
    }
}

The Sender is extremely simple: It allows you to call the put function, but only once. This is enforced by the fact that this send function doesn't take a reference to self, but instead consumes it. After using send, the Sender's lifetime has ended and it can't be used again. Calling put once is safe because when creating the Sender using split, we ensured there was no message stored.

impl<'a, T> Sender<'a, T> {
    pub fn send(self, message: T) {
        unsafe { self.0.put(message) };
    }
}

The future

The receiver we get from split is a future that can be awaited for the message being sent. This is the last tricky bit in our implementation. A Future is a thing with a poll function which will be called by the executor. In it, it receives a pinned reference to itself and a context containing a waker. In our case, we don't care that the reference is pinned. The poll can return one of two things: Poll::Ready(message) indicates the future is done and the .await will return with the message. Poll::Pending means the future is not done yet and Rust will handle control back to the executor that called poll so it can find another task to run.

The Future will be polled for the first time when it is first awaited. After that, in principle it won't run again until we use the waker we received in the previous poll to wake it again. This means we need to store the waker somewhere where the sending task can use it to wake the receiving task again. This is of course in the waker field of the Oneshot. Thanks to using an atomic waker, we don't need to worry if it's safe to store the waker, we can do it at any time.

In practice, some executors poll the Future even if it hasn't been woken so it is always important to check if we're really done. In this case, we can safely call take since we have the unique reference to the Receiver and the Sender will never call take. Note that the documentation of AtomicWaker states that consumers should call register before checking the result of a computation so we can't call it only in case the message is not ready yet.

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

impl<'a, T> Future for Receiver<'a, T> {
    type Output = T;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
        self.0.waker.register(cx.waker());
        if let Some(message) = unsafe { self.0.take() } {
            Poll::Ready(message)
        } else {
            Poll::Pending
        }
    }
}

If we drop the receiver, any message in the oneshot will not be used anymore. It will already be dropped eventually when the oneshot itself is dropped or reused, but we might save resources by dropping it early, so we drop the message here if it exists. Note that it is safe to do so because we have the unique reference to the receiver and the sender will never call take so it won't run concurrently.

impl<'a, T> Drop for Receiver<'a, T> {
    fn drop(&mut self) {
        unsafe { self.0.take() };
    }
}

That's it

And that wraps up the implementation of our oneshot. Not too bad, right? Hopefully you have a good idea now about what goes into a concurrency primitive. The code for this blog is available at https://github.com/tweedegolf/async-heapless so please create an issue or pull request if there's something wrong with it. Also see https://github.com/tweedegolf/async-spi/blob/main/src/common.rs where I use (the unsafe methods of) the oneshot to build an async SPI driver.

I want to encourage you to try writing your own abstractions. For starters, you could create a channel with a capacity of one message where the sender and receiver can be used multiple times and the sender can block until there is space in the channel. After that, you could look into channels that can store multiple messages or allow multiple concurrent senders or receivers. Or maybe you want to build a primitive where both tasks wait for each other and then exchange messages at the same time. Just make sure you don't try to add everything at once, each of those things is hard enough on its own.

Finally, remember to document every unsafe function with the conditions needed to use it safely!