Is it valid to wake a Rust future while it's being polled?

Question

I'd like to be able to sleep my future for a single "frame" so that other work can happen. Is this a valid implementation of this idea?

use std::future::Future;
use std::task::{Context, Poll};
use std::pin::Pin;

struct Yield {
    yielded: bool,
}

impl Future for Yield {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<()> {
        if self.yielded {
            Poll::Ready(())
        } else {
            self.yielded = true;

            // This is the part I'm concerned about
            ctx.waker().wake_by_ref();

            Poll::Pending
        }
    }
}

Specifically, my concern is that the context won't "notice" the wake_by_ref call if it's made before the poll returns Pending. Does the interface contract of poll make any guarantees about this task being immediately re-polled when executed in this way?

Bournemouth

Please allow me to say what a great question this is. I just stumbled upon a use case where I needed to know this. – Gilligan

Answer 1

TL;DR: Your code is valid.

Based on the contract for the waker, it has to poll your future one more time. Otherwise, it is possible to have a race condition between the Future::poll call and the counterpart of the future which actually does some work.

Let's take a look at an example:

impl Future for Foo {
    type Output = ();
    fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<()> {
        let result = communicate_with_worker(ctx); // returns false

        // <-- Time point (1)

        return match result {
            true => Poll::Pending,
            false => Poll::Ready(()),
        };
    }
}

At time point (1), the future has decided that it is not ready, but it's possible that the polling thread is paused here and the worker thread was scheduled and finished its work.

The worker thread will then call the waker and request the future be polled again. If the waker decided to not poll the future again since it's polling the future right now, then the waker will never receive a wake up request again.

This means that the waker may discard wake up requests which come before poll was called, but it's not allowed to discard wake up requests which came during the future's poll call.

---

The only question I have: why would you like to reschedule polling for one more frame?

Since your actual work has to be done in a separate thread (not inside fn poll) then it doesn't make any sense to reschedule polling.

Clausen

To answer your question: this is for a dataloader implementation, which is originally a facebook library that does "gather a series of keys over the course of an async frame, then issue a batched request to the API". In this case, the "work" is "wait for other tasks to have the chance to add their keys to this batch before executing it. It makes more sense in a single-threaded executor; I'm still working out the specific design. – Bournemouth