tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-25 00:00:18 +02:00
Code Issues Packages Projects Releases Wiki Activity

rt: move scheduler ctxs to runtime::context (#5727)

Browse Source 
This commit eliminates the current_thread::CURRENT and multi_thread::current
thread-local variables in favor of using `runtime::context`. This is another step
towards reducing the total number of thread-local variables used by Tokio.
This commit is contained in:
Carl Lerche
2023-05-26 19:07:20 -07:00
committed by GitHub
parent d274ef3748
commit 9f9db7da63
8 changed files with 204 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/macros/mod.rs
-4
View File
@@ -23,10 +23,6 @@ cfg_trace! {
mod trace;
}
#[macro_use]
#[cfg(feature = "rt")]
pub(crate) mod scoped_tls;
cfg_macros! {
#[macro_use]
mod select;
tokio/src/macros/scoped_tls.rs
-77
View File
@@ -1,77 +0,0 @@
use crate::loom::thread::LocalKey;
use std::cell::Cell;
use std::marker;
/// Sets a reference as a thread-local.
macro_rules! scoped_thread_local {
($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => (
$(#[$attrs])*
$vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty>
= $crate::macros::scoped_tls::ScopedKey {
inner: {
tokio_thread_local!(static FOO: ::std::cell::Cell<*const ()> = const {
std::cell::Cell::new(::std::ptr::null())
});
&FOO
},
_marker: ::std::marker::PhantomData,
};
)
}
/// Type representing a thread local storage key corresponding to a reference
/// to the type parameter `T`.
pub(crate) struct ScopedKey<T> {
pub(crate) inner: &'static LocalKey<Cell<*const ()>>,
pub(crate) _marker: marker::PhantomData<T>,
}
unsafe impl<T> Sync for ScopedKey<T> {}
impl<T> ScopedKey<T> {
/// Inserts a value into this scoped thread local storage slot for a
/// duration of a closure.
pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R
where
F: FnOnce() -> R,
{
struct Reset {
key: &'static LocalKey<Cell<*const ()>>,
val: *const (),
}
impl Drop for Reset {
fn drop(&mut self) {
self.key.with(|c| c.set(self.val));
}
}
let prev = self.inner.with(|c| {
let prev = c.get();
c.set(t as *const _ as *const ());
prev
});
let _reset = Reset {
key: self.inner,
val: prev,
};
f()
}
/// Gets a value out of this scoped variable.
pub(crate) fn with<F, R>(&'static self, f: F) -> R
where
F: FnOnce(Option<&T>) -> R,
{
let val = self.inner.with(|c| c.get());
if val.is_null() {
f(None)
} else {
unsafe { f(Some(&*(val as *const T))) }
}
}
}
tokio/src/runtime/context.rs
+21
View File
@@ -7,6 +7,9 @@ use std::cell::Cell;
use crate::util::rand::{FastRand, RngSeed};
cfg_rt! {
mod scoped;
use scoped::Scoped;
use crate::runtime::{scheduler, task::Id, Defer};
use std::cell::RefCell;
@@ -27,6 +30,10 @@ struct Context {
#[cfg(feature = "rt")]
handle: RefCell<Option<scheduler::Handle>>,
/// Handle to the scheduler's internal "context"
#[cfg(feature = "rt")]
scheduler: Scoped<scheduler::Context>,
#[cfg(feature = "rt")]
current_task_id: Cell<Option<Id>>,
@@ -70,6 +77,11 @@ tokio_thread_local! {
/// accessing drivers, etc...
#[cfg(feature = "rt")]
handle: RefCell::new(None),
/// Tracks the current scheduler internal context
#[cfg(feature = "rt")]
scheduler: Scoped::new(),
#[cfg(feature = "rt")]
current_task_id: Cell::new(None),
@@ -287,6 +299,15 @@ cfg_rt! {
})
}
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();
tokio/src/runtime/context/scoped.rs
+58
View File
@@ -0,0 +1,58 @@
use std::cell::Cell;
use std::ptr;
/// Scoped thread-local storage
pub(super) struct Scoped<T> {
pub(super) inner: Cell<*const T>,
}
unsafe impl<T> Sync for Scoped<T> {}
impl<T> Scoped<T> {
pub(super) fn new() -> Scoped<T> {
Scoped {
inner: Cell::new(ptr::null()),
}
}
/// Inserts a value into the scoped cell for the duration of the closure
pub(super) fn set<F, R>(&self, t: &T, f: F) -> R
where
F: FnOnce() -> R,
{
struct Reset<'a, T> {
cell: &'a Cell<*const T>,
prev: *const T,
}
impl<T> Drop for Reset<'_, T> {
fn drop(&mut self) {
self.cell.set(self.prev);
}
}
let prev = self.inner.get();
self.inner.set(t as *const _);
let _reset = Reset {
cell: &self.inner,
prev,
};
f()
}
/// Gets the value out of the scoped cell;
pub(super) fn with<F, R>(&self, f: F) -> R
where
F: FnOnce(Option<&T>) -> R,
{
let val = self.inner.get();
if val.is_null() {
f(None)
} else {
unsafe { f(Some(&*(val as *const T))) }
}
}
}
tokio/src/runtime/scheduler/current_thread.rs
+76 -49
View File
@@ -86,7 +86,9 @@ struct Shared {
}
/// Thread-local context.
struct Context {
///
/// pub(crate) to store in `runtime::context`.
pub(crate) struct Context {
/// Scheduler handle
handle: Arc<Handle>,
@@ -100,9 +102,6 @@ type Notified = task::Notified<Arc<Handle>>;
/// Initial queue capacity.
const INITIAL_CAPACITY: usize = 64;
// Tracks the current CurrentThread.
scoped_thread_local!(static CURRENT: Context);
impl CurrentThread {
pub(crate) fn new(
driver: Driver,
@@ -185,10 +184,10 @@ impl CurrentThread {
let core = self.core.take()?;
Some(CoreGuard {
context: Context {
context: scheduler::Context::CurrentThread(Context {
handle: handle.clone(),
core: RefCell::new(Some(core)),
},
}),
scheduler: self,
})
}
@@ -205,41 +204,60 @@ impl CurrentThread {
None => panic!("Oh no! We never placed the Core back, this is a bug!"),
};
core.enter(|mut core, _context| {
// Drain the OwnedTasks collection. This call also closes the
// collection, ensuring that no tasks are ever pushed after this
// call returns.
handle.shared.owned.close_and_shutdown_all();
// Check that the thread-local is not being destroyed
let tls_available = context::with_current(|_| ()).is_ok();
// Drain local queue
// We already shut down every task, so we just need to drop the task.
while let Some(task) = core.next_local_task(handle) {
drop(task);
}
if tls_available {
core.enter(|core, _context| {
let core = shutdown2(core, handle);
(core, ())
});
} else {
// Shutdown without setting the context. `tokio::spawn` calls will
// fail, but those will fail either way because the thread-local is
// not available anymore.
let context = core.context.expect_current_thread();
let core = context.core.borrow_mut().take().unwrap();
// Close the injection queue
handle.shared.inject.close();
// Drain remote queue
while let Some(task) = handle.shared.inject.pop() {
drop(task);
}
assert!(handle.shared.owned.is_empty());
// Submit metrics
core.submit_metrics(handle);
// Shutdown the resource drivers
if let Some(driver) = core.driver.as_mut() {
driver.shutdown(&handle.driver);
}
(core, ())
});
let core = shutdown2(core, handle);
*context.core.borrow_mut() = Some(core);
}
}
}
fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
// Drain the OwnedTasks collection. This call also closes the
// collection, ensuring that no tasks are ever pushed after this
// call returns.
handle.shared.owned.close_and_shutdown_all();
// Drain local queue
// We already shut down every task, so we just need to drop the task.
while let Some(task) = core.next_local_task(handle) {
drop(task);
}
// Close the injection queue
handle.shared.inject.close();
// Drain remote queue
while let Some(task) = handle.shared.inject.pop() {
drop(task);
}
assert!(handle.shared.owned.is_empty());
// Submit metrics
core.submit_metrics(handle);
// Shutdown the resource drivers
if let Some(driver) = core.driver.as_mut() {
driver.shutdown(&handle.driver);
}
core
}
impl fmt::Debug for CurrentThread {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("CurrentThread").finish()
@@ -425,10 +443,10 @@ impl Handle {
let mut traces = vec![];
// todo: how to make this work outside of a runtime context?
CURRENT.with(|maybe_context| {
context::with_scheduler(|maybe_context| {
// drain the local queue
let context = if let Some(context) = maybe_context {
context
context.expect_current_thread()
} else {
return;
};
@@ -522,8 +540,10 @@ impl Schedule for Arc<Handle> {
}
fn schedule(&self, task: task::Notified<Self>) {
CURRENT.with(|maybe_cx| match maybe_cx {
Some(cx) if Arc::ptr_eq(self, &cx.handle) => {
use scheduler::Context::CurrentThread;
context::with_scheduler(|maybe_cx| match maybe_cx {
Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
let mut core = cx.core.borrow_mut();
// If `None`, the runtime is shutting down, so there is no need
@@ -552,11 +572,14 @@ impl Schedule for Arc<Handle> {
// Do nothing
}
UnhandledPanic::ShutdownRuntime => {
use scheduler::Context::CurrentThread;
// This hook is only called from within the runtime, so
// `CURRENT` should match with `&self`, i.e. there is no
// opportunity for a nested scheduler to be called.
CURRENT.with(|maybe_cx| match maybe_cx {
Some(cx) if Arc::ptr_eq(self, &cx.handle) => {
// `context::with_scheduler` should match with `&self`, i.e.
// there is no opportunity for a nested scheduler to be
// called.
context::with_scheduler(|maybe_cx| match maybe_cx {
Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
let mut core = cx.core.borrow_mut();
// If `None`, the runtime is shutting down, so there is no need to signal shutdown
@@ -590,7 +613,7 @@ impl Wake for Handle {
/// Used to ensure we always place the `Core` value back into its slot in
/// `CurrentThread`, even if the future panics.
struct CoreGuard<'a> {
context: Context,
context: scheduler::Context,
scheduler: &'a CurrentThread,
}
@@ -672,13 +695,15 @@ impl CoreGuard<'_> {
where
F: FnOnce(Box<Core>, &Context) -> (Box<Core>, R),
{
let context = self.context.expect_current_thread();
// Remove `core` from `context` to pass into the closure.
let core = self.context.core.borrow_mut().take().expect("core missing");
let core = context.core.borrow_mut().take().expect("core missing");
// Call the closure and place `core` back
let (core, ret) = CURRENT.set(&self.context, || f(core, &self.context));
let (core, ret) = context::set_scheduler(&self.context, || f(core, context));
*self.context.core.borrow_mut() = Some(core);
*context.core.borrow_mut() = Some(core);
ret
}
@@ -686,7 +711,9 @@ impl CoreGuard<'_> {
impl Drop for CoreGuard<'_> {
fn drop(&mut self) {
if let Some(core) = self.context.core.borrow_mut().take() {
let context = self.context.expect_current_thread();
if let Some(core) = context.core.borrow_mut().take() {
// Replace old scheduler back into the state to allow
// other threads to pick it up and drive it.
self.scheduler.core.set(core);
tokio/src/runtime/scheduler/mod.rs
+29
View File
@@ -25,6 +25,14 @@ pub(crate) enum Handle {
Disabled,
}
#[cfg(feature = "rt")]
pub(super) enum Context {
CurrentThread(current_thread::Context),
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
MultiThread(multi_thread::Context),
}
impl Handle {
#[cfg_attr(not(feature = "full"), allow(dead_code))]
pub(crate) fn driver(&self) -> &driver::Handle {
@@ -184,6 +192,27 @@ cfg_rt! {
}
}
}
impl Context {
#[track_caller]
pub(crate) fn expect_current_thread(&self) -> &current_thread::Context {
match self {
Context::CurrentThread(context) => context,
#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
_ => panic!("expected `CurrentThread::Context`")
}
}
cfg_rt_multi_thread! {
#[track_caller]
pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context {
match self {
Context::MultiThread(context) => context,
_ => panic!("expected `MultiThread::Context`")
}
}
}
}
}
cfg_not_rt! {
tokio/src/runtime/scheduler/multi_thread/mod.rs
+1 -1
View File
@@ -15,7 +15,7 @@ pub(crate) use park::{Parker, Unparker};
pub(crate) mod queue;
mod worker;
pub(crate) use worker::Launch;
pub(crate) use worker::{Context, Launch};
pub(crate) use worker::block_in_place;
tokio/src/runtime/scheduler/multi_thread/worker.rs
+19 -10
View File
@@ -170,7 +170,7 @@ struct Remote {
}
/// Thread-local context
struct Context {
pub(crate) struct Context {
/// Worker
worker: Arc<Worker>,
@@ -192,9 +192,6 @@ type Task = task::Task<Arc<Handle>>;
/// A notified task handle
type Notified = task::Notified<Arc<Handle>>;
// Tracks thread-local state
scoped_thread_local!(static CURRENT: Context);
/// Value picked out of thin-air. Running the LIFO slot a handful of times
/// seemms sufficient to benefit from locality. More than 3 times probably is
/// overweighing. The value can be tuned in the future with data that shows
@@ -277,7 +274,7 @@ where
impl Drop for Reset {
fn drop(&mut self) {
CURRENT.with(|maybe_cx| {
with_current(|maybe_cx| {
if let Some(cx) = maybe_cx {
let core = cx.worker.core.take();
let mut cx_core = cx.core.borrow_mut();
@@ -294,7 +291,7 @@ where
let mut had_entered = false;
let setup_result = CURRENT.with(|maybe_cx| {
let setup_result = with_current(|maybe_cx| {
match (
crate::runtime::context::current_enter_context(),
maybe_cx.is_some(),
@@ -414,12 +411,14 @@ fn run(worker: Arc<Worker>) {
let _enter = crate::runtime::context::enter_runtime(&handle, true);
// Set the worker context.
let cx = Context {
let cx = scheduler::Context::MultiThread(Context {
worker,
core: RefCell::new(None),
};
});
context::set_scheduler(&cx, || {
let cx = cx.expect_multi_thread();
CURRENT.set(&cx, || {
// This should always be an error. It only returns a `Result` to support
// using `?` to short circuit.
assert!(cx.run(core).is_err());
@@ -868,7 +867,7 @@ impl task::Schedule for Arc<Handle> {
impl Handle {
pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) {
CURRENT.with(|maybe_cx| {
with_current(|maybe_cx| {
if let Some(cx) = maybe_cx {
// Make sure the task is part of the **current** scheduler.
if self.ptr_eq(&cx.worker.handle) {
@@ -1008,6 +1007,16 @@ fn wake_deferred_tasks() {
context::with_defer(|deferred| deferred.wake());
}
#[track_caller]
fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R {
use scheduler::Context::MultiThread;
context::with_scheduler(|ctx| match ctx {
Some(MultiThread(ctx)) => f(Some(ctx)),
_ => f(None),
})
}
cfg_metrics! {
impl Shared {
pub(super) fn injection_queue_depth(&self) -> usize {