Repository: tokio-rs/tokio
File: tokio/src/runtime/context.rs (Rust)

Code before:
        let reset = CONTEXT.with(|c| {
            if let EnterRuntime::Entered {
                allow_block_in_place: true,
            } = c.runtime.get()
            {
                c.runtime.set(EnterRuntime::Entered {
                    allow_block_in_place: false,
                });
                true
            } else {
                false
            }
        });

        DisallowBlockInPlaceGuard(reset)
    }

    #[track_caller]
    pub(crate) fn defer(waker: &Waker) {
        with_scheduler(|maybe_scheduler| {
            if let Some(scheduler) = maybe_scheduler {
                scheduler.defer(waker);
            } else {
                waker.wake_by_ref();
            }
        });
    }

    pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R {
        CONTEXT.with(|c| c.scheduler.set(v, f))
    }

    #[track_caller]
    pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R {
        CONTEXT.with(|c| c.scheduler.with(f))
    }

    impl Context {
        fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard {
            let rng_seed = handle.seed_generator().next_seed();

            let old_handle = self.handle.borrow_mut().replace(handle.clone());
            let old_seed = self.rng.replace_seed(rng_seed);

            SetCurrentGuard {
                old_handle,
                old_seed,
            }
        }
    }


Review comment (suggestion):
We should also `wake_by_ref` if the thread local is destroyed due to thread shutdown.

Code after:
        let reset = CONTEXT.with(|c| {
            if let EnterRuntime::Entered {
                allow_block_in_place: true,
            } = c.runtime.get()
            {
                c.runtime.set(EnterRuntime::Entered {
                    allow_block_in_place: false,
                });
                true
            } else {
                false
            }
        });

        DisallowBlockInPlaceGuard(reset)
    }

    #[track_caller]
    pub(crate) fn defer(waker: &Waker) {
        with_scheduler(|maybe_scheduler| {
            if let Some(scheduler) = maybe_scheduler {
                scheduler.defer(waker);
            } else {
                // Called from outside of the runtime, immediately wake the
                // task.
                waker.wake_by_ref();
            }
        });
    }

    pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R {
        CONTEXT.with(|c| c.scheduler.set(v, f))
    }

    #[track_caller]
    pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R {
        CONTEXT.with(|c| c.scheduler.with(f))
    }

    impl Context {
        fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard {
            let rng_seed = handle.seed_generator().next_seed();

            let old_handle = self.handle.borrow_mut().replace(handle.clone());
            let old_seed = self.rng.replace_seed(rng_seed);

            SetCurrentGuard {
                old_handle,
                old_seed,
            }
        }