ronantakizawa/codereview-bench huggingface dataset:
Repository: tokio-rs/tokio
File: tokio/src/task/consume_budget.rs (Rust)

Code before:
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task went out ouf budget.
///
/// The task will only yield if it ran out of its coop budget.
/// It can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to runtime each time.
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub async fn consume_budget() {
    struct ConsumeBudget {
        status: Poll<()>,
    }

    impl Future for ConsumeBudget {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.status.is_ready() {
                return self.status;
            }
            self.status = crate::coop::poll_proceed(cx).map(|restore| {
                restore.made_progress();
            });
            self.status
        }
    }

    ConsumeBudget { status: Poll::Pending }.await
}


Review comment (suggestion):
one note about the documentation for this is that this is the first actual mention of the cooperative scheduling budget in the API documentation. I think it is probably worth adding a high-level summary of the task budget somewhere in the `tokio::task` documentation and having this API documentation reference that? Currently, the task budget is basically only documented in [this blog post](https://tokio.rs/blog/2020-04-preemption) which seems maybe not ideal...

Code after:
use std::task::Poll;

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task's coop budget was exhausted.
///
/// The task will only yield if its entire coop budget has been exhausted.
/// This function can can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to the runtime each time.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// # Examples
///
/// Make sure that a function which returns a sum of (potentially lots of)
/// iterated values is cooperative.
///
/// ```
/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
///     let mut sum: i64 = 0;
///     while let Some(i) = input.next() {
///         sum += i;
///         tokio::task::consume_budget().await
///     }
///     sum
/// }
/// ```
/// [unstable]: crate#unstable-features
#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
pub async fn consume_budget() {
    let mut status = Poll::Pending;

    crate::future::poll_fn(move |cx| {
        if status.is_ready() {
Repository: tokio-rs/tokio
File: tokio/src/task/consume_budget.rs (Rust)

Code before:
///
/// The task will only yield if it ran out of its coop budget.
/// It can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to runtime each time.
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub async fn consume_budget() {
    struct ConsumeBudget {
        status: Poll<()>,
    }

    impl Future for ConsumeBudget {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.status.is_ready() {
                return self.status;
            }
            self.status = crate::coop::poll_proceed(cx).map(|restore| {
                restore.made_progress();
            });
            self.status
        }
    }

    ConsumeBudget { status: Poll::Pending }.await
}


Review comment (suggestion):
nit, take it or leave it: i think we could _probably_ implement this a bit more simply using [`poll_fn`](https://github.com/tokio-rs/tokio/blob/ac0f894dd9be2a3a7e86405607ae240e9bc0a770/tokio/src/future/poll_fn.rs):
```suggestion
    let mut status = Poll::Pending;
    
    crate::future::poll_fn(move |cx| {
        if status.is_ready() {
            return status;
        }

        status = crate::coop::poll_proceed(cx).map(|restore| {
            restore.made_progress();
        });

        status
    }).await
```
i believe this should work identically, and avoids a little bit of boilerplate by not having to `impl Future` for a type.

Code after:
/// This function can can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to the runtime each time.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// # Examples
///
/// Make sure that a function which returns a sum of (potentially lots of)
/// iterated values is cooperative.
///
/// ```
/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
///     let mut sum: i64 = 0;
///     while let Some(i) = input.next() {
///         sum += i;
///         tokio::task::consume_budget().await
///     }
///     sum
/// }
/// ```
/// [unstable]: crate#unstable-features
#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
pub async fn consume_budget() {
    let mut status = Poll::Pending;

    crate::future::poll_fn(move |cx| {
        if status.is_ready() {
            return status;
        }
        status = crate::coop::poll_proceed(cx).map(|restore| {
            restore.made_progress();
        });
        status
    })
    .await
}
Repository: tokio-rs/tokio
File: tokio/src/task/consume_budget.rs (Rust)

Code before:
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task went out ouf budget.
///
/// The task will only yield if it ran out of its coop budget.
/// It can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to runtime each time.
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub async fn consume_budget() {
    struct ConsumeBudget {
        status: Poll<()>,
    }

    impl Future for ConsumeBudget {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.status.is_ready() {
                return self.status;
            }
            self.status = crate::coop::poll_proceed(cx).map(|restore| {
                restore.made_progress();
            });
            self.status
        }
    }


Review comment (suggestion):
nit: how about
```suggestion
/// runtime *if* the task's coop budget was exhausted.

Code after:
use std::task::Poll;

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task's coop budget was exhausted.
///
/// The task will only yield if its entire coop budget has been exhausted.
/// This function can can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to the runtime each time.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// # Examples
///
/// Make sure that a function which returns a sum of (potentially lots of)
/// iterated values is cooperative.
///
/// ```
/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
///     let mut sum: i64 = 0;
///     while let Some(i) = input.next() {
///         sum += i;
///         tokio::task::consume_budget().await
///     }
///     sum
/// }
/// ```
/// [unstable]: crate#unstable-features
#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
Repository: tokio-rs/tokio
File: tokio/src/task/consume_budget.rs (Rust)

Code before:
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task went out ouf budget.
///
/// The task will only yield if it ran out of its coop budget.
/// It can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to runtime each time.
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub async fn consume_budget() {
    struct ConsumeBudget {
        status: Poll<()>,
    }

    impl Future for ConsumeBudget {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.status.is_ready() {
                return self.status;
            }
            self.status = crate::coop::poll_proceed(cx).map(|restore| {
                restore.made_progress();
            });
            self.status
        }
    }

    ConsumeBudget { status: Poll::Pending }.await
}


Review comment (suggestion):
nit: maybe
```suggestion
/// The task will only yield if its entire coop budget has been exhausted.
/// This function can be used in order to insert optional yield points into long
```

Code after:
use std::task::Poll;

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task's coop budget was exhausted.
///
/// The task will only yield if its entire coop budget has been exhausted.
/// This function can can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to the runtime each time.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// # Examples
///
/// Make sure that a function which returns a sum of (potentially lots of)
/// iterated values is cooperative.
///
/// ```
/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
///     let mut sum: i64 = 0;
///     while let Some(i) = input.next() {
///         sum += i;
///         tokio::task::consume_budget().await
///     }
///     sum
/// }
/// ```
/// [unstable]: crate#unstable-features
#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
pub async fn consume_budget() {
    let mut status = Poll::Pending;
Repository: tokio-rs/tokio
File: tokio/src/task/consume_budget.rs (Rust)

Code before:
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task went out ouf budget.
///
/// The task will only yield if it ran out of its coop budget.
/// It can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to runtime each time.
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub async fn consume_budget() {
    struct ConsumeBudget {
        status: Poll<()>,
    }

    impl Future for ConsumeBudget {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
            if self.status.is_ready() {
                return self.status;
            }
            self.status = crate::coop::poll_proceed(cx).map(|restore| {
                restore.made_progress();
            });
            self.status
        }
    }

    ConsumeBudget { status: Poll::Pending }.await
}


Review comment (suggestion):
nit:
```suggestion
/// without redundantly yielding to the runtime each time.
```

Code after:
use std::task::Poll;

/// Consumes a unit of budget and returns the execution back to the Tokio
/// runtime *if* the task's coop budget was exhausted.
///
/// The task will only yield if its entire coop budget has been exhausted.
/// This function can can be used in order to insert optional yield points into long
/// computations that do not use Tokio resources like sockets or semaphores,
/// without redundantly yielding to the runtime each time.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// # Examples
///
/// Make sure that a function which returns a sum of (potentially lots of)
/// iterated values is cooperative.
///
/// ```
/// async fn sum_iterator(input: &mut impl std::iter::Iterator<Item=i64>) -> i64 {
///     let mut sum: i64 = 0;
///     while let Some(i) = input.next() {
///         sum += i;
///         tokio::task::consume_budget().await
///     }
///     sum
/// }
/// ```
/// [unstable]: crate#unstable-features
#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]
pub async fn consume_budget() {
    let mut status = Poll::Pending;

    crate::future::poll_fn(move |cx| {
        if status.is_ready() {
Repository: tokio-rs/tokio
File: tokio/src/runtime/metrics/io.rs (Rust)

Code before:
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

#[derive(Default)]
pub(crate) struct IoDriverMetrics {
    pub(super) fd_count: AtomicU64,
    pub(super) ready_count: AtomicU64,
}

impl IoDriverMetrics {
    pub(crate) fn incr_fd_count(&self) {
        self.fd_count.fetch_add(1, Relaxed);
    }

    pub(crate) fn dec_fd_count(&self) {
        self.fd_count.fetch_sub(1, Relaxed);
    }

    pub(crate) fn incr_ready_count(&self) {
        self.ready_count.fetch_add(1, Relaxed);
    }
}


Review comment (suggestion):
Since there is only one "producer" of metrics instead of `fetch_add`, i would do a load / store combo. It is cheaper.

Code after:
#![cfg_attr(not(feature = "net"), allow(dead_code))]

use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

#[derive(Default)]
pub(crate) struct IoDriverMetrics {
    pub(super) fd_registered_count: AtomicU64,
    pub(super) fd_deregistered_count: AtomicU64,
    pub(super) ready_count: AtomicU64,
}

impl IoDriverMetrics {
    pub(crate) fn incr_fd_count(&self) {
        let prev = self.fd_registered_count.load(Relaxed);
        let new = prev.wrapping_add(1);
        self.fd_registered_count.store(new, Relaxed);
    }

    pub(crate) fn dec_fd_count(&self) {
        let prev = self.fd_deregistered_count.load(Relaxed);
        let new = prev.wrapping_add(1);
        self.fd_deregistered_count.store(new, Relaxed);
    }

    pub(crate) fn incr_ready_count_by(&self, amt: u64) {
        let prev = self.ready_count.load(Relaxed);
        let new = prev.wrapping_add(amt);
        self.ready_count.store(new, Relaxed);
    }
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/metrics/io.rs (Rust)

Code before:
#![cfg_attr(not(feature = "net"), allow(dead_code))]

use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

#[derive(Default)]
pub(crate) struct IoDriverMetrics {
    pub(super) fd_count: AtomicU64,
    pub(super) ready_count: AtomicU64,
}

impl IoDriverMetrics {
    pub(crate) fn incr_fd_count(&self) {
        let prev = self.fd_count.load(Relaxed);
        let new = prev.wrapping_add(1);
        self.fd_count.store(new, Relaxed);
    }

    pub(crate) fn dec_fd_count(&self) {
        let prev = self.fd_count.load(Relaxed);
        let new = prev.wrapping_sub(1);
        self.fd_count.store(new, Relaxed);
    }

    pub(crate) fn incr_ready_count_by(&self, amt: u64) {
        let prev = self.fd_count.load(Relaxed);
        let new = prev.wrapping_add(amt);
        self.ready_count.store(new, Relaxed);
    }
}


Review comment (suggestion):
Is it worth tracking registered & deregistered separately? Right now, if an FD is registered then immediately deregistered, this will show up as the counter not moving. What is the value in tracking outstanding FDs? If registered & deregistered ops are tracked separately, churn will be detectable.

Separately, do we want to track `fd_count` or punt that to `tracing`? I am not sure. Thoughts @hawkw 

Code after:
#![cfg_attr(not(feature = "net"), allow(dead_code))]

use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

#[derive(Default)]
pub(crate) struct IoDriverMetrics {
    pub(super) fd_registered_count: AtomicU64,
    pub(super) fd_deregistered_count: AtomicU64,
    pub(super) ready_count: AtomicU64,
}

impl IoDriverMetrics {
    pub(crate) fn incr_fd_count(&self) {
        let prev = self.fd_registered_count.load(Relaxed);
        let new = prev.wrapping_add(1);
        self.fd_registered_count.store(new, Relaxed);
    }

    pub(crate) fn dec_fd_count(&self) {
        let prev = self.fd_deregistered_count.load(Relaxed);
        let new = prev.wrapping_add(1);
        self.fd_deregistered_count.store(new, Relaxed);
    }

    pub(crate) fn incr_ready_count_by(&self, amt: u64) {
        let prev = self.ready_count.load(Relaxed);
        let new = prev.wrapping_add(amt);
        self.ready_count.store(new, Relaxed);
    }
}
Repository: tokio-rs/tokio
File: tokio/src/sync/broadcast.rs (Rust)

Code before:
            let waker = waiter.waker.take().unwrap();
            waker.wake();
        }
    }
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Sender<T> {
        let shared = self.shared.clone();
        shared.num_tx.fetch_add(1, SeqCst);

        Sender { shared }
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        if 1 == self.shared.num_tx.fetch_sub(1, SeqCst) {
            let _ = self.send2(None);
        }
    }
}

impl<T> Receiver<T> {
    /// Returns the number of messages that were sent into the channel and that
    /// this [`Receiver`] has yet to receive.
    ///
    /// [`Receiver`]: crate::sync::broadcast::Receiver
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, mut rx1) = broadcast::channel(16);
    ///
    ///     tx.send(10).unwrap();
    ///     tx.send(20).unwrap();
    ///
    ///     assert_eq!(rx1.num_msgs(), 2);
    ///     assert_eq!(rx1.recv().await.unwrap(), 10);
    ///     assert_eq!(rx1.num_msgs(), 1);
    ///     assert_eq!(rx1.recv().await.unwrap(), 20);     
    ///     assert_eq!(rx1.num_msgs(), 0);
    /// }
    /// ```
    pub fn num_msgs(&self) -> u64 {
        let next_send_pos = self.shared.tail.lock().pos;
        next_send_pos - self.next

Review comment (suggestion):
This should document what happens when the receiver gets so far behind that it will receive a `Lagged` error.

Code after:
            let waker = waiter.waker.take().unwrap();
            waker.wake();
        }
    }
}

impl<T> Clone for Sender<T> {
    fn clone(&self) -> Sender<T> {
        let shared = self.shared.clone();
        shared.num_tx.fetch_add(1, SeqCst);

        Sender { shared }
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        if 1 == self.shared.num_tx.fetch_sub(1, SeqCst) {
            let _ = self.send2(None);
        }
    }
}

impl<T> Receiver<T> {
    /// Returns the number of messages that were sent into the channel and that
    /// this [`Receiver`] has yet to receive.
    ///
    /// If the returned value from `len` is larger than the next largest power of 2
    /// of the capacity of the channel any call to [`recv`] will return an
    /// `Err(RecvError::Lagged)` and any call to [`try_recv`] will return an
    /// `Err(TryRecvError::Lagged)`, e.g. if the capacity of the channel is 10,
    /// [`recv`] will start to return `Err(RecvError::Lagged)` once `len` returns
    /// values larger than 16.
    ///
    /// [`Receiver`]: crate::sync::broadcast::Receiver
    /// [`recv`]: crate::sync::broadcast::Receiver::recv
    /// [`try_recv`]: crate::sync::broadcast::Receiver::try_recv
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, mut rx1) = broadcast::channel(16);
    ///
    ///     tx.send(10).unwrap();
    ///     tx.send(20).unwrap();
    ///
    ///     assert_eq!(rx1.len(), 2);
Repository: tokio-rs/tokio
File: tokio-util/src/io/stream_reader.rs (Rust)

Code before:
        Self {
            inner: stream,
            chunk: None,
        }
    }

    /// Do we have a chunk and is it non-empty?
    fn has_chunk(self: Pin<&mut Self>) -> bool {
        if let Some(chunk) = self.project().chunk {
            chunk.remaining() > 0
        } else {
            false
        }
    }

    /// Consumes this `StreamReader`, returning a Tuple consisting
    /// of the underlying stream and an Option of the interal buffer,
    /// which is Some in case the buffer contains elements.
    pub fn into_inner_with_chunk(self) -> (S, Option<B>) {
        if let Some(chunk) = self.chunk {
            if chunk.remaining() > 0 {
                return (self.inner, Some(chunk));
            }
        }

        (self.inner, None)
    }
}

impl<S, B> StreamReader<S, B> {
    /// Gets a reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Gets a mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Gets a pinned mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut S> {
        self.project().inner
    }


Review comment (suggestion):
```suggestion
        if self.has_chunk() {
            (self.inner, self.chunk)
        } else {
            (self.inner, None)
        }
```

Code after:
        Self {
            inner: stream,
            chunk: None,
        }
    }

    /// Do we have a chunk and is it non-empty?
    fn has_chunk(&self) -> bool {
        if let Some(ref chunk) = self.chunk {
            chunk.remaining() > 0
        } else {
            false
        }
    }

    /// Consumes this `StreamReader`, returning a Tuple consisting
    /// of the underlying stream and an Option of the interal buffer,
    /// which is Some in case the buffer contains elements.
    pub fn into_inner_with_chunk(self) -> (S, Option<B>) {
        if self.has_chunk() {
            (self.inner, self.chunk)
        } else {
            (self.inner, None)
        }
    }
}

impl<S, B> StreamReader<S, B> {
    /// Gets a reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Gets a mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Gets a pinned mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut S> {
        self.project().inner
    }

    /// Consumes this `BufWriter`, returning the underlying stream.
    ///