actix_rt/
system.rs

1use std::{
2    cell::RefCell,
3    collections::HashMap,
4    future::Future,
5    io,
6    pin::Pin,
7    sync::atomic::{AtomicUsize, Ordering},
8    task::{Context, Poll},
9};
10
11use futures_core::ready;
12use tokio::sync::{mpsc, oneshot};
13
14use crate::{arbiter::ArbiterHandle, Arbiter};
15
16static SYSTEM_COUNT: AtomicUsize = AtomicUsize::new(0);
17
18thread_local!(
19    static CURRENT: RefCell<Option<System>> = const { RefCell::new(None) };
20);
21
22/// A manager for a per-thread distributed async runtime.
23#[derive(Clone, Debug)]
24pub struct System {
25    id: usize,
26    sys_tx: mpsc::UnboundedSender<SystemCommand>,
27
28    /// Handle to the first [Arbiter] that is created with the System.
29    arbiter_handle: ArbiterHandle,
30}
31
32#[cfg(not(feature = "io-uring"))]
33impl System {
34    /// Create a new system.
35    ///
36    /// # Panics
37    /// Panics if underlying Tokio runtime can not be created.
38    #[allow(clippy::new_ret_no_self)]
39    pub fn new() -> SystemRunner {
40        Self::with_tokio_rt(|| {
41            crate::runtime::default_tokio_runtime()
42                .expect("Default Actix (Tokio) runtime could not be created.")
43        })
44    }
45
46    /// Create a new System using the [Tokio Runtime](tokio-runtime) returned from a closure.
47    ///
48    /// [tokio-runtime]: tokio::runtime::Runtime
49    pub fn with_tokio_rt<F>(runtime_factory: F) -> SystemRunner
50    where
51        F: FnOnce() -> tokio::runtime::Runtime,
52    {
53        let (stop_tx, stop_rx) = oneshot::channel();
54        let (sys_tx, sys_rx) = mpsc::unbounded_channel();
55
56        let rt = crate::runtime::Runtime::from(runtime_factory());
57        let sys_arbiter = rt.block_on(async { Arbiter::in_new_system() });
58        let system = System::construct(sys_tx, sys_arbiter.clone());
59
60        system
61            .tx()
62            .send(SystemCommand::RegisterArbiter(usize::MAX, sys_arbiter))
63            .unwrap();
64
65        // init background system arbiter
66        let sys_ctrl = SystemController::new(sys_rx, stop_tx);
67        rt.spawn(sys_ctrl);
68
69        SystemRunner { rt, stop_rx }
70    }
71}
72
73#[cfg(feature = "io-uring")]
74impl System {
75    /// Create a new system.
76    ///
77    /// # Panics
78    /// Panics if underlying Tokio runtime can not be created.
79    #[allow(clippy::new_ret_no_self)]
80    pub fn new() -> SystemRunner {
81        SystemRunner
82    }
83
84    /// Create a new System using the [Tokio Runtime](tokio-runtime) returned from a closure.
85    ///
86    /// [tokio-runtime]: tokio::runtime::Runtime
87    #[doc(hidden)]
88    pub fn with_tokio_rt<F>(_: F) -> SystemRunner
89    where
90        F: FnOnce() -> tokio::runtime::Runtime,
91    {
92        unimplemented!("System::with_tokio_rt is not implemented for io-uring feature yet")
93    }
94}
95
96impl System {
97    /// Constructs new system and registers it on the current thread.
98    pub(crate) fn construct(
99        sys_tx: mpsc::UnboundedSender<SystemCommand>,
100        arbiter_handle: ArbiterHandle,
101    ) -> Self {
102        let sys = System {
103            sys_tx,
104            arbiter_handle,
105            id: SYSTEM_COUNT.fetch_add(1, Ordering::SeqCst),
106        };
107
108        System::set_current(sys.clone());
109
110        sys
111    }
112
113    /// Get current running system.
114    ///
115    /// # Panics
116    /// Panics if no system is registered on the current thread.
117    pub fn current() -> System {
118        CURRENT.with(|cell| match *cell.borrow() {
119            Some(ref sys) => sys.clone(),
120            None => panic!("System is not running"),
121        })
122    }
123
124    /// Try to get current running system.
125    ///
126    /// Returns `None` if no System has been started.
127    ///
128    /// Unlike [`current`](Self::current), this never panics.
129    pub fn try_current() -> Option<System> {
130        CURRENT.with(|cell| cell.borrow().clone())
131    }
132
133    /// Get handle to a the System's initial [Arbiter].
134    pub fn arbiter(&self) -> &ArbiterHandle {
135        &self.arbiter_handle
136    }
137
138    /// Check if there is a System registered on the current thread.
139    pub fn is_registered() -> bool {
140        CURRENT.with(|sys| sys.borrow().is_some())
141    }
142
143    /// Register given system on current thread.
144    #[doc(hidden)]
145    pub fn set_current(sys: System) {
146        CURRENT.with(|cell| {
147            *cell.borrow_mut() = Some(sys);
148        })
149    }
150
151    /// Numeric system identifier.
152    ///
153    /// Useful when using multiple Systems.
154    pub fn id(&self) -> usize {
155        self.id
156    }
157
158    /// Stop the system (with code 0).
159    pub fn stop(&self) {
160        self.stop_with_code(0)
161    }
162
163    /// Stop the system with a given exit code.
164    pub fn stop_with_code(&self, code: i32) {
165        let _ = self.sys_tx.send(SystemCommand::Exit(code));
166    }
167
168    pub(crate) fn tx(&self) -> &mpsc::UnboundedSender<SystemCommand> {
169        &self.sys_tx
170    }
171}
172
173/// Runner that keeps a [System]'s event loop alive until stop message is received.
174#[cfg(not(feature = "io-uring"))]
175#[must_use = "A SystemRunner does nothing unless `run` is called."]
176#[derive(Debug)]
177pub struct SystemRunner {
178    rt: crate::runtime::Runtime,
179    stop_rx: oneshot::Receiver<i32>,
180}
181
182#[cfg(not(feature = "io-uring"))]
183impl SystemRunner {
184    /// Starts event loop and will return once [System] is [stopped](System::stop).
185    pub fn run(self) -> io::Result<()> {
186        let exit_code = self.run_with_code()?;
187
188        match exit_code {
189            0 => Ok(()),
190            nonzero => Err(io::Error::other(format!("Non-zero exit code: {nonzero}"))),
191        }
192    }
193
194    /// Runs the event loop until [stopped](System::stop_with_code), returning the exit code.
195    pub fn run_with_code(self) -> io::Result<i32> {
196        let SystemRunner { rt, stop_rx, .. } = self;
197
198        // run loop
199        rt.block_on(stop_rx).map_err(io::Error::other)
200    }
201
202    /// Retrieves a reference to the underlying [Actix runtime](crate::Runtime) associated with this
203    /// `SystemRunner` instance.
204    ///
205    /// The Actix runtime is responsible for managing the event loop for an Actix system and
206    /// executing asynchronous tasks. This method provides access to the runtime, allowing direct
207    /// interaction with its features.
208    ///
209    /// In a typical use case, you might need to share the same runtime between different
210    /// parts of your project. For example, some components might require a [`Runtime`] to spawn
211    /// tasks on the same runtime.
212    ///
213    /// Read more in the documentation for [`Runtime`].
214    ///
215    /// # Examples
216    ///
217    /// ```
218    /// let system_runner = actix_rt::System::new();
219    /// let actix_runtime = system_runner.runtime();
220    ///
221    /// // Use the runtime to spawn an async task or perform other operations
222    /// ```
223    ///
224    /// # Note
225    ///
226    /// While this method provides an immutable reference to the Actix runtime, which is safe to
227    /// share across threads, be aware that spawning blocking tasks on the Actix runtime could
228    /// potentially impact system performance. This is because the Actix runtime is responsible for
229    /// driving the system, and blocking tasks could delay other tasks in the run loop.
230    ///
231    /// [`Runtime`]: crate::Runtime
232    pub fn runtime(&self) -> &crate::runtime::Runtime {
233        &self.rt
234    }
235
236    /// Runs the provided future, blocking the current thread until the future completes.
237    #[track_caller]
238    #[inline]
239    pub fn block_on<F: Future>(&self, fut: F) -> F::Output {
240        self.rt.block_on(fut)
241    }
242}
243
244/// Runner that keeps a [System]'s event loop alive until stop message is received.
245#[cfg(feature = "io-uring")]
246#[must_use = "A SystemRunner does nothing unless `run` is called."]
247#[derive(Debug)]
248pub struct SystemRunner;
249
250#[cfg(feature = "io-uring")]
251impl SystemRunner {
252    /// Starts event loop and will return once [System] is [stopped](System::stop).
253    pub fn run(self) -> io::Result<()> {
254        unimplemented!("SystemRunner::run is not implemented for io-uring feature yet");
255    }
256
257    /// Runs the event loop until [stopped](System::stop_with_code), returning the exit code.
258    pub fn run_with_code(self) -> io::Result<i32> {
259        unimplemented!("SystemRunner::run_with_code is not implemented for io-uring feature yet");
260    }
261
262    /// Runs the provided future, blocking the current thread until the future completes.
263    #[inline]
264    pub fn block_on<F: Future>(&self, fut: F) -> F::Output {
265        tokio_uring::start(async move {
266            let (stop_tx, stop_rx) = oneshot::channel();
267            let (sys_tx, sys_rx) = mpsc::unbounded_channel();
268
269            let sys_arbiter = Arbiter::in_new_system();
270            let system = System::construct(sys_tx, sys_arbiter.clone());
271
272            system
273                .tx()
274                .send(SystemCommand::RegisterArbiter(usize::MAX, sys_arbiter))
275                .unwrap();
276
277            // init background system arbiter
278            let sys_ctrl = SystemController::new(sys_rx, stop_tx);
279            tokio_uring::spawn(sys_ctrl);
280
281            let res = fut.await;
282            drop(stop_rx);
283            res
284        })
285    }
286}
287
288#[derive(Debug)]
289pub(crate) enum SystemCommand {
290    Exit(i32),
291    RegisterArbiter(usize, ArbiterHandle),
292    DeregisterArbiter(usize),
293}
294
295/// There is one `SystemController` per [System]. It runs in the background, keeping track of
296/// [Arbiter]s and is able to distribute a system-wide stop command.
297#[derive(Debug)]
298pub(crate) struct SystemController {
299    stop_tx: Option<oneshot::Sender<i32>>,
300    cmd_rx: mpsc::UnboundedReceiver<SystemCommand>,
301    arbiters: HashMap<usize, ArbiterHandle>,
302}
303
304impl SystemController {
305    pub(crate) fn new(
306        cmd_rx: mpsc::UnboundedReceiver<SystemCommand>,
307        stop_tx: oneshot::Sender<i32>,
308    ) -> Self {
309        SystemController {
310            cmd_rx,
311            stop_tx: Some(stop_tx),
312            arbiters: HashMap::with_capacity(4),
313        }
314    }
315}
316
317impl Future for SystemController {
318    type Output = ();
319
320    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
321        // process all items currently buffered in channel
322        loop {
323            match ready!(self.cmd_rx.poll_recv(cx)) {
324                // channel closed; no more messages can be received
325                None => return Poll::Ready(()),
326
327                // process system command
328                Some(cmd) => match cmd {
329                    SystemCommand::Exit(code) => {
330                        // stop all arbiters
331                        for arb in self.arbiters.values() {
332                            arb.stop();
333                        }
334
335                        // stop event loop
336                        // will only fire once
337                        if let Some(stop_tx) = self.stop_tx.take() {
338                            let _ = stop_tx.send(code);
339                        }
340                    }
341
342                    SystemCommand::RegisterArbiter(id, arb) => {
343                        self.arbiters.insert(id, arb);
344                    }
345
346                    SystemCommand::DeregisterArbiter(id) => {
347                        self.arbiters.remove(&id);
348                    }
349                },
350            }
351        }
352    }
353}
actix_rt/system.rs

actix_rt/
system.rs
