AsyncWrite: A Tale of Four Implementations
The "When" and "Why" of Using pin_project
In Rust, we often don’t need to work directly with futures or implement low-level async APIs, but sometimes they are inevitable. In this post, I report and reflect on my experience implementing an async API. I am going to compare four different implementations, and walk you through the use of futures in a non-trivial example. I will discuss the “when” and “why” of using pin_project.
The following table summarizes the four alternative implementations, and includes references to the corresponding code in the associated GitHub repository.
| Variant | Path | Is struct stateful? | Calls an async function? | Custom Future? | Pinning of the nested Future |
|---|---|---|---|---|---|
| Impl-1 | pinned.rs | No | Yes | No | Pinned field |
| Impl-2 | stateful.rs | Yes | Yes | No | Pinned field |
| Impl-3 | pinless.rs | Yes | No | Yes (Unpin) |
Safe on-the-fly pinning |
| Impl-4 | asyncless.rs | No (but could be) | No | Yes (!Unpin) |
pin_project or unsafe |
| 🚫 | - | - | Yes | - | pin_project |
In each alternative, we will implement AsyncWrite for a struct.
In the first alternative, the struct does not hold any internal state.
Inside the AsyncWrite implementation, we call an async function, which, as we will see, requires pinning.
In the second implementation, we do the same for a struct that maintains internal state.
We will see how some refactoring can help us address a class of lifetime issues.
The third variant wraps the async call inside a custom Future.
Thus, the AsyncWrite implementation does not have to deal directly with the async call.
Here, we will observe the difference between Unpin and !Unpin types.
Finally, in the fourth variant, we will completely remove the async call, and will additionally discuss the use of pin_project.
These four scenarios are not the only imaginable combinations. However, not all combinations are interesting or feasible. For example, the last row in the table represents a combination that won’t work without resorting to unsafe code, potentially introducing unsoundness.
Why custom implement AsyncWrite?
Having a rough idea about how futures work is different from actually using them to build a complex working system!
A while ago, I had to work with futures at a low level for the first time.
With a lot of cursory readings here and there, particularly about pinning, I ended up with a working solution 🎉
Looking at that solution now, after a few months, I see that my solution could have been simplified, if I had a deeper understanding of some concepts.
Notably, in my solution, I had used pin_project unnecessarily!
Before diving into implementing AsyncWrite, let me tell you about my use case.
In my project, I needed to compress large amounts of data and upload the compressed content to AWS S3.
For S3, I used the aws_sdk_s3 crate, which is async and relies on tokio.
For compression, I used ZstdEncoder`, from async_compression.
Async was chosen first and foremost to achieve better performance, but also because I could not find a reliable alternative for aws_sdk_s3.
ZstdEncoder has a suspiciously simple API. It took me a while to understand how it works.
Looking at the examples was necessary for an aha moment!
The API for ZstdEncoder says:
This structure implements an AsyncWrite interface and will take in uncompressed data and write it compressed to an underlying stream.
ZstdEncoder provides a few methods for creating instances of it (e.g., new).
These methods take an AsyncWrite object (inner) as input! This is the underlying stream that the compressed data is written into.
To feed uncompressed data into an instance of ZstdEncoder, the most convenient and efficient way is to wrap it in a BufWriter and call write_all, or any other method from tokio::io::AsyncWriteExt on it.
To retrieve the compressed data from ZstdEncoder, you must find a way to access the underlying inner stream that you provided at the creation time. ZstdEncoder takes ownership of the inner stream, and during the operation, only gives you immutable access to it. With this restricted access, and considering the available AsyncWrite types (e.g., Vec) as the underlying stream, you cannot consume the compressed data, in chunks, during the operation. In other words, to consume the compressed data, you have to wait until the end of the operation, and call into_inner on the ZstdEncoder instance, to gain ownership of the inner stream and access the compressed data. The following snippet shows this scenario.
async fn compress(in_data: &[u8]) -> Result<Vec<u8>> {
let mut encoder = ZstdEncoder::new(Vec::new());
encoder.write_all(in_data).await?;
encoder.shutdown().await?;
Ok(encoder.into_inner())
}
For a small example this works fine, but if you are using async, particularly in a streaming manner, you are most likely working with large amounts of data. You cannot simply wait until the end of the process to examine or use the results.
This was exactly the case in my project. I was processing terabytes of data, and needed to upload the compressed results to S3. More specifically, I wanted an end-to-end in-memory streaming solution. As the input stream of uncompressed data came in, it would be compressed, split into chunks and uploaded to S3 (using its multi-part upload API, with create, upload, and complete methods).
Despite my best efforts to avoid implementing AsyncWrite, it proved inevitable as no combination of existing structs and APIs could fulfill the task.
Implementing AsyncWrite
Let’s start with a quick look at AsyncWrite’s API:
pub trait AsyncWrite {
// Required methods
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8]
) -> Poll<Result<usize, Error>>;
fn poll_flush(
self: Pin<&mut Self>,
cx: &mut Context<'_>
) -> Poll<Result<(), Error>>;
fn poll_shutdown(
self: Pin<&mut Self>,
cx: &mut Context<'_>
) -> Poll<Result<(), Error>>;
// Provided methods
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[IoSlice<'_>]
) -> Poll<Result<usize, Error>> { ... }
fn is_write_vectored(&self) -> bool { ... }
}
The first thing you notice about this API is that you need to work with pinned objects!
A large part of this post is dedicated to discussing pinning options. From this API, methods that need to be implemented are poll_write, poll_flush, and poll_shutdown.
My focus will be on poll_write.
In this article, I am using tokio::io::AsyncWrite; futures::io::AsyncWrite is very similar.
AsyncWrite for a stateless struct
Let’s say we have an AsyncUploader struct that can asynchronously upload content to a remote server.
We want to use this uploader in contexts where an AsyncWrite is needed, and will implement AsyncWrite for it.
This uploader could be as simple as the following struct:
// Impl-1.a: pinned.rs
struct AsyncUploader {}
impl AsyncUploader {
fn new() -> Self {
Self {}
}
async fn upload(&self, _data: &[u8]){
// TODO: Upload the given bytes.
}
}
Inside upload, I am going to call some async functions (for instance, functions from aws_sdk_s3), so it has to be async. But for now, we will leave the body of upload empty.
To see AsyncUploader in action, I will use the following simple caller code.
// Impl-1.b: pinned.rs
#[tokio::main]
async fn main() {
let mut uploader = AsyncUploader::new();
// In reality, we will read from an input stream instead of this for loop!
for i in 0..10 {
let msg = format!("{i}: Hello, world!\n");
uploader.write_all(msg.as_bytes()).await.unwrap();
}
uploader.shutdown().await.unwrap();
}
Futures under the hood
The most important part of the code above is the call to write_all (from AsyncWriteExt).
Internally, write_all calls poll_write (from AsyncWrite).
The following sequence diagram shows the objects that are involved in executing write_all, including uploader and a Future on which poll_write and poll are called, respectively.
This diagram is simplified and not very accurate.
I have left out the details of tokio workers, and tasks.
This is because in our simple example, we have only one task, and one worker.
Adding them to the diagram is not worth the added clutter.
The sequence diagram does not cover all different paths corresponding to the return value of poll_write, or the termination condition for the loop.
But I hope it helps with understanding the big picture.
As shown in the diagram, in the main function, the call to write_all on the uploader results in the creation of a Future of type WriteAll.
This Future is then executed by a worker in the tokio runtime.
The runtime periodically calls poll on the Future, until it is ready.
Each call to poll on WriteAll results in a call to poll_write (in more general cases where a buffer is involved, it may result in multiple calls to poll_write).
This continues until poll_write returns Ready.
This completes one iteration of the for loop in the main function. Note that the call to poll_write may sometimes return Ready immediately.
In that case, we don’t really enter the loop in the sequence diagram.
Now that we’ve seen the big picture, let’s implement AsyncWrite for AsyncUploader, with a focus on poll_write.
Inside poll_write, one way or another, we have to call AsyncUploader::upload, which is an async function.
Normally, we’d call .await on an async function (e.g., self.upload(buf).await), but poll_write is not an async function.
So inside it, we cannot use .await! Instead, we have to create a Future from the call to upload, and call poll on it.
This is possible, because the call to an async function (in this case upload) returns a Future.
Futures and pinning
Recall that Future is a trait with the following API:
pub trait Future {
type Output;
// Required method
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}
Note that poll takes a pinned mutable pointer to Self as its receiver.
Pin is a smart pointer that wraps another pointer in itself, and expresses a guarantee about the immovability of the object (pointee in the following diagram) that the inner pointer points to. Pin also restricts mutable access to the pointee.
To be able to call poll on the Future returned from the call to upload, we have to first pin a pointer to that Future. Here is an attempt at it:
// Impl-1.c: pinned.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
Let mut fut = self.upload(buf);
let mut future = Pin::new(&mut fut);
// TODO: Use the future.
Poll::Pending
}
}
Compiling this results in the following error:
error[E0277]: `{async fn body@src/bin/main.rs:20:41: 22:6}` cannot be unpinned
--> src/bin/pinned.rs:45:35
|
20 | async fn upload(&self, _buf: &[u8]) {
| ----------------------------------- within this `impl futures::Future<Output = ()>`
...
45 | let mut future = Pin::new(&mut fut);
| -------- ^^^^^^^^ within `impl futures::Future<Output = ()>`, the trait `Unpin` is not implemented for `{async fn body@src/bin/pinned.rs:20:41: 22:6}`
| |
| required by a bound introduced by this call
|
= note: consider using the `pin!` macro
consider using `Box::pin` if you need to access the pinned value outside of the current scope
note: required because it appears within the type `impl Future<Output = ()>`
--> src/bin/pinned.rs:20:5
|
20 | async fn upload(&self, _buf: &[u8]) {
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: required by a bound in `Pin::<P>::new`
--> /Users/razieh/.rustup/toolchains/nightly-x86_64-apple-darwin/lib/rustlib/src/rust/library/core/src/pin.rs:484:23
|
484 | impl<P: Deref<Target: Unpin>> Pin<P> {
| ^^^^^ required by this bound in `Pin::<P>::new`
...
503 | pub const fn new(pointer: P) -> Pin<P> {
| --- required by a bound in this associated function
The compiler is telling us that the Future created from the call to upload cannot be unpinned!
Meaning that this Future does not implement Unpin.
Unpin is a marker trait that tells the compiler that an object does not care about its location in memory.
Most types are Unpin.
A prominent exception are Future types created from calling async functions.
We are getting this compiler error because the input to Pin::new has to be a pointer to a pointee that is Unpin, but the Future we have is not Unpin.
We have a few solutions here.
If we look at Pin’s documentation, we see another method, new_unchecked, for creating Pin that does not require Unpin.
However, this method is unsafe. Let’s avoid that for now.
The other alternative is Box::pin, which is suggested by the compiler.
// Impl-1.d: pinned.rs
use futures::Future;
...
impl AsyncWrite for AsyncUploader {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
let mut fut = Box::pin(self.upload(buf));
match fut.as_mut().poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(_) => Poll::Ready(Ok(buf.len())),
}
}
}
Now that we have a pinned Future, we can call poll on it.
Note that we have to call as_mut to get a pinned mutable reference to the Future object.
This is because fut is a Pin<Box<...>>; not exactly the type we can call poll on.
This implementation creates a new Future every time poll_write is invoked.
This is not at all useful! For a Future to advance, one must regularly invoke poll on it.
In this setup, however, we only call poll once on each Future created inside poll_write!
These Futures go out of scope and get destroyed when we return from poll_write.
We never get to poll on them again or make any progress with them!
In fact, if you try this implementation with an upload function that implements a long-running task (as below), you will see that the program never terminates, and never produces any output!
// Impl-1.e: pinned.rs
struct AsyncUploader {}
impl AsyncUploader {
...
async fn upload(&self, data: &[u8]) -> Result<(), Error> {
// Sleep for a bit to simulate a long-running operation.
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
println!(
"Uploaded {}",
String::from_utf8(data.to_vec())
.map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?
);
Ok(())
}
}
For simplicity, and to avoid getting into complications of interacting with a remote server, in this implementation of upload, I have used sleep (the non-blocking version from tokio) to simulate a long-running task.
The execution of main gets stuck, because with this implementation of upload, poll_write never returns anything other than Poll::Pending!
A long-living Future
To fix this, we have to change our design to give the Future created from the call to upload a longer lifetime, so that we can keep track of its progress, and repeatedly call poll on it.
To do this, we have to bind the Future to the AsyncUploader instance, by declaring the Future as a field.
The following is the updated sequence diagram.
This diagram is again not very precise. It only shows the scenario where the first call to poll_write returns Poll::Pending, which is not always the case. In some cases, the call to write_all may complete with a single call to poll_write that immediately returns Poll::Ready.
To implement this design, we have to add a new field (which we call fut) to AsyncUploader.
But what exactly should be the type of this field?
We know the following:
- It is a
Future, andpollwill be called on it. Sincepoll’s receiver is aPin, one way or another, we have to pin thisFuture. - The
Futureis the result of calling an async function. We don’t know the concrete type of thisFuture, and have to refer to it asdyn Future<Output = T>, whereTis the return type of theuploadmethod. - When we create an instance of
AsyncUploader, we don’t yet have an instance of thisFuture. To handle this case,futhas to be anOption. - Since
dyn Future<Output = T>is not sized, we have to keep a pointer to it. It is more convenient, and more ergonomic to useBoxand keep theFutureon the heap.
There is a type alias for Pin<Box<dyn Future<Output = T>>>, called BoxFuture, which is exactly the type we need here. In addition to T, BoxFuture takes a lifetime parameter, which we can set to ‘static.
Putting all this together, we update AsyncUploader as follows:
// Impl-1.f: pinned.rs
struct AsyncUploader {
fut: Option<BoxFuture<'static, Result<(), Error>>>,
}
With this new field, inside poll_write, we first check if self.fut is None. If that is the case, we (1) call upload, passing the content of the input buf to it, (2) update the value of self.fut with the resulting Future, and (3) return Poll::Ready. Since we have consumed the incoming bytes in buf, it is important to return Poll::Ready. If we already have a Future in fut, then we poll on it. If the result of poll is Poll::Ready with success, we do as we did above: call upload again, update fut, and return Poll::Ready. Otherwise, we either return Poll::Pending, or Poll::Ready with an error. This implementation is shown below.
// Impl-1.g: pinned.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
if self.fut.is_none() {
self.fut = Some(Box::pin(self.upload(buf)));
// The future we created is not necessarily completed yet, but
// we successfully used all the input bytes, and are now ready
// to accept new bytes. The poll_write contract requires us to
// return Ready.
return Poll::Ready(Ok(buf.len()));
}
// We have an on-going Future. We can only upload the given bytes
// in buf, and create a new Future, if the previous one is
// complete. Otherwise, we return Pending or Ready with an error.
match self.fut.as_mut().unwrap().as_mut().poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(())) => {
self.fut = Some(Box::pin(self.upload(buf)));
Poll::Ready(Ok(buf.len()))
}
Poll::Ready(Err(err)) => {
let _ = self.fut.take();
Poll::Ready(Err(err))
}
}
}
}
If we run this, we get the following error:
error: future cannot be sent between threads safely
--> src/bin/pinned.rs:61:29
|
61 | self.fut = Some(Box::pin(self.upload(buf)));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^ future returned by `upload` is not `Send`
|
= help: the trait `Sync` is not implemented for `(dyn futures::Future<Output = Result<(), std::io::Error>> + std::marker::Send + 'static)`
note: captured value is not `Send` because `&` references cannot be sent unless their referent is `Sync`
--> src/bin/pinned.rs:35:9
|
35 | &self,
| ^^^^^ has type `&AsyncUploader` which is not `Send`, because `AsyncUploader` is not `Sync`
= note: required for the cast from `Pin<Box<impl futures::Future<Output = Result<(), std::io::Error>>>>` to `Pin<Box<dyn futures::Future<Output = Result<(), std::io::Error>> + std::marker::Send>>`
The error says that the Future that we have wrapped inside BoxFuture is not Send. One might be tempted to fix this by declaring a new type based on BoxFuture, but without the Send trait bound:
pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
If we replace BoxFuture with BoxFut in the declaration of fut, we make the compiler very unhappy:
error[E0502]: cannot borrow `self` as mutable because it is also borrowed as immutable
--> src/bin/pinned.rs:61:13
|
61 | self.fut = Some(Box::pin(self.upload(buf)));
| ^^^^ ---------------------------
| | | |
| | | immutable borrow occurs here
| | cast requires that `self` is borrowed for `'static`
| mutable borrow occurs here
error: lifetime may not live long enough
--> src/bin/pinned.rs:61:29
|
53 | mut self: std::pin::Pin<&mut Self>,
| - let's call the lifetime of this reference `'1`
...
61 | self.fut = Some(Box::pin(self.upload(buf)));
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^ cast requires that `'1` must outlive `'static`
error[E0597]: `self` does not live long enough
--> src/bin/pinned.rs:61:38
|
53 | mut self: std::pin::Pin<&mut Self>,
| -------- binding `self` declared here
...
61 | self.fut = Some(Box::pin(self.upload(buf)));
| ---------^^^^--------------
| | |
| | borrowed value does not live long enough
| cast requires that `self` is borrowed for `'static`
...
84 | }
| - `self` dropped here while still borrowed
error[E0502]: cannot borrow `self` as mutable because it is also borrowed as immutable
--> src/bin/pinned.rs:76:17
|
76 | self.fut = Some(Box::pin(self.upload(buf)));
| ^^^^ ---------------------------
| | | |
| | | immutable borrow occurs here
| | cast requires that `self` is borrowed for `'static`
| mutable borrow occurs here
error[E0597]: `self` does not live long enough
--> src/bin/pinned.rs:76:42
|
53 | mut self: std::pin::Pin<&mut Self>,
| -------- binding `self` declared here
...
76 | self.fut = Some(Box::pin(self.upload(buf)));
| ---------^^^^--------------
| | |
| | borrowed value does not live long enough
| cast requires that `self` is borrowed for `'static`
...
84 | }
| - `self` dropped here while still borrowed
The main problem is that in our implementation, we are violating one of Rust’s most important memory management rules. We are borrowing self, simultaneously, as both mutable and immutable. Additionally, we have a lifetime issue. This comes from the fact that (1) we declared fut to have a static lifetime, then (2) bound its lifetime to that of self (by setting it to be the result of self.upload), which (3) is dropped at the end of poll_write.
Let’s backtrack the change from BoxFuture to BoxFut. In fact, the Send trait bound is required on BoxFuture because in a multithreaded context, calls to Future::poll may run on different threads. But the borrowing and lifetime issues are genuine problems that we have to fix.
Refactoring to the rescue
We are getting the “future cannot be sent between threads safely” error because the Future we are creating by the call to upload is capturing a reference to self, which is not Sync.
Luckily, inside upload, we never use self. So we can make upload an ordinary function, and remove self from its signature.
Let’s do this for now, and call the new function upload_bytes.
I’ve also changed the type of data from &[u8] to Vec<u8> to solve a similar lifetime problem.
Note that this kind of refactoring might not be as straightforward in more complex situations, but is often the cleanest way to address lifetime issues when working with Futures from async calls.
Here is our upload_bytes function:
// Impl-1.h: pinned.rs
/// This is the slow upload function, which we want to use inside the
/// AsyncWrite implementation.
async fn upload_bytes(data: Vec<u8>) -> Result<(), Error> {
// Sleep for a bit to simulate a long-running operation.
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
println!(
"Uploaded {}",
String::from_utf8(data)
.map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?
);
Ok(())
}
With this change, we can implement poll_write as follows.
// Impl-1.i: pinned.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
if self.fut.is_none() {
// We need a vec to avoid lifetime issues.
let data = buf.to_vec();
self.fut = Some(Box::pin(upload_bytes(data)));
// The future we created is not necessarily completed yet, but
// we successfully used all the input bytes, and are now ready
// to accept new bytes. The poll_write contract requires us to
// return Ready.
return Poll::Ready(Ok(buf.len()));
}
// We have an on-going Future. We can only upload the given bytes
// in buf, and create a new Future, if the previous one is
// complete. Otherwise, we return Pending or Ready with an error.
match self.fut.as_mut().unwrap().as_mut().poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(())) => {
self.fut = Some(Box::pin(upload_bytes(buf.to_vec())));
Poll::Ready(Ok(buf.len()))
}
Poll::Ready(Err(err)) => {
// Clean up fut to avoid a "resumed after completion"
// panic, if for whatever reason the caller decides to
// ignore this error and proceed with another write or a
// call to shutdown.
let _ = self.fut.take();
Poll::Ready(Err(err))
}
}
}
}
Note that I have made the first argument of poll_write (self: std::pin::Pin<&mut Self>) mut! This is allowed, and not a violation of the AsyncWrite::poll_write’s signature, because we are consuming self, and can mutate it anyway.
Equivalently, we could redeclare self as mutable in the body of the method.
We need self to be mutable, because, we are modifying its fut, and polling on it, which requires mutable access to fut and therefore self.
For this implementation to work correctly, we need a similar code in poll_shutdown, except that in poll_shutdown, we don’t call upload_bytes, as we don’t have any bytes to upload! See the final code on github.
If we run the main function, we can see that it prints out ten “Uploaded {i}: Hello World!” messages, with one-second delays between them!
Now that we have this simple AsyncUploader working, let’s make it more interesting.
Specifically, the AsyncUploader currently does not have any internal state. In reality, we need to work with structs that have some internal state.
AsyncWrite for a stateful struct
In my project, I had to upload content to S3.
To speed up the process, I used S3’s multi-part upload API, which allows splitting a file (or object) into parts, and uploading them concurrently.
S3 multipart upload requires each part to be at least 5M bytes, except for the last part.
Addressing this requirement in the implementation can be achieved in various ways.
One approach is to task AsyncUploader with ensuring that this condition is met.
In such a case, AsyncUploader will invoke the upload_bytes method only if it has a minimum number of bytes to upload.
Otherwise, it will simply append the bytes to an internal buffer, which is the internal state we are going to add to AsyncUploader:
// Impl-2.a: stateful.rs
struct AsyncUploader {
/// Minimum size of the chunks to upload.
min_size: usize,
/// Internal buffer.
buf: Vec<u8>,
/// Future corresponding to the long-running upload operation.
fut: Option<BoxFuture<'static, Result<(), Error>>>,
}
impl AsyncUploader {
fn new(min_size: usize) -> Self {
Self {
min_size,
buf: Vec::with_capacity(2 * min_size),
fut: None,
}
}
...
}
I’ve initialized the buffer with a capacity twice the minimum size. This is just a heuristic to avoid reallocating the buffer too many times, as we add bytes to it.
To make poll_write more readable after this change, I refactored it slightly, and moved the main logic to an upload method, which (1) adds the incoming bytes to the inner buffer if it does not have enough bytes, (2) calls uploade_bytes if it has enough bytes, and (3) updates fut.
// Impl-2.b: stateful.rs
impl AsyncUploader {
...
/// Precondition: This method should only be called if self.fut is None
/// or contains a completed future.
fn upload(&mut self, buf: &[u8]) -> usize {
// Here we could check that the stated precondition about self.fut
// holds. Checking the completion of the Future in self.fut is only
// possible via a call to poll, which requires passing to it a
// Context object. Doing that requires changing the signature of
// the method, and passing to it an additional Context object.
let mut size = 0;
if self.buf.len() < self.min_size {
self.buf.extend_from_slice(buf);
size = buf.len();
}
if self.buf.len() >= self.min_size {
let data = self.buf.drain(..).collect();
self.fut = Some(Box::pin(upload_bytes(data)));
} else {
// It is safe to set self.fut to None because of the
// precondition. Without setting fut to none, we'd continue
// polling on it after completion, which results in a panic.
self.fut = None;
}
size
}
}
There are a few things to note about this implementation:
- The receiver of
uploadis&mut self; as simple as that! We did not have to pin it, because we are not callingpollin this method. - There is a precondition for calling
upload, as stated in the doc comment. According to this precondition,uploadcan only be called ifself.futisNone, or its innerFutureisReady. I added this condition to keep the logic simple, without violating the contract ofpoll_write.
The interested reader is encouraged to modify the code, remove the precondition, and make the implementation more performant. But make sure not to violate the contract of poll_write when you call self.upload. Namely, if bytes are used and added to the inner buffer, poll_write must not return Pending.
Here is the updated poll_write, which calls self.upload, after checking that the precondition holds:
// Impl-2.c: stateful.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
if self.fut.is_none() {
let size = self.upload(buf);
// We successfully added all input bytes to the inner buffer,
// and are now ready to accept new bytes. The poll_write
// contract requires us to return Ready.
return Poll::Ready(Ok(size));
}
// We have an on-going future. We can only upload the given bytes
// in buf, and create a new Future, if the previous future
// is complete. Otherwise, we return Pending or an error.
match self.fut.as_mut().unwrap().as_mut().poll(cx) {
Poll::Pending => {
// For simplicity, we don't use the incoming bytes if we
// have an on-going future. For a more performant
// implementation, one should add incoming bytes to the
// inner buffer as long as its capacity allows, and only
// return Pending if the buffer is full, and the Future in
// fut is not yet complete.
Poll::Pending
}
Poll::Ready(Ok(())) => {
let size = self.upload(buf);
Poll::Ready(Ok(size))
}
Poll::Ready(Err(err)) => {
// Clean up fut to avoid a "resumed after completion"
// panic, if for whatever reason the caller decides to
// ignore this error and proceed with another write or a
// call to shutdown.
let _ = self.fut.take();
Poll::Ready(Err(err))
}
}
}
}
In this version, we are invoking self.upload inside poll_write without a problem.
Inside self.upload, the Future in self.fut is updated, but without holding a reference to self.
We achieved this by draining the inner buffer to a separate vector before calling upload_bytes.
This way, we removed the restriction on the lifetime of self imposed by that of the Future inside self.fut.
Now that we have an inner buffer, we have to update poll_shutdown to upload any remaining bytes in the internal buffer, before termination.
Normally, with an inner buffer, you’d have to flush the buffer in poll_flush too.
But in our case, flushing the inner buffer, if we don’t have enough bytes, will violate the minimum upload size condition.
So we leave the implementation of poll_flush empty.
If we again run the main function after these changes, we see that this time our “Uploaded {i}: Hello world!” messages are written in batches of 8, due to the internal buffering.
Custom Future to hide Future from async
In both of the examples above, the main logic happens inside an async function (i.e., upload_bytes).
The difficulty with async functions is that they result in unnamed Future types, whereas it is often easier to work with concrete known types.
In this example, we are going to introduce a custom Future type, and push the async function call away from our AsyncUploader!
A custom Future type
In the following snippet, I have changed upload_bytes to return a custom Future named UploadFuture:
// Impl-3.a: pinless.rs
struct UploadFuture {
/// Inner Future, resulted from an async call.
fut: BoxFuture<'static, Result<(), Error>>,
}
impl Future for UploadFuture {
type Output = Result<(), Error>;
fn poll(...) -> Poll<Self::Output> {...}
}
/// This is the slow upload function, which we want to use inside the
/// AsyncWrite implementation.
fn upload_bytes(data: Vec<u8>) -> UploadFuture {
let fut = async {
// Sleep for a bit to simulate a long-running operation.
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
println!(
"Uploaded {}",
String::from_utf8(data).map_err(|e| {
Error::new(ErrorKind::Other, format!("{e}"))
})?
);
Ok(())
};
UploadFuture { fut: Box::pin(fut) }
}
I have simply wrapped the entire logic in an async-block, and have wrapped the Future from the async-block in an instance of UploadFuture, which is what is returned from upload_bytes.
So far, I have only introduced some nesting in the structure!
To complete the nesting, we have to update the declaration of fut in AsyncUploader:
// Impl-3.b: pinless.rs
struct AsyncUploader {
/// Minimum size of the chunks to upload.
min_size: usize,
/// Internal buffer.
buf: Vec<u8>,
/// Future corresponding to the long-running upload operation.
fut: Option<UploadFuture>,
}
Additionally, we have to tweak a few lines in AsyncUploader::upload.
In particular, we have to change self.fut = Some(Box::pin(upload_bytes(data))); to self.fut = Some(upload_bytes(data));.
Disentangled lifetimes
One interesting thing we can do now is to make upload_bytes a method in AsyncUploader, with &mut self as its receiver:
// Impl-3.c: pinless.rs
impl AsyncUploader {
...
fn upload_bytes(&mut self, data: Vec<u8>) -> UploadFuture {
let fut = async {/* same as before */ };
UploadFuture { fut: Box::pin(fut) }
}
}
Recall that in Impl-2, achieving this was not possible with the async version of upload_bytes.
Our attempt led to an error stating, “future cannot be sent between threads safely”.
By encapsulating the result of the async block in a custom Future type, we have effectively decoupled self from the Future that we have to poll on, thus eliminating the error.
To complete the code, we have to implement Future for UploadFuture:
// Example 3.d: pinless.rs
impl Future for UploadFuture {
type Output = Result<(), Error>;
fn poll(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Self::Output> {
match self.fut.as_mut().poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(o) => Poll::Ready(o),
}
}
}
Safe on-the-fly pinning of Unpin objects
Finally, we have to fix poll_write, which calls poll on self.fut.
For that, we need pinned access to the Future in self.fut.
But in the new version of AsyncUploader, we have not declared fut as Pin.
While we could have accomplished this by declaring fut as Option<Pin<UploadFuture>>, such a declaration is unnecessary.
Instead, we can dynamically create a Pin<&mut UploadFuture> from self.fut when invoking poll, as demonstrated below:
// Impl-3.e: pinless.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
if self.fut.is_none() {
let size = self.upload(buf);
return Poll::Ready(Ok(size));
}
// We have to pin fut.
let pinned = Pin::new(self.fut.as_mut().unwrap());
match pinned.poll(cx) {
Poll::Pending => {
Poll::Pending
}
Poll::Ready(Ok(())) => {
let size = self.upload(buf);
Poll::Ready(Ok(size))
}
Poll::Ready(Err(err)) => {
let _ = self.fut.take();
Poll::Ready(Err(err))
}
}
}
}
An attentive reader may wonder why we found it unnecessary to declare AsyncUploader::fut as Pin, whereas in all previous examples, the Future was pinned via BoxFuture. The reason is that UploadFuture is Unpin. So to satisfy the Pin’s contract, we don’t have to keep the pinning promise about UploadFuture uninterrupted for the entirety of its lifetime. However, an unnamed Future that we get from a call to an async function is !Unpin (not Unpin). For such a Future to work correctly, it must persist in the same location in memory throughout its entire lifetime. Using BoxFuture is the only way we can guarantee this safely.
Unsafe pinning and pin-projection
Currently, inside the body of upload_bytes, we still have an async block. Interestingly, we can remove it!
The only source of async in upload_bytes is the call to tokio::sleep.
The return type of tokio::sleep is a Future named Sleep.
Instead of awaiting on it in an async-block, we can directly wrap it inside UpdateFuture:
// Impl-4.a: asyncless.rs
struct UploadFuture {
content: String,
fut: Sleep,
}
...
impl AsyncUploader {
...
/// This is the slow upload function, which we want to use inside the
/// AsyncWrite implementation.
fn upload_bytes(&self, data: &[u8]) -> Result<UploadFuture, Error> {
let content = String::from_utf8(data.to_vec())
.map_err(|e| Error::new(ErrorKind::Other, format!("{e}")))?;
Ok(UploadFuture {
content,
fut: tokio::time::sleep(tokio::time::Duration::from_secs(1)),
})
}
}
Note that I did not wrap Sleep in a BoxFuture. I could, but let’s see what happens if we deviate from that pattern.
A little bit of unsafe code
When implementing Future for UploadFuture we need pinned access to Sleep, because that is what we are going to call poll on.
However, Sleep is !Unpin and we cannot use Pin::new to pin it.
Instead, I am going to use the unsafe method Pin::map_unchecked_mut, to map self to a pinned mutable reference to its inner Sleep:
// Impl-4.b: asyncless.rs
impl Future for UploadFuture {
type Output = ();
fn poll(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Self::Output> {
// Make a copy of content before moving out of self.
let content = self.content.clone();
let pinned = unsafe { self.map_unchecked_mut(|s| &mut s.fut) };
match pinned.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(_) => {
println!("Uploaded {}", content);
Poll::Ready(())
}
}
}
}
Owned !Unpin calls for pin_project
Since Sleep is !Unpin, after changing UploadFuture::fut from a BoxFuture to Sleep, UploadFuture too has become !Unpin.
As a result, we can no longer simply pin it using Pin::new, as we did in poll_write in Impl-3. Here is the complete error message:
error[E0277]: `PhantomPinned` cannot be unpinned
--> src/bin/asyncless.rs:82:31
|
82 | let pinned = Pin::new(self.fut.as_mut().unwrap());
| -------- ^^^^^^^^^^^^^^^^^^^^^^^^^^ within `(PhantomData<&()>, PhantomPinned)`, the trait `Unpin` is not implemented for `PhantomPinned`
| |
| required by a bound introduced by this call
|
= note: consider using the `pin!` macro
consider using `Box::pin` if you need to access the pinned value outside of the current scope
= note: required because it appears within the type `(PhantomData<&()>, PhantomPinned)`
= note: required for `Sleep` to implement `Unpin`
note: required because it appears within the type `UploadFuture`
--> src/bin/asyncless.rs:39:8
|
39 | struct UploadFuture {
| ^^^^^^^^^^^^
note: required by a bound in `Pin::<P>::new`
--> /Users/razieh/.rustup/toolchains/nightly-x86_64-apple-darwin/lib/rustlib/src/rust/library/core/src/pin.rs:484:23
|
484 | impl<P: Deref<Target: Unpin>> Pin<P> {
| ^^^^^ required by this bound in `Pin::<P>::new`
...
503 | pub const fn new(pointer: P) -> Pin<P> {
| --- required by a bound in this associated function
To fix this, we need some more gymnastics similar to what we did above.
This time, instead of writing unsafe code, I am going to use pin_project.
The main benefit of using pin_project is that it saves us from writing unsafe code.
The first step is to annotate AsyncUploader with the pin_project::pin_project attribute (note that, for simplicity, I’ve based this version on Impl-1, without the internal buffer in AsyncUploader):
// Impl-4.c: asyncless.rs
#[pin_project::pin_project]
struct AsyncUploader {
#[pin]
fut: Option<UploadFuture>,
}
The pin_project::pin_project attribute declares a new struct, which is very similar to AsyncUploader but instead of fut has a pinned field of type Pin<&'pin mut (Option<UploadFuture>)>.
The declaration of this field is instructed by the #[pin] attribute that we added to fut.
Moreover, pin_project::pin_project adds a project(Pin<&’a mut Self>) method to AsyncUploader, which returns an instance of the newly declared struct.
Inside poll_write, we use self.project().fut to get access to a pinned version of UploadFuture.
// Impl-4.d: asyncless.rs
impl AsyncWrite for AsyncUploader {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
// Method `project` consumes `self`. Before calling it on self,
// let's create the future that we might need later.
let fut = self.upload_bytes(buf);
let mut this = self.project();
if this.fut.is_none() {
this.fut.set(fut.ok());
// The future we created is not necessarily completed yet, but
// we successfully used all the input bytes, and are now ready
// to accept new bytes. The poll_write contract requires us to
// return Ready.
return Poll::Ready(Ok(buf.len()));
}
// We have an on-going future. We can only upload the given bytes
// in buf, and create a new Future, if the previous future is
// complete. Otherwise, we return Pending or Ready with an error.
match this.fut.as_mut().as_pin_mut().unwrap().poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(()) => match fut {
Ok(fut) => {
this.fut.set(Some(fut));
Poll::Ready(Ok(buf.len()))
}
Err(err) => Poll::Ready(Err(err)),
},
}
}
}
A few things to note:
- We could have used
pin_projectforUploadFuture(in Impl-4.a) in a similar manner to avoid the unsafe code inUploadFuture::poll. - When you have your fields pinned, as in Impl-1, Impl-2, and Impl-3, adding
pin_projectwon’t offer much more.
As a result of Sleep and UploadFuture being !Unpin, now AsyncUploader is !Unpin too. This means that we cannot easily get mutable access to it. To account for this, I had to change the receiver of upload_bytes from &mut self to &self. If mutable access is really needed, a good solution is to refactor the code.
See the full implementation of this version on GitHub.
There is also a more complex variant with a FileUploader, which is still simplified, but closer to what you may need to implement in a real-world project.
When to use pin_project?
In the four implementations we discussed in this article, we used different approaches in the design of AsyncUploader to pin the Future we needed to poll on.
In each case, the choice was made based on the type of the Future wrapped in AsyncUploader.
The following diagram summarizes these choices.
The focus is on designing the Outer data structure (e.g., AsyncUploader), considering that it will be used as a pinned value in its context (e.g., poll_write). The question is, what is needed to gain a pinned access to the inner, in that context, so that we can call poll on it.
If the inner is Unpin, we can simply access and pin it on-the-fly, on the call site, without having to write any unsafe code.
Otherwise, if Outer owns the inner, we can achieve our goal with some unsafe code (e.g., Impl-4). In this case, it is easy to reason about the soundness of the unsafe operation. Since outer is pinned, and since it owns inner, we can guarantee that inner does not change its location in memory during its lifetime. Because outer, being pinned and !Unpin, does not change its location in memory during its lifetime. In this case, to avoid writing unsafe code, we can use pin_project.
If Outer does not own the inner (e.g., Impl-1), then we cannot reason about the immovability of the inner so easily. Such situations often arise when working with Futures originated from invocations of async functions. Such a Future does not have a known concrete type, and thus neither a known size, at compile time. As a result, Outer cannot directly own the Future, but can only have a reference to it. To ensure that the referenced Future remains immovable and is not unsoundly mutated, we need to keep a pinned pointer to it for the entirety of its lifetime. Therefore, we need to use BoxFuture when declaring such an inner in the Outer. The pinning requirement cannot be achieved using pin_project in this case.
However, there is another scenario where it might be beneficial to annotate the outer with pin_project.
Suppose that you have an Outer that will be used as pinned, but you don’t necessarily need pinned access to its fields.
In these cases, especially if your struct is generic, it is advised to add pin_project to your struct to avoid potential unsound implementations of Unpin for your struct.
This is because in addition to providing pin-projection functionalities, pin_project adds a conditional Unpin implementation to your struct.
If I add impl Unpin for AsyncUploader {} in Impl-4 above, I get the following error:
error[E0119]: conflicting implementations of trait `Unpin` for type `AsyncUploader`
--> src/bin/asyncless.rs:16:1
|
16 | #[pin_project::pin_project]
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^ conflicting implementation for `AsyncUploader`
...
22 | impl Unpin for AsyncUploader {}
| ---------------------------- first implementation here
|
= note: this error originates in the derive macro `::pin_project::__private::__PinProjectInternalDerive` (in Nightly builds, run with -Z macro-backtrace for more info)
Interestingly, if you search for pin_project in GitHub, you can find many uses of pin_project without any fields being declared #[pin], or without a call to project. Some of these don’t have any generic types, and use BoxFuture (meaning that the Outer type is already unconditionally Unpin). I honestly don’t know what is the point of using pin_project in those cases. Some could be sloppy copy-pasting. But I hope some are intentional and have a good justification. Without any why-documentation it is hard to tell which one is the case!
Summary of Errors and Fixes!
Error: “future cannot be sent between threads safely”
- Likely root-cause:
Futurecapturesself, which is neitherSendnorSync. - Fix: Refactor your struct so that the
Futureyou have topollon does not depend onself. - Example: See Impl-1.g
Error: “lifetime may not live long enough”
- Likely root-cause:
Futurecaptures references, potentially causing a cyclic ordering among lifetimes! - Fix: If possible, refactor your code so that the
Futureowns everything it needs, or declare a customFuturewith well-thought generic lifetimes. - Example: A hypothetical variant of Impl-3.a, with
&[u8]as the type of data passed toupload_bytes.
Error: “T cannot be unpinned”
- Fix: Use
BoxFutureorpin_project. - Example: See Impl-1.g, and Impl-4.b
Error: “cannot borrow data in dereference of
Pin<&mut T>as mutable”
- Likely root-cause: Pinned object is not
Unpin. - Fix: Refactor to avoid the need for a mutable receiver, or make your type
Unpin. - Example: A hypothetical version of Impl-4.a, with
&mut selfas the receiver.
There is more!
While crafting this article, I realized I needed a deeper understanding of the concept of pinning. This realization led me to an extensive exploration of the topic, and a write-up that I plan to publish as a follow-up of this post. Stay tuned for a deep dive into pinning!
Related Articles
- The Async Rust book
- Pin and Suffering: A very good step-by-step intro to futures and pinning
- Comprehensive-Rust; the section and exercises on Async Rust