Pinning Rust async data types for memory safety

Implementing asynchronous programming can help your application run faster and handle high traffic. But, the challenge of async programming in Rust is ensuring that data accessed by multiple tasks is consistent and safe.

Concurrency is one strategy to prevent this issue, as it provides benefits in memory usage. However, for memory safety with asynchronous data types in Rust, pinning is the best solution.

Pinning is a feature in Rust that allows an async data type to be secured or “pinned” to a specific location in memory so that it cannot be moved while it’s being used by multiple tasks.

In this article, we’ll investigate different methods for implementing pinning in Rust. We’ll also look at how to safely access pinned data.

Jump ahead:

• Using the Future trait for async data types in Rust
• Understanding the problem with async data types and memory safety
• Implementing pinning in Rust
  • Pinning to the stack
  • Pinning to the heap
• Accessing data values with pin projections

Using the Future trait for async data types in Rust

When working with async data in Rust, we’ve got this thing called the Future trait that helps us out. Any function or struct that uses this trait will be in one of two states: Pending or Ready.

We use the await keyword to wait for the Future trait to finish and give us a value. It’s pretty handy!

Here’s how to implement the Future trait into a struct:

impl Future for MyFutureStruct {
    type Output = ();
    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
        Poll::Ready(())
    }
}

Here’s how to implement the Future trait into a function:

impl MyFutureStruct {
    fn get_something_async(&self) -> impl Future<Output = String> + '_ {
        async {
            "Something".to_string()
        }
    }
}

Any Rust data type can be used in async as long as it implements the Future trait.

Understanding the problem with async data types and memory safety

Let’s take a look at how we do async in Rust. The language itself doesn’t provide us with a runner to do async programming, but fortunately, there is a community package, tokio, that can help us to handle the async runner. Let’s install it in our Rust project:

cargo add tokio --features full

There are two ways to create an async runner with tokio, with a macro or by using the tokio runtime builder.

Here’s how to create an async runtime with a macro:

#[tokio::main(flavor = "current_thread")]
async fn main() {
   // your code here
}

And here’s how to create an async runtime using the tokio runtime builder:

fn main() {
    let rt = tokio::runtime::Builder::new_current_thread()
        .build()
        .unwrap();
    rt.block_on(/* your async code or function here */);
}

For this tutorial, we’ll use the macro version since it’s much simpler.

In Rust, we have the struct data type that we can use to store a custom value — just like a class in Java.

We can create a struct like so:

#[derive(Debug)]
struct Post {
    name: String,
    slug: *const String
}

But, we can‘t just use this struct directly to perform an asynch process. First, we need to implement a Future trait to this struct. That will tell our async runtime that the value of slug from this struct will be filled later, enabling us to run the async process to get the value of this struct.

Let’s add a method to our struct so that we can fill the value for the name and slug properties:

impl Post {
    pub fn new(name: String) -> Self {
        Self { 
            name,
            slug: std::ptr::null()
        }
    }
    fn collect_slug(&mut self) {
        let ref_slug = &self.name as *const _;
        self.slug = ref_slug;
    }
    fn get_slug(&self) -> String {
        unsafe { &*(self.slug) }.replace(" ", "-").to_lowercase()
    }
}

Now, let’s implement a Future trait to our struct:

use std::future::Future;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

impl Future for Post {
    type Output = ();
    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
        println!("Post:::poll {}", self);
        Poll::Ready(())
    }
}

Let’s also implement a Display to make it easier for us to read the value inside our struct:

impl Display for Post {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "name:`{}` slug:`{}`",self.name self.get_slug()).unwrap();
        Ok(())
    }
}

Now let’s try working with this struct:

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let mut post1 = Post::new("Article First".to_string());
    let mut post2 = Post::new("Article Second".to_string());
    post1.collect_slug();
    post2.collect_slug();
    post1.await;
    post2.await;
}

We got everything that we wanted! Here’s the output:

Post:::poll name:`Article First` slug:`article-first`
Post:::poll name:`Article Second` slug:`article-second`

But, there’s one issue we encounter when moving the value of the Post struct. Take a look at the below code:

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let mut post1 = Post::new("Article First".to_string());
    let mut post2 = Post::new("Article Second".to_string());
    post1.collect_slug();
    post2.collect_slug();
    std::mem::swap(&mut post1, &mut post2);
    post1.await;
    post2.await;
}

With this memory swap, we expect to get a value like this:

Post:::poll name:`Article Second` slug:`article-second`
Post:::poll name:`Article First` slug:`article-first`

Instead, here’s our result:

Post:::poll name:`Article Second` slug:`article-first`
Post:::poll name:`Article First` slug:``

Why did this happen? To investigate further, let’s add some simple tracing using println:

println!("Before swap");
println!("post1: {:?} => {:p}, {:?}", post1.name, &post1.name, post1.get_slug());
println!("post2: {:?} => {:p}, {:?}\n", post2.name, &post2.name, post2.get_slug());
std::mem::swap(&mut post1, &mut post2);
println!("After swap");
println!("post1: {:?} => {:p}, {:?}", post1.name, &post1.name, post1.get_slug());
println!("post2: {:?} => {:p}, {:?}\n", post2.name, &post2.name, post2.get_slug());

If we print the values in memory we get the below output. What’s strange is that even though we used a Future trait, we didn’t get a swapped slug value for post2:

Before swap
post1: "Article Firstd" => 0x16af320b0, "article-firstd"
post2: "Article Second" => 0x16af320d0, "article-second"

After swap
post1: "Article Second" => 0x16af320b0, "article-firstd"
post2: "Article Firstd" => 0x16af320d0, "article-second"

We successfully swapped the memory for the name variable, but the original slug address is still the same — nothing changed there.

The below diagram may be helpful for understanding how the swap works:

In Rust, if a type contains pointers that refer to itself, it is not safe to move the value of that type because the pointers will not be updated and will still point to the old memory location.

For example, if we try to move the value of post1 to post2, the slug pointer in post1 will still point to the old location in memory, and the slug pointer in post2 will not be updated.

This will cause both post1 and post2 to have invalid pointers, which is not safe. Swapping the value of post1 and post2 would require changing the slug field, and it can’t be done because it’s self-referential.

Let’s see how pinning can help us address this issue.

Understanding and implementing pinning in Rust

[Pin](https://github.com/rust-lang/rfcs/blob/master/text/2349-pin.md) is a reference type in Rust; it was introduced to handle data that should not be moved or that can not be moved safely. For example, asynchronously accessed data can be pinned in memory to prevent it from being moved or modified by the Rust runtime.

With pinning, we can ensure that async data is accessed predictably and consistently, avoiding concurrency issues like data races. There are two ways to implement pinning in Rust; we can pin to the stack or we can pin to the heap.

Let’s take a closer look!

Pinning to the stack

Adding the PhantomPinned trait to the Post struct allows it to be pinned in memory, meaning that its location in memory will be fixed and will not be moved by the program.

To achieve this, we first need to import PhantomPinned to our project and then add it to the Post struct “!Unpin”:

use std::marker::PhantomPinned;
#[derive(Debug)]
struct Post {
    name: String,
    slug: *const String,
    _marker: PhantomPinned,
}

Now, we can create a new instance of the Post struct, like so:

impl Post {
    pub fn new(name: String) -> Self {
        Self {
            name,
            slug: std::ptr::null(),
            _marker: PhantomPinned,
        }
    }
}

Next, we need to change the way we mutate the value of our struct:

impl Post {
    pub fn collect_slug(self: Pin<&mut Self>) {
        let ref_slug = &self.name as *const String;
        let this = unsafe {
            self.get_unchecked_mut()
        };
        this.slug = ref_slug;
    }
    pub fn get_slug(self: Pin<&Self>) -> String {
        unsafe { &*(self.slug) }.replace(" ", "-").to_lowercase()
    }
    fn get_name(&self) -> &str {
        &self.name
    }
}

We also need to change to the Display trait implementation:

impl Display for Post {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let this = unsafe { Pin::new_unchecked(self) };
        write!(f, "name:`{}` slug:`{}`", this.as_ref().get_name(), Post::get_slug(this)).unwrap();
        Ok(())
    }
}

Finally, we can work with the pinned value of our struct:

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let mut post1 = Post::new("Article Firstd".into());
    let mut post2 = Post::new("Article Second".into());
    let mut post1 = unsafe { Pin::new_unchecked(&mut post1) };
    let mut post2= unsafe { Pin::new_unchecked(&mut post2) };
    Post::collect_slug(post1.as_mut());
    Post::collect_slug(post2.as_mut());
    post1.await;
    post2.await;
}

Here’s the result of the operations. It works and it gives us the value that we expected:

Post:::poll name:`Article Second` slug:`article-second`
Post:::poll name:`Article Firstd` slug:`article-firstd`

But, when we use the swap this way, the Rust compiler will “complain” as this is not memory safe:

std::mem::swap(post1.get_mut(), post2.get_mut());

N.B.,be sure to use .get_mut() when your data implements PhantomPinned if you are going to mutate the data. This trait guarantees that the struct will not be moved in memory during the lifetime of the mutable reference and can prevent certain types of bugs that can occur when working with pinned data

Pinning to the heap

Data that is stored on the stack can be moved around in memory. But when data is “pinned” to the heap, it means that its location in memory is fixed and will not be moved by the program. Data that is pinned to the heap will be in one location for as long as the program is running and will be easy to access and use.

To store data in the heap, we’ll use the Box data type:

impl Post {
    pub fn new(name: String) -> Pin<Box<Self>> {
        let post = Self {
            name,
            slug: std::ptr::null(),
            _marker: PhantomPinned,
        };
        let mut boxed = Box::pin(post);
        let self_ptr = &boxed.name as *const String;
        unsafe { boxed.as_mut().get_unchecked_mut().slug = self_ptr }
        boxed
    }
    fn get_slug(&self) -> String {
        unsafe { &*(self.slug) }.replace(" ", "-").to_lowercase()
    }
    fn get_name(&self) -> &str {
        &self.name
    }
}

To display the value of the data, we’ll use the as_ref function from Box:

impl Display for Post {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, 
          "name:`{}` slug:`{}`", 
          self.as_ref().get_name(), self.as_ref().get_slug()
        ).unwrap();
        Ok(())
    }
}

Here’s how we access the value:

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let post1 = Post::new("Article Firstd".into());
    let post2 = Post::new("Article Second".into());
    post1.await;
    post2.await;
}

// Post:::poll name:`Article Firstd` slug:`article-firstd`
// Post:::poll name:`Article Second` slug:`article-second`

Accessing data values with pin projections

Pin projections enable us to access the inner value of a Pin<P> (i.e., the value of type P) while keeping it pinned in memory.

To access the inner value of a Pin<P>, we can use the .as_ref() or .as_mut() method to get a reference or mutable reference to the inner value. This allows us to work with the inner value as if it were a regular value while ensuring that it will not be moved.

For example, we can create a Pin, like so:

let mut post1 = Post::new("Article Firstd".into());

Then, we can access the inner value using as_ref:

println!("post1: {:?} => {:?}", post1.as_ref().get_name(), post1.as_ref().get_slug());

This is useful when we are working with a nested struct that implements the Future trait. Let’s say we want to implement some tracing for our post struct:

struct TraceDuration<Fut: Future> {
    start: Option<time::Instant>,
    child: Fut
}

We can’t use the child directly like this:

impl<Fut: Future> Future for TraceDuration<Fut> {
    type Output = (Fut::Output, Duration);
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let start = self.start.get_or_insert(time::Instant::now());
        let post_poll = self.child.poll(cx);
        let elapsed = start.elapsed();
        match post_poll {
            Poll::Ready(res) => Poll::Ready((res, elapsed)),
            Poll::Pending => Poll::Pending,
        }
    }
}

Instead, we need to extract that field from the self object, like so:

impl<Fut: Future> Future for TraceDuration<Fut> {
    type Output = (Fut::Output, Duration);
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let (mut start, child) = unsafe {
            let this = self.get_unchecked_mut();
            (
                Pin::new_unchecked(&mut this.start),
                Pin::new_unchecked(&mut this.child),
            )
        };
        let start = start.get_or_insert(time::Instant::now());
        let post_poll = child.poll(cx);
        let elapsed = start.elapsed();
        match post_poll {
            Poll::Ready(res) => Poll::Ready((res, elapsed)),
            Poll::Pending => Poll::Pending,
        }
    }
}

Luckily there’s a crate called pin-project that makes our code even cleaner so that we don’t need to do the pin projection manually:

cargo add pin-project

Here’s how our code looks when we access the child using pin-project:

impl<Fut: Future> Future for TraceDuration<Fut> {
    type Output = (Fut::Output, Duration);
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let start = this.start.get_or_insert(time::Instant::now());
        let post_poll = this.child.poll(cx);
        let elapsed = start.elapsed();
        match post_poll {
            Poll::Ready(res) => Poll::Ready((res, elapsed)),
            Poll::Pending => Poll::Pending,
        }
    }
}

Conclusion

By using pinning with async programming, we can ensure that the data accessed by multiple tasks is always consistent and that there are no data race conditions.

When working with self-referential structs and pinning in Rust, be aware that sometimes it’s necessary to use unsafe code. One example is giving deference to the pointer stored in the slug variable. Using points in this way can be dangerous. If it is not used correctly, it can lead to memory safety issues such as data races or null pointer dereferences.

It’s very important to understand the level of safety of your code, and use unsafe code cautiously. Additionally, it is crucial to ensure that your struct does not move in memory when accessed by single or multiple tasks, otherwise it can cause data races or undefined behavior.