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

Code before:

    pub(crate) fn set_current_task_id(id: Option<Id>) {
        CONTEXT.with(|ctx| ctx.current_task_id.replace(id));
    }

    #[track_caller]
    pub(crate) fn current_task_id() -> Id {
        match CONTEXT.try_with(|ctx| ctx.current_task_id.get()) {
            Ok(Some(id)) => id,
            _ => panic!("can't get a task id when not inside a task"),
        }
    }

    pub(crate) fn try_current() -> Result<scheduler::Handle, TryCurrentError> {
        match CONTEXT.try_with(|ctx| ctx.scheduler.borrow().clone()) {
            Ok(Some(handle)) => Ok(handle),
            Ok(None) => Err(TryCurrentError::new_no_context()),
            Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()),
        }
    }

    /// Sets this [`Handle`] as the current active [`Handle`].
    ///
    /// [`Handle`]: crate::runtime::scheduler::Handle
    pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> {
        let rng_seed = handle.seed_generator().next_seed();

        CONTEXT.try_with(|ctx| {
            let old_handle = ctx.scheduler.borrow_mut().replace(handle.clone());
            let old_seed = ctx.rng.replace_seed(rng_seed);

            SetCurrentGuard {
                old_handle,
                old_seed,
            }
        }).ok()
    }


    /// Marks the current thread as being within the dynamic extent of an
    /// executor.
    #[track_caller]
    pub(crate) fn enter_runtime(handle: &scheduler::Handle, allow_block_in_place: bool) -> EnterRuntimeGuard {
        if let Some(enter) = try_enter_runtime(allow_block_in_place) {
            // Set the current runtime handle. This should not fail. A later
            // cleanup will remove the unwrap().
            try_set_current(handle).unwrap();
            return enter;
        }

        panic!(

Review comment (bug):
This will panic if it runs during the destructor of a thread-local. I would silently do nothing here so that destructors can still run even if thread-locals are not available.

Some sort of test for this would be nice as well, but there are already existing bugs related to this (see #5162), and you don't have to fix those in this PR, so you should only add such a test if it doesn't hit an existing bug about this.

Code after:
    }

    #[derive(Debug)]
    #[must_use]
    pub(crate) struct SetCurrentGuard {
        old_handle: Option<scheduler::Handle>,
        old_seed: RngSeed,
    }

    /// Guard tracking that a caller has entered a runtime context.
    #[must_use]
    pub(crate) struct EnterRuntimeGuard {
        pub(crate) blocking: BlockingRegionGuard,
        #[allow(dead_code)] // Only tracking the guard.
        pub(crate) handle: SetCurrentGuard,
    }

    /// Guard tracking that a caller has entered a blocking region.
    #[must_use]
    pub(crate) struct BlockingRegionGuard {
        _p: PhantomData<RefCell<()>>,
    }

    pub(crate) struct DisallowBlockInPlaceGuard(bool);

    pub(crate) fn set_current_task_id(id: Option<Id>) -> Option<Id> {
        CONTEXT.try_with(|ctx| ctx.current_task_id.replace(id)).unwrap_or(None)
    }

    pub(crate) fn current_task_id() -> Option<Id> {
        CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None)
    }

    pub(crate) fn try_current() -> Result<scheduler::Handle, TryCurrentError> {
        match CONTEXT.try_with(|ctx| ctx.handle.borrow().clone()) {
            Ok(Some(handle)) => Ok(handle),
            Ok(None) => Err(TryCurrentError::new_no_context()),
            Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()),
        }
    }

    /// Sets this [`Handle`] as the current active [`Handle`].
    ///
    /// [`Handle`]: crate::runtime::scheduler::Handle
    pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> {
        CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok()
    }


    /// Marks the current thread as being within the dynamic extent of an
    /// executor.