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

Code before:
/// Tasks are pinned.
unsafe impl<S> linked_list::Link for Task<S> {
    type Handle = Task<S>;
    type Target = Header;

    fn as_raw(handle: &Task<S>) -> NonNull<Header> {
        handle.raw.header_ptr()
    }

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::from_raw(ptr)
    }

    unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
        self::core::Trailer::addr_of_owned(Header::get_trailer(target))
    }
}

/// # Safety
///
/// Tasks are pinned.
unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        let task_id = unsafe { Header::get_id(target) };
        // safety: in 32-bits machine, the lower 32 bits will be reserved
        task_id.0 as usize
    }
}


Review comment (suggestion):
```suggestion
/// # Safety
///
/// The id of a task is never changed after creation of the task, so the return value of
/// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but
/// the shard still won't change from call to call.)
unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        // SAFETY: The caller guarantees that `target` points at a valid task.
        let task_id = unsafe { Header::get_id(target) };
        task_id.0 as usize
```


Code after:
/// Tasks are pinned.
unsafe impl<S> linked_list::Link for Task<S> {
    type Handle = Task<S>;
    type Target = Header;

    fn as_raw(handle: &Task<S>) -> NonNull<Header> {
        handle.raw.header_ptr()
    }

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::from_raw(ptr)
    }

    unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
        self::core::Trailer::addr_of_owned(Header::get_trailer(target))
    }
}

/// # Safety
///
/// The id of a task is never changed after creation of the task, so the return value of
/// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but
/// the shard id still won't change from call to call.)
unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        // SAFETY: The caller guarantees that `target` points at a valid task.
        let task_id = unsafe { Header::get_id(target) };
        task_id.0 as usize
    }
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/mod.rs (Rust)

Code before:

/// # Safety
///
/// Tasks are pinned.
unsafe impl<S> linked_list::Link for Task<S> {
    type Handle = Task<S>;
    type Target = Header;

    fn as_raw(handle: &Task<S>) -> NonNull<Header> {
        handle.raw.header_ptr()
    }

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::from_raw(ptr)
    }

    unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
        self::core::Trailer::addr_of_owned(Header::get_trailer(target))
    }
}

/// # Safety
///
/// The id of a task is never changed after creation of the task, so the return value of
/// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but
/// the shard still won't change from call to call.)
unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        // SAFETY: The caller guarantees that `target` points at a valid task.
        let task_id = unsafe { Header::get_id(target) };
        task_id.0 as usize
    }
}


Review comment (suggestion):
```suggestion
/// the shard id still won't change from call to call.)
```

Code after:

/// # Safety
///
/// Tasks are pinned.
unsafe impl<S> linked_list::Link for Task<S> {
    type Handle = Task<S>;
    type Target = Header;

    fn as_raw(handle: &Task<S>) -> NonNull<Header> {
        handle.raw.header_ptr()
    }

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::from_raw(ptr)
    }

    unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
        self::core::Trailer::addr_of_owned(Header::get_trailer(target))
    }
}

/// # Safety
///
/// The id of a task is never changed after creation of the task, so the return value of
/// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but
/// the shard id still won't change from call to call.)
unsafe impl<S> sharded_list::ShardedListItem for Task<S> {
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize {
        // SAFETY: The caller guarantees that `target` points at a valid task.
        let task_id = unsafe { Header::get_id(target) };
        task_id.0 as usize
    }
}
Repository: tokio-rs/tokio
File: tokio/src/task/local.rs (Rust)

Code before:
    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) {
            None => panic!("`spawn_local` called from outside of a `task::LocalSet`"),
            Some(cx) => cx.spawn(future, name)
       }
    }
}

/// Initial queue capacity.
const INITIAL_CAPACITY: usize = 64;

/// Max number of tasks to poll per tick.
const MAX_TASKS_PER_TICK: usize = 61;

/// How often it check the remote queue first.
const REMOTE_FIRST_INTERVAL: u8 = 31;

/// Context guard for LocalSet
pub struct LocalEnterGuard {
    ctx: Option<Rc<Context>>,

    /// Distinguish whether the context was entered or polled.
    entered: bool,
}

impl Drop for LocalEnterGuard {
    fn drop(&mut self) {
        CURRENT.with(|LocalData { ctx, entered }| {
            ctx.set(self.ctx.take());
            entered.set(self.entered);
        })
    }
}

impl fmt::Debug for LocalEnterGuard {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LocalEnterGuard").finish()
    }
}

impl LocalSet {
    /// Returns a new local task set.
    pub fn new() -> LocalSet {
        let owner = context::thread_id().expect("cannot create LocalSet during thread shutdown");

        LocalSet {
            tick: Cell::new(0),

Review comment (suggestion):
```suggestion
    /// Distinguishes whether the context was entered or polled.
```

Code after:
    {
        spawn_local_inner(future, None)
    }


    #[track_caller]
    pub(super) fn spawn_local_inner<F>(future: F, name: Option<&str>) -> JoinHandle<F::Output>
    where F: Future + 'static,
          F::Output: 'static
    {
        match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) {
            None => panic!("`spawn_local` called from outside of a `task::LocalSet`"),
            Some(cx) => cx.spawn(future, name)
       }
    }
}

/// Initial queue capacity.
const INITIAL_CAPACITY: usize = 64;

/// Max number of tasks to poll per tick.
const MAX_TASKS_PER_TICK: usize = 61;

/// How often it check the remote queue first.
const REMOTE_FIRST_INTERVAL: u8 = 31;

/// Context guard for LocalSet
pub struct LocalEnterGuard {
    ctx: Option<Rc<Context>>,

    /// Distinguishes whether the context was entered or being polled.
    /// When we enter it, the value `wake_on_schedule` is set. In this case
    /// `spawn_local` refers the context, whereas it is not being polled now.
    wake_on_schedule: bool,
}

impl Drop for LocalEnterGuard {
    fn drop(&mut self) {
        CURRENT.with(
            |LocalData {
                 ctx,
                 wake_on_schedule,
             }| {
                ctx.set(self.ctx.take());
                wake_on_schedule.set(self.wake_on_schedule);
            },
        )
    }
}

impl fmt::Debug for LocalEnterGuard {
Repository: tokio-rs/tokio
File: tokio/tests/task_local_set.rs (Rust)

Code before:
        pending::<()>().await;
        tx1.send(()).unwrap();
    });

    // Tick the loop
    local
        .run_until(async {
            task::yield_now().await;
        })
        .await;

    // Drop the LocalSet
    drop(local);
}

#[tokio::test]
async fn spawn_wakes_localset() {
    let local = LocalSet::new();
    futures::select! {
        _ = local.run_until(pending::<()>()).fuse() => unreachable!(),
        ret = async { local.spawn_local(ready(())).await.unwrap()}.fuse() => ret
    }
}

/// Checks that the task wakes up with `enter`.
/// <https://github.com/tokio-rs/tokio/issues/5020>.
#[tokio::test]
async fn sleep_with_local_enter_guard() {
    let local = LocalSet::new();
    let _guard = local.enter();

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

    local
        .run_until(async move {
            tokio::task::spawn_local(async move {
                time::sleep(Duration::ZERO).await;

                tx.send(()).expect("failed to send");
            });
            assert_eq!(rx.await, Ok(()));
        })
        .await;
}

#[test]
fn store_local_set_in_thread_local_with_runtime() {
    use tokio::runtime::Runtime;

    thread_local! {
        static CURRENT: RtAndLocalSet = RtAndLocalSet::new();

Review comment (suggestion):
nit, not important: I might add:
```suggestion
/// Reproduces <https://github.com/tokio-rs/tokio/issues/5020>.
```

Code after:
        pending::<()>().await;
        tx1.send(()).unwrap();
    });

    // Tick the loop
    local
        .run_until(async {
            task::yield_now().await;
        })
        .await;

    // Drop the LocalSet
    drop(local);
}

#[tokio::test]
async fn spawn_wakes_localset() {
    let local = LocalSet::new();
    futures::select! {
        _ = local.run_until(pending::<()>()).fuse() => unreachable!(),
        ret = async { local.spawn_local(ready(())).await.unwrap()}.fuse() => ret
    }
}

/// Checks that the task wakes up with `enter`.
/// Reproduces <https://github.com/tokio-rs/tokio/issues/5020>.
#[tokio::test]
async fn sleep_with_local_enter_guard() {
    let local = LocalSet::new();
    let _guard = local.enter();

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

    local
        .run_until(async move {
            tokio::task::spawn_local(async move {
                time::sleep(Duration::ZERO).await;

                tx.send(()).expect("failed to send");
            });
            assert_eq!(rx.await, Ok(()));
        })
        .await;
}

#[test]
fn store_local_set_in_thread_local_with_runtime() {
    use tokio::runtime::Runtime;

    thread_local! {
        static CURRENT: RtAndLocalSet = RtAndLocalSet::new();
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/trace/mod.rs (Rust)

Code before:
        while let Some(notified) = injection.pop(&mut synced.inject) {
            drop(notified);
        }

        drop(synced);

        // notify each task
        let mut tasks = vec![];
        owned.for_each(|task| {
            // set the notified bit
            task.as_raw().state().transition_to_notified_for_tracing();
            // store the raw tasks into a vec
            tasks.push(task.as_raw());
            // do NOT poll `task` here, since we hold a lock on `owned` and the task
            // may complete and need to remove itself from `owned`.
        });

        tasks
            .into_iter()
            .map(|task| {
                // trace the task
                let ((), trace) = Trace::capture(|| task.poll());

                // reschedule the task
                let _ = task.state().transition_to_notified_by_ref();
                task.schedule();

                trace
            })
            .collect()
    }
}


Review comment (bug):
I'm a few months removed from when I last hacked on this, so my memory of Tokio's internals is a bit hazy. 

These two lines are now the only difference between how owned tasks are handled on the current thread scheduler, and how they're handled on the multi thread scheduler. The current thread scheduler doesn't have these two lines, and adding them causes Tokio to panic.

It *feels* wrong for them to be missing from the current-thread scheduler: We *need* to re-schedule tasks after tracing them, right? 

Any insights here?

Code after:
        // precondition: we have drained the tasks from the local and injection
        // queues.
        trace_owned(owned)
    }
}

/// Trace the `OwnedTasks`.
///
/// # Preconditions
///
/// This helper presumes exclusive access to each task. The tasks must not exist
/// in any other queue.
fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>) -> Vec<Trace> {
    // notify each task
    let mut tasks = vec![];
    owned.for_each(|task| {
        // notify the task (and thus make it poll-able) and stash it
        tasks.push(task.notify_for_tracing());
        // we do not poll it here since we hold a lock on `owned` and the task
        // may complete and need to remove itself from `owned`.
    });

    tasks
        .into_iter()
        .map(|task| {
            let local_notified = owned.assert_owner(task);
            let ((), trace) = Trace::capture(|| local_notified.run());
            trace
        })
        .collect()
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/trace/mod.rs (Rust)

Code before:
    pub(in crate::runtime) unsafe fn trace_multi_thread(
        owned: &OwnedTasks<Arc<multi_thread::Handle>>,
        local: &mut multi_thread::queue::Local<Arc<multi_thread::Handle>>,
        synced: &Mutex<Synced>,
        injection: &Shared<Arc<multi_thread::Handle>>,
    ) -> Vec<Trace> {
        // clear the local queue
        while let Some(notified) = local.pop() {
            drop(notified);
        }

        // clear the injection queue
        let mut synced = synced.lock();
        while let Some(notified) = injection.pop(&mut synced.inject) {
            drop(notified);
        }

        drop(synced);

        // notify each task
        let mut tasks = vec![];
        owned.for_each(|task| {
            // set the notified bit
            task.as_raw().state().transition_to_notified_for_tracing();
            // store the raw tasks into a vec
            tasks.push(task.as_raw());
            // do NOT poll `task` here, since we hold a lock on `owned` and the task
            // may complete and need to remove itself from `owned`.
        });

        tasks
            .into_iter()
            .map(|task| {
                // trace the task
                let ((), trace) = Trace::capture(|| task.poll());

                // reschedule the task
                let _ = task.state().transition_to_notified_by_ref();
                task.schedule();

                trace
            })
            .collect()
    }
}


Review comment (suggestion):
We should be taking a refcount to the task instead of using `as_raw`. 

Code after:
        // clear the local queue
        while let Some(notified) = local.pop() {
            drop(notified);
        }

        // clear the injection queue
        let mut synced = synced.lock();
        while let Some(notified) = injection.pop(&mut synced.inject) {
            drop(notified);
        }

        drop(synced);

        // precondition: we have drained the tasks from the local and injection
        // queues.
        trace_owned(owned)
    }
}

/// Trace the `OwnedTasks`.
///
/// # Preconditions
///
/// This helper presumes exclusive access to each task. The tasks must not exist
/// in any other queue.
fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>) -> Vec<Trace> {
    // notify each task
    let mut tasks = vec![];
    owned.for_each(|task| {
        // notify the task (and thus make it poll-able) and stash it
        tasks.push(task.notify_for_tracing());
        // we do not poll it here since we hold a lock on `owned` and the task
        // may complete and need to remove itself from `owned`.
    });

    tasks
        .into_iter()
        .map(|task| {
            let local_notified = owned.assert_owner(task);
            let ((), trace) = Trace::capture(|| local_notified.run());
            trace
        })
        .collect()
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/mod.rs (Rust)

Code before:
        Task::new(RawTask::from_raw(ptr))
    }

    #[cfg(all(
        tokio_unstable,
        tokio_taskdump,
        feature = "rt",
        target_os = "linux",
        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
    ))]
    pub(super) fn as_raw(&self) -> RawTask {
        self.raw
    }

    fn header(&self) -> &Header {
        self.raw.header()
    }

    fn header_ptr(&self) -> NonNull<Header> {
        self.raw.header_ptr()
    }

    cfg_taskdump! {
        pub(super) fn notify_for_tracing(&self) -> Notified<S> {
            self.as_raw().state().transition_to_notified_for_tracing();
            Notified(self.clone())
        }
    }
}

impl<S: 'static> Notified<S> {
    fn header(&self) -> &Header {
        self.0.header()
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
        Notified(Task::new(ptr))
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) fn into_raw(self) -> RawTask {
        let raw = self.0.raw;
        mem::forget(self);
        raw
    }
}

impl<S: Schedule> Task<S> {

Review comment (performance):
Since `transition_to_notified_for_tracing` already increments the refcount, you do not also need to clone it. You probably have a memory leak right now.
```suggestion
            Notified(self)
```

Code after:
        Task::new(RawTask::from_raw(ptr))
    }

    #[cfg(all(
        tokio_unstable,
        tokio_taskdump,
        feature = "rt",
        target_os = "linux",
        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
    ))]
    pub(super) fn as_raw(&self) -> RawTask {
        self.raw
    }

    fn header(&self) -> &Header {
        self.raw.header()
    }

    fn header_ptr(&self) -> NonNull<Header> {
        self.raw.header_ptr()
    }

    cfg_taskdump! {
        pub(super) fn notify_for_tracing(&self) -> Notified<S> {
            self.as_raw().state().transition_to_notified_for_tracing();
            // SAFETY: `transition_to_notified_for_tracing` increments the refcount.
            unsafe { Notified(Task::new(self.raw)) }
        }
    }
}

impl<S: 'static> Notified<S> {
    fn header(&self) -> &Header {
        self.0.header()
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
        Notified(Task::new(ptr))
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) fn into_raw(self) -> RawTask {
        let raw = self.0.raw;
        mem::forget(self);
        raw
    }
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/mod.rs (Rust)

Code before:

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::new(RawTask::from_raw(ptr))
    }

    #[cfg(all(
        tokio_unstable,
        tokio_taskdump,
        feature = "rt",
        target_os = "linux",
        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
    ))]
    pub(super) fn as_raw(&self) -> RawTask {
        self.raw
    }

    fn header(&self) -> &Header {
        self.raw.header()
    }

    fn header_ptr(&self) -> NonNull<Header> {
        self.raw.header_ptr()
    }

    cfg_taskdump! {
        pub(super) fn notify_for_tracing(self) -> Notified<S> {
            self.as_raw().state().transition_to_notified_for_tracing();
            Notified(self)
        }
    }
}

impl<S: 'static> Notified<S> {
    fn header(&self) -> &Header {
        self.0.header()
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
        Notified(Task::new(ptr))
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) fn into_raw(self) -> RawTask {
        let raw = self.0.raw;
        mem::forget(self);
        raw
    }
}

Review comment (suggestion):
You should be able to fix that like this:
```suggestion
        pub(super) fn notify_for_tracing(&self) -> Notified<S> {
```

Code after:

    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
        Task::new(RawTask::from_raw(ptr))
    }

    #[cfg(all(
        tokio_unstable,
        tokio_taskdump,
        feature = "rt",
        target_os = "linux",
        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
    ))]
    pub(super) fn as_raw(&self) -> RawTask {
        self.raw
    }

    fn header(&self) -> &Header {
        self.raw.header()
    }

    fn header_ptr(&self) -> NonNull<Header> {
        self.raw.header_ptr()
    }

    cfg_taskdump! {
        pub(super) fn notify_for_tracing(&self) -> Notified<S> {
            self.as_raw().state().transition_to_notified_for_tracing();
            // SAFETY: `transition_to_notified_for_tracing` increments the refcount.
            unsafe { Notified(Task::new(self.raw)) }
        }
    }
}

impl<S: 'static> Notified<S> {
    fn header(&self) -> &Header {
        self.0.header()
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
        Notified(Task::new(ptr))
    }
}

impl<S: 'static> Notified<S> {
    pub(crate) fn into_raw(self) -> RawTask {
        let raw = self.0.raw;
        mem::forget(self);
        raw
    }
Repository: tokio-rs/tokio
File: tokio/src/runtime/io/registration_set.rs (Rust)

Code before:
    }

    pub(super) fn release(&self, synced: &mut Synced) {
        for io in synced.pending_release.drain(..) {
            // safety: the registration is part of our list
            let _ = unsafe { synced.registrations.remove(io.as_ref().into()) };
        }

        self.num_pending_release.store(0, Release);
    }
}

// Safety: `Arc` pins the inner data
unsafe impl linked_list::Link for Arc<ScheduledIo> {
    type Handle = Arc<ScheduledIo>;
    type Target = ScheduledIo;

    fn as_raw(handle: &Self::Handle) -> NonNull<ScheduledIo> {
        // safety: Arc::as_ptr never returns null
        unsafe { NonNull::new_unchecked(Arc::as_ptr(handle) as *mut _) }
    }

    unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Arc<ScheduledIo> {
        // safety: the linked list currently owns a ref count
        Arc::from_raw(ptr.as_ptr() as *const _)
    }

    unsafe fn pointers(
        target: NonNull<Self::Target>,
    ) -> NonNull<linked_list::Pointers<ScheduledIo>> {
        NonNull::new_unchecked(target.as_ref().linked_list_pointers.get())
    }
}


Review comment (suggestion):
This is not a change I want. Rather, I agree with [the unsafe-block-in-unsafe-fn RFC][1]. I would rather enable `#[warn(unsafe_op_in_unsafe_fn)]` and fix anywhere that is missing an unsafe block inside unsafe functions.

[1]: https://rust-lang.github.io/rfcs/2585-unsafe-block-in-unsafe-fn.html

Code after:
    }

    pub(super) fn release(&self, synced: &mut Synced) {
        for io in synced.pending_release.drain(..) {
            // safety: the registration is part of our list
            let _ = unsafe { synced.registrations.remove(io.as_ref().into()) };
        }

        self.num_pending_release.store(0, Release);
    }
}

// Safety: `Arc` pins the inner data
unsafe impl linked_list::Link for Arc<ScheduledIo> {
    type Handle = Arc<ScheduledIo>;
    type Target = ScheduledIo;

    fn as_raw(handle: &Self::Handle) -> NonNull<ScheduledIo> {
        // safety: Arc::as_ptr never returns null
        unsafe { NonNull::new_unchecked(Arc::as_ptr(handle) as *mut _) }
    }

    unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Arc<ScheduledIo> {
        // safety: the linked list currently owns a ref count
        unsafe { Arc::from_raw(ptr.as_ptr() as *const _) }
    }

    unsafe fn pointers(
        target: NonNull<Self::Target>,
    ) -> NonNull<linked_list::Pointers<ScheduledIo>> {
        NonNull::new_unchecked(target.as_ref().linked_list_pointers.get())
    }
}