task: fix LocalSet having a single shared task budget (#2462)

## Motivation

Currently, an issue exists where a `LocalSet` has a single cooperative
task budget that's shared across all futures spawned on the `LocalSet`
_and_ by any future passed to `LocalSet::run_until` or
`LocalSet::block_on`. Because these methods will poll the `run_until`
future before polling spawned tasks, it is possible for that task to
_always_ deterministically starve the entire `LocalSet` so that no local
tasks can proceed. When the completion of that future _itself_ depends
on other tasks on the `LocalSet`, this will then result in a deadlock,
as in issue #2460.

A detailed description of why this is the case, taken from [this 
comment][1]:

`LocalSet` wraps each time a local task is run in `budget`:
https://github.com/tokio-rs/tokio/blob/947045b9445f15fb9314ba0892efa2251076ae73/tokio/src/task/local.rs#L406

This is identical to what tokio's other schedulers do when running
tasks, and in theory should give each task its own budget every time
it's polled. 

_However_, `LocalSet` is different from other schedulers. Unlike the
runtime schedulers, a `LocalSet` is itself a future that's run on
another scheduler, in `block_on`.  `block_on` _also_ sets a budget:
https://github.com/tokio-rs/tokio/blob/947045b9445f15fb9314ba0892efa2251076ae73/tokio/src/runtime/basic_scheduler.rs#L131

The docs for `budget` state that:
https://github.com/tokio-rs/tokio/blob/947045b9445f15fb9314ba0892efa2251076ae73/tokio/src/coop.rs#L73

This means that inside of a `LocalSet`, the calls to `budget` are
no-ops. Instead, each future polled by the `LocalSet` is subtracting
from a single global budget.

`LocalSet`'s `RunUntil` future polls the provided future before polling
any other tasks spawned on the local set:
https://github.com/tokio-rs/tokio/blob/947045b9445f15fb9314ba0892efa2251076ae73/tokio/src/task/local.rs#L525-L535

In this case, the provided future is `JoinAll`. Unfortunately, every
time a `JoinAll` is polled, it polls _every_ joined future that has not
yet completed. When the number of futures in the `JoinAll` is >= 128,
this means that the `JoinAll` immediately exhausts the task budget. This
would, in theory, be a _good_ thing --- if the `JoinAll` had a huge
number of `JoinHandle`s in it and none of them are ready, it would limit
the time we spend polling those join handles. 

However, because the `LocalSet` _actually_ has a single shared task
budget, this means polling the `JoinAll` _always_ exhausts the entire
budget. There is now no budget remaining to poll any other tasks spawned
on the `LocalSet`, and they are never able to complete.

[1]: https://github.com/tokio-rs/tokio/issues/2460#issuecomment-621403122

## Solution

This branch solves this issue by resetting the task budget when polling
a `LocalSet`. I've added a new function to `coop` for resetting the task
budget to `UNCONSTRAINED` for the duration of a closure, and thus
allowing the `budget` calls in `LocalSet` to _actually_ create a new
budget for each spawned local task. Additionally, I've changed
`LocalSet` to _also_ ensure that a separate task budget is applied to
any future passed to `block_on`/`run_until`.

Additionally, I've added a test reproducing the issue described in
#2460. This test fails prior to this change, and passes after it.

Fixes #2460

Signed-off-by: Eliza Weisman <eliza@buoyant.io>

1 files changed, 41 insertions, 8 deletions

diff --git a/tokio/src/coop.rs b/tokio/src/coop.rs
index 1d624591..606ba3a7 100644
--- a/tokio/src/coop.rs
+++ b/tokio/src/coop.rs
@@ -85,15 +85,11 @@ where
             return f();
         }
 
-        struct Guard<'a>(&'a Cell<usize>);
-        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, R>(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<usize>,
+    prev: usize,
+}
+
 impl<F: Future> Future for CoopFuture<F> {
     type Output = F::Output;
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
@@ -327,6 +354,12 @@ cfg_sync! {
     impl<F> 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::*;