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

Code before:
    ///
    /// ```
    /// use std::pin::Pin;
    /// use std::io::Result as IoResult;
    /// use std::task::{ready, Context, Poll};
    /// use tokio::task::coop;
    /// use tokio::io::{AsyncRead, ReadBuf};
    ///
    /// struct Zero;
    ///
    /// impl AsyncRead for Zero {
    ///     fn poll_read(
    ///         self: Pin<&mut Self>,
    ///         cx: &mut Context<'_>,
    ///         buf: &mut ReadBuf<'_>
    ///     ) -> Poll<IoResult<()>> {
    ///        const ZEROES: [u8; 64] = [0; 64];
    ///
    ///        let coop = ready!(coop::poll_proceed(cx));
    ///        if buf.remaining() == 0 {
    ///            Poll::Ready(Ok(()))
    ///        } else {
    ///            while buf.remaining() > 0 {
    ///                let len = ZEROES.len().min(buf.remaining());
    ///                buf.put_slice(&ZEROES[..len]);
    ///            }
    ///            // we have made progress, so don't restore the budget
    ///            // when `coop` is dropped.
    ///            coop.made_progress();
    ///            Poll::Pending
    ///        }
    ///     }
    /// }
    /// ```
    ///
    /// [`/dev/zero`]: https://man7.org/linux/man-pages/man4/zero.4.html
    #[inline]
    pub fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
        context::budget(|cell| {
            let mut budget = cell.get();

            let decrement = budget.decrement();

            if decrement.success {
                let restore = RestoreOnPending::new(cell.get());
                cell.set(budget);

                // avoid double counting
                if decrement.hit_zero {
                    inc_budget_forced_yield_count();
                }

Review comment (suggestion):
You can improve the implementation a bit using [the `BufMut` impl](https://docs.rs/tokio/latest/tokio/io/struct.ReadBuf.html#impl-BufMut-for-ReadBuf%3C'a%3E).
```suggestion
    ///            buf.put_bytes(0, buf.remaining());
```

Code after:
    /// is consumed. If no budget is available, the task yields to the scheduler.
    ///
    /// ```
    /// use std::pin::Pin;
    /// use std::task::{ready, Context, Poll};
    /// use tokio::task::coop;
    /// use futures::stream::{Stream, StreamExt};
    /// use futures::channel::mpsc::UnboundedReceiver;
    ///
    /// struct CoopUnboundedReceiver<T> {
    ///    receiver: UnboundedReceiver<T>,
    /// }
    ///
    /// impl<T> Stream for CoopUnboundedReceiver<T> {
    ///     type Item = T;
    ///     fn poll_next(
    ///         mut self: Pin<&mut Self>,
    ///         cx: &mut Context<'_>
    ///     ) -> Poll<Option<T>> {
    ///         let coop = ready!(coop::poll_proceed(cx));
    ///         match self.receiver.poll_next_unpin(cx) {
    ///             Poll::Ready(v) => {
    ///                 // We received a value, so consume budget.
    ///                 coop.made_progress();
    ///                 Poll::Ready(v)
    ///             }
    ///             Poll::Pending => Poll::Pending,
    ///        }
    ///     }
    /// }
    /// ```
    #[inline]
    pub fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
        context::budget(|cell| {
            let mut budget = cell.get();

            let decrement = budget.decrement();

            if decrement.success {
                let restore = RestoreOnPending::new(cell.get());
                cell.set(budget);

                // avoid double counting
                if decrement.hit_zero {
                    inc_budget_forced_yield_count();
                }

                Poll::Ready(restore)
            } else {
                register_waker(cx);
                Poll::Pending