ronantakizawa/codereview-bench huggingface dataset:
Repository: tokio-rs/tokio
File: tokio/src/runtime/blocking/schedule.rs (Rust)

Code before:
use crate::runtime::task::{self, Task};

/// `task::Schedule` implementation that does nothing (except some bookkeeping
/// in test-util builds). This is unique to the blocking scheduler as tasks
/// scheduled are not really futures but blocking operations.
///
/// We avoid storing the task by forgetting it in `bind` and re-materializing it
/// in `release`.
pub(crate) struct BlockingSchedule;

impl task::Schedule for BlockingSchedule {
    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
        #[cfg(feature = "test-util")]
        crate::time::allow_auto_advance();
        None
    }

    fn schedule(&self, _task: task::Notified<Self>) {
        unreachable!();
    }
}


Review comment (suggestion):
I would prefer to have it store a reference to the clock.
```Rust
pub(crate) struct BlockingSchedule {
    #[cfg(feature = "test-util")]
    clock: Clock,
}

impl task::Schedule for BlockingSchedule {
    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
        #[cfg(feature = "test-util")]
        self.clock.allow_auto_advance();
        None
    }
}
```

Code after:
#[cfg(feature = "test-util")]
use crate::runtime::scheduler;
use crate::runtime::task::{self, Task};
use crate::runtime::Handle;

/// `task::Schedule` implementation that does nothing (except some bookkeeping
/// in test-util builds). This is unique to the blocking scheduler as tasks
/// scheduled are not really futures but blocking operations.
///
/// We avoid storing the task by forgetting it in `bind` and re-materializing it
/// in `release`.
pub(crate) struct BlockingSchedule {
    #[cfg(feature = "test-util")]
    handle: Handle,
}

impl BlockingSchedule {
    #[cfg_attr(not(feature = "test-util"), allow(unused_variables))]
    pub(crate) fn new(handle: &Handle) -> Self {
        #[cfg(feature = "test-util")]
        {
            match &handle.inner {
                scheduler::Handle::CurrentThread(handle) => {
                    handle.driver.clock.inhibit_auto_advance();
                }
                #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
                scheduler::Handle::MultiThread(_) => {}
            }
        }
        BlockingSchedule {
            #[cfg(feature = "test-util")]
            handle: handle.clone(),
        }
    }
}

impl task::Schedule for BlockingSchedule {
    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
        #[cfg(feature = "test-util")]
Repository: tokio-rs/tokio
File: tokio-test/src/task.rs (Rust)

Code before:
        self.task.enter(|cx| fut.poll(cx))
    }
}

impl<T: Stream> Spawn<T> {
    /// Polls a stream
    pub fn poll_next(&mut self) -> Poll<Option<T::Item>> {
        let stream = self.future.as_mut();
        self.task.enter(|cx| stream.poll_next(cx))
    }
}

impl<T: Future> Future for Spawn<T> {
    type Output = T::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.future.as_mut().poll(cx)
    }
}

impl<T: Stream> Stream for Spawn<T> {
    type Item = T::Item;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.future.as_mut().poll_next(cx)
    }
}

impl MockTask {
    /// Creates new mock task
    fn new() -> Self {
        MockTask {
            waker: Arc::new(ThreadWaker::new()),
        }
    }

    /// Runs a closure from the context of the task.
    ///
    /// Any wake notifications resulting from the execution of the closure are
    /// tracked.
    fn enter<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut Context<'_>) -> R,
    {
        self.waker.clear();
        let waker = self.waker();
        let mut cx = Context::from_waker(&waker);

        f(&mut cx)
    }


Review comment (suggestion):
a little misleading—kinda makes me think that the code would be

`if t.is_future() { poll } else { do nothing }`

Code after:
    pub fn enter<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut Context<'_>, Pin<&mut T>) -> R,
    {
        let fut = self.future.as_mut();
        self.task.enter(|cx| f(cx, fut))
    }
}

impl<T: Unpin> ops::Deref for Spawn<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.future
    }
}

impl<T: Unpin> ops::DerefMut for Spawn<T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.future
    }
}

impl<T: Future> Spawn<T> {
    /// If `T` is a [`Future`] then poll it. This will handle pinning and the context
    /// type for the future.
    pub fn poll(&mut self) -> Poll<T::Output> {
        let fut = self.future.as_mut();
        self.task.enter(|cx| fut.poll(cx))
    }
}

impl<T: Stream> Spawn<T> {
    /// If `T` is a [`Stream`] then poll_next it. This will handle pinning and the context
    /// type for the stream.
    pub fn poll_next(&mut self) -> Poll<Option<T::Item>> {
        let stream = self.future.as_mut();
        self.task.enter(|cx| stream.poll_next(cx))
    }
}

impl<T: Future> Future for Spawn<T> {
    type Output = T::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.future.as_mut().poll(cx)
    }
}

impl<T: Stream> Stream for Spawn<T> {
    type Item = T::Item;
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/harness.rs (Rust)

Code before:
            core: &'a CoreStage<T>,
        }
        impl<'a, T: Future> Drop for Guard<'a, T> {
            fn drop(&mut self) {
                // If the future panics on poll, we drop it inside the panic
                // guard.
                self.core.drop_future_or_output();
            }
        }
        struct TaskIdGuard {}
        impl TaskIdGuard {
            fn new(id: Id) -> Self {
                context::set_current_task_id(Some(id));
                TaskIdGuard {}
            }
        }
        impl Drop for TaskIdGuard {
            fn drop(&mut self) {
                context::set_current_task_id(None);
            }
        }
        let _task_id_guard = TaskIdGuard::new(id);
        let guard = Guard { core };
        let res = guard.core.poll(cx);
        mem::forget(guard);
        res
    }));

    // Prepare output for being placed in the core stage.
    let output = match output {
        Ok(Poll::Pending) => return Poll::Pending,
        Ok(Poll::Ready(output)) => Ok(output),
        Err(panic) => {
            scheduler.unhandled_panic();
            Err(JoinError::panic(id, panic))
        }
    };

    // Catch and ignore panics if the future panics on drop.
    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        core.store_output(output);
    }));

    if res.is_err() {
        scheduler.unhandled_panic();
    }

    Poll::Ready(())
}


Review comment (refactor):
We might want to move this outside the `catch_unwind` so that the `store_output` call is also covered by the id, as this runs the future's destructor when the future completes normally.

Code after:
    let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        struct Guard<'a, T: Future, S: Schedule> {
            core: &'a Core<T, S>,
        }
        impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> {
            fn drop(&mut self) {
                // If the future panics on poll, we drop it inside the panic
                // guard.
                self.core.drop_future_or_output();
            }
        }
        let guard = Guard { core };
        let res = guard.core.poll(cx);
        mem::forget(guard);
        res
    }));

    // Prepare output for being placed in the core stage.
    let output = match output {
        Ok(Poll::Pending) => return Poll::Pending,
        Ok(Poll::Ready(output)) => Ok(output),
        Err(panic) => {
            core.scheduler.unhandled_panic();
            Err(JoinError::panic(core.task_id, panic))
        }
    };

    // Catch and ignore panics if the future panics on drop.
    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
        core.store_output(output);
    }));

    if res.is_err() {
        core.scheduler.unhandled_panic();
    }

    Poll::Ready(())
}
Repository: tokio-rs/tokio
File: tokio/tests/task_local.rs (Rust)

Code before:

    struct MyFuture {
        tx: Option<oneshot::Sender<u32>>,
    }
    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Ready(())
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(KEY.get());
        }
    }

    let (tx, rx) = oneshot::channel();

    let h = tokio::spawn(KEY.scope(42, MyFuture { tx: Some(tx) }));

    assert_eq!(rx.await.unwrap(), 42);
    h.await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn task_id() {
    use tokio::task;

    let handle = tokio::spawn(async { println!("task id: {}", task::id()) });

    handle.await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_collision_multi_thread() {
    use tokio::task;

    let handle1 = tokio::spawn(async { task::id() });
    let handle2 = tokio::spawn(async { task::id() });

    let (id1, id2) = tokio::join!(handle1, handle2);
    assert_ne!(id1.unwrap(), id2.unwrap());
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_id_collision_current_thread() {
    use tokio::task;


Review comment (suggestion):
Please add the following tests:

 * Can get id from `spawn_blocking` task.
 * Can get id from `spawn_local` task on a `LocalSet`.
 * Can't get id in `block_on` call.
 * Can get id in future destructor when it completes normally.
 * Can get id in future destructor when it is aborted with `JoinHandle::abort`.
 * Can get id in destructor of output when `JoinHandle` dropped before task completion.
 * Can get id in destructor of output when `JoinHandle` dropped after task completion.
 * If you use a `LocalSet` to run a local task within another task, the original task's id is not erased once we return.
 * If you use `block_in_place` + `spawn` + `block_on` on a current-thread runtime to execute a spawned task within another task, the original task's id is not erased.

Code after:
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<u32>>,
    }
    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Ready(())
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(KEY.get());
        }
    }

    let (tx, rx) = oneshot::channel();

    let h = tokio::spawn(KEY.scope(42, MyFuture { tx: Some(tx) }));

    assert_eq!(rx.await.unwrap(), 42);
    h.await.unwrap();
}
Repository: tokio-rs/tokio
File: tokio/src/task/mod.rs (Rust)

Code before:
        pub use consume_budget::consume_budget;
    }

    mod local;
    pub use local::{spawn_local, LocalSet, LocalEnterGuard};

    mod task_local;
    pub use task_local::LocalKey;

    mod unconstrained;
    pub use unconstrained::{unconstrained, Unconstrained};

    #[doc(inline)]
    pub use join_set::JoinSet;
    pub use crate::runtime::task::AbortHandle;

    // Uses #[cfg(...)] instead of macro since the macro adds docsrs annotations.
    #[cfg(not(tokio_unstable))]
    mod join_set;
    #[cfg(tokio_unstable)]
    pub mod join_set;

    cfg_unstable! {
        pub use crate::runtime::task::Id;
        pub use crate::runtime::task::id;
        pub use crate::runtime::task::try_id;
    }

    cfg_trace! {
        mod builder;
        pub use builder::Builder;
    }

    /// Task-related futures.
    pub mod futures {
        pub use super::task_local::TaskLocalFuture;
    }
}


Review comment (suggestion):
```suggestion
        pub use crate::runtime::task::{Id, id, try_id};
```

Code after:
        pub use consume_budget::consume_budget;
    }

    mod local;
    pub use local::{spawn_local, LocalSet, LocalEnterGuard};

    mod task_local;
    pub use task_local::LocalKey;

    mod unconstrained;
    pub use unconstrained::{unconstrained, Unconstrained};

    #[doc(inline)]
    pub use join_set::JoinSet;
    pub use crate::runtime::task::AbortHandle;

    // Uses #[cfg(...)] instead of macro since the macro adds docsrs annotations.
    #[cfg(not(tokio_unstable))]
    mod join_set;
    #[cfg(tokio_unstable)]
    pub mod join_set;

    cfg_unstable! {
        pub use crate::runtime::task::{Id, id, try_id};
    }

    cfg_trace! {
        mod builder;
        pub use builder::Builder;
    }

    /// Task-related futures.
    pub mod futures {
        pub use super::task_local::TaskLocalFuture;
    }
}
Repository: tokio-rs/tokio
File: tokio/tests/task_id.rs (Rust)

Code before:
#![warn(rust_2018_idioms)]
#![allow(clippy::declare_interior_mutable_const)]
#![cfg(all(feature = "full", not(tokio_wasi)))]

#[cfg(tokio_unstable)]
use std::error::Error;
#[cfg(tokio_unstable)]
use std::future::Future;
#[cfg(tokio_unstable)]
use std::pin::Pin;
#[cfg(tokio_unstable)]
use std::task::{Context, Poll};
#[cfg(tokio_unstable)]
use tokio::runtime::{Builder, Runtime};
#[cfg(tokio_unstable)]
use tokio::sync::oneshot;
#[cfg(tokio_unstable)]
use tokio::task::{self, LocalSet};

#[cfg(tokio_unstable)]
mod support {
    pub mod panic;
}
#[cfg(tokio_unstable)]
use support::panic::test_panic;

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]

Review comment (suggestion):
If you do this, then you can remove all of the other `#[cfg(tokio_unstable)]` annotations.
```suggestion
#![cfg(all(feature = "full", tokio_unstable, not(tokio_wasi)))]
```

Code after:
#![warn(rust_2018_idioms)]
#![allow(clippy::declare_interior_mutable_const)]
#![cfg(all(feature = "full", tokio_unstable))]

#[cfg(not(tokio_wasi))]
use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
#[cfg(not(tokio_wasi))]
use tokio::runtime::{Builder, Runtime};
use tokio::sync::oneshot;
use tokio::task::{self, Id, LocalSet};

#[cfg(not(tokio_wasi))]
mod support {
    pub mod panic;
}
#[cfg(not(tokio_wasi))]
use support::panic::test_panic;

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

Code before:
    });
    tx.send(handle.id()).unwrap();
    handle.await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_completion() {
    struct MyFuture;

    impl Future for MyFuture {
        type Output = ();

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

    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    tokio::spawn(MyFuture).await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_abort() {
    struct MyFuture;

    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    tokio::spawn(MyFuture).abort();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_before_completion() {

Review comment (bug):
I'm not sure that this test will fail even if a panic happens. You should send a notification back to the main function and assert the success there.

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

    impl Future for MyFuture {
        type Output = ();

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

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

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

#[cfg(not(tokio_wasi))]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_abort() {
    struct MyFuture {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    let (tx, rx) = oneshot::channel();
    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);
Repository: tokio-rs/tokio
File: tokio/tests/task_id.rs (Rust)

Code before:
#![warn(rust_2018_idioms)]
#![allow(clippy::declare_interior_mutable_const)]
#![cfg(all(feature = "full", tokio_unstable, not(tokio_wasi)))]

use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::runtime::{Builder, Runtime};
use tokio::sync::oneshot;
use tokio::task::{self, Id, LocalSet};

mod support {
    pub mod panic;
}
use support::panic::test_panic;

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn() {
    tokio::spawn(async { println!("task id: {}", task::id()) })
        .await
        .unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn_blocking() {
    task::spawn_blocking(|| println!("task id: {}", task::id()))
        .await

Review comment (suggestion):
Except for the panic-location tests, the tests here should work under wasm. You can disable specific tests under wasm with `#[cfg_attr(tokio_wasm, ignore)]`.
```suggestion
#![cfg(all(feature = "full", tokio_unstable))]
```

Code after:
#![warn(rust_2018_idioms)]
#![allow(clippy::declare_interior_mutable_const)]
#![cfg(all(feature = "full", tokio_unstable))]

#[cfg(not(tokio_wasi))]
use std::error::Error;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
#[cfg(not(tokio_wasi))]
use tokio::runtime::{Builder, Runtime};
use tokio::sync::oneshot;
use tokio::task::{self, Id, LocalSet};

#[cfg(not(tokio_wasi))]
mod support {
    pub mod panic;
}
#[cfg(not(tokio_wasi))]
use support::panic::test_panic;

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

Code before:
}

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

    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    let id = handle.id();
    handle.abort();
    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() })
        }
    }

Review comment (suggestion):
```suggestion
    handle.abort();
    assert!(handle.await.unwrap_err().is_cancelled());
    assert_eq!(rx.await.unwrap(), id);
```

Code after:
        }
    }

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

#[cfg(not(tokio_wasi))]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_abort() {
    struct MyFuture {
        tx: Option<oneshot::Sender<Id>>,
    }

    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(task::id());
        }
    }

    let (tx, rx) = oneshot::channel();
    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 {