actix-rt 2.12.0

use std::{
    cell::RefCell,
    collections::HashMap,
    future::Future,
    io,
    pin::Pin,
    sync::atomic::{AtomicUsize, Ordering},
    task::{Context, Poll},
};

use futures_core::ready;
use tokio::sync::{mpsc, watch};

use crate::{arbiter::ArbiterHandle, Arbiter};

static SYSTEM_COUNT: AtomicUsize = AtomicUsize::new(0);

thread_local!(
    static CURRENT: RefCell<Option<System>> = const { RefCell::new(None) };
);

/// A manager for a per-thread distributed async runtime.
#[derive(Clone, Debug)]
pub struct System {
    id: usize,
    sys_tx: mpsc::UnboundedSender<SystemCommand>,

    /// Handle to the first [Arbiter] that is created with the System.
    arbiter_handle: ArbiterHandle,
}

impl System {
    /// Create a new system.
    ///
    /// # Panics
    /// Panics if underlying Tokio runtime can not be created.
    #[allow(clippy::new_ret_no_self)]
    pub fn new() -> SystemRunner {
        Self::with_tokio_rt(|| {
            crate::runtime::default_tokio_runtime()
                .expect("Default Actix (Tokio) runtime could not be created.")
        })
    }

    /// Create a new System using the [Tokio Runtime](tokio-runtime) returned from a closure.
    ///
    /// The closure may return any type that can be converted into [`Runtime`], such as
    /// `tokio::runtime::Runtime`, `Arc<tokio::runtime::Runtime>`, or
    /// `&'static tokio::runtime::Runtime`.
    ///
    /// [tokio-runtime]: tokio::runtime::Runtime
    /// [`Runtime`]: crate::Runtime
    pub fn with_tokio_rt<F, R>(runtime_factory: F) -> SystemRunner
    where
        F: FnOnce() -> R,
        R: Into<crate::runtime::Runtime>,
    {
        let (stop_tx, stop_rx) = watch::channel(None);
        let (sys_tx, sys_rx) = mpsc::unbounded_channel();

        let rt = runtime_factory().into();
        let sys_arbiter = rt.block_on(async { Arbiter::in_new_system() });
        let system = System::construct(sys_tx, sys_arbiter.clone());

        system
            .tx()
            .send(SystemCommand::RegisterArbiter(usize::MAX, sys_arbiter))
            .unwrap();

        // init background system arbiter
        let sys_ctrl = SystemController::new(sys_rx, stop_tx);
        rt.spawn(sys_ctrl);

        SystemRunner { rt, stop_rx }
    }
}

impl System {
    /// Constructs new system and registers it on the current thread.
    pub(crate) fn construct(
        sys_tx: mpsc::UnboundedSender<SystemCommand>,
        arbiter_handle: ArbiterHandle,
    ) -> Self {
        let sys = System {
            sys_tx,
            arbiter_handle,
            id: SYSTEM_COUNT.fetch_add(1, Ordering::SeqCst),
        };

        System::set_current(sys.clone());

        sys
    }

    /// Get current running system.
    ///
    /// # Panics
    /// Panics if no system is registered on the current thread.
    pub fn current() -> System {
        CURRENT.with(|cell| match *cell.borrow() {
            Some(ref sys) => sys.clone(),
            None => panic!("System is not running"),
        })
    }

    /// Try to get current running system.
    ///
    /// Returns `None` if no System has been started.
    ///
    /// Unlike [`current`](Self::current), this never panics.
    pub fn try_current() -> Option<System> {
        CURRENT.with(|cell| cell.borrow().clone())
    }

    /// Get handle to a the System's initial [Arbiter].
    pub fn arbiter(&self) -> &ArbiterHandle {
        &self.arbiter_handle
    }

    /// Check if there is a System registered on the current thread.
    pub fn is_registered() -> bool {
        CURRENT.with(|sys| sys.borrow().is_some())
    }

    /// Register given system on current thread.
    #[doc(hidden)]
    pub fn set_current(sys: System) {
        CURRENT.with(|cell| {
            *cell.borrow_mut() = Some(sys);
        })
    }

    /// Numeric system identifier.
    ///
    /// Useful when using multiple Systems.
    pub fn id(&self) -> usize {
        self.id
    }

    /// Stop the system (with code 0).
    pub fn stop(&self) {
        self.stop_with_code(0)
    }

    /// Stop the system with a given exit code.
    pub fn stop_with_code(&self, code: i32) {
        let _ = self.sys_tx.send(SystemCommand::Exit(code));
    }

    pub(crate) fn tx(&self) -> &mpsc::UnboundedSender<SystemCommand> {
        &self.sys_tx
    }
}

/// Runner that keeps a [System]'s event loop alive until stop message is received.
#[must_use = "A SystemRunner does nothing unless `run` is called."]
#[derive(Debug)]
pub struct SystemRunner {
    rt: crate::runtime::Runtime,
    stop_rx: watch::Receiver<Option<i32>>,
}

impl SystemRunner {
    /// Starts event loop and will return once [System] is [stopped](System::stop).
    pub fn run(self) -> io::Result<()> {
        let exit_code = self.run_with_code()?;

        match exit_code {
            0 => Ok(()),
            nonzero => Err(io::Error::other(format!("Non-zero exit code: {nonzero}"))),
        }
    }

    /// Runs the event loop until [stopped](System::stop_with_code), returning the exit code.
    pub fn run_with_code(self) -> io::Result<i32> {
        let SystemRunner { rt, stop_rx, .. } = self;

        // run loop
        rt.block_on(wait_for_stop(stop_rx))
    }

    /// Retrieves a reference to the underlying [Actix runtime](crate::Runtime) associated with this
    /// `SystemRunner` instance.
    ///
    /// The Actix runtime is responsible for managing the event loop for an Actix system and
    /// executing asynchronous tasks. This method provides access to the runtime, allowing direct
    /// interaction with its features.
    ///
    /// In a typical use case, you might need to share the same runtime between different
    /// parts of your project. For example, some components might require a [`Runtime`] to spawn
    /// tasks on the same runtime.
    ///
    /// Read more in the documentation for [`Runtime`].
    ///
    /// # Examples
    ///
    /// ```
    /// let system_runner = actix_rt::System::new();
    /// let actix_runtime = system_runner.runtime();
    ///
    /// // Use the runtime to spawn an async task or perform other operations
    /// ```
    ///
    /// # Note
    ///
    /// While this method provides an immutable reference to the Actix runtime, which is safe to
    /// share across threads, be aware that spawning blocking tasks on the Actix runtime could
    /// potentially impact system performance. This is because the Actix runtime is responsible for
    /// driving the system, and blocking tasks could delay other tasks in the run loop.
    ///
    /// [`Runtime`]: crate::Runtime
    pub fn runtime(&self) -> &crate::runtime::Runtime {
        &self.rt
    }

    /// Returns a future that resolves with the system's exit code when it is stopped.
    ///
    /// This can be used to react to a system stop signal while running a future with
    /// [`SystemRunner::block_on`], such as when coordinating shutdown with `tokio::select!`.
    ///
    /// # Examples
    /// ```no_run
    /// use std::process::ExitCode;
    /// use actix_rt::System;
    ///
    /// let sys = System::new();
    /// let stop = sys.stop_future();
    ///
    /// let exit = sys.block_on(async move {
    ///     actix_rt::spawn(async {
    ///         System::current().stop_with_code(0);
    ///     });
    ///
    ///     let code = stop.await.unwrap_or(1);
    ///     ExitCode::from(code as u8)
    /// });
    ///
    /// # drop(exit);
    /// ```
    pub fn stop_future(&self) -> SystemStop {
        SystemStop::new(self.stop_rx.clone())
    }

    /// Splits this runner into its runtime and a future that resolves when the system stops.
    ///
    /// After calling this method, [`SystemRunner::run`] and [`SystemRunner::run_with_code`] can no
    /// longer be used.
    pub fn into_parts(self) -> (crate::runtime::Runtime, SystemStop) {
        let SystemRunner { rt, stop_rx } = self;
        (rt, SystemStop::new(stop_rx))
    }

    /// Runs the provided future, blocking the current thread until the future completes.
    #[track_caller]
    #[inline]
    pub fn block_on<F: Future>(&self, fut: F) -> F::Output {
        self.rt.block_on(fut)
    }
}

/// Future that resolves with the exit code when a [`System`] is stopped.
#[must_use = "SystemStop does nothing unless polled or awaited."]
pub struct SystemStop {
    inner: Pin<Box<dyn Future<Output = io::Result<i32>> + 'static>>,
}

impl SystemStop {
    fn new(stop_rx: watch::Receiver<Option<i32>>) -> Self {
        Self {
            inner: Box::pin(wait_for_stop(stop_rx)),
        }
    }
}

impl Future for SystemStop {
    type Output = io::Result<i32>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.inner.as_mut().poll(cx)
    }
}

async fn wait_for_stop(mut stop_rx: watch::Receiver<Option<i32>>) -> io::Result<i32> {
    loop {
        if let Some(code) = *stop_rx.borrow() {
            return Ok(code);
        }

        stop_rx.changed().await.map_err(io::Error::other)?;
    }
}

#[derive(Debug)]
pub(crate) enum SystemCommand {
    Exit(i32),
    RegisterArbiter(usize, ArbiterHandle),
    DeregisterArbiter(usize),
}

/// There is one `SystemController` per [System]. It runs in the background, keeping track of
/// [Arbiter]s and is able to distribute a system-wide stop command.
#[derive(Debug)]
pub(crate) struct SystemController {
    stop_tx: Option<watch::Sender<Option<i32>>>,
    cmd_rx: mpsc::UnboundedReceiver<SystemCommand>,
    arbiters: HashMap<usize, ArbiterHandle>,
}

impl SystemController {
    pub(crate) fn new(
        cmd_rx: mpsc::UnboundedReceiver<SystemCommand>,
        stop_tx: watch::Sender<Option<i32>>,
    ) -> Self {
        SystemController {
            cmd_rx,
            stop_tx: Some(stop_tx),
            arbiters: HashMap::with_capacity(4),
        }
    }
}

impl Future for SystemController {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // process all items currently buffered in channel
        loop {
            match ready!(self.cmd_rx.poll_recv(cx)) {
                // channel closed; no more messages can be received
                None => return Poll::Ready(()),

                // process system command
                Some(cmd) => match cmd {
                    SystemCommand::Exit(code) => {
                        // stop all arbiters
                        for arb in self.arbiters.values() {
                            arb.stop();
                        }

                        // stop event loop
                        // will only fire once
                        if let Some(stop_tx) = self.stop_tx.take() {
                            let _ = stop_tx.send(Some(code));
                        }
                    }

                    SystemCommand::RegisterArbiter(id, arb) => {
                        self.arbiters.insert(id, arb);
                    }

                    SystemCommand::DeregisterArbiter(id) => {
                        self.arbiters.remove(&id);
                    }
                },
            }
        }
    }
}