use crate::loom::thread::AccessError;
use crate::runtime::coop;
use std::cell::Cell;
#[cfg(any(feature = "rt", feature = "macros"))]
use crate::util::rand::{FastRand, RngSeed};
cfg_rt! {
use crate::runtime::{scheduler, task::Id, Defer};
use std::cell::RefCell;
use std::marker::PhantomData;
use std::time::Duration;
}
struct Context {
#[cfg(feature = "rt")]
handle: RefCell<Option<scheduler::Handle>>,
#[cfg(feature = "rt")]
current_task_id: Cell<Option<Id>>,
#[cfg(feature = "rt")]
runtime: Cell<EnterRuntime>,
#[cfg(feature = "rt")]
defer: RefCell<Option<Defer>>,
#[cfg(any(feature = "rt", feature = "macros"))]
rng: FastRand,
budget: Cell<coop::Budget>,
}
tokio_thread_local! {
static CONTEXT: Context = {
Context {
#[cfg(feature = "rt")]
handle: RefCell::new(None),
#[cfg(feature = "rt")]
current_task_id: Cell::new(None),
#[cfg(feature = "rt")]
runtime: Cell::new(EnterRuntime::NotEntered),
#[cfg(feature = "rt")]
defer: RefCell::new(None),
#[cfg(any(feature = "rt", feature = "macros"))]
rng: FastRand::new(RngSeed::new()),
budget: Cell::new(coop::Budget::unconstrained()),
}
}
}
#[cfg(feature = "macros")]
pub(crate) fn thread_rng_n(n: u32) -> u32 {
CONTEXT.with(|ctx| ctx.rng.fastrand_n(n))
}
pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> {
CONTEXT.try_with(|ctx| f(&ctx.budget))
}
cfg_rt! {
use crate::runtime::TryCurrentError;
use std::fmt;
#[derive(Debug, Clone, Copy)]
#[must_use]
pub(crate) enum EnterRuntime {
#[cfg_attr(not(feature = "rt"), allow(dead_code))]
Entered { allow_block_in_place: bool },
NotEntered,
}
#[derive(Debug)]
#[must_use]
pub(crate) struct SetCurrentGuard {
old_handle: Option<scheduler::Handle>,
old_seed: RngSeed,
}
#[must_use]
pub(crate) struct EnterRuntimeGuard {
pub(crate) blocking: BlockingRegionGuard,
#[allow(dead_code)] pub(crate) handle: SetCurrentGuard,
is_root: bool,
}
#[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()),
}
}
pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> {
CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok()
}
#[track_caller]
pub(crate) fn enter_runtime(handle: &scheduler::Handle, allow_block_in_place: bool) -> EnterRuntimeGuard {
if let Some(enter) = try_enter_runtime(handle, allow_block_in_place) {
return enter;
}
panic!(
"Cannot start a runtime from within a runtime. This happens \
because a function (like `block_on`) attempted to block the \
current thread while the thread is being used to drive \
asynchronous tasks."
);
}
fn try_enter_runtime(handle: &scheduler::Handle, allow_block_in_place: bool) -> Option<EnterRuntimeGuard> {
CONTEXT.with(|c| {
if c.runtime.get().is_entered() {
None
} else {
c.runtime.set(EnterRuntime::Entered { allow_block_in_place });
let mut defer = c.defer.borrow_mut();
let is_root = if defer.is_none() {
*defer = Some(Defer::new());
true
} else {
false
};
Some(EnterRuntimeGuard {
blocking: BlockingRegionGuard::new(),
handle: c.set_current(handle),
is_root,
})
}
})
}
pub(crate) fn try_enter_blocking_region() -> Option<BlockingRegionGuard> {
CONTEXT.try_with(|c| {
if c.runtime.get().is_entered() {
None
} else {
Some(BlockingRegionGuard::new())
}
}).unwrap_or_else(|_| Some(BlockingRegionGuard::new()))
}
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)
}
pub(crate) fn with_defer<R>(f: impl FnOnce(&mut Defer) -> R) -> Option<R> {
CONTEXT.with(|c| {
let mut defer = c.defer.borrow_mut();
defer.as_mut().map(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,
}
}
}
impl Drop for SetCurrentGuard {
fn drop(&mut self) {
CONTEXT.with(|ctx| {
*ctx.handle.borrow_mut() = self.old_handle.take();
ctx.rng.replace_seed(self.old_seed.clone());
});
}
}
impl fmt::Debug for EnterRuntimeGuard {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Enter").finish()
}
}
impl Drop for EnterRuntimeGuard {
fn drop(&mut self) {
CONTEXT.with(|c| {
assert!(c.runtime.get().is_entered());
c.runtime.set(EnterRuntime::NotEntered);
if self.is_root {
*c.defer.borrow_mut() = None;
}
});
}
}
impl BlockingRegionGuard {
fn new() -> BlockingRegionGuard {
BlockingRegionGuard { _p: PhantomData }
}
pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError>
where
F: std::future::Future,
{
use crate::runtime::park::CachedParkThread;
let mut park = CachedParkThread::new();
park.block_on(f)
}
pub(crate) fn block_on_timeout<F>(&mut self, f: F, timeout: Duration) -> Result<F::Output, ()>
where
F: std::future::Future,
{
use crate::runtime::park::CachedParkThread;
use std::task::Context;
use std::task::Poll::Ready;
use std::time::Instant;
let mut park = CachedParkThread::new();
let waker = park.waker().map_err(|_| ())?;
let mut cx = Context::from_waker(&waker);
pin!(f);
let when = Instant::now() + timeout;
loop {
if let Ready(v) = crate::runtime::coop::budget(|| f.as_mut().poll(&mut cx)) {
return Ok(v);
}
let now = Instant::now();
if now >= when {
return Err(());
}
with_defer(|defer| defer.wake());
park.park_timeout(when - now);
}
}
}
impl Drop for DisallowBlockInPlaceGuard {
fn drop(&mut self) {
if self.0 {
CONTEXT.with(|c| {
if let EnterRuntime::Entered {
allow_block_in_place: false,
} = c.runtime.get()
{
c.runtime.set(EnterRuntime::Entered {
allow_block_in_place: true,
});
}
})
}
}
}
impl EnterRuntime {
pub(crate) fn is_entered(self) -> bool {
matches!(self, EnterRuntime::Entered { .. })
}
}
}
cfg_rt_multi_thread! {
pub(crate) fn current_enter_context() -> EnterRuntime {
CONTEXT.with(|c| c.runtime.get())
}
pub(crate) fn exit_runtime<F: FnOnce() -> R, R>(f: F) -> R {
struct Reset(EnterRuntime);
impl Drop for Reset {
fn drop(&mut self) {
CONTEXT.with(|c| {
assert!(!c.runtime.get().is_entered(), "closure claimed permanent executor");
c.runtime.set(self.0);
});
}
}
let was = CONTEXT.with(|c| {
let e = c.runtime.get();
assert!(e.is_entered(), "asked to exit when not entered");
c.runtime.set(EnterRuntime::NotEntered);
e
});
let _reset = Reset(was);
f()
}
}