FuturesUnordered: Respect yielding from future

[taiki-e]

If the future yields, FuturesUnordered respects it and yields too.
In this patch, if the future waken before it finishes polling, we assume the future yields.

This is the approach suggested by @carllerche in discord (if I understand correctly).
And I noticed that stjepang had also suggested this approach before: rust-lang/rust#74335 (comment)

cc @cramertj @jonhoo
cc #2053

---

The following is a result of an example in #2526 with the addition of a counter for the number of times the futures were polled.

code

use futures::stream::futures_unordered::FuturesUnordered;
use futures::stream::StreamExt;
use futures_util::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::Mutex;

static COUNT: AtomicUsize = AtomicUsize::new(0);

#[pin_project::pin_project]
struct PollCounter<F>(#[pin] F);

impl<F: Future> Future for PollCounter<F> {
    type Output = F::Output;
    fn poll(
        self: Pin<&mut Self>,
        cx: &mut futures_task::Context<'_>,
    ) -> futures_task::Poll<Self::Output> {
        COUNT.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        self.project().0.poll(cx)
    }
}

async fn do_task(mutex: Arc<Mutex<()>>) {
    mutex.lock().await;
}

async fn benchmark(n: usize) {
    COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
    let start_time: std::time::Instant = std::time::Instant::now();

    let mutex: Arc<Mutex<()>> = Arc::new(Mutex::new(()));
    let mutex_guard = mutex.lock().await;

    let mut futs = Vec::new();
    for _ in 0..n {
        futs.push(PollCounter(do_task(mutex.clone())));
    }
    let mut futs_unordered: FuturesUnordered<_> = futs.into_iter().collect();

    std::mem::drop(mutex_guard);

    for _ in 0..n {
        futs_unordered.select_next_some().await;
    }

    println!(
        "n: {}, time: {}ms, polled: {}",
        n,
        start_time.elapsed().as_millis(),
        COUNT.load(std::sync::atomic::Ordering::SeqCst)
    );
}

#[tokio::main]
async fn main() {
    for n in [10_000, 20_000, 40_000, 80_000, 160_000] {
        benchmark(n).await;
    }
}

before:

n: 10000, time: 6ms, polled: 12481
n: 20000, time: 13ms, polled: 24992
n: 40000, time: 25ms, polled: 49984
n: 80000, time: 50ms, polled: 100000
n: 160000, time: 96ms, polled: 200000

after:

n: 10000, time: 6ms, polled: 10078
n: 20000, time: 13ms, polled: 20156
n: 40000, time: 25ms, polled: 40312
n: 80000, time: 48ms, polled: 80625
n: 160000, time: 93ms, polled: 161250

[Darksonn]

This seems good to me.

[jonhoo]

I'm okay with the change as-is, but see my one concern around over-eager yielding.

[jonhoo]

I worry a little about this being too eager to yield. We're saying here that if any future is awoken while we're polling it (even if it's due to something else legitimately waking it) then we're going to yield from the entire FuturesUnordered. I wonder if it makes more sense to require observing, say, two such yields before concluding we should yield? That buffers us slightly against yielding unnecessarily while still making us yield fairly quickly when everything inside of us will yield going forward.

[taiki-e]

I adopted this heuristic (FuturesUnordered needs to support no-std, which cannot use thread ID). Thanks for the suggestion.

[Darksonn]

Another possible implementation would be to check if the thread ID in the waker is equal to the thread currently polling the FuturesUnordered, if any.

[Darksonn]

Any updates on this?

[taiki-e]

@jonhoo @Darksonn Thanks both of you for the reviews!

I've addressed all reviews, so this should be ready to merge now.

[taiki-e]

Published in 0.3.20.