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

Code before:
#[must_use]
pub(crate) struct BlockingRegionGuard {
    _p: PhantomData<NotSendOrSync>,
}

pub(crate) struct DisallowBlockInPlaceGuard(bool);

pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> {
    CONTEXT
        .try_with(|c| {
            if c.runtime.get().is_entered() {
                None
            } else {
                Some(BlockingRegionGuard::new())
            }
            // If accessing the thread-local fails, the thread is terminating
            // and thread-locals are being destroyed. Because we don't know if
            // we are currently in a runtime or not, we default to being
            // permissive.
        })
        .unwrap_or_else(|_| Some(BlockingRegionGuard::new()))
}

/// Disallows blocking in the current runtime context until the guard is dropped.
pub(crate) fn disallow_block_in_place() -> DisallowBlockInPlaceGuard {
    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)
}

impl BlockingRegionGuard {
    pub(super) fn new() -> BlockingRegionGuard {
        BlockingRegionGuard { _p: PhantomData }
    }

    /// Blocks the thread on the specified future, returning the value with
    /// which that future completes.
    pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError>
    where

Review comment (suggestion):
It might be worth aligning here too?
https://github.com/tokio-rs/tokio/blob/ab8d7b82a1252b41dc072f641befb6d2afcb3373/tokio/src/runtime/context/blocking.rs#L109

I think drop still gets panic if `CONTEXT.try_with` returns Err when constructing a `DisallowBlockInPlaceGuard`.

Code after:
#[must_use]
pub(crate) struct BlockingRegionGuard {
    _p: PhantomData<NotSendOrSync>,
}

pub(crate) struct DisallowBlockInPlaceGuard(bool);

pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> {
    CONTEXT
        .try_with(|c| {
            if c.runtime.get().is_entered() {
                None
            } else {
                Some(BlockingRegionGuard::new())
            }
            // If accessing the thread-local fails, the thread is terminating
            // and thread-locals are being destroyed. Because we don't know if
            // we are currently in a runtime or not, we default to being
            // permissive.
        })
        .unwrap_or_else(|_| Some(BlockingRegionGuard::new()))
}

/// Disallows blocking in the current runtime context until the guard is dropped.
pub(crate) fn disallow_block_in_place() -> DisallowBlockInPlaceGuard {
    let reset = CONTEXT.try_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.unwrap_or(false))
}

impl BlockingRegionGuard {
    pub(super) fn new() -> BlockingRegionGuard {
        BlockingRegionGuard { _p: PhantomData }
    }

    /// Blocks the thread on the specified future, returning the value with
    /// which that future completes.
    pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError>
    where