diff --git a/tokio/src/coop.rs b/tokio/src/coop.rs index 1d6245916..606ba3a73 100644 --- a/tokio/src/coop.rs +++ b/tokio/src/coop.rs @@ -85,15 +85,11 @@ where return f(); } - struct Guard<'a>(&'a Cell); - impl<'a> Drop for Guard<'a> { - fn drop(&mut self) { - self.0.set(UNCONSTRAINED); - } - } - hits.set(BUDGET); - let _guard = Guard(hits); + let _guard = ResetGuard { + hits, + prev: UNCONSTRAINED, + }; f() }) } @@ -114,6 +110,32 @@ cfg_blocking_impl! { } } +cfg_rt_core! { + cfg_rt_util! { + /// Run the given closure with a new task budget, resetting the previous + /// budget when the closure finishes. + /// + /// This is intended for internal use by `LocalSet` and (potentially) other + /// similar schedulers which are themselves futures, and need a fresh budget + /// for each of their children. + #[inline(always)] + pub(crate) fn reset(f: F) -> R + where + F: FnOnce() -> R, + { + HITS.with(move |hits| { + let prev = hits.get(); + hits.set(UNCONSTRAINED); + let _guard = ResetGuard { + hits, + prev, + }; + f() + }) + } + } +} + /// Invoke `f` with a subset of the remaining budget. /// /// This is useful if you have sub-futures that you need to poll, but that you want to restrict @@ -289,6 +311,11 @@ pin_project_lite::pin_project! { } } +struct ResetGuard<'a> { + hits: &'a Cell, + prev: usize, +} + impl Future for CoopFuture { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { @@ -327,6 +354,12 @@ cfg_sync! { impl CoopFutureExt for F where F: Future {} } +impl<'a> Drop for ResetGuard<'a> { + fn drop(&mut self) { + self.hits.set(self.prev); + } +} + #[cfg(all(test, not(loom)))] mod test { use super::*; diff --git a/tokio/src/macros/cfg.rs b/tokio/src/macros/cfg.rs index 0679aa737..85f95cbd3 100644 --- a/tokio/src/macros/cfg.rs +++ b/tokio/src/macros/cfg.rs @@ -35,6 +35,23 @@ macro_rules! cfg_blocking_impl { } } +/// Enables blocking API internals +macro_rules! cfg_blocking_impl_or_task { + ($($item:item)*) => { + $( + #[cfg(any( + feature = "blocking", + feature = "fs", + feature = "dns", + feature = "io-std", + feature = "rt-threaded", + feature = "task", + ))] + $item + )* + } +} + /// Enables enter::block_on macro_rules! cfg_block_on { ($($item:item)*) => { diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 9af50cee4..346fe437f 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -454,20 +454,24 @@ impl Future for LocalSet { // Register the waker before starting to work self.context.shared.waker.register_by_ref(cx.waker()); - if self.with(|| self.tick()) { - // If `tick` returns true, we need to notify the local future again: - // there are still tasks remaining in the run queue. - cx.waker().wake_by_ref(); - Poll::Pending - } else if self.context.tasks.borrow().owned.is_empty() { - // If the scheduler has no remaining futures, we're done! - Poll::Ready(()) - } else { - // There are still futures in the local set, but we've polled all the - // futures in the run queue. Therefore, we can just return Pending - // since the remaining futures will be woken from somewhere else. - Poll::Pending - } + // Reset any previous task budget while polling tasks spawned on the + // `LocalSet`, ensuring that each has its own separate budget. + crate::coop::reset(|| { + if self.with(|| self.tick()) { + // If `tick` returns true, we need to notify the local future again: + // there are still tasks remaining in the run queue. + cx.waker().wake_by_ref(); + Poll::Pending + } else if self.context.tasks.borrow().owned.is_empty() { + // If the scheduler has no remaining futures, we're done! + Poll::Ready(()) + } else { + // There are still futures in the local set, but we've polled all the + // futures in the run queue. Therefore, we can just return Pending + // since the remaining futures will be woken from somewhere else. + Poll::Pending + } + }) } } @@ -521,18 +525,23 @@ impl Future for RunUntil<'_, T> { .register_by_ref(cx.waker()); let _no_blocking = crate::runtime::enter::disallow_blocking(); + // Reset any previous task budget so that the future passed to + // `run_until` and any tasks spawned on the `LocalSet` have their + // own budgets. + crate::coop::reset(|| { + let f = me.future; + if let Poll::Ready(output) = crate::coop::budget(|| f.poll(cx)) { + return Poll::Ready(output); + } - if let Poll::Ready(output) = me.future.poll(cx) { - return Poll::Ready(output); - } + if me.local_set.tick() { + // If `tick` returns `true`, we need to notify the local future again: + // there are still tasks remaining in the run queue. + cx.waker().wake_by_ref(); + } - if me.local_set.tick() { - // If `tick` returns `true`, we need to notify the local future again: - // there are still tasks remaining in the run queue. - cx.waker().wake_by_ref(); - } - - Poll::Pending + Poll::Pending + }) }) } } diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index 1a10fefa6..38c7c9392 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -312,28 +312,17 @@ fn drop_cancels_tasks() { assert_eq!(1, Rc::strong_count(&rc1)); } -#[test] -fn drop_cancels_remote_tasks() { - // This test reproduces issue #1885. +/// Runs a test function in a separate thread, and panics if the test does not +/// complete within the specified timeout, or if the test function panics. +/// +/// This is intended for running tests whose failure mode is a hang or infinite +/// loop that cannot be detected otherwise. +fn with_timeout(timeout: Duration, f: impl FnOnce() + Send + 'static) { use std::sync::mpsc::RecvTimeoutError; let (done_tx, done_rx) = std::sync::mpsc::channel(); let thread = std::thread::spawn(move || { - let (tx, mut rx) = mpsc::channel::<()>(1024); - - let mut rt = rt(); - - let local = LocalSet::new(); - local.spawn_local(async move { while let Some(_) = rx.recv().await {} }); - local.block_on(&mut rt, async { - time::delay_for(Duration::from_millis(1)).await; - }); - - drop(tx); - - // This enters an infinite loop if the remote notified tasks are not - // properly cancelled. - drop(local); + f(); // Send a message on the channel so that the test thread can // determine if we have entered an infinite loop: @@ -349,10 +338,11 @@ fn drop_cancels_remote_tasks() { // // Note that it should definitely complete in under a minute, but just // in case CI is slow, we'll give it a long timeout. - match done_rx.recv_timeout(Duration::from_secs(60)) { + match done_rx.recv_timeout(timeout) { Err(RecvTimeoutError::Timeout) => panic!( - "test did not complete within 60 seconds, \ - we have (probably) entered an infinite loop!" + "test did not complete within {:?} seconds, \ + we have (probably) entered an infinite loop!", + timeout, ), // Did the test thread panic? We'll find out for sure when we `join` // with it. @@ -366,6 +356,49 @@ fn drop_cancels_remote_tasks() { thread.join().expect("test thread should not panic!") } +#[test] +fn drop_cancels_remote_tasks() { + // This test reproduces issue #1885. + with_timeout(Duration::from_secs(60), || { + let (tx, mut rx) = mpsc::channel::<()>(1024); + + let mut rt = rt(); + + let local = LocalSet::new(); + local.spawn_local(async move { while let Some(_) = rx.recv().await {} }); + local.block_on(&mut rt, async { + time::delay_for(Duration::from_millis(1)).await; + }); + + drop(tx); + + // This enters an infinite loop if the remote notified tasks are not + // properly cancelled. + drop(local); + }); +} + +#[test] +fn local_tasks_wake_join_all() { + // This test reproduces issue #2460. + with_timeout(Duration::from_secs(60), || { + use futures::future::join_all; + use tokio::task::LocalSet; + + let mut rt = rt(); + let set = LocalSet::new(); + let mut handles = Vec::new(); + + for _ in 1..=128 { + handles.push(set.spawn_local(async move { + tokio::task::spawn_local(async move {}).await.unwrap(); + })); + } + + rt.block_on(set.run_until(join_all(handles))); + }); +} + #[tokio::test] async fn local_tasks_are_polled_after_tick() { // Reproduces issues #1899 and #1900