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

Code before:
                    .and_then(|queue| queue.pop_front())
            })
        };

        task.map(|task| unsafe {
            // Safety: because the `LocalSet` itself is `!Send`, we know we are
            // on the same thread if we have access to the `LocalSet`, and can
            // therefore access the local run queue.
            self.context.shared.local_state.assert_owner(task)
        })
    }

    fn pop_local(&self) -> Option<task::Notified<Arc<Shared>>> {
        unsafe {
            // Safety: because the `LocalSet` itself is `!Send`, we know we are
            // on the same thread if we have access to the `LocalSet`, and can
            // therefore access the local run queue.
            self.context.shared.local_state.task_pop_front()
        }
    }

    fn with<T>(&self, f: impl FnOnce() -> T) -> T {
        CURRENT.with(|LocalData { ctx, entered, .. }| {
            struct Reset<'a> {
                ctx_ref: &'a RcCell<Context>,
                entered_ref: &'a Cell<bool>,
                ctx_val: Option<Rc<Context>>,
                entered_val: bool,
            }
            impl<'a> Drop for Reset<'a> {
                fn drop(&mut self) {
                    self.ctx_ref.set(self.ctx_val.take());
                    self.entered_ref.set(self.entered_val)
                }
            }
            let ctx_old = ctx.replace(Some(self.context.clone()));
            let entered_old = entered.replace(false);

            let _reset = Reset {
                ctx_ref: ctx,
                entered_ref: entered,
                ctx_val: ctx_old,
                entered_val: entered_old,
            };

            f()
        })
    }

    /// This method is like `with`, but it just calls `f` without setting the thread-local if that
    /// fails.

Review comment (suggestion):
nit, not a big deal: both `ctx_ref` and `entered_ref` reference fields of the `LocalData`. we could make this struct a word smaller by changing these to a single `&'a LocalData` reference, and just using that to access both `ctx` and `entered`.

Code after:

    /// Ticks the scheduler, returning whether the local future needs to be
    /// notified again.
    fn tick(&self) -> bool {
        for _ in 0..MAX_TASKS_PER_TICK {
            // Make sure we didn't hit an unhandled panic
            if self.context.unhandled_panic.get() {
                panic!("a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic");
            }

            match self.next_task() {
                // Run the task
                //
                // Safety: As spawned tasks are `!Send`, `run_unchecked` must be
                // used. We are responsible for maintaining the invariant that
                // `run_unchecked` is only called on threads that spawned the
                // task initially. Because `LocalSet` itself is `!Send`, and
                // `spawn_local` spawns into the `LocalSet` on the current
                // thread, the invariant is maintained.
                Some(task) => crate::runtime::coop::budget(|| task.run()),
                // We have fully drained the queue of notified tasks, so the
                // local future doesn't need to be notified again — it can wait
                // until something else wakes a task in the local set.
                None => return false,
            }
        }

        true
    }

    fn next_task(&self) -> Option<task::LocalNotified<Arc<Shared>>> {
        let tick = self.tick.get();
        self.tick.set(tick.wrapping_add(1));

        let task = if tick % REMOTE_FIRST_INTERVAL == 0 {
            self.context
                .shared
                .queue
                .lock()
                .as_mut()
                .and_then(|queue| queue.pop_front())
                .or_else(|| self.pop_local())
        } else {
            self.pop_local().or_else(|| {
                self.context
                    .shared
                    .queue
                    .lock()
                    .as_mut()
                    .and_then(|queue| queue.pop_front())
            })