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

Code before:
}

impl LocalSet {
    /// Returns a new local task set.
    pub fn new() -> LocalSet {
        LocalSet {
            tick: Cell::new(0),
            context: Context {
                owned: LocalOwnedTasks::new(),
                queue: VecDequeCell::with_capacity(INITIAL_CAPACITY),
                shared: Arc::new(Shared {
                    queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))),
                    waker: AtomicWaker::new(),
                }),
            },
            _not_send: PhantomData,
        }
    }

    /// Enter current LocalSet context
    pub fn enter(&self) -> LocalEnterGuard<'_> {
        CURRENT.inner.with(|c| {
            c.set(&self.context as *const _ as *const ());
            LocalEnterGuard { _guard: &self }
        })
    }

    /// Spawns a `!Send` task onto the local task set.
    ///
    /// This task is guaranteed to be run on the current thread.
    ///
    /// Unlike the free function [`spawn_local`], this method may be used to
    /// spawn local tasks when the task set is _not_ running. For example:
    /// ```rust
    /// use tokio::task;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let local = task::LocalSet::new();
    ///
    ///     // Spawn a future on the local set. This future will be run when
    ///     // we call `run_until` to drive the task set.
    ///     local.spawn_local(async {
    ///        // ...
    ///     });
    ///
    ///     // Run the local task set.
    ///     local.run_until(async move {
    ///         // ...
    ///     }).await;
    ///

Review comment (suggestion):
I don't think I was paying attention when I first looked at this. This solution is not going to work. This never unsets the thread-local when the `LocalEnterGuard` is dropped, and to make matters worse, even if you add a destructor that unsets it, you can use `std::mem::forget` to not run that destructor. Since we are using raw pointers to set it, this results in the raw pointer still existing when the `LocalSet` is destroyed - this is a use-after-free.

In fact, I no longer think it is clear how we can solve this issue. We would probably have to heap-allocate the context.

Code after:
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(Option<Rc<Context>>);

impl Drop for LocalEnterGuard {
    fn drop(&mut self) {
        CURRENT.with(|ctx| {
            ctx.replace(self.0.take());
        })
    }
}

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 {
        LocalSet {
            tick: Cell::new(0),
            context: Rc::new(Context {
                owned: LocalOwnedTasks::new(),
                queue: VecDequeCell::with_capacity(INITIAL_CAPACITY),
                shared: Arc::new(Shared {
                    queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))),
                    waker: AtomicWaker::new(),
                    #[cfg(tokio_unstable)]
                    unhandled_panic: crate::runtime::UnhandledPanic::Ignore,
                }),
                unhandled_panic: Cell::new(false),
            }),
            _not_send: PhantomData,
        }
    }

    /// Enters the context of this `LocalSet`.
    ///
    /// The [`spawn_local`] method will spawn tasks on the `LocalSet` whose
    /// context you are inside.
    ///
    /// [`spawn_local`]: fn@crate::task::spawn_local