TakeUntil + Tokio task budgeting + infinite stream fails

[Nemo157]

When using take_until with a Tokio derived future as the condition, wrapping an always ready stream, Tokio's task budgeting will stop the condition from ever completing (playground)

use std::time::Duration;
use futures::stream::StreamExt as _;

#[tokio::main]
async fn main() {
    futures::stream::repeat(0)
        .take_until(tokio::time::delay_for(Duration::from_secs(1)))
        .for_each(|_| async {}).await;
}

What happens is that for the first 128 times the ForEach polls the stream (or whatever the task budget is) the TakeUntil will first poll the Delay, which will return Pending as it is not complete, then poll the Repeat and return an element. After the task budget is exhausted, essentially the same thing will continue happening, but now Delay will be returning Pending because the budget is exhausted. After 1 second, when the Delay should complete, it continues returning Pending because of the budget.

[Darksonn]

Unlike some other issues we've seen, this is not the fault of Tokio's task budgeting, rather this is one of the issues that Tokio's task budgeting fixes. Of course, types that don't participate in the budgeting don't get the benefit, such as not blocking the executor in an example as yours.

To be specific, this is because Tokio's time driver runs on the same thread as where the future is polled, and since you don't yield to the executor, the time driver does not get a chance to notify delay_for.

[Darksonn]

With a closer look, this is related to coop, but it is a different kind of issue from #2047 and #2130, and is more comparable to using executor::block_on. The first two issues would be solved by e.g. notifying the coop task after the call to poll. In this case, it will keep calling poll on the timer in a hot loop until it completes, and it will never yield back to the executor, even if coop doesn't immediately wake it.