Crate drive_async_iterator 

• GitHub repo
• rendered docs

HEADS UP: This experimental library crate currently requires a fork of rustc to build, because it depends on the proposed AsyncIterator::poll_progress method.

With the introduction of AsyncIterator::poll_progress, the .next() method on Stream/StreamExt is probably not going to work, because there’s no way for it to fulfill the poll_progress contract. The most common use case of .next() can be replaced with a for await loop, but more complicated use cases are hard to translate. The drive_next function in this crate supports .next() in a poll_progress-compatible form, by taking ownership of an AsyncIterator and providing a handle with a .next() method. It calls poll_progress concurrently with its body when an .await other than .next().await is pending, so the new AsyncIterator contract is satisfied.

A caller that previously used Stream/StreamExt this way…

let mut my_stream = pin!(some_stream());
assert_eq!(my_stream.next().await, Some(42));
assert_eq!(my_stream.next().await, None);

…can use AsyncIterator + poll_progress this way (note that pin! is gone):

use drive_async_iterator::drive_next;

drive_next(some_stream(), async |my_stream| {
    assert_eq!(my_stream.next().await, Some(42));
    assert_eq!(my_stream.next().await, None);
}).await;

drive_next takes ownership of an AsyncIterator and provides a handle that has a next method (and a couple others, see below).

Fixing deadlocks

Besides helping with migration, drive_next solves a class of deadlocks that present-day .next() loops are vulnerable to. Or rather, poll_progress solves them, and drive_next wires .next() into poll_progress. Consider this example (playground link):

// This function acquires a static `Mutex` and does a brief sleep,
// simulating some sort of IO with a shared resource.
async fn do_work() {
    static LOCK: Mutex<()> = Mutex::const_new(());
    let _guard = LOCK.lock().await;
    sleep(Duration::from_millis(10)).await;
}

let mut futures = FuturesUnordered::new();
futures.push(do_work());
futures.push(do_work());
while let Some(_) = futures.next().await {
    do_work().await; // Deadlock!
}

That example deadlocks because one of the do_work futures in the FuturesUnordered is holding the LOCK, but it’s not making progress. Using drive_next the same loop runs smoothly, because FuturesUnordered can poll its contents concurrently with the loop body:

let mut futures = FuturesUnordered::new();
futures.push(do_work());
futures.push(do_work());
drive_next(futures, async |futures| {
    while let Some(_) = futures.next().await {
        do_work().await; // Not a deadlock!
    }
}).await;

More complicated cases

The example above could work with a standard for await loop, so it doesn’t necessarily need the drive_next helper. However, one of the powerful features of FuturesUnordered (and also for example StreamMap) is that you can add more work to it while it’s running (playground link):

let mut futures = FuturesUnordered::new();
loop {
    select! {
        Some(_) = futures.next() => {
            // Do something with the result...
        }
        job = more_work() => futures.push(job),
    }
}

That example isn’t going to work with for await, for two reasons:

1. We need access to futures in the loop body, but for await takes ownership of it.
2. When FuturesUnordered is empty, its poll_next method returns Done. But that makes for await and also drive_next drop it immediately. We need it to return Pending instead.

For the first problem, the handle provided by drive_next supports the with_pin_mut and (for Unpin types) with_mut methods. For the second problem, this crate provides a NeverDone async iterator adapter. Putting those two things together, we can implement the example above while still calling poll_progress correctly under the hood:

drive_next(NeverDone::new(FuturesUnordered::new()), async |futures| {
    loop {
        select! {
            Some(_) = futures.next() => {
                // Do something with the result...
            }
            job = more_work() => {
                futures.with_mut(|maybe_futures: Option<_>| {
                    let futures = maybe_futures.expect("never dropped");
                    futures.push(job);
                });
            }
        }
    }
}).await;

There’s a tricky implementation detail here. FuturesUnordered doesn’t trigger a wakeup when push takes it from empty to non-empty, and it also doesn’t poll the future you pushed until you call poll_next. To make sure we keep looping and poll the new future, drive_next re-polls both its iterator and its body anytime with_mut or with_pin_mut is called. This is a bit of a hack, and it’s possible it could cause accidental hot loops in some cases. It would be better if FuturesUnordered managed these wakeups internally, but that’s not how it works today.

Note also that calling next on a NeverDone async iterator will never return None. Instead, it’ll block potentially forever waiting for the iterator to yield more items. That’s different from the behavior of FuturesUnordered::next today, which returns None immediately in the empty case. The blocking behavior might avoid accidental hot loops in some cases, but it can also cause some existing callers to block forever unexpectedly.

Structs

DrivenAsyncIterator

The AsyncIterator wrapper type that drive_next provides to its closure.

NeverDone

An AsyncIterator wrapper that never returns Done

Functions

drive_next

Take ownership of an AsyncIterator, create a DrivenAsyncIterator handle for it, and call the user’s closure with that handle as an argument. The handle has an async .next() method.