ronantakizawa/codereview-bench huggingface dataset:
Repository: tokio-rs/tokio
File: tokio/tests/task_id.rs (Rust)

Code before:
            })
        }
    }

    let (tx_output, rx_output) = oneshot::channel();
    let (tx_future, rx_future) = oneshot::channel();
    let handle = tokio::spawn(MyFuture {
        tx_output: Some(tx_output),
        tx_future: Some(tx_future),
    });
    let id = handle.id();
    rx_future.await.unwrap();
    drop(handle);
    assert_eq!(rx_output.await.unwrap(), id);
}

#[test]
fn task_try_id_outside_task() {
    assert_eq!(None, task::try_id());
}

#[test]
fn task_try_id_inside_block_on() {
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        assert_eq!(None, task::try_id());
    });
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn_local() {
    LocalSet::new()
        .run_until(async {
            task::spawn_local(async { println!("task id: {}", task::id()) })
                .await
                .unwrap();
        })
        .await
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_nested_spawn_local() {
    LocalSet::new()
        .run_until(async {
            task::spawn_local(async {
                let parent_id = task::id();
                LocalSet::new()
                    .run_until(async {
                        task::spawn_local(async move {
                            assert_ne!(parent_id, task::id());
                        })

Review comment (suggestion):
We can verify in the test that the output has not yet been dropped when we drop the handle.
```suggestion
    assert!(rx_output.try_recv().is_err());
    drop(handle);
    assert_eq!(rx_output.await.unwrap(), id);
```

Code after:
        tx_output: Option<oneshot::Sender<Id>>,
        tx_future: Option<oneshot::Sender<()>>,
    }

    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            let _ = self.tx_future.take().unwrap().send(());
            Poll::Ready(MyOutput {
                tx: self.tx_output.take(),
            })
        }
    }

    let (tx_output, mut rx_output) = oneshot::channel();
    let (tx_future, rx_future) = oneshot::channel();
    let handle = tokio::spawn(MyFuture {
        tx_output: Some(tx_output),
        tx_future: Some(tx_future),
    });
    let id = handle.id();
    rx_future.await.unwrap();
    assert!(rx_output.try_recv().is_err());
    drop(handle);
    assert_eq!(rx_output.await.unwrap(), id);
}

#[test]
fn task_try_id_outside_task() {
    assert_eq!(None, task::try_id());
}

#[cfg(not(tokio_wasi))]
#[test]
fn task_try_id_inside_block_on() {
    let rt = Runtime::new().unwrap();
    rt.block_on(async {
        assert_eq!(None, task::try_id());
    });
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn_local() {
    LocalSet::new()
        .run_until(async {
            task::spawn_local(async { println!("task id: {}", task::id()) })
                .await
                .unwrap();
        })
        .await
Repository: tokio-rs/tokio
File: tokio/tests/task_id.rs (Rust)

Code before:
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Drop for MyOutput {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            Poll::Ready(MyOutput { tx: self.tx.take() })
        }
    }

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    let id = handle.id();
    drop(handle);
    assert_eq!(rx.await.unwrap(), id);
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_after_completion() {
    struct MyOutput {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Drop for MyOutput {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    struct MyFuture {
        tx_output: Option<oneshot::Sender<Id>>,
        tx_future: Option<oneshot::Sender<()>>,
    }

    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            let _ = self.tx_future.take().unwrap().send(());
            Poll::Ready(MyOutput {

Review comment (suggestion):
```suggestion
    drop(handle);
    assert!(rx.try_recv().is_err());
    assert_eq!(rx.await.unwrap(), id);
```

Code after:
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    let id = handle.id();
    handle.abort();
    assert!(handle.await.unwrap_err().is_cancelled());
    assert_eq!(rx.await.unwrap(), id);
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_before_completion() {
    struct MyOutput {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Drop for MyOutput {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            Poll::Ready(MyOutput { tx: self.tx.take() })
        }
    }

    let (tx, mut rx) = oneshot::channel();
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    let id = handle.id();
    drop(handle);
    assert!(rx.try_recv().is_err());
    assert_eq!(rx.await.unwrap(), id);
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_after_completion() {
    struct MyOutput {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Drop for MyOutput {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }
Repository: tokio-rs/tokio
File: tokio/src/runtime/builder.rs (Rust)

Code before:
        ///
        /// let rt = runtime::Builder::new_multi_thread()
        ///     .enable_io()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_io(&mut self) -> &mut Self {
            self.enable_io = true;
            self
        }

        /// Enables the I/O driver and configures the capacity of events to be 
        /// processed per tick.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new_current_thread()
        ///     .enable_io()
        ///     .with_events_capacity(1024)
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn with_events_capacity(&mut self, capacity: usize) -> &mut Self {
            self.nevents = capacity;
            self
        }
    }
}

cfg_time! {
    impl Builder {
        /// Enables the time driver.
        ///
        /// Doing this enables using `tokio::time` on the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new_multi_thread()
        ///     .enable_time()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_time(&mut self) -> &mut Self {
            self.enable_time = true;
            self

Review comment (suggestion):
we don't use `with`, maybe try something like `max_io_events_per_tick` or something. Not perfect, but be descriptive.

Code after:
        ///
        /// let rt = runtime::Builder::new_multi_thread()
        ///     .enable_io()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_io(&mut self) -> &mut Self {
            self.enable_io = true;
            self
        }

        /// Enables the I/O driver and configures the max number of events to be
        /// processed per tick.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new_current_thread()
        ///     .enable_io()
        ///     .max_io_events_per_tick(1024)
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn max_io_events_per_tick(&mut self, capacity: usize) -> &mut Self {
            self.nevents = capacity;
            self
        }
    }
}

cfg_time! {
    impl Builder {
        /// Enables the time driver.
        ///
        /// Doing this enables using `tokio::time` on the runtime.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime;
        ///
        /// let rt = runtime::Builder::new_multi_thread()
        ///     .enable_time()
        ///     .build()
        ///     .unwrap();
        /// ```
        pub fn enable_time(&mut self) -> &mut Self {
            self.enable_time = true;
            self
Repository: tokio-rs/tokio
File: tokio/src/loom/std/atomic_u64_static_once_cell.rs (Rust)

Code before:
use super::AtomicU64;
use crate::loom::sync::{Mutex, atomic::Ordering};
use crate::util::once_cell::OnceCell;

pub(crate) struct StaticAtomicU64 {
    init: u64,
    cell: OnceCell<Mutex<u64>>,
}

impl AtomicU64 {
    pub(crate) fn new(val: u64) -> Self {
        Self {
            inner: Mutex::new(val),
        }
    }
}

impl StaticAtomicU64 {
    pub(crate) const fn new(val: u64) -> StaticAtomicU64 {
        StaticAtomicU64 {
            init: val,
            cell: OnceCell::new(),
        }
    }

    pub(crate) fn fetch_add(&self, val: u64, order: Ordering) -> u64 {
        let mut lock = self.inner().lock();
        let prev = *lock;
        *lock = prev + val;
        prev
    }

Review comment (suggestion):
It is a bit annoying to store the initial value forever, but this is *only* for platforms that do not have `AtomicU64` **and** no `const fn Mutex::new()` (i.e. old versions of Rust), so it isn't a big deal IMO.

Code after:
use super::AtomicU64;
use crate::loom::sync::{atomic::Ordering, Mutex};
use crate::util::once_cell::OnceCell;

pub(crate) struct StaticAtomicU64 {
    init: u64,
    cell: OnceCell<Mutex<u64>>,
}

impl AtomicU64 {
    pub(crate) fn new(val: u64) -> Self {
        Self {
            inner: Mutex::new(val),
        }
    }
}

impl StaticAtomicU64 {
    pub(crate) const fn new(val: u64) -> StaticAtomicU64 {
        StaticAtomicU64 {
            init: val,
            cell: OnceCell::new(),
        }
    }

    pub(crate) fn fetch_add(&self, val: u64, order: Ordering) -> u64 {
        let mut lock = self.inner().lock();
        let prev = *lock;
        *lock = prev + val;
        prev
    }
Repository: tokio-rs/tokio
File: tokio/src/sync/mpsc/block.rs (Rust)

Code before:
            crate::loom::thread::yield_now();
        }
    }
}

/// Returns `true` if the specified slot has a value ready to be consumed.
fn is_ready(bits: usize, slot: usize) -> bool {
    let mask = 1 << slot;
    mask == mask & bits
}

/// Returns `true` if the closed flag has been set.
fn is_tx_closed(bits: usize) -> bool {
    TX_CLOSED == bits & TX_CLOSED
}

impl<T> Values<T> {
    unsafe fn initialize(_value: NonNull<Values<T>>) {
        // When fuzzing, `UnsafeCell` needs to be initialized.
        if_loom! {
            let p = _value.as_ptr() as *mut UnsafeCell<MaybeUninit<T>>;
            for i in 0..BLOCK_CAP {
                p.add(i)
                    .write(UnsafeCell::new(MaybeUninit::uninit()));
            }
        }
    }
}

impl<T> ops::Index<usize> for Values<T> {
    type Output = UnsafeCell<MaybeUninit<T>>;

    fn index(&self, index: usize) -> &Self::Output {
        self.0.index(index)
    }
}

#[cfg(all(test, not(loom)))]
#[test]
fn assert_no_stack_overflow() {
    // https://github.com/tokio-rs/tokio/issues/5293

    struct Foo {
        _a: [u8; 2_000_000],
    }

    assert_eq!(
        Layout::new::<MaybeUninit<Block<Foo>>>(),
        Layout::new::<Block<Foo>>()
    );


Review comment (suggestion):
```suggestion
    /// # Safety
    /// The raw pointer must be valid for writing a `Values<T>`.
    unsafe fn initialize(_value: NonNull<Values<T>>) {
```

Code after:
                Err(curr) => curr,
            };

            crate::loom::thread::yield_now();
        }
    }
}

/// Returns `true` if the specified slot has a value ready to be consumed.
fn is_ready(bits: usize, slot: usize) -> bool {
    let mask = 1 << slot;
    mask == mask & bits
}

/// Returns `true` if the closed flag has been set.
fn is_tx_closed(bits: usize) -> bool {
    TX_CLOSED == bits & TX_CLOSED
}

impl<T> Values<T> {
    /// Initialize a `Values` struct from a pointer.
    ///
    /// # Safety
    ///
    /// The raw pointer must be valid for writing a `Values<T>`.
    unsafe fn initialize(_value: NonNull<Values<T>>) {
        // When fuzzing, `UnsafeCell` needs to be initialized.
        if_loom! {
            let p = _value.as_ptr() as *mut UnsafeCell<MaybeUninit<T>>;
            for i in 0..BLOCK_CAP {
                p.add(i)
                    .write(UnsafeCell::new(MaybeUninit::uninit()));
            }
        }
    }
}

impl<T> ops::Index<usize> for Values<T> {
    type Output = UnsafeCell<MaybeUninit<T>>;

    fn index(&self, index: usize) -> &Self::Output {
        self.0.index(index)
    }
}

#[cfg(all(test, not(loom)))]
#[test]
fn assert_no_stack_overflow() {
    // https://github.com/tokio-rs/tokio/issues/5293

    struct Foo {
Repository: tokio-rs/tokio
File: tokio/src/runtime/io/registration.rs (Rust)

Code before:
        // cycle would remain.
        //
        // See tokio-rs/tokio#3481 for more details.
        self.shared.clear_wakers();
    }
}

fn gone() -> io::Error {
    io::Error::new(
        io::ErrorKind::Other,
        crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR,
    )
}

cfg_io_readiness! {
    impl Registration {
        pub(crate) async fn readiness(&self, interest: Interest) -> io::Result<ReadyEvent> {
            use std::future::Future;
            use std::pin::Pin;

            let fut = self.shared.readiness(interest);
            pin!(fut);

            let ev = crate::future::poll_fn(|cx| {
                Pin::new(&mut fut).poll(cx)
            }).await;

            if ev.is_shutdown {
                return Err(gone())
            }

            Ok(ev)
        }

        pub(crate) async fn async_io<R>(&self, interest: Interest, mut f: impl FnMut() -> io::Result<R>) -> io::Result<R> {
            loop {
                let event = self.readiness(interest).await?;

                match f() {
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        self.clear_readiness(event);
                    }
                    x => return x,
                }
            }
        }
    }
}


Review comment (suggestion):
```Rust
let ev = self.shared.readiness(interest).await;
```

Code after:
        // cycle would remain.
        //
        // See tokio-rs/tokio#3481 for more details.
        self.shared.clear_wakers();
    }
}

fn gone() -> io::Error {
    io::Error::new(
        io::ErrorKind::Other,
        crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR,
    )
}

cfg_io_readiness! {
    impl Registration {
        pub(crate) async fn readiness(&self, interest: Interest) -> io::Result<ReadyEvent> {
            let ev = self.shared.readiness(interest).await;

            if ev.is_shutdown {
                return Err(gone())
            }

            Ok(ev)
        }

        pub(crate) async fn async_io<R>(&self, interest: Interest, mut f: impl FnMut() -> io::Result<R>) -> io::Result<R> {
            loop {
                let event = self.readiness(interest).await?;

                match f() {
                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                        self.clear_readiness(event);
                    }
                    x => return x,
                }
            }
        }
    }
}
Repository: tokio-rs/tokio
File: tokio/tests/sync_broadcast.rs (Rust)

Code before:
    assert_recv!(rx1);

    assert_eq!(tx.len(), 3);
    assert!(!tx.is_empty());

    assert_recv!(rx2);

    assert_eq!(tx.len(), 2);
    assert!(!tx.is_empty());

    tx.send(4).unwrap();
    tx.send(5).unwrap();
    tx.send(6).unwrap();

    assert_eq!(tx.len(), 4);
    assert!(!tx.is_empty());
}

#[test]
#[cfg(not(tokio_wasm_not_wasi))]
fn sender_len_random() {
    let (tx, mut rx1) = broadcast::channel(16);
    let mut rx2 = tx.subscribe();

    for _ in 0..1000 {
        match rand::thread_rng().gen_range(0..3) {
            0 => {
                let _ = rx1.try_recv();
            }
            1 => {
                let _ = rx2.try_recv();
            }
            _ => {
                tx.send(0).unwrap();
            }
        }

        let expected_len = usize::min(usize::max(rx1.len(), rx2.len()), 16);
        assert_eq!(tx.len(), expected_len);
    }
}


Review comment (suggestion):
We should probably give recv and send the same probability.
```suggestion
        match rand::thread_rng().gen_range(0..4) {
```

Code after:
    assert_eq!(tx.len(), 3);
    assert!(!tx.is_empty());

    assert_recv!(rx2);

    assert_eq!(tx.len(), 2);
    assert!(!tx.is_empty());

    tx.send(4).unwrap();
    tx.send(5).unwrap();
    tx.send(6).unwrap();

    assert_eq!(tx.len(), 4);
    assert!(!tx.is_empty());
}

#[test]
#[cfg(not(tokio_wasm_not_wasi))]
fn sender_len_random() {
    use rand::Rng;

    let (tx, mut rx1) = broadcast::channel(16);
    let mut rx2 = tx.subscribe();

    for _ in 0..1000 {
        match rand::thread_rng().gen_range(0..4) {
            0 => {
                let _ = rx1.try_recv();
            }
            1 => {
                let _ = rx2.try_recv();
            }
            _ => {
                tx.send(0).unwrap();
            }
        }

        let expected_len = usize::min(usize::max(rx1.len(), rx2.len()), 16);
        assert_eq!(tx.len(), expected_len);
    }
}
Repository: tokio-rs/tokio
File: tokio/src/sync/broadcast.rs (Rust)

Code before:
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, _rx) = broadcast::channel(16);
    ///
    ///     // Will not be seen
    ///     tx.send(10).unwrap();
    ///
    ///     let mut rx = tx.subscribe();
    ///
    ///     tx.send(20).unwrap();
    ///
    ///     let value = rx.recv().await.unwrap();
    ///     assert_eq!(20, value);
    /// }
    /// ```
    pub fn subscribe(&self) -> Receiver<T> {
        let shared = self.shared.clone();
        new_receiver(shared)
    }

    /// Returns the number of queued values.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, mut rx1) = broadcast::channel(16);
    ///     let mut rx2 = tx.subscribe();
    ///
    ///     tx.send(10).unwrap();
    ///     tx.send(20).unwrap();
    ///     tx.send(30).unwrap();
    ///
    ///     assert_eq!(tx.len(), 3);
    ///
    ///     rx1.recv().await.unwrap();
    ///
    ///     // The len is still 3 since rx2 hasn't seen the first value yet.
    ///     assert_eq!(tx.len(), 3);
    ///
    ///     rx2.recv().await.unwrap();
    ///
    ///     assert_eq!(tx.len(), 2);

Review comment (suggestion):
The documentation should explain that this is the number of messages that at least one channel have not yet received. (And also note that lagged messages are not counted, since the receiver len includes those.)

Code after:
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, _rx) = broadcast::channel(16);
    ///
    ///     // Will not be seen
    ///     tx.send(10).unwrap();
    ///
    ///     let mut rx = tx.subscribe();
    ///
    ///     tx.send(20).unwrap();
    ///
    ///     let value = rx.recv().await.unwrap();
    ///     assert_eq!(20, value);
    /// }
    /// ```
    pub fn subscribe(&self) -> Receiver<T> {
        let shared = self.shared.clone();
        new_receiver(shared)
    }

    /// Returns the number of queued values.
    ///
    /// A value is queued until it has either been seen by all receivers that were alive at the time
    /// it was sent, or has been evicted from the queue by subsequent sends that exceeded the
    /// queue's capacity.
    ///
    /// # Note
    ///
    /// In contrast to [`Receiver::len`], this method only reports queued values and not values that
    /// have been evicted from the queue before being seen by all receivers.
    ///
    /// # Examples
    ///
    /// ```
    /// use tokio::sync::broadcast;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let (tx, mut rx1) = broadcast::channel(16);
    ///     let mut rx2 = tx.subscribe();
    ///
    ///     tx.send(10).unwrap();
    ///     tx.send(20).unwrap();
    ///     tx.send(30).unwrap();
    ///
    ///     assert_eq!(tx.len(), 3);
Repository: tokio-rs/tokio
File: tokio-stream/tests/watch.rs (Rust)

Code before:
    let mut stream = WatchStream::from(rx);

    assert_eq!(stream.next().await.unwrap(), "hello");

    tx.send("bye").unwrap();

    assert_eq!(stream.next().await.unwrap(), "bye");
}

#[tokio::test]
async fn watch_stream_new_on_change() {
    let (tx, rx) = watch::channel("hello");

    let mut stream = WatchStream::new_on_changed(rx);

    let task = tokio::spawn(async move {
        let value = stream.next().await.unwrap();
        if value != "bye" {
            panic!("received unexpected value: {:?}", value);
        }
    });

    tx.send("bye").unwrap();

    task.await.unwrap();
}


Review comment (suggestion):
Since tests run on the current-thread runtime, this test is actually _guaranteed_ to call `send(bye)` before it calls `stream.next()`. This means that the test would also pass when using `new`.

Please use our tokio-test utilities to instead assert that `next` returns pending before the send, and then assert that it returns `"bye"` after the send.

Code after:
#![cfg(feature = "sync")]

use tokio::sync::watch;
use tokio_stream::wrappers::WatchStream;
use tokio_stream::StreamExt;
use tokio_test::assert_pending;
use tokio_test::task::spawn;

#[tokio::test]
async fn watch_stream_message_not_twice() {
    let (tx, rx) = watch::channel("hello");

    let mut counter = 0;
    let mut stream = WatchStream::new(rx).map(move |payload| {
        println!("{}", payload);
        if payload == "goodbye" {
            counter += 1;
        }
        if counter >= 2 {
            panic!("too many goodbyes");
        }
    });

    let task = tokio::spawn(async move { while stream.next().await.is_some() {} });

    // Send goodbye just once
    tx.send("goodbye").unwrap();

    drop(tx);
    task.await.unwrap();
}

#[tokio::test]
async fn watch_stream_from_rx() {
    let (tx, rx) = watch::channel("hello");

    let mut stream = WatchStream::from(rx);

    assert_eq!(stream.next().await.unwrap(), "hello");

    tx.send("bye").unwrap();

    assert_eq!(stream.next().await.unwrap(), "bye");
}

#[tokio::test]
async fn watch_stream_from_changes() {
    let (tx, rx) = watch::channel("hello");

    let mut stream = WatchStream::from_changes(rx);

    assert_pending!(spawn(&mut stream).poll_next());

    tx.send("bye").unwrap();

    assert_eq!(stream.next().await.unwrap(), "bye");
}