diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 17935831708..25fc7bab672 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -893,7 +893,7 @@ impl Future for RunUntil<'_, T> { let _no_blocking = crate::runtime::enter::disallow_block_in_place(); let f = me.future; - if let Poll::Ready(output) = crate::runtime::coop::budget(|| f.poll(cx)) { + if let Poll::Ready(output) = f.poll(cx) { return Poll::Ready(output); } diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index a0ade93412e..271afb8f5cf 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -587,6 +587,56 @@ mod unstable { }) .await; } + + // This test compares that, when the task driving `run_until` has already + // consumed budget, the `run_until` future has less budget than a "spawned" + // task. + // + // "Budget" is a fuzzy metric as the Tokio runtime is able to change values + // internally. This is why the test uses indirection to test this. + #[tokio::test] + async fn run_until_does_not_get_own_budget() { + // Consume some budget + tokio::task::consume_budget().await; + + LocalSet::new() + .run_until(async { + let spawned = tokio::spawn(async { + let mut spawned_n = 0; + + { + let mut spawned = tokio_test::task::spawn(async { + loop { + spawned_n += 1; + tokio::task::consume_budget().await; + } + }); + // Poll once + assert!(!spawned.poll().is_ready()); + } + + spawned_n + }); + + let mut run_until_n = 0; + { + let mut run_until = tokio_test::task::spawn(async { + loop { + run_until_n += 1; + tokio::task::consume_budget().await; + } + }); + // Poll once + assert!(!run_until.poll().is_ready()); + } + + let spawned_n = spawned.await.unwrap(); + assert_ne!(spawned_n, 0); + assert_ne!(run_until_n, 0); + assert!(spawned_n > run_until_n); + }) + .await + } } fn rt() -> runtime::Runtime {