How Async Works in Rust

So I recently hit the recursion limit in a function I was writing and I couldn't remember where the recursion limit thing even came from - all I remember(ed) is something something async something something coroutines - so I decided to relearn about it and write about it.

#How Async works in rust (abbreviated)

If you've read any of my previous writing you've probably seen me talk about "Zero cost abstractions". The simplest way to understand this is that the compiler generates all the code needed at compile time - there isn't an external runtime that does stuff that you have to "pay for" in terms of memory or CPU.

However in every other language with async/await, there is a runtime that manages the tasks, so how can you have zero cost and async await?

Rust does this by turning async await calls into a state machine wrapped with a Future implementation.

For example say you have a simple function like

async fn get_users(db: DatabaseConnection) -> Result<Vec<User>, DbError> {
	println!("About to get users");
	db.execute::<Users>("select * from users limit 10").await
}

A few things happen.

First - async functions don't actually exist! It's sugar for fn(...) -> impl Future<Output = T> where T is your original function's output type.

Secondly - and this is where the explanation is abbreviated - it makes a state machine enum.

#Async State Machine

Now, if you're reading this I'm assuming you don't already understand rust compiler internals, so this is a simplified explanation. The below is an approximation because while this would work, the real rust compiler uses unstable/nightly only features that you don't have access to.

The intuition about how this works is pretty straightforward. You can represent a thing for the computer to do as a "Task", and this "Task" can either be something that's going to be started in the future, something that's in progress or something that's complete. Naturally these steps transition Pending -> Started -> Done, which you could model as a simple state machine.

So going back to our example

async fn get_users(db: DatabaseConnection) -> Result<Vec<User>, DbError> {
	println!("About to get users");
	db.execute::<Users>("select * from users limit 10").await
}

Becomes

enum GetUsersFuture {
	State0 {
	   db: DatabaseConnection
	}, 
	State1 {
		execute_future: ExecuteFuture
	},
	Done
}

enum ExecuteFuture {
	State0 {
	  query: String
	},
	Done
}

fn get_users(db: DatabaseConnection) -> impl Future<Output = Result<Vec<User>, DbError>> {
	GetUsersFuture::State0 { db: db }
}

Now if you're like me, the first time you see this you'll be like wait, what just happened? Where'd my code go?

#impl Future

This is where the magic happens. If you look at the Future Trait you'll see that it only requires two things - 1) the Output type and 2) a function called poll.

Poll is what is responsible for driving the state machine between the transitions. Whenever the async runtime wants to know what work can be done it calls poll. Poll can either return Pending - which means, "I'm not done, check in later¹" or Ready(T) with the Output value specified (this is why you normally return Result<T, E> from future, since typically things can go wrong in the background)

So how would you implement this function? Again the high level idea is pretty straightforward, we need to implement this function:

fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output>{}

and since we'd impl Future for GetUsersFuture, self is the enum state machine we talked about earlier. Context is important too but for now, pretend it doesn't exist - we'll come back to it.

So now that our code is represented as a state machine, the logical thing to do is to check which state we're in (ie, match self) and decide what to do

loop {
	match self {
		State0 { db } => {
		
		} 
	}
}

Ok in the initial state we have a db connection, so any code between here and the next .await we can run immediately (synchronously)

loop { 
	match self {
		State0 { db } => {
			println!("About to get users"); // This is where that code went
			let db_future = db.execute(...); // run the data query, it returns a future too
			// once we have the future, we can transition to the next state
			*self = State1 { execute_future: db_future };
			// we don't return anything here
		} 
	}
}

So the first "State" essentially is just the first line of this example function (the example is purposefully not very complicated here) but you can see once it's done you transition to the next state. This is why async code can "yield". We can simply not execute the next state until other synchronous tasks are done - that is we won't call "poll" on any other Futures while the synchronous code in another future is running. This is called cooperative multitasking - each task voluntarily yields control back to the runtime rather than being preempted.

So to flesh out this example - how do we handle State1? Essentially it's the same song and dance but with a twist

loop { 
	match self {
		State0 { db } => {
			println!("About to get users"); // This is where that code went
			let db_future = db.execute(...); // run the data query, it returns a future too
			// once we have the future, we can transition to the next state
			*self = State1 { execute_future: db_future };
			// we don't return anything here
		}
		State1 { execute_future } => {
			// Remember, futures only run when the poll method (like the one we're in)
			// are called on their futures, so we're responsible for kicking off our subtasks
			match execute_future.poll(cx) {
				// Remember the state machine
				Pending => return Poll::Pending,
				Ready(output) => {
					// Transition the State Machine to done
					*self = Done;
					return Poll::Ready(output);
				}
			}
		}
		Done => panic!("future polled after completion"),
	}
}

As you can see, we're responsible for driving the tasks we're waiting for forward. Whenever the runtime calls poll on our function we delegate that to the poll of the database call we're waiting for (and so on the more nested async fns you have)

#Sidenote: Nested Async type explosion

In the beginning I mentioned that I wrote this because of hitting the recursion limit and wracking my brain for what that means. Now I can explain how this happens.

Recall that for each .await we generate this state machine. Then recall that each state machine step stores the Future of its sub-tasks, because just like synchronous code, each step enum stores the parameters you'll need for that step

enum FutureA {
  Stage0 { inner: FutureB }, // FutureB { inner: FutureC { inner: FutureD } }
  Stage1 { inner: FutureE }, // FutureE { inner: FutureF { inner: FutureG } }
}

The problem is that this type resolution process can only go so many steps deep (128 by default) so when there are too many .awaits in a function body, or nested, you'll hit this limit²

#Who watches the watchmen?

The last thing you might wonder is if each future is responsible for calling its child future isn't all the code essentially synchronous?

Well yes, but no. Typically the reason you'd use async at all is that you expect there to be a point in which you're waiting for something you don't have control of - typically IO. In our database example, when the request eventually gets sent to the database how do we know when the database replies? We could loop over read_socket or something but that'd block the runtime and defeat the purpose of being async as well.

The short answer is that tokio & other async runtimes have an asynchronous IO implementation that allows the runtime to say "I'm going to tell the operating system to tell me when I should check back for new data" and then only "poll" when that happens. Remember how I glossed over the Context parameter earlier? This is where it comes into play. Context allows you to register a callback with the operating system that will tell your async runtime when to call poll again, thus preventing it from polling in a busy loop and blocking the runtime.

The details of how that works are pretty interesting (and if you've heard of io_uring you might know it's rather controversial in rust) but that's a topic for next time.