coop: expose an unconstrained() opt-out

tokio/src/coop.rs

@@ -1,6 +1,6 @@
 #![cfg_attr(not(feature = "full"), allow(dead_code))]
 
-//! Opt-in yield points for improved cooperative scheduling.
+//! 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
@@ -21,37 +21,41 @@
 //!
 //! 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:
+//! yield, and will starve other tasks and resources on the same executor.
 //!
-//! ```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.
+//! To account for this, Tokio has explicit yield points in a number of library
+//! functions, which force tasks to return to the executor periodically.
 //!
 //! [`poll`]: method@std::future::Future::poll
 
-// NOTE: The doctests in this module are ignored since the whole module is (currently) private.
-
+// ```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.
+
+use pin_project_lite::pin_project;
 use std::cell::Cell;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
 
 thread_local! {
     static CURRENT: Cell<Budget> = Cell::new(Budget::unconstrained());
@@ -147,9 +151,57 @@ cfg_rt! {
     }
 }
 
-cfg_coop! {
-    use std::task::{Context, Poll};
+pin_project! {
+    /// Future for the [`unconstrained`](unconstrained) method.
+    #[must_use = "Unconstrained does nothing unless polled"]
+    pub struct Unconstrained<F> {
+        #[pin]
+        inner: F,
+    }
+}
+
+impl<F> Future for Unconstrained<F>
+where
+    F: Future,
+{
+    type Output = <F as Future>::Output;
 
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let inner = self.project().inner;
+        with_budget(Budget::unconstrained(), || inner.poll(cx))
+    }
+}
+
+/// Turn off cooperative scheduling for a future. The future will never be forced to yield by
+/// Tokio. Using this exposes your service to starvation if the unconstrained future never yields
+/// otherwise.
+///
+/// # Examples
+///
+/// ```
+/// # #[tokio::main]
+/// # async fn main() {
+/// use tokio::{coop, sync::mpsc};
+///
+/// let fut = async {
+///     let (tx, mut rx) = mpsc::unbounded_channel();
+///
+///     for i in 0..1000 {
+///         let _ = tx.send(());
+///         // This will always be ready. If coop was in effect, this code would be forced to yield
+///         // periodically. However, if left unconstrained, then this code will never yield.
+///         rx.recv().await;
+///     }
+/// };
+///
+/// coop::unconstrained(fut).await;
+/// # }
+/// ```
+pub fn unconstrained<F>(inner: F) -> Unconstrained<F> {
+    Unconstrained { inner }
+}
+
+cfg_coop! {
     #[must_use]
     pub(crate) struct RestoreOnPending(Cell<Budget>);
 

tokio/src/lib.rs

@@ -372,7 +372,7 @@ cfg_rt! {
     pub mod runtime;
 }
 
-pub(crate) mod coop;
+pub mod coop;
 
 cfg_signal! {
     pub mod signal;

tokio/tests/rt_common.rs

@@ -1017,6 +1017,32 @@ rt_test! {
         });
     }
 
+    #[test]
+    fn coop_unconstrained() {
+        use std::task::Poll::Ready;
+
+        let rt = rt();
+
+        rt.block_on(async {
+            // Create a bunch of tasks
+            let mut tasks = (0..1_000).map(|_| {
+                tokio::spawn(async { })
+            }).collect::<Vec<_>>();
+
+            // Hope that all the tasks complete...
+            time::sleep(Duration::from_millis(100)).await;
+
+            tokio::coop::unconstrained(poll_fn(|cx| {
+                // All the tasks should be ready
+                for task in &mut tasks {
+                    assert!(Pin::new(task).poll(cx).is_ready());
+                }
+
+                Ready(())
+            })).await;
+        });
+    }
+
     // Tests that the "next task" scheduler optimization is not able to starve
     // other tasks.
     #[test]