Repository: tokio-rs/tokio
File: tokio/src/task/coop/mod.rs (Rust)

Code before:
            #[pin]
            pub(crate) fut: F,
        }
    }

    impl<F: Future> Future for Coop<F> {
        type Output = F::Output;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            let coop = ready!(poll_proceed(cx));
            let me = self.project();
            if let Poll::Ready(ret) = me.fut.poll(cx) {
                coop.made_progress();
                Poll::Ready(ret)
            } else {
                Poll::Pending
            }
        }
    }

    /// Run a future with a budget constraint for cooperative scheduling.
    /// If the future exceeds its budget while being polled, control is yielded back to the
    /// runtime.
    #[inline]
    pub fn cooperative<F: Future>(fut: F) -> Coop<F> {
        Coop { fut }
    }
}

#[cfg(all(test, not(loom)))]
mod test {
    use super::*;

    #[cfg(all(target_family = "wasm", not(target_os = "wasi")))]
    use wasm_bindgen_test::wasm_bindgen_test as test;

    fn get() -> Budget {
        context::budget(|cell| cell.get()).unwrap_or(Budget::unconstrained())
    }

    #[test]
    fn budgeting() {
        use std::future::poll_fn;
        use tokio_test::*;

        assert!(get().0.is_none());

        let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));

        assert!(get().0.is_none());
        drop(coop);

Review comment (suggestion):
Could we have an example or something. This is really short.

It could worth to add something like "this should be used to make futures from external crates cooperative". And you shouldn't use it with Tokio futures that are already cooperative. Maybe an example with the mpsc channel from futures and making `recv` cooperative.

Code after:

    impl<F: Future> Future for Coop<F> {
        type Output = F::Output;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            let coop = ready!(poll_proceed(cx));
            let me = self.project();
            if let Poll::Ready(ret) = me.fut.poll(cx) {
                coop.made_progress();
                Poll::Ready(ret)
            } else {
                Poll::Pending
            }
        }
    }

    /// Creates a wrapper future that makes the inner future cooperate with the Tokio scheduler.
    ///
    /// When polled, the wrapper will first call [`poll_proceed`] to consume task budget, and
    /// immediately yield if the budget has been depleted. If budget was available, the inner future
    /// is polled. The budget consumption will be made final using [`RestoreOnPending::made_progress`]
    /// if the inner future resolves to its final value.
    ///
    /// # Examples
    ///
    /// When you call `recv` on the `Receiver` of a [`tokio::sync::mpsc`](crate::sync::mpsc)
    /// channel, task budget will automatically be consumed when the next value is returned.
    /// This makes tasks that use Tokio mpsc channels automatically cooperative.
    ///
    /// If you're using [`futures::channel::mpsc`](https://docs.rs/futures/latest/futures/channel/mpsc/index.html)
    /// instead, automatic task budget consumption will not happen. This example shows how can use
    /// `cooperative` to make `futures::channel::mpsc` channels cooperate with the scheduler in the
    /// same way Tokio channels do.
    ///
    /// ```
    /// use tokio::task::coop::cooperative;
    /// use futures::channel::mpsc::Receiver;
    /// use futures::stream::StreamExt;
    ///
    /// async fn receive_next<T>(receiver: &mut Receiver<T>) -> Option<T> {
    ///     // Use `StreamExt::next` to obtain a `Future` that resolves to the next value
    ///     let recv_future = receiver.next();
    ///     // Wrap it a cooperative wrapper
    ///     let coop_future = cooperative(recv_future);
    ///     // And await
    ///     coop_future.await
    /// }
    #[inline]
    pub fn cooperative<F: Future>(fut: F) -> Coop<F> {
        Coop { fut }
    }