Add cooperative task yielding

A single call to `poll` on a top-level task may potentially do a lot of
work before it returns `Poll::Pending`. If a task runs for a long period
of time without yielding back to the executor, it can starve other tasks
waiting on that executor to execute them, or drive underlying resources.
See for example rust-lang/futures-rs#2047, rust-lang/futures-rs#1957,
and rust-lang/futures-rs#869. Since Rust does not have a runtime, it is
difficult to forcibly preempt a long-running task.

Consider a future like this one:

```rust
use tokio::stream::StreamExt;
async fn drop_all<I: Stream>(input: I) {
    while let Some(_) = input.next().await {}
}
```

It may look harmless, but consider what happens under heavy load if the
input stream is _always_ ready. If we spawn `drop_all`, the task will
never yield, and will starve other tasks and resources on the same
executor.

This patch adds a `coop` module that provides an opt-in mechanism for
futures to cooperate with the executor to avoid starvation. This
alleviates the problem above:

```
use tokio::stream::StreamExt;
async fn drop_all<I: Stream>(input: I) {
    while let Some(_) = input.next().await {
        tokio::coop::proceed().await;
    }
}
```

The call to [`proceed`] will coordinate with the executor to make sure
that every so often control is yielded back to the executor so it can
run other tasks.

The implementation uses a thread-local counter that simply counts how
many "cooperation points" we have passed since the task was first
polled. Once the "budget" has been spent, any subsequent points will
return `Poll::Pending`, eventually making the top-level task yield. When
it finally does yield, the executor resets the budget before
running the next task.

The budget per task poll is currently hard-coded to 128. Eventually, we
may want to make it dynamic as more cooperation points are added. The
number 128 was chosen more or less arbitrarily to balance the cost of
yielding unnecessarily against the time an executor may be "held up".

At the moment, all the tokio leaf futures ("resources") call into coop,
but external futures have no way of doing so. We probably want to
continue limiting coop points to leaf futures in the future, but may
want to also enable third-party leaf futures to cooperate to benefit the
ecosystem as a whole. This is reflected in the methods marked as `pub`
in `mod coop` (even though the module is only `pub(crate)`). We will
likely also eventually want to expose `coop::limit`, which enables
sub-executors and manual `impl Future` blocks to avoid one sub-task
spending all of their poll budget.

Benchmarks (see #2160) suggest that the overhead of `coop`
is marginal.

tokio/src/coop.rs

@@ -0,0 +1,302 @@
+//! Opt-in yield points for improved cooperative scheduling.
+//!
+//! A single call to [`poll`] on a top-level task may potentially do a lot of work before it
+//! returns `Poll::Pending`. If a task runs for a long period of time without yielding back to the
+//! executor, it can starve other tasks waiting on that executor to execute them, or drive
+//! underlying resources. Since Rust does not have a runtime, it is difficult to forcibly preempt a
+//! long-running task. Instead, this module provides an opt-in mechanism for futures to collaborate
+//! with the executor to avoid starvation.
+//!
+//! Consider a future like this one:
+//!
+//! ```
+//! # use tokio::stream::{Stream, StreamExt};
+//! async fn drop_all<I: Stream + Unpin>(mut input: I) {
+//!     while let Some(_) = input.next().await {}
+//! }
+//! ```
+//!
+//! It may look harmless, but consider what happens under heavy load if the input stream is
+//! _always_ ready. If we spawn `drop_all`, the task will never yield, and will starve other tasks
+//! and resources on the same executor. With opt-in yield points, this problem is alleviated:
+//!
+//! ```ignore
+//! # use tokio::stream::{Stream, StreamExt};
+//! async fn drop_all<I: Stream + Unpin>(mut input: I) {
+//!     while let Some(_) = input.next().await {
+//!         tokio::coop::proceed().await;
+//!     }
+//! }
+//! ```
+//!
+//! The `proceed` future will coordinate with the executor to make sure that every so often control
+//! is yielded back to the executor so it can run other tasks.
+//!
+//! # Placing yield points
+//!
+//! Voluntary yield points should be placed _after_ at least some work has been done. If they are
+//! not, a future sufficiently deep in the task hierarchy may end up _never_ getting to run because
+//! of the number of yield points that inevitably appear before it is reached. In general, you will
+//! want yield points to only appear in "leaf" futures -- those that do not themselves poll other
+//! futures. By doing this, you avoid double-counting each iteration of the outer future against
+//! the cooperating budget.
+//!
+//!   [`poll`]: https://doc.rust-lang.org/std/future/trait.Future.html#tymethod.poll
+
+// NOTE: The doctests in this module are ignored since the whole module is (currently) private.
+
+use std::cell::Cell;
+use std::task::{Context, Poll};
+
+/// Constant used to determine how much "work" a task is allowed to do without yielding.
+///
+/// The value itself is chosen somewhat arbitrarily. It needs to be high enough to amortize wakeup
+/// and scheduling costs, but low enough that we do not starve other tasks for too long. The value
+/// also needs to be high enough that particularly deep tasks are able to do at least some useful
+/// work at all.
+///
+/// Note that as more yield points are added in the ecosystem, this value will probably also have
+/// to be raised.
+const BUDGET: usize = 128;
+
+/// Constant used to determine if budgeting has been disabled.
+const UNCONSTRAINED: usize = usize::max_value();
+
+thread_local! {
+    static HITS: Cell<usize> = Cell::new(UNCONSTRAINED);
+}
+
+/// Run the given closure with a cooperative task budget.
+///
+/// Enabling budgeting when it is already enabled is a no-op.
+#[inline(always)]
+pub(crate) fn budget<F, R>(f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    HITS.with(move |hits| {
+        if hits.get() != UNCONSTRAINED {
+            // We are already being budgeted.
+            //
+            // Arguably this should be an error, but it can happen "correctly"
+            // such as with block_on + LocalSet, so we make it a no-op.
+            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);
+        f()
+    })
+}
+
+cfg_blocking_impl! {
+    /// Forcibly remove the budgeting constraints early.
+    pub(crate) fn stop() {
+        HITS.with(|hits| {
+            hits.set(UNCONSTRAINED);
+        });
+    }
+}
+
+/// 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
+/// from using up your entire budget. For example, imagine the following future:
+///
+/// ```rust
+/// # use std::{future::Future, pin::Pin, task::{Context, Poll}};
+/// use futures::stream::FuturesUnordered;
+/// struct MyFuture<F1, F2> {
+///     big: FuturesUnordered<F1>,
+///     small: F2,
+/// }
+///
+/// use tokio::stream::Stream;
+/// impl<F1, F2> Future for MyFuture<F1, F2>
+///   where F1: Future, F2: Future
+/// # , F1: Unpin, F2: Unpin
+/// {
+///     type Output = F2::Output;
+///
+///     // fn poll(...)
+/// # fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F2::Output> {
+/// #   let this = &mut *self;
+///     let mut big = // something to pin self.big
+/// #                 Pin::new(&mut this.big);
+///     let small = // something to pin self.small
+/// #             Pin::new(&mut this.small);
+///
+///     // see if any of the big futures have finished
+///     while let Some(e) = futures::ready!(big.as_mut().poll_next(cx)) {
+///         // do something with e
+/// #       let _ = e;
+///     }
+///
+///     // see if the small future has finished
+///     small.poll(cx)
+/// }
+/// # }
+/// ```
+///
+/// It could be that every time `poll` gets called, `big` ends up spending the entire budget, and
+/// `small` never gets polled. That would be sad. If you want to stick up for the little future,
+/// that's what `limit` is for. It lets you portion out a smaller part of the yield budget to a
+/// particular segment of your code. In the code above, you would write
+///
+/// ```rust,ignore
+/// # use std::{future::Future, pin::Pin, task::{Context, Poll}};
+/// # use futures::stream::FuturesUnordered;
+/// # struct MyFuture<F1, F2> {
+/// #     big: FuturesUnordered<F1>,
+/// #     small: F2,
+/// # }
+/// #
+/// # use tokio::stream::Stream;
+/// # impl<F1, F2> Future for MyFuture<F1, F2>
+/// #   where F1: Future, F2: Future
+/// # , F1: Unpin, F2: Unpin
+/// # {
+/// # type Output = F2::Output;
+/// # fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F2::Output> {
+/// #   let this = &mut *self;
+/// #   let mut big = Pin::new(&mut this.big);
+/// #   let small = Pin::new(&mut this.small);
+/// #
+///     // see if any of the big futures have finished
+///     while let Some(e) = futures::ready!(tokio::coop::limit(64, || big.as_mut().poll_next(cx))) {
+/// #       // do something with e
+/// #       let _ = e;
+/// #   }
+/// #   small.poll(cx)
+/// # }
+/// # }
+/// ```
+///
+/// Now, even if `big` spends its entire budget, `small` will likely be left with some budget left
+/// to also do useful work. In particular, if the remaining budget was `N` at the start of `poll`,
+/// `small` will have at least a budget of `N - 64`. It may be more if `big` did not spend its
+/// entire budget.
+///
+/// Note that you cannot _increase_ your budget by calling `limit`. The budget provided to the code
+/// inside the buget is the _minimum_ of the _current_ budget and the bound.
+///
+#[allow(unreachable_pub, dead_code)]
+pub fn limit<R>(bound: usize, f: impl FnOnce() -> R) -> R {
+    HITS.with(|hits| {
+        let budget = hits.get();
+        // with_bound cannot _increase_ the remaining budget
+        let bound = std::cmp::min(budget, bound);
+        // When f() exits, how much should we add to what is left?
+        let floor = budget.saturating_sub(bound);
+        // Make sure we restore the remaining budget even on panic
+        struct RestoreBudget<'a>(&'a Cell<usize>, usize);
+        impl<'a> Drop for RestoreBudget<'a> {
+            fn drop(&mut self) {
+                let left = self.0.get();
+                self.0.set(self.1 + left);
+            }
+        }
+        // Time to restrict!
+        hits.set(bound);
+        let _restore = RestoreBudget(&hits, floor);
+        f()
+    })
+}
+
+/// Returns `Poll::Pending` if the current task has exceeded its budget and should yield.
+#[allow(unreachable_pub, dead_code)]
+#[inline]
+pub fn poll_proceed(cx: &mut Context<'_>) -> Poll<()> {
+    HITS.with(|hits| {
+        let n = hits.get();
+        if n == UNCONSTRAINED {
+            // opted out of budgeting
+            Poll::Ready(())
+        } else if n == 0 {
+            cx.waker().wake_by_ref();
+            Poll::Pending
+        } else {
+            hits.set(n.saturating_sub(1));
+            Poll::Ready(())
+        }
+    })
+}
+
+/// Resolves immediately unless the current task has already exceeded its budget.
+///
+/// This should be placed after at least some work has been done. Otherwise a future sufficiently
+/// deep in the task hierarchy may end up never getting to run because of the number of yield
+/// points that inevitably appear before it is even reached. For example:
+///
+/// ```ignore
+/// # use tokio::stream::{Stream, StreamExt};
+/// async fn drop_all<I: Stream + Unpin>(mut input: I) {
+///     while let Some(_) = input.next().await {
+///         tokio::coop::proceed().await;
+///     }
+/// }
+/// ```
+#[allow(unreachable_pub, dead_code)]
+#[inline]
+pub async fn proceed() {
+    use crate::future::poll_fn;
+    poll_fn(|cx| poll_proceed(cx)).await;
+}
+
+#[cfg(all(test, not(loom)))]
+mod test {
+    use super::*;
+
+    fn get() -> usize {
+        HITS.with(|hits| hits.get())
+    }
+
+    #[test]
+    fn bugeting() {
+        use tokio_test::*;
+
+        assert_eq!(get(), UNCONSTRAINED);
+        assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+        assert_eq!(get(), UNCONSTRAINED);
+        budget(|| {
+            assert_eq!(get(), BUDGET);
+            assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+            assert_eq!(get(), BUDGET - 1);
+            assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+            assert_eq!(get(), BUDGET - 2);
+        });
+        assert_eq!(get(), UNCONSTRAINED);
+
+        budget(|| {
+            limit(3, || {
+                assert_eq!(get(), 3);
+                assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+                assert_eq!(get(), 2);
+                limit(4, || {
+                    assert_eq!(get(), 2);
+                    assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+                    assert_eq!(get(), 1);
+                });
+                assert_eq!(get(), 1);
+                assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+                assert_eq!(get(), 0);
+                assert_pending!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+                assert_eq!(get(), 0);
+                assert_pending!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+                assert_eq!(get(), 0);
+            });
+            assert_eq!(get(), BUDGET - 3);
+            assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+            assert_eq!(get(), BUDGET - 4);
+            assert_ready!(task::spawn(proceed()).poll());
+            assert_eq!(get(), BUDGET - 5);
+        });
+    }
+}

tokio/src/io/registration.rs

@@ -139,6 +139,9 @@ impl Registration {
     ///
     /// This function will panic if called from outside of a task context.
     pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<mio::Ready>> {
+        // Keep track of task budget
+        ready!(crate::coop::poll_proceed(cx));
+
         let v = self.poll_ready(Direction::Read, Some(cx))?;
         match v {
             Some(v) => Poll::Ready(Ok(v)),
@@ -190,6 +193,9 @@ impl Registration {
     ///
     /// This function will panic if called from outside of a task context.
     pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<mio::Ready>> {
+        // Keep track of task budget
+        ready!(crate::coop::poll_proceed(cx));
+
         let v = self.poll_ready(Direction::Write, Some(cx))?;
         match v {
             Some(v) => Poll::Ready(Ok(v)),

tokio/src/lib.rs

@@ -321,6 +321,8 @@ cfg_process! {
 
 pub mod runtime;
 
+pub(crate) mod coop;
+
 cfg_signal! {
     pub mod signal;
 }

tokio/src/process/mod.rs

@@ -700,6 +700,9 @@ where
     type Output = Result<T, E>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // Keep track of task budget
+        ready!(crate::coop::poll_proceed(cx));
+
         let ret = Pin::new(&mut self.inner).poll(cx);
 
         if let Poll::Ready(Ok(_)) = ret {

tokio/src/runtime/basic_scheduler.rs

@@ -128,7 +128,7 @@ where
             pin!(future);
 
             'outer: loop {
-                if let Ready(v) = future.as_mut().poll(&mut cx) {
+                if let Ready(v) = crate::coop::budget(|| future.as_mut().poll(&mut cx)) {
                     return v;
                 }
 

tokio/src/runtime/blocking/task.rs

@@ -27,6 +27,14 @@ where
             .take()
             .expect("[internal exception] blocking task ran twice.");
 
+        // This is a little subtle:
+        // For convenience, we'd like _every_ call tokio ever makes to Task::poll() to be budgeted
+        // using coop. However, the way things are currently modeled, even running a blocking task
+        // currently goes through Task::poll(), and so is subject to budgeting. That isn't really
+        // what we want; a blocking task may itself want to run tasks (it might be a Worker!), so
+        // we want it to start without any budgeting.
+        crate::coop::stop();
+
         Poll::Ready(func())
     }
 }

tokio/src/runtime/enter.rs

@@ -98,7 +98,7 @@ cfg_blocking_impl! {
             let mut f = unsafe { Pin::new_unchecked(&mut f) };
 
             loop {
-                if let Ready(v) = f.as_mut().poll(&mut cx) {
+                if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) {
                     return Ok(v);
                 }
 
@@ -130,7 +130,7 @@ cfg_blocking_impl! {
             let when = Instant::now() + timeout;
 
             loop {
-                if let Ready(v) = f.as_mut().poll(&mut cx) {
+                if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) {
                     return Ok(v);
                 }
 

tokio/src/runtime/shell.rs

@@ -46,7 +46,7 @@ impl Shell {
         let mut cx = Context::from_waker(&self.waker);
 
         loop {
-            if let Ready(v) = f.as_mut().poll(&mut cx) {
+            if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) {
                 return v;
             }