Skip to content

Commit

Permalink
rt: split runtime::context into multiple files (#5768)
Browse files Browse the repository at this point in the history 
This PR restructures `runtime::context` into multiple files by component and feature flag. The goal is to reduce code defined in macros and make each context component more manageable.

There should be no behavior changes except tweaking how the RNG seed is set. Instead of putting it in `set_current`, we set it when entering the runtime. This aligns better with the feature's original intent, enabling users to make a runtime's RNG deterministic. The seed should not be changed by `Handle::enter()`, so there is no need to have the code in `context::set_current`.
  • Loading branch information
@carllerche
carllerche committed Jun 6, 2023
1 parent e75ca93 commit 1204da7
Show file tree
Hide file tree
Showing 10 changed files with 388 additions and 351 deletions.
307 changes: 15 additions & 292 deletions tokio/src/runtime/context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,33 @@ use std::cell::Cell;
use crate::util::rand::FastRand;

cfg_rt! {
mod blocking;
pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard};

mod current;
pub(crate) use current::{with_current, try_set_current, SetCurrentGuard};

mod runtime;
pub(crate) use runtime::{EnterRuntime, enter_runtime};

mod scoped;
use scoped::Scoped;

use crate::runtime::{scheduler, task::Id};

use std::cell::RefCell;
use std::marker::PhantomData;
use std::task::Waker;
use std::time::Duration;

cfg_taskdump! {
use crate::runtime::task::trace;
}
}

cfg_rt_multi_thread! {
mod runtime_mt;
pub(crate) use runtime_mt::{current_enter_context, exit_runtime};
}

struct Context {
/// Uniquely identifies the current thread
#[cfg(feature = "rt")]
Expand Down Expand Up @@ -125,10 +137,7 @@ pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, A
}

cfg_rt! {
use crate::runtime::{ThreadId, TryCurrentError};
use crate::util::rand::RngSeed;

use std::fmt;
use crate::runtime::ThreadId;

pub(crate) fn thread_id() -> Result<ThreadId, AccessError> {
CONTEXT.try_with(|ctx| {
Expand All @@ -143,45 +152,6 @@ cfg_rt! {
})
}

#[derive(Debug, Clone, Copy)]
#[must_use]
pub(crate) enum EnterRuntime {
/// Currently in a runtime context.
#[cfg_attr(not(feature = "rt"), allow(dead_code))]
Entered { allow_block_in_place: bool },

/// Not in a runtime context **or** a blocking region.
NotEntered,
}

#[derive(Debug)]
#[must_use]
pub(crate) struct SetCurrentGuard {
old_handle: Option<scheduler::Handle>,
old_seed: RngSeed,
// Should not be `Send` since it must be *dropped* on the same thread as
// created, but there is no issue with sync access.
_p: PhantomData<crate::util::markers::SyncNotSend>,
}

/// Guard tracking that a caller has entered a runtime context.
#[must_use]
pub(crate) struct EnterRuntimeGuard {
/// Tracks that the current thread has entered a blocking function call.
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)
}
Expand All @@ -190,93 +160,6 @@ cfg_rt! {
CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None)
}

pub(crate) fn with_current<F, R>(f: F) -> Result<R, TryCurrentError>
where
F: FnOnce(&scheduler::Handle) -> R,
{

match CONTEXT.try_with(|ctx| ctx.handle.borrow().as_ref().map(f)) {
Ok(Some(ret)) => Ok(ret),
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.
#[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."
);
}

/// Tries to enter a runtime context, returns `None` if already in a runtime
/// context.
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 {
// Set the entered flag
c.runtime.set(EnterRuntime::Entered { allow_block_in_place });

Some(EnterRuntimeGuard {
blocking: BlockingRegionGuard::new(),
handle: c.set_current(handle),
})
}
})
}

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)
}

#[track_caller]
pub(crate) fn defer(waker: &Waker) {
with_scheduler(|maybe_scheduler| {
Expand All @@ -299,127 +182,6 @@ cfg_rt! {
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 mut rng = self.rng.get().unwrap_or_else(FastRand::new);
let old_seed = rng.replace_seed(rng_seed);
self.rng.set(Some(rng));

SetCurrentGuard {
old_handle,
old_seed,
_p: PhantomData,
}
}
}

impl Drop for SetCurrentGuard {
fn drop(&mut self) {
CONTEXT.with(|ctx| {
*ctx.handle.borrow_mut() = self.old_handle.take();

let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new);
rng.replace_seed(self.old_seed.clone());
ctx.rng.set(Some(rng));
});
}
}

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);
});
}
}

impl BlockingRegionGuard {
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
F: std::future::Future,
{
use crate::runtime::park::CachedParkThread;

let mut park = CachedParkThread::new();
park.block_on(f)
}

/// Blocks the thread on the specified future for **at most** `timeout`
///
/// If the future completes before `timeout`, the result is returned. If
/// `timeout` elapses, then `Err` is returned.
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(());
}

park.park_timeout(when - now);
}
}
}

impl Drop for DisallowBlockInPlaceGuard {
fn drop(&mut self) {
if self.0 {
// XXX: Do we want some kind of assertion here, or is "best effort" okay?
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_taskdump! {
/// SAFETY: Callers of this function must ensure that trace frames always
/// form a valid linked list.
Expand All @@ -428,42 +190,3 @@ cfg_rt! {
}
}
}

// Forces the current "entered" state to be cleared while the closure
// is executed.
//
// # Warning
//
// This is hidden for a reason. Do not use without fully understanding
// executors. Misusing can easily cause your program to deadlock.
cfg_rt_multi_thread! {
/// Returns true if in a runtime context.
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 {
// Reset in case the closure panics
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);
// dropping _reset after f() will reset ENTERED
f()
}
}
Loading

0 comments on commit 1204da7

Please sign in to comment.