nuclei/rt/runtime/
enter.rs

1//! A runtime set on a thread without `block_on` ([`Runtime::enter`]), the
2//! handle of the runtime a thread is inside without falling back to the
3//! default one ([`RuntimeHandle::try_current`]), runtime ids, the worker a
4//! thread is, and a shutdown that does not wait.
5
6use core::fmt;
7use core::marker::PhantomData;
8use std::sync::Arc;
9
10use super::{Inner, Runtime, RuntimeHandle};
11use crate::rt::context;
12
13/// A runtime's id, distinct among the runtimes of the process.
14#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
15pub struct RuntimeId(u64);
16
17impl RuntimeId {
18    pub(crate) fn from_raw(id: u64) -> Self {
19        Self(id)
20    }
21}
22
23impl fmt::Display for RuntimeId {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        self.0.fmt(f)
26    }
27}
28
29/// The runtime [`Runtime::enter`] set on this thread, until it drops: the
30/// free functions ([`spawn`](crate::spawn), ...) use that runtime meanwhile.
31/// It stays on the thread that made it.
32#[must_use = "the runtime is set only while the guard lives"]
33pub struct EnterGuard<'a> {
34    _set: context::RuntimeSet,
35    _runtime: PhantomData<&'a ()>,
36    _not_send: PhantomData<*const ()>,
37}
38
39impl fmt::Debug for EnterGuard<'_> {
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        f.debug_struct("EnterGuard").finish_non_exhaustive()
42    }
43}
44
45/// [`RuntimeHandle::try_current`] found no runtime set on this thread.
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub struct TryCurrentError {
48    _private: (),
49}
50
51impl TryCurrentError {
52    /// Whether no runtime was set on the thread (the one reason today).
53    pub fn is_missing_context(&self) -> bool {
54        true
55    }
56}
57
58impl fmt::Display for TryCurrentError {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        f.write_str(
61            "no runtime is set on this thread: call it inside `block_on`, a task, or `enter`",
62        )
63    }
64}
65
66impl std::error::Error for TryCurrentError {}
67
68fn enter(inner: &Arc<Inner>) -> EnterGuard<'_> {
69    EnterGuard {
70        _set: context::set_runtime(inner),
71        _runtime: PhantomData,
72        _not_send: PhantomData,
73    }
74}
75
76impl Runtime {
77    /// Sets this runtime on the calling thread until the guard drops, so the
78    /// free functions use it outside `block_on`.
79    ///
80    /// ```
81    /// let runtime = nuclei::Runtime::builder().worker_threads(1).build().unwrap();
82    /// let _guard = runtime.enter();
83    /// let task = nuclei::spawn(async { 2 + 2 });
84    /// assert_eq!(runtime.block_on(task), 4);
85    /// ```
86    pub fn enter(&self) -> EnterGuard<'_> {
87        enter(&self.inner)
88    }
89
90    /// The runtime's id.
91    pub fn id(&self) -> RuntimeId {
92        RuntimeId(self.inner.id)
93    }
94
95    /// Stops the workers and returns at once: the worker threads drop their
96    /// tasks and exit on their own, without being waited for, and what the
97    /// handles of those tasks left behind is delivered on threads of its
98    /// own.
99    pub fn shutdown_background(self) {
100        let lingering = self.inner.lingering();
101        lingering.close();
102        self.inner.workers.stop_all();
103        // Detached: dropping the handles does not wait. `Drop` finds none
104        // left to join, and no lingering to wait for.
105        drop(self.inner.workers.take_threads());
106        lingering.end();
107    }
108}
109
110impl RuntimeHandle {
111    /// Sets the handle's runtime on the calling thread until the guard
112    /// drops.
113    pub fn enter(&self) -> EnterGuard<'_> {
114        enter(&self.inner)
115    }
116
117    /// The runtime's id.
118    pub fn id(&self) -> RuntimeId {
119        RuntimeId(self.inner.id)
120    }
121
122    /// The handle of the runtime set on this thread (inside its `block_on`,
123    /// on one of its workers, or under [`enter`](Self::enter)); unlike
124    /// [`current`](Self::current) it never falls back to the default
125    /// runtime.
126    pub fn try_current() -> Result<Self, TryCurrentError> {
127        context::runtime()
128            .map(|inner| Self { inner })
129            .ok_or(TryCurrentError { _private: () })
130    }
131}
132
133/// The index of the worker this code runs on (0, 1, ... among the workers
134/// of its runtime), `None` off the workers (in `block_on`'s own thread, on
135/// the blocking pool, outside any runtime).
136pub fn worker_index() -> Option<usize> {
137    let runtime = context::runtime()?;
138    context::current(|shard| shard.and_then(|shard| shard.worker_index(runtime.id)))
139}
nuclei/rt/runtime/enter.rs

nuclei/rt/runtime/
enter.rs
