Understanding Rust's Polling with Custom Futures

Introduction

Asynchronous programming has become an indispensable paradigm for building performant and responsive applications, especially in domains like networking, I/O-bound tasks, and high-concurrency systems. Rust, with its strong type system and ownership model, offers a powerful and safe approach to async programming built upon its Future trait. While you often interact with futures through the async/await syntax, truly comprehending how these abstractions work under the hood is crucial for debugging, optimizing, and even designing custom asynchronous components. This deep dive into writing a custom Future will demystify the polling mechanism, revealing the fundamental dance between your async tasks and the executor, ultimately empowering you to wield Rust's async capabilities with greater confidence and precision.

The Heart of Asynchronous Execution: Polling

Before we construct our custom future, let's establish a clear understanding of the core concepts involved:

• Future Trait: In Rust, a Future is a trait that represents an asynchronous computation that may or may not have completed yet. It has a single method, poll, which an executor repeatedly calls to check the future's progress.
• Executor: An executor is responsible for taking Futures and driving them to completion by repeatedly calling their poll method. It manages the lifecycle of futures, schedules tasks, and handles waking up tasks when they are ready to make progress. Popular executors include tokio and async-std.
• Polling: This is the act of the executor calling the poll method on an uncompleted Future. When poll is called, the future attempts to make progress.
• Poll Enum: The poll method returns a Poll enum, which has two variants:
  • Poll::Ready(T): Indicates that the future has completed successfully, and T is the result of the computation.
  • Poll::Pending: Indicates that the future is not yet complete. When Pending is returned, the future must ensure that it arranges to be woken up (via Waker) when it's ready to make further progress.
• Waker: A Waker is a low-level mechanism provided by the executor to allow a Future to signal that it's ready to be polled again. When a future returns Poll::Pending, it captures and clones the Waker from the Context. Later, when an event occurs that might unblock the future (e.g., data arrives on a socket, a timer expires), the future calls waker.wake_by_ref() to notify the executor to re-poll it.
• Context: The Context passed to the poll method contains a Waker and other information useful for the future to interact with the executor.

Building a Custom Future: A Simple Delay

Let's create a custom Future that introduces a non-blocking delay. This will allow us to observe the polling mechanism directly.

We'll define a Delay struct that holds a deadline (when it should complete) and an optional Waker to wake up the task.

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::thread;

// Represents the state of our delay future
struct Delay {
    deadline: Instant,
    // We need to store the waker to wake up the future when the deadline passes.
    // Arc<Mutex<Option<Waker>>> allows us to share and safely modify the waker across threads.
    waker_storage: Arc<Mutex<Option<Waker>>>,
    // A flag to ensure we only spawn the timer thread once.
    timer_thread_spawned: bool,
}

impl Delay {
    fn new(duration: Duration) -> Self {
        Delay {
            deadline: Instant::now() + duration,
            waker_storage: Arc::new(Mutex::new(None)),
            timer_thread_spawned: false,
        }
    }
}

// Implement the Future trait for our Delay struct
impl Future for Delay {
    // The output type of our future is unit, as it just represents a delay.
    type Output = ();

    // The core of the future: the poll method
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // If the deadline has already passed, the future is ready.
        if Instant::now() >= self.deadline {
            println!("Delay future: Deadline reached. Returning Ready.");
            return Poll::Ready(());
        }

        // --- Store the Waker and Set up the Timer (only once) ---
        // If the timer thread hasn't been spawned yet, set it up.
        if !self.timer_thread_spawned {
            println!("Delay future: First poll. Storing waker and spawning timer thread.");

            // Store the current waker from the context.
            // This waker will be used by the timer thread to wake up this task.
            let mut waker_guard = self.waker_storage.lock().unwrap();
            *waker_guard = Some(cx.waker().clone());
            drop(waker_guard); // Release the lock early

            // Clone the Arc to pass it to the new thread.
            let waker_storage_clone = self.waker_storage.clone();
            let duration_until_deadline = self.deadline.duration_since(Instant::now());

            // Spawn a new thread that will sleep until the deadline
            // and then wake up the original task.
            thread::spawn(move || {
                thread::sleep(duration_until_deadline);
                println!("Delay timer thread: Deadline passed. Waking up the task.");
                // Retrieve the waker and wake the task
                if let Some(waker) = waker_storage_clone.lock().unwrap().take() {
                    waker.wake();
                }
            });

            // Mark that the timer thread is spawned to avoid re-spawning
            self.timer_thread_spawned = true;
        } else {
            // This part handles subsequent polls if the timer thread is already running.
            // It's crucial to update the waker if the executor decides to move the task
            // or re-schedule it. If the waker isn't updated, the previous waker might become
            // stale, leading to the task never being woken up.
            let mut waker_guard = self.waker_storage.lock().unwrap();
            if waker_guard.as_ref().map_or(true, |w| !w.will_wake(cx.waker())) {
                println!("Delay future: Waker changed. Updating.");
                *waker_guard = Some(cx.waker().clone());
            }
        }


        // If the deadline has not passed yet, the future is pending.
        // It will be re-polled when the `waker.wake()` is called by the timer thread.
        println!("Delay future: Deadline not yet reached. Returning Pending.");
        Poll::Pending
    }
}

#[tokio::main]
async fn main() {
    println!("Main: Starting program.");

    let delay_future = Delay::new(Duration::from_secs(2)); // Create a 2-second delay
    println!("Main: Awaiting delay future...");
    delay_future.await; // Await our custom future

    println!("Main: Delay completed. Program finished.");
}

Explanation of the Delay Future:

1. struct Delay:

   • deadline: An Instant representing when the delay should finish.
   • waker_storage: An Arc<Mutex<Option<Waker>>> is essential. The Waker needs to be shared between the Future (which owns self.waker_storage) and the separate thread::spawn (which will call wake). Arc enables shared ownership, and Mutex provides safe interior mutability to store and retrieve the Waker. Option is used because the Waker might not be available on the very first poll before it's stored.
   • timer_thread_spawned: A simple boolean flag to ensure we only spawn our "timer" thread once.

2. impl Future for Delay:

   • type Output = ();: Our delay future simply completes, producing no meaningful value.
   • poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>: This is the core.
     • if Instant::now() >= self.deadline: On every poll, we check if the deadline has passed. If so, we're Ready and return Poll::Ready(()).
     • if !self.timer_thread_spawned: This conditional block ensures that we only set up the actual timer (the thread::spawn part) once.
       • let mut waker_guard = self.waker_storage.lock().unwrap(); *waker_guard = Some(cx.waker().clone());: We acquire a lock on our waker_storage, clone the Waker from the current Context, and store it. This Waker points back to this specific task that is currently being polled.
       • thread::spawn(...): We launch a standard Rust thread. This thread will sleep() for the remaining duration. This is a blocking sleep from the perspective of this helper thread, but it doesn't block the executor thread because it's in a separate OS thread.
       • Inside the spawned thread, after sleeping, it retrieves the stored Waker and calls waker.wake(). This wake() call tells the async runtime (Tokio in our main) that the task associated with this Waker is now ready to be re-polled.
       • self.timer_thread_spawned = true;: We set the flag to true to prevent spawning multiple timer threads for the same Delay future.
     • else { ... }: If the timer thread has been spawned (i.e., this is a subsequent poll of an already pending future), we still need to check if the Waker in the Context has changed (!w.will_wake(cx.waker())). If it has, we update our stored Waker. This is important because executors can sometimes move or re-schedule tasks, requiring a new Waker to correctly notify the task.
     • Poll::Pending: If the deadline hasn't passed and the timer is set up, the future is still waiting. We return Poll::Pending. The executor will stop polling this future until waker.wake() is called.

How it operates with tokio::main and await:

1. Delay::new(Duration::from_secs(2)): A Delay instance is created.
2. delay_future.await: This is where the magic happens.
   • Tokio's executor receives delay_future.
   • First Poll: The executor calls delay_future.poll(cx).
     • The deadline is not met.
     • timer_thread_spawned is false.
     • The Waker from cx is cloned and stored in delay_future.waker_storage.
     • A new thread::spawn is created. This thread starts sleeping for 2 seconds.
     • timer_thread_spawned is set to true.
     • poll returns Poll::Pending.
   • Executor's action after Poll::Pending: The executor now knows delay_future is not ready. It puts this task aside and starts polling other ready tasks (if any) or waits for waker.wake() calls. Crucially, the Tokio runtime thread is not blocked by our thread::spawn's thread::sleep.
   • After 2 seconds: The thread::spawn completes its thread::sleep.
     • It retrieves the stored Waker and calls waker.wake().
   • Executor's action after waker.wake(): The executor receives the wake-up signal for the task associated with delay_future. It schedules delay_future to be polled again.
   • Second (or later) Poll: The executor calls delay_future.poll(cx) again.
     • Now, Instant::now() >= self.deadline is true.
     • poll returns Poll::Ready(()).
   • Completion: The delay_future.await expression finally completes, and the main function continues.

Conclusion

By implementing a custom Delay future, we've gained a hands-on understanding of Rust's asynchronous polling mechanism. We've seen how Future::poll is repeatedly invoked by an executor, how Poll::Pending signals an incomplete state, and critically, how the Waker acts as the bridge, allowing an awaiting task to signal the executor to recommence polling when progress can be made. This explicit dance between the Future and the Executor via the Waker is the bedrock of Rust's efficient, non-blocking asynchronous programming, enabling high-performance and scalable applications without the overhead of blocking threads. Mastering custom Future implementations is an advanced skill that unlocks deeper insights into Rust's powerful async ecosystem.