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

Code before:
    ctx: RcCell::new(),
    entered: Cell::new(false),
} });

struct LocalData {
    ctx: RcCell<Context>,
    entered: Cell<bool>,
}

impl LocalData {
    fn enter(&self, ctx: Rc<Context>) -> LocalDataEnterGuard<'_> {
        let ctx = self.ctx.replace(Some(ctx));
        let entered = self.entered.replace(true);

        LocalDataEnterGuard {
            local_data_ref: self,
            ctx,
            entered,
        }
    }
}

struct LocalDataEnterGuard<'a> {
    local_data_ref: &'a LocalData,
    ctx: Option<Rc<Context>>,
    entered: bool,
}

impl<'a> Drop for LocalDataEnterGuard<'a> {
    fn drop(&mut self) {
        self.local_data_ref.ctx.set(self.ctx.take());
        self.local_data_ref.entered.set(self.entered)
    }
}

cfg_rt! {
    /// Spawns a `!Send` future on the current [`LocalSet`].
    ///
    /// The spawned future will run on the same thread that called `spawn_local`.
    ///
    /// The provided future will start running in the background immediately
    /// when `spawn_local` is called, even if you don't await the returned
    /// `JoinHandle`.
    ///
    /// # Panics
    ///
    /// This function panics if called outside of a [`LocalSet`].
    ///
    /// Note that if [`tokio::spawn`] is used from within a `LocalSet`, the
    /// resulting new task will _not_ be inside the `LocalSet`, so you must use
    /// `spawn_local` if you want to stay within the `LocalSet`.

Review comment (suggestion):
could we maybe add a comment on this explaining how it differs from the public `LocalEnterGuard`?  we now have two types in this module with similar names and similar behavior, one of which is private and the other of which is public. So, it would be nice to make the difference clear to readers. We may want to consider giving this type a different name, as well, if you can come up with something more descriptive.

Code after:
    ctx: RcCell::new(),
    wake_on_schedule: Cell::new(false),
} });

struct LocalData {
    ctx: RcCell<Context>,
    wake_on_schedule: Cell<bool>,
}

impl LocalData {
    /// Should be called except when we call `LocalSet::enter`.
    /// Especially when we poll a LocalSet.
    #[must_use = "dropping this guard will reset the entered state"]
    fn enter(&self, ctx: Rc<Context>) -> LocalDataEnterGuard<'_> {
        let ctx = self.ctx.replace(Some(ctx));
        let wake_on_schedule = self.wake_on_schedule.replace(false);
        LocalDataEnterGuard {
            local_data_ref: self,
            ctx,
            wake_on_schedule,
        }
    }
}

/// A guard for `LocalData::enter()`
struct LocalDataEnterGuard<'a> {
    local_data_ref: &'a LocalData,
    ctx: Option<Rc<Context>>,
    wake_on_schedule: bool,
}

impl<'a> Drop for LocalDataEnterGuard<'a> {
    fn drop(&mut self) {
        self.local_data_ref.ctx.set(self.ctx.take());
        self.local_data_ref
            .wake_on_schedule
            .set(self.wake_on_schedule)
    }
}

cfg_rt! {
    /// Spawns a `!Send` future on the current [`LocalSet`].
    ///
    /// The spawned future will run on the same thread that called `spawn_local`.
    ///
    /// The provided future will start running in the background immediately
    /// when `spawn_local` is called, even if you don't await the returned
    /// `JoinHandle`.
    ///
    /// # Panics
    ///