Jump to content

Time-travel in Rust: How Async Functions Resumes Where it Left off!

From JOHNWICK

Have you ever imagined how a Rust async function resumes from the middle when an async I/O is complete? It feels like pure magic what rust does under the hood, for example:

async fn get_user(id: i32) -> Result<User, Error> {
    let url = format!("https://api.example.com/users/{}", id);
    // wait until the api call completes
    let response = reqwest::get(url).await?;
    // then wait till the json stream to User struct, 
    let user = response.json::<User>().await?;
    Ok(user)
}

The above function pauses at each await, first it make an API GET call then pauses it till the API completes, in the meantime the thread is released to do other task, then when the API completes it resumes from the exact point. Then it converts the raw body into JSON using the same principles.

In Rust, this magic lies in Futures — types that represent asynchronous computations, pausing and resuming seamlessly without blocking threads. The secret? A compiler-generated state machine, a polling system, that signals when to continue.

pub trait Future {
    type Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

Lets explore the Rust’s zero-cost async model in this blog. Rust Basics Recap

In Rust, synchronous programming executes tasks sequentially, blocking the thread until each completes. Asynchronous programming, however, lets the tasks pause and yield control, allowing the thread to handle other task. The Future trait’s poll method is key here: it checks if a Future is ready (Poll::Ready) or still waiting (Poll::Pending). Executors drive these Futures by repeatedly polling them until completion with the help of Wakers, managing tasks efficiently. I have delved deep into this topic, you can check my blog below. But I should warn you that it’s a heavy read! So, for now assume executor call poll multiple times and drive this async functions to completion.

Building a MiniTokio in Rust: How does Tokio Work ! Rust’s async-await mechanism allows us to write asynchronous code that efficiently manages I/O operations. To… medium.com

The Magic Behind Async/Await: Compiler Transformations Rust’s async functions appear to run seamlessly, pausing at await and resuming later, but this is powered by the compiler’s transformation of async fn into a state machine. This store the current state of the async function that does different task depending on the current state.

When you write an async fn, the Rust compiler translate it into a enum/struct implementing the Future trait, with an internal state machine that tracks execution progress. Each await becomes a yield point, allowing the function to pause and resume without losing its state.

// simple async function 
async fn get_user(id: i32) -> User {
    let resp = read_data(id).await;
    let processed_resp = process1(resp).await;
    let user = process2(processed_resp).await;
    user
}

Is Translated into a state machine:

// type alias for readabilility
type PinnedFuture<T> = Pin<Box<dyn Future<Output = T>>>;

// state machine to make the get_user resumable
enum StateMachine {
    Start { id: i32 },
    AfterStart { resp_fut: PinnedFuture<String> },
    AfterRead { resp_fut: PinnedFuture<String> },
    AfterProcess { user_fut: PinnedFuture<User> },
}

// impl state machine for the state machine
impl Future for StateMachine { 
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
      ...
    }
}

// the original async funtion becomes sync 
// returning a future that can be polled
fn get_user(id: i32) -> StateMachine {
    Start { id }
}

By implementing the Future trait for this enum, Rust knows exactly where the async function stopped. If the work isn’t done yet, it returns Poll::Pending, when the User is finally ready, it returns Poll::Ready.


#[tokio::main]
async fn main() {
    // this .await is converted to .poll() by rust
    let user = get_user(32).await;
    println!("{}", user.name);
}

Polling and the Future Trait in Depth This code is a hand-written asynchronous state machine that implements the Future trait. The async fn, is translated into this state machine by Rust under the hood.


impl Future for StateMachine {
    // The final value produced by this state-machine future
    // after all steps complete successfully.
    type Output = User;

    // poll() drives the state machine forward. Each call attempts to make
    // progress by polling the currently active inner future. If the inner
    // future is not ready, we return Poll::Pending
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this: &mut StateMachine = self.get_mut();
        // this loop is required if any of the future directly return Poll::Ready, 
        // e.g 
        // async fn read_data(id: i32) -> String {
        //     String::from("hello")
        // }
        // 
        loop {
            // Drive the current state. On readiness, transition to the next
            // state by replacing `this` and continue looping.
            match this {
                Start { id } => {
                    // Transition: Start -> AfterStart by creating the
                    // read_data future for the given id.
                    println!("Start is called");
                    let fut = read_data(*id);
                    *this = AfterStart {
                        resp_fut: Box::pin(fut),
                    };
                    // Use the Loop to poll the future via the next state
                    continue;
                }
                AfterStart { resp_fut } => {
                    // Poll the read_data future. If ready, transform the
                    // response with process1 and move to AfterRead.
                    println!("After Start is called");
                    match resp_fut.as_mut().poll(cx) {
                        Poll::Ready(resp) => {
                            let fut = process1(resp);
                            *this = AfterRead {
                                resp_fut: Box::pin(fut),
                            };
                            // Use the Loop to poll via the next state
                            continue;
                        }
                        // Not ready yet - return Pending so the executor
                        // can wake us when progress can be made.
                        Poll::Pending => return Poll::Pending,
                    }
                }
                AfterRead { resp_fut } => {
                    // Poll the process1 future. If ready, build the user via
                    // process2 and move to AfterProcess.
                    println!("AfterRead is called");
                    match resp_fut.as_mut().poll(cx) {
                        Poll::Ready(resp) => {
                            let fut = process2(resp);
                            *this = AfterProcess {
                                user_fut: Box::pin(fut),
                            };
                            // Use the Loop to poll via the next state
                            continue;
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                },
                AfterProcess { user_fut } => {
                    // Final state: poll the future producing the User. If
                    // ready, return Poll::Ready(User) and finish the state
                    // machine.
                    println!("AfterProcess is called");
                    return user_fut.as_mut().poll(cx);
                }
            }
        }
    }
}

The type StateMachine has multiple states (Start, AfterStart, AfterRead, AfterProcess), each representing progress / resumable point in an async function. Whenever the poll function is called it starts polling its inner futures.

The Loop ensures that if any sub-future resolves immediately with Poll::Ready, the state transitions and execution continues within the same poll call instead of stalling or waiting for another poll call. The flow is like this,

  • The sync get_user function return the start state which is then polled by the main method .await .
  • In the Start state, it creates a read_data future and transitions to AfterStart.
  • In AfterStart, it polls read_data. If ready, it runs process1 and transitions to AfterRead.
  • In AfterRead, it polls process1. If ready, it runs process2 and transitions to AfterProcess.
  • In AfterProcess, it polls the final future and, once ready, returns Poll::Ready(User).

If any of the future has returned pending it will return pending back to the main, and let the thread, execute other task, the future that has returned pending will use a waker to put the current function back in the event loop and so that this get_user is polled again. I have explained this in detail in my other blog in detail : Building a MiniTokio in Rust: How does Tokio Work ! Rust’s async-await mechanism allows us to write asynchronous code that efficiently manages I/O operations. To… medium.com


main
└── let state_machine = get_user("id123");   // Creates state machine (Future)
    └── state_machine.poll(cx)               // First poll
    │   ├── state == Start
    │   │   └── let resp_fut = read_data(id) → Inner future created
    │   │   └── Transition to AfterStart { resp_fut: Box::pin(read_data("id123")) }
    │   │   └── // Continue loop to poll new state
    │   │
    │   └── state == AfterStart
    │       └── resp_fut.poll(cx)
    │           └── Returns Poll::Pending   // read_data not ready, yield to executor
    │
    │── ... later, waker wakes ...           // Executor wakes the future
    │── state_machine.poll(cx)               // Second poll
    │   ├── state == AfterStart
    │   │   └── resp_fut.poll(cx)
    │   │       ├── Returns Poll::Ready(resp)   // read_data completes with resp (e.g., String)
    │   │       └── let fut = process1(resp)    // Create process1 future
    │   │       └── Transition to AfterRead { resp_fut: Box::pin(fut) }
    │   │       └── // Continue loop to poll new state
    │   │
    │   └── state == AfterRead
    │       └── resp_fut.poll(cx)
    │           ├── Returns Poll::Pending   // process1 not ready, yield to executor
    │
    ├── ... later, waker wakes ...           // Executor wakes the future again
    └── state_machine.poll(cx)               // Third poll
        ├── state == AfterRead
        │   └── resp_fut.poll(cx)
        │       ├── Returns Poll::Ready(resp)   // process1 completes with resp
        │       └── let fut = process2(resp)    // Create process2 future
        │       └── Transition to AfterProcess { user_fut: Box::pin(fut) }
        │       └── Continue loop to poll new state
        │
        └── state == AfterProcess
            └── user_fut.poll(cx)
                ├── user_fut returns Poll::Ready(user)   // process2 completes with User
                └── Return Poll::Ready(user)    // State machine completes

Transformation with If-Else Suppose you have a if else in the async function, how do you think Rust would transform it, e.g.

async fn get_user(id: i32) -> User {
    let resp = read_data(id).await;
    let resp = if id % 2 == 0 {
        process1(resp).await
    } else {
        resp
    };
    let user = process2(resp).await;
    user
}

If the id is even then we need to do process1, or else we need to return the data to do process2. The state machine would look something like this, all the state needs to be captured but depending on the branch the state may/may not execute.

enum GetUserFunctionFuture {
    Start { id: i32 },
    AfterStart { id: i32, resp_fut: PinnedFuture<String> },
    AfterRead { id: i32, resp: String },
    AfterProcess1 { resp_fut: PinnedFuture<String> },
    AfterProcess2 { user_fut: PinnedFuture<User> },
}

the if condition will be placed after the AfterRead then it move to next based on the correct state.

impl Future for GetUserFunctionFuture {
    type Output = User;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        loop {
            match this {
                // same as above 
                GetUserFunctionFuture::AfterRead { id, resp } => {
                    println!("AfterRead is called");
                    if *id % 2 == 0 {
                        // Even ID: call process1
                        let fut = process1(resp.clone());
                        *this = GetUserFunctionFuture::AfterProcess1 {
                            resp_fut: Box::pin(fut),
                        };
                    } else {
                        // Odd ID: skip to process2
                        let fut = process2(resp.clone());
                        *this = GetUserFunctionFuture::AfterProcess2 {
                            user_fut: Box::pin(fut),
                        };
                    }
                    continue;
                }
                GetUserFunctionFuture::AfterProcess1 { resp_fut } => {
                    println!("AfterProcess1Waiting is called");
                    match resp_fut.as_mut().poll(cx) {
                        Poll::Ready(resp) => {
                            // After process1, call process2
                            let fut = process2(resp);
                            *this = GetUserFunctionFuture::AfterProcess2 {
                                user_fut: Box::pin(fut),
                            };
                            continue;
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                GetUserFunctionFuture::AfterProcess2 { user_fut } => {
                    println!("AfterReadDirect is called");
                    return user_fut.as_mut().poll(cx);
                }
            }
        }
    }
}

Transformation with Loop

Lets see how a async function with a loop is transformed, the loop need to repeatedly call another async function and pause inside the loop, lets see now how it is transformed.

async fn get_user(id: i32) -> User {
    let mut resp = read_data(id).await;
    for _ in 0..3 { // Loop 3 times
        resp = process1(resp).await; // Append " world" each time
    }
    let user = process2(resp).await;
    user
}

The above function will capture a state called loop that store the index the of the loop and the data that need to be passed to the inner function, with this the loop can resume the function at exact iteration.

enum GetUserFunctionFuture {
    Start { id: i32 },
    AfterStart { id: i32, resp_fut: PinnedFuture<String> },
    /// Loop with saving the i varaible state
    Loop { i: i32, resp: String, resp_fut: Option<PinnedFuture<String>> }, 
    AfterLoop { user_fut: PinnedFuture<User> },
}

Most of the transformation is same except the loop state that reuses the outer loop with the state management to iterate and forward the loop, at each iteration it store a new state and continues or return pending.

impl Future for GetUserFunctionFuture {
    type Output = User;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        loop {
            match this {
                GetUserFunctionFuture::Start { id } => {
                    println!("Start is called");
                    let fut = read_data(*id);
                    *this = GetUserFunctionFuture::AfterStart {
                        id: *id,
                        resp_fut: Box::pin(fut),
                    };
                    continue;
                }
                GetUserFunctionFuture::AfterStart { id: _, resp_fut } => {
                    println!("AfterStart is called");
                    match resp_fut.as_mut().poll(cx) {
                        Poll::Ready(resp) => {
                            // Initialize loop with i=0
                            *this = GetUserFunctionFuture::Loop {
                                i: 0,
                                resp,
                                resp_fut: None, // No active future yet
                            };
                            continue;
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                GetUserFunctionFuture::Loop { i, resp, resp_fut } => {
                    println!("Loop is called (i={})", i);
                    if *i < 3 {
                        match resp_fut {
                            None => {
                                // Start a new iteration
                                let fut = process1(resp.clone());
                                *this = GetUserFunctionFuture::Loop {
                                    i: *i,
                                    resp: resp.clone(),
                                    resp_fut: Some(Box::pin(fut)),
                                };
                                continue;
                            }
                            Some(fut) => {
                                match fut.as_mut().poll(cx) {
                                    Poll::Ready(new_resp) => {
                                        // Iteration complete, increment i
                                        *this = GetUserFunctionFuture::Loop {
                                            i: *i + 1,
                                            resp: new_resp,
                                            resp_fut: None, // Reset for next iteration
                                        };
                                        continue;
                                    }
                                    Poll::Pending => return Poll::Pending,
                                }
                            }
                        }
                    } else {
                        // Loop done, move to process2
                        let fut = process2(resp.clone());
                        *this = GetUserFunctionFuture::AfterLoop {
                            user_fut: Box::pin(fut),
                        };
                        continue;
                    }
                }
                GetUserFunctionFuture::AfterLoop { user_fut } => {
                    println!("AfterLoop is called");
                    return user_fut.as_mut().poll(cx);
                }
            }
        }
    }
}

Conclusion So, what actually makes a Rust Future resume where it left off?
It’s the combination of two key pieces working together:

  • Compiler-Generated State Machine — Every async fn is lowered into a state machine that remembers its current position and local variables at each await.
  • The Polling System — The Future trait’s poll method drives progress. Executors keep calling poll until the state machine reaches Poll::Ready.

I hope now you have a good understanding of how rust makes its function resumable.

Read the full article here: https://medium.com/@souravdas08/time-travel-in-rust-how-async-functions-resumes-where-it-left-off-6d50b41f0e2c