Add cooperative task yielding (#2160)

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 tokio-rs/tokio#2160) suggest that the overhead of `coop`
is marginal.

7 files changed, 22 insertions, 6 deletions

diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs
index a494f9e3..0419c209 100644
--- a/tokio/src/runtime/basic_scheduler.rs
+++ b/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;
                 }
 
diff --git a/tokio/src/runtime/blocking/task.rs b/tokio/src/runtime/blocking/task.rs
index 0553c9bd..f98b8549 100644
--- a/tokio/src/runtime/blocking/task.rs
+++ b/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())
     }
 }
diff --git a/tokio/src/runtime/enter.rs b/tokio/src/runtime/enter.rs
index 64648b71..afdb67a3 100644
--- a/tokio/src/runtime/enter.rs
+++ b/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);
                 }
 
diff --git a/tokio/src/runtime/shell.rs b/tokio/src/runtime/shell.rs
index 992244de..efe77f20 100644
--- a/tokio/src/runtime/shell.rs
+++ b/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;
             }
 
diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs
index 43e6b471..dee55a54 100644
--- a/tokio/src/runtime/task/core.rs
+++ b/tokio/src/runtime/task/core.rs
@@ -161,7 +161,7 @@ impl<T: Future, S: Schedule> Core<T, S> {
                 let waker_ref = waker_ref::<T, S>(header);
                 let mut cx = Context::from_waker(&*waker_ref);
 
-                future.poll(&mut cx)
+                crate::coop::budget(|| future.poll(&mut cx))
             })
         };
 
diff --git a/tokio/src/runtime/task/join.rs b/tokio/src/runtime/task/join.rs
index ed893a35..fdcc346e 100644
--- a/tokio/src/runtime/task/join.rs
+++ b/tokio/src/runtime/task/join.rs
@@ -101,6 +101,9 @@ impl<T> Future for JoinHandle<T> {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut ret = Poll::Pending;
 
+        // Keep track of task budget
+        ready!(crate::coop::poll_proceed(cx));
+
         // Raw should always be set. If it is not, this is due to polling after
         // completion
         let raw = self
diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs
index cc79530a..53c7e5b6 100644
--- a/tokio/src/runtime/thread_pool/worker.rs
+++ b/tokio/src/runtime/thread_pool/worker.rs
@@ -187,7 +187,12 @@ cfg_blocking! {
 
             // Get the worker core. If none is set, then blocking is fine!
             let core = match cx.core.borrow_mut().take() {
-                Some(core) => core,
+                Some(core) => {
+                    // We are effectively leaving the executor, so we need to
+                    // forcibly end budgeting.
+                    crate::coop::stop();
+                    core
+                },
                 None => return,
             };