coop: expose coop as a public module

tokio/src/io/util/copy.rs

@@ -94,7 +94,7 @@ impl CopyBuffer {
             feature = "time",
         ))]
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
         loop {
             // If there is some space left in our buffer, then we try to read some
             // data to continue, thus maximizing the chances of a large write.

tokio/src/io/util/mem.rs

@@ -332,7 +332,7 @@ impl AsyncRead for SimplexStream {
             buf: &mut ReadBuf<'_>,
         ) -> Poll<std::io::Result<()>> {
             ready!(crate::trace::trace_leaf(cx));
-            let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+            let coop = ready!(crate::task::coop::poll_proceed(cx));
 
             let ret = self.poll_read_internal(cx, buf);
             if ret.is_ready() {
@@ -362,7 +362,7 @@ impl AsyncWrite for SimplexStream {
             buf: &[u8],
         ) -> Poll<std::io::Result<usize>> {
             ready!(crate::trace::trace_leaf(cx));
-            let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+            let coop = ready!(crate::task::coop::poll_proceed(cx));
 
             let ret = self.poll_write_internal(cx, buf);
             if ret.is_ready() {
@@ -390,7 +390,7 @@ impl AsyncWrite for SimplexStream {
             bufs: &[std::io::IoSlice<'_>],
         ) -> Poll<Result<usize, std::io::Error>> {
             ready!(crate::trace::trace_leaf(cx));
-            let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+            let coop = ready!(crate::task::coop::poll_proceed(cx));
 
             let ret = self.poll_write_vectored_internal(cx, bufs);
             if ret.is_ready() {

tokio/src/io/util/mod.rs

@@ -88,7 +88,7 @@ cfg_io_util! {
 
     cfg_coop! {
         fn poll_proceed_and_make_progress(cx: &mut std::task::Context<'_>) -> std::task::Poll<()> {
-            let coop = std::task::ready!(crate::runtime::coop::poll_proceed(cx));
+            let coop = std::task::ready!(crate::task::coop::poll_proceed(cx));
             coop.made_progress();
             std::task::Poll::Ready(())
         }

tokio/src/process/mod.rs

@@ -1048,7 +1048,7 @@ where
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         ready!(crate::trace::trace_leaf(cx));
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         let ret = Pin::new(&mut self.inner).poll(cx);
 

tokio/src/runtime/blocking/task.rs

@@ -37,7 +37,7 @@ where
         // 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::runtime::coop::stop();
+        crate::task::coop::stop();
 
         Poll::Ready(func())
     }

tokio/src/runtime/context.rs

@@ -1,5 +1,5 @@
 use crate::loom::thread::AccessError;
-use crate::runtime::coop;
+use crate::task::coop;
 
 use std::cell::Cell;
 
@@ -135,7 +135,7 @@ pub(crate) fn thread_rng_n(n: u32) -> u32 {
     })
 }
 
-pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> {
+pub(crate) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> {
     CONTEXT.try_with(|ctx| f(&ctx.budget))
 }
 

tokio/src/runtime/context/blocking.rs

@@ -87,7 +87,7 @@ impl BlockingRegionGuard {
         let when = Instant::now() + timeout;
 
         loop {
-            if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) {
+            if let Ready(v) = crate::task::coop::budget(|| f.as_mut().poll(&mut cx)) {
                 return Ok(v);
             }
 

tokio/src/runtime/io/registration.rs

@@ -148,7 +148,7 @@ impl Registration {
     ) -> Poll<io::Result<ReadyEvent>> {
         ready!(crate::trace::trace_leaf(cx));
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
         let ev = ready!(self.shared.poll_readiness(cx, direction));
 
         if ev.is_shutdown {
@@ -219,7 +219,7 @@ impl Registration {
         loop {
             let event = self.readiness(interest).await?;
 
-            let coop = std::future::poll_fn(crate::runtime::coop::poll_proceed).await;
+            let coop = std::future::poll_fn(crate::task::coop::poll_proceed).await;
 
             match f() {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {

tokio/src/runtime/mod.rs

@@ -321,8 +321,6 @@ mod tests;
 
 pub(crate) mod context;
 
-pub(crate) mod coop;
-
 pub(crate) mod park;
 
 mod driver;

tokio/src/runtime/park.rs

@@ -281,7 +281,7 @@ impl CachedParkThread {
         pin!(f);
 
         loop {
-            if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) {
+            if let Ready(v) = crate::task::coop::budget(|| f.as_mut().poll(&mut cx)) {
                 return Ok(v);
             }
 

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -361,7 +361,7 @@ impl Context {
     /// thread-local context.
     fn run_task<R>(&self, mut core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
         core.metrics.start_poll();
-        let mut ret = self.enter(core, || crate::runtime::coop::budget(f));
+        let mut ret = self.enter(core, || crate::task::coop::budget(f));
         ret.0.metrics.end_poll();
         ret
     }
@@ -726,7 +726,7 @@ impl CoreGuard<'_> {
 
                 if handle.reset_woken() {
                     let (c, res) = context.enter(core, || {
-                        crate::runtime::coop::budget(|| future.as_mut().poll(&mut cx))
+                        crate::task::coop::budget(|| future.as_mut().poll(&mut cx))
                     });
 
                     core = c;

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -63,10 +63,9 @@ use crate::runtime::scheduler::multi_thread::{
 };
 use crate::runtime::scheduler::{inject, Defer, Lock};
 use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks};
-use crate::runtime::{
-    blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics,
-};
+use crate::runtime::{blocking, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics};
 use crate::runtime::{context, TaskHooks};
+use crate::task::coop;
 use crate::util::atomic_cell::AtomicCell;
 use crate::util::rand::{FastRand, RngSeedGenerator};
 

tokio/src/runtime/scheduler/multi_thread_alt/worker.rs

@@ -64,8 +64,9 @@ use crate::runtime::scheduler::multi_thread_alt::{
 };
 use crate::runtime::scheduler::{self, inject, Lock};
 use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks};
-use crate::runtime::{blocking, coop, driver, task, Config, SchedulerMetrics, WorkerMetrics};
+use crate::runtime::{blocking, driver, task, Config, SchedulerMetrics, WorkerMetrics};
 use crate::runtime::{context, TaskHooks};
+use crate::task::coop;
 use crate::util::atomic_cell::AtomicCell;
 use crate::util::rand::{FastRand, RngSeedGenerator};
 

tokio/src/runtime/task/join.rs

@@ -322,7 +322,7 @@ impl<T> Future for JoinHandle<T> {
         let mut ret = Poll::Pending;
 
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         // Try to read the task output. If the task is not yet complete, the
         // waker is stored and is notified once the task does complete.

tokio/src/sync/batch_semaphore.rs

@@ -591,11 +591,11 @@ impl Future for Acquire<'_> {
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let coop = ready!(trace_poll_op!(
             "poll_acquire",
-            crate::runtime::coop::poll_proceed(cx),
+            crate::task::coop::poll_proceed(cx),
         ));
 
         #[cfg(not(all(tokio_unstable, feature = "tracing")))]
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         let result = match semaphore.poll_acquire(cx, needed, node, *queued) {
             Poll::Pending => {

tokio/src/sync/broadcast.rs

@@ -119,7 +119,7 @@
 use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::atomic::{AtomicBool, AtomicUsize};
 use crate::loom::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard};
-use crate::runtime::coop::cooperative;
+use crate::task::coop::cooperative;
 use crate::util::linked_list::{self, GuardedLinkedList, LinkedList};
 use crate::util::WakeList;
 

tokio/src/sync/mpsc/chan.rs

@@ -292,7 +292,7 @@ impl<T, S: Semaphore> Rx<T, S> {
         ready!(crate::trace::trace_leaf(cx));
 
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         self.inner.rx_fields.with_mut(|rx_fields_ptr| {
             let rx_fields = unsafe { &mut *rx_fields_ptr };
@@ -354,7 +354,7 @@ impl<T, S: Semaphore> Rx<T, S> {
         ready!(crate::trace::trace_leaf(cx));
 
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         if limit == 0 {
             coop.made_progress();

tokio/src/sync/oneshot.rs

@@ -794,7 +794,7 @@ impl<T> Sender<T> {
         ready!(crate::trace::trace_leaf(cx));
 
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         let inner = self.inner.as_ref().unwrap();
 
@@ -1142,7 +1142,7 @@ impl<T> Inner<T> {
     fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
         ready!(crate::trace::trace_leaf(cx));
         // Keep track of task budget
-        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        let coop = ready!(crate::task::coop::poll_proceed(cx));
 
         // Load the state
         let mut state = State::load(&self.state, Acquire);

tokio/src/sync/watch.rs

@@ -111,8 +111,8 @@
 //! [`Sender::closed`]: crate::sync::watch::Sender::closed
 //! [`Sender::subscribe()`]: crate::sync::watch::Sender::subscribe
 
-use crate::runtime::coop::cooperative;
 use crate::sync::notify::Notify;
+use crate::task::coop::cooperative;
 
 use crate::loom::sync::atomic::AtomicUsize;
 use crate::loom::sync::atomic::Ordering::{AcqRel, Relaxed};

tokio/src/task/consume_budget.rs → tokio/src/task/coop/consume_budget.rs

@@ -1,5 +1,3 @@
-use std::task::{ready, Poll};
-
 /// Consumes a unit of budget and returns the execution back to the Tokio
 /// runtime *if* the task's coop budget was exhausted.
 ///
@@ -25,14 +23,14 @@ use std::task::{ready, Poll};
 /// ```
 #[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
 pub async fn consume_budget() {
-    let mut status = Poll::Pending;
+    let mut status = std::task::Poll::Pending;
 
     std::future::poll_fn(move |cx| {
-        ready!(crate::trace::trace_leaf(cx));
+        std::task::ready!(crate::trace::trace_leaf(cx));
         if status.is_ready() {
             return status;
         }
-        status = crate::runtime::coop::poll_proceed(cx).map(|restore| {
+        status = crate::task::coop::poll_proceed(cx).map(|restore| {
             restore.made_progress();
         });
         status