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

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

            let _no_blocking = crate::runtime::enter::disallow_blocking();
            let f = me.future;

            if let Poll::Ready(output) = crate::coop::budget(|| 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(|maybe_cx| match maybe_cx.borrow().as_ref() {
            Some(cx) if cx.shared.ptr_eq(self) => {
                cx.queue.push_back(task);
            }
            _ => {
                // 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();

                if let Some(queue) = lock.as_mut() {
                    queue.push_back(task);
                    drop(lock);
                    self.waker.wake();
                }
            }
        });
    }

    fn ptr_eq(&self, other: &Shared) -> bool {
        std::ptr::eq(self, other)
    }
}

impl task::Schedule for Arc<Shared> {
    fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
        CURRENT.with(|maybe_cx| match maybe_cx.borrow().as_ref() {
            None => panic!("scheduler context missing"),

Review comment (refactor):
Hmm, I may have changed my opinion on whether `Cell` is a good idea. Regardless, I am a bit worried about having us borrow the `RefCell` for too long. For example, in this case, the `RefCell` is borrowed until after the entire `match`, so the arbitrary user code inside the call to `self.waker.wake()` will run before the `RefCell` is released. Thus, if waking the waker involves calling `LocalSet::enter`, then that will result in a panic.

I would recommend that you define a function like this:
```
use std::cell::RefCell;
use std::rc::Rc;

fn clone_rc<T>(rc: &RefCell<Option<Rc<T>>>) -> Option<Rc<T>> {
    rc.borrow().clone()
}
```
And use it when matching on the context.

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

impl Future for LocalSet {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
        // Register the waker before starting to work
        self.context.shared.waker.register_by_ref(cx.waker());

        if self.with(|| self.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
        } else if self.context.owned.is_empty() {
            // If the scheduler has no remaining futures, we're done!
            Poll::Ready(())
        } else {
            // There are still futures in the local set, but we've polled all the
            // futures in the run queue. Therefore, we can just return Pending
            // since the remaining futures will be woken from somewhere else.
            Poll::Pending
        }
    }
}

impl Default for LocalSet {
    fn default() -> LocalSet {
        LocalSet::new()
    }
}

impl Drop for LocalSet {
    fn drop(&mut self) {
        self.with(|| {
            // Shut down all tasks in the LocalOwnedTasks and close it to
            // prevent new tasks from ever being added.
            self.context.owned.close_and_shutdown_all();

            // We already called shutdown on all tasks above, so there is no
            // need to call shutdown.
            for task in self.context.queue.take() {
                drop(task);
            }

            // Take the queue from the Shared object to prevent pushing
            // notifications to it in the future.