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

Code before:
                .register_by_ref(cx.waker());

            let _no_blocking = crate::runtime::context::disallow_block_in_place();
            let f = me.future;

            if let Poll::Ready(output) = f.poll(cx) {
                return Poll::Ready(output);
            }

            if me.local_set.tick() {
                // If `tick` returns `true`, we need to notify the local future again:
                // there are still tasks remaining in the run queue.
                cx.waker().wake_by_ref();
            }

            Poll::Pending
        })
    }
}

impl Shared {
    /// Schedule the provided task on the scheduler.
    fn schedule(&self, task: task::Notified<Arc<Self>>) {
        CURRENT.with(|localdata| {
            match localdata.ctx.get() {
                Some(cx) if cx.shared.ptr_eq(self) && !localdata.entered.get() => unsafe {
                    // Safety: if the current `LocalSet` context points to this
                    // `LocalSet`, then we are on the thread that owns it.
                    cx.shared.local_state.task_push_back(task);
                },

                // We are on the thread that owns the `LocalSet`, so we can
                // wake to the local queue.
                _ if context::thread_id().ok() == Some(self.local_state.owner) => {
                    unsafe {
                        // Safety: we just checked that the thread ID matches
                        // the localset's owner, so this is safe.
                        self.local_state.task_push_back(task);
                    }
                    // We still have to wake the `LocalSet`, because it isn't
                    // currently being polled.
                    self.waker.wake();
                }

                // We are *not* on the thread that owns the `LocalSet`, so we
                // have to wake to the remote queue.
                _ => {
                    // First, check whether the queue is still there (if not, the
                    // LocalSet is dropped). Then push to it if so, and if not,
                    // do nothing.
                    let mut lock = self.queue.lock();

Review comment (suggestion):
it might be worth adding a comment here explaining why we only choose this branch if `entered` is false? although, if we expand the doc comment on `entered`, that could also be fine. really, i just think there ought to be an explicit comment somewhere stating that the `LocalSet` is woken if `entered` is true.

Code after:

        handle
    }
}

// === impl LocalFuture ===

impl<T: Future> Future for RunUntil<'_, T> {
    type Output = T::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
        let me = self.project();

        me.local_set.with(|| {
            me.local_set
                .context
                .shared
                .waker
                .register_by_ref(cx.waker());

            let _no_blocking = crate::runtime::context::disallow_block_in_place();
            let f = me.future;

            if let Poll::Ready(output) = f.poll(cx) {
                return Poll::Ready(output);
            }

            if me.local_set.tick() {
                // If `tick` returns `true`, we need to notify the local future again:
                // there are still tasks remaining in the run queue.
                cx.waker().wake_by_ref();
            }

            Poll::Pending
        })
    }
}

impl Shared {
    /// Schedule the provided task on the scheduler.
    fn schedule(&self, task: task::Notified<Arc<Self>>) {
        CURRENT.with(|localdata| {
            match localdata.ctx.get() {
                // If the current `LocalSet` is being polled, we don't need to wake it.
                // When we `enter` it, then the value `wake_on_schedule` is set to be true.
                // In this case it is not being polled, so we need to wake it.
                Some(cx) if cx.shared.ptr_eq(self) && !localdata.wake_on_schedule.get() => unsafe {
                    // Safety: if the current `LocalSet` context points to this
                    // `LocalSet`, then we are on the thread that owns it.
                    cx.shared.local_state.task_push_back(task);
                },