Skip to content

Commit

Permalink
rt: rm coop::budget from LocalSet::run_until (#5155)
Browse files Browse the repository at this point in the history
The `LocalSet::run_until` future is just a "plain" future that should
run on a runtime that already has a coop budget. In other words, the
`run_until` future should not get its own budget but should inherit the
calling task's budget. Getting this behavior is done by removing the
call to `budget` in `run_until`.
  • Loading branch information
carllerche authored Nov 1, 2022
1 parent a051ed7 commit ca8e176
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 1 deletion.
2 changes: 1 addition & 1 deletion tokio/src/task/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -893,7 +893,7 @@ impl<T: Future> 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);
}

Expand Down
50 changes: 50 additions & 0 deletions tokio/tests/task_local_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit ca8e176

Please sign in to comment.