ntex_rt/

arbiter.rs

#![allow(clippy::missing_panics_doc)]
use std::sync::{Arc, atomic::AtomicBool, atomic::AtomicUsize, atomic::Ordering};
use std::{any::Any, any::TypeId, cell::RefCell, fmt, mem, panic, rc::Rc, thread};

use async_channel::{Receiver, Sender, unbounded};
use parking_lot::Mutex;

use crate::metrics::{LatencyMetrics, LatencyStats, start_latency_sampler};
use crate::{Handle, HashMap, Id, System};

thread_local!(
    static ADDR: RefCell<Option<Arbiter>> = const { RefCell::new(None) };
    static STORAGE: RefCell<HashMap<TypeId, Rc<dyn Any>>> = RefCell::new(HashMap::default());
    static ON_SHUTDOWN: RefCell<Vec<Box<dyn FnOnce()>>> = const { RefCell::new(Vec::new()) };
);

pub(super) static COUNT: AtomicUsize = AtomicUsize::new(99);

pub(super) enum ArbiterCommand {
    Stop,
}

/// An asynchronous execution environment running on one OS thread.
///
/// Creating an arbiter starts a thread with its own local event loop. Futures
/// spawned on that event loop are not required to implement `Send`.
/// The system's primary arbiter instead runs on the thread driving its
/// [`SystemRunner`](crate::SystemRunner).
///
/// Clones refer to the same arbiter and may be sent to other threads. Use
/// [`handle`](Self::handle) to submit `Send` futures from another thread.
pub struct Arbiter(pub(crate) Arc<ArbiterInner>);

type OnCloseStorage = Arc<Mutex<Vec<Box<dyn Fn() + Send + Sync>>>>;

pub(crate) struct ArbiterInner {
    id: usize,
    name: Arc<String>,
    sys_id: usize,
    hnd: Option<Handle>,
    pub(crate) sender: Sender<ArbiterCommand>,
    thread_handle: Mutex<Option<thread::JoinHandle<()>>>,
    on_stop: OnCloseStorage,
    running: AtomicBool,
    pub(crate) latency: Arc<LatencyStats>,
    #[cfg(target_os = "linux")]
    tid: i32,
    #[cfg(windows)]
    thread: Option<crate::watchdog::ThreadHandle>,
}

impl fmt::Debug for Arbiter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Arbiter({:?})", self.0.name.as_ref())
    }
}

impl Clone for Arbiter {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

impl Default for Arbiter {
    fn default() -> Self {
        Self::new()
    }
}

impl Arbiter {
    #[allow(clippy::borrowed_box)]
    pub(super) fn new_system(id: usize, name: String) -> (Self, ArbiterController) {
        let (tx, rx) = unbounded();

        let aid = COUNT.fetch_add(1, Ordering::Relaxed);
        let arb = Arbiter::with_sender(id, aid, Arc::new(name), tx, Arc::default());
        ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
        clear_storage();

        (
            arb,
            ArbiterController {
                rx,
                sys: None,
                stop: None,
            },
        )
    }

    /// Returns the arbiter running on the current thread.
    ///
    /// # Panics
    ///
    /// Panics if no arbiter is running on the current thread.
    pub fn current() -> Arbiter {
        ADDR.with(|cell| match *cell.borrow() {
            Some(ref addr) => addr.clone(),
            None => panic!("Arbiter is not running"),
        })
    }

    /// Requests that the arbiter stop its event loop.
    ///
    /// This returns without waiting for shutdown; use [`join`](Self::join) to
    /// wait for a spawned arbiter's thread. To stop the primary arbiter's event
    /// loop, stop its [`System`] or complete the future driving the system.
    pub fn stop(&self) {
        let _ = self.0.sender.try_send(ArbiterCommand::Stop);
    }

    /// Starts an arbiter on a new thread with an automatically generated name.
    ///
    /// Uses the current system's runner and thread configuration.
    ///
    /// # Panics
    ///
    /// Panics if no current system is installed or the arbiter thread cannot
    /// be started.
    pub fn new() -> Arbiter {
        let id = COUNT.fetch_add(1, Ordering::Relaxed);
        Arbiter::start(id, format!("{}:arb:{}", System::current().name(), id), None)
    }

    /// Starts an arbiter on a new thread with the specified name.
    ///
    /// Uses the current system's runner and thread configuration.
    ///
    /// # Panics
    ///
    /// Panics if no current system is installed or the arbiter thread cannot
    /// be started.
    pub fn with_name(name: String) -> Arbiter {
        Arbiter::start(COUNT.fetch_add(1, Ordering::Relaxed), name, None)
    }

    /// Starts an arbiter on a new thread with the specified name.
    ///
    /// `init` runs on the new thread before the arbiter's runtime is created,
    /// for example to set the thread's CPU affinity, so that the runtime's
    /// memory is allocated on the NUMA node of that CPU.
    /// The new thread has no current system or arbiter during `init`, and its
    /// arbiter storage is cleared after `init` returns.
    ///
    /// # Panics
    ///
    /// Panics if no current system is installed, the thread cannot be started,
    /// or initialization fails, including a panic in `init`.
    pub fn with_init<F>(name: String, init: F) -> Arbiter
    where
        F: FnOnce() + Send + 'static,
    {
        Arbiter::start(
            COUNT.fetch_add(1, Ordering::Relaxed),
            name,
            Some(Box::new(init)),
        )
    }

    fn start(id: usize, name: String, init: Option<Box<dyn FnOnce() + Send>>) -> Arbiter {
        let sys = System::current();
        let name2 = Arc::new(name.clone());
        let config = sys.config();
        let (arb_tx, arb_rx) = unbounded();

        let builder = if sys.config().stack_size > 0 {
            thread::Builder::new()
                .name(name)
                .stack_size(sys.config().stack_size)
        } else {
            thread::Builder::new().name(name)
        };

        let name = name2.clone();
        let sys_id = sys.id();
        let (arb_hnd_tx, arb_hnd_rx) = oneshot::channel();

        let handle = builder
            .spawn(move || {
                if let Some(init) = init {
                    init();
                }
                let name3 = name2.clone();
                log::info!("Starting {name3:?} arbiter");

                let sys2 = sys.clone();
                let (stop, stop_rx) = oneshot::channel();
                clear_storage();

                // `block_on` finalizes the arbiter once its event loop exits
                let result = crate::driver::block_on(config.runner.as_ref(), async move {
                    let arb = Arbiter::with_sender(sys_id.0, id, name2, arb_tx, Arc::default());
                    sys.register_arbiter(arb.clone());
                    // `Arbiter::latency()` is available once the arbiter is returned
                    start_latency_sampler(&arb.0.latency, config.latency_interval);
                    arb_hnd_tx
                        .send(arb.clone())
                        .expect("Controller thread has gone");

                    // start arbiter controller
                    crate::spawn_detached(
                        ArbiterController {
                            sys: None,
                            stop: Some(stop),
                            rx: arb_rx,
                        }
                        .run(sys),
                    );
                    ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));

                    // run loop
                    let _ = stop_rx.await;
                });

                let cleanup = || {
                    // unregister arbiter
                    sys2.unregister_arbiter(Id(id));
                    // skipped by `block_on` if the event loop panicked
                    run_shutdown_callbacks();
                    // SAFETY: The event loop and shutdown callbacks have ended, so storage borrows are released.
                    unsafe {
                        remove_all_items();
                    }
                };
                if result.is_err() {
                    crate::abort_on_panic(cleanup);
                } else {
                    cleanup();
                }

                if let Err(e) = result {
                    log::error!("Arbiter {name3:?} has panicked.");
                    panic::resume_unwind(e);
                }
                log::info!("Arbiter {name3:?} has stopped");
            })
            .unwrap_or_else(|err| panic!("Cannot spawn an arbiter's thread {name:?}: {err:?}"));

        let arb = arb_hnd_rx.recv().expect("Could not start new arbiter");
        *arb.0.thread_handle.lock() = Some(handle);
        arb
    }

    fn with_sender(
        sys_id: usize,
        id: usize,
        name: Arc<String>,
        sender: Sender<ArbiterCommand>,
        on_stop: OnCloseStorage,
    ) -> Self {
        #[cfg(feature = "tokio")]
        let hnd = { Handle::new() };

        #[cfg(feature = "compio")]
        let hnd = { Handle::new() };

        #[cfg(all(not(feature = "compio"), not(feature = "tokio")))]
        let hnd = { Handle::current() };

        Self(Arc::new(ArbiterInner {
            id,
            sys_id,
            name,
            sender,
            on_stop,
            hnd: Some(hnd),
            thread_handle: Mutex::new(None),
            running: AtomicBool::new(true),
            latency: Arc::default(),
            #[cfg(target_os = "linux")]
            #[allow(clippy::cast_possible_truncation)]
            // SAFETY: `SYS_gettid` takes no pointer arguments and returns this thread's ID.
            tid: unsafe { libc::syscall(libc::SYS_gettid) } as i32,
            #[cfg(windows)]
            thread: crate::watchdog::ThreadHandle::current(),
        }))
    }

    /// Returns the arbiter identifier.
    pub fn id(&self) -> Id {
        Id(self.0.id)
    }

    #[cfg(target_os = "linux")]
    /// TID of the arbiter
    pub(crate) fn tid(&self) -> i32 {
        self.0.tid
    }

    #[cfg(windows)]
    /// Handle of the arbiter thread
    pub(crate) fn thread(&self) -> Option<&crate::watchdog::ThreadHandle> {
        self.0.thread.as_ref()
    }

    /// Returns the arbiter name.
    pub fn name(&self) -> &str {
        self.0.name.as_ref()
    }

    #[inline]
    /// Returns a handle to the arbiter's runtime.
    pub fn handle(&self) -> &Handle {
        self.0.hnd.as_ref().unwrap()
    }

    /// Returns the arbiter's scheduling latency.
    ///
    /// Returns `None` if latency sampling is disabled, see
    /// [`Builder::latency_interval`](crate::Builder::latency_interval).
    pub fn latency(&self) -> Option<LatencyMetrics> {
        self.0.latency.metrics()
    }

    #[inline]
    /// Returns whether the arbiter is running.
    ///
    /// A stop request does not change this immediately. The flag is cleared
    /// after the event loop exits, before [`on_stop`](Self::on_stop) callbacks.
    pub fn is_running(&self) -> bool {
        self.0.running.load(Ordering::Relaxed)
    }

    /// Returns a value from thread-local arbiter storage, inserting it if absent.
    ///
    /// Values are keyed by type and shared with [`get_item`], [`set_item`], and
    /// [`with_item`] on the calling thread. Existing values are cloned.
    /// The initializer runs while storage is mutably borrowed and must not
    /// reenter these storage APIs.
    ///
    /// If the storage has already been destroyed because the thread is
    /// exiting, the value returned by `f` is not stored.
    pub fn get_value<T, F>(f: F) -> T
    where
        T: Clone + 'static,
        F: FnOnce() -> T,
    {
        let mut f = Some(f);
        STORAGE
            .try_with(|cell| {
                let mut st = cell.borrow_mut();
                if let Some(boxed) = st.get(&TypeId::of::<T>())
                    && let Some(val) = (&**boxed as &(dyn Any + 'static)).downcast_ref::<T>()
                {
                    return val.clone();
                }
                let val = (f.take().unwrap())();
                st.insert(TypeId::of::<T>(), Rc::new(val.clone()));
                val
            })
            .unwrap_or_else(|_| (f.take().unwrap())())
    }

    /// Registers a callback to run when the current thread's arbiter shuts down.
    ///
    /// Callbacks run once, in registration order, on the arbiter's thread when
    /// its stop is requested, by [`Arbiter::stop()`] or [`System::stop()`],
    /// while the event loop is still running. An arbiter that ends without a
    /// stop request, such as one driven by
    /// [`SystemRunner::block_on()`](crate::SystemRunner::block_on), runs them
    /// once its event loop has exited instead. A callback registered by
    /// another callback runs in the same shutdown.
    ///
    /// If the thread is exiting and its thread-local storage has already been
    /// destroyed, `f` is dropped without running. If `f` panics after the
    /// event loop panicked, the process is aborted.
    pub fn on_shutdown<F>(f: F)
    where
        F: FnOnce() + 'static,
    {
        let f: Box<dyn FnOnce()> = Box::new(f);
        let _ = ON_SHUTDOWN.try_with(move |cell| cell.borrow_mut().push(f));
    }

    #[must_use]
    /// Adds a callback to run after the arbiter stops.
    ///
    /// The callback runs on the arbiter thread after its event loop exits and
    /// its running flag is cleared. Register before shutdown; registration after
    /// finalization does not invoke the callback.
    ///
    /// If `f` panics after the event loop panicked, the process is aborted.
    pub fn on_stop<F>(self, f: F) -> Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.0.on_stop.lock().push(Box::new(f));
        self
    }

    /// Returns the arbiter of the current thread, if any.
    pub(crate) fn try_current() -> Option<Arbiter> {
        ADDR.try_with(|cell| cell.borrow().clone()).ok().flatten()
    }

    /// Marks the arbiter as stopped and runs its `on_stop` callbacks,
    /// called once its event loop has exited.
    pub(crate) fn finalize(&self) {
        self.0.running.store(false, Ordering::Relaxed);
        let on_stop = mem::take(&mut *self.0.on_stop.lock());
        // a panic while unwinding aborts the process
        for f in on_stop {
            f();
        }
    }

    /// Waits for the arbiter's thread to stop.
    ///
    /// This returns immediately for an arbiter that does not own a thread
    /// handle, including the system's primary arbiter. The join handle is shared
    /// by all clones and consumed by the first call. This does not request a stop.
    ///
    /// # Errors
    ///
    /// Returns the panic payload if the arbiter thread panicked.
    ///
    /// # Panics
    ///
    /// May panic when attempting to join the current thread.
    pub fn join(&mut self) -> thread::Result<()> {
        if let Some(thread_handle) = self.0.thread_handle.lock().take() {
            thread_handle.join()
        } else {
            Ok(())
        }
    }
}

impl Eq for Arbiter {}

impl PartialEq for Arbiter {
    fn eq(&self, other: &Self) -> bool {
        self.0.id == other.0.id && self.0.sys_id == other.0.sys_id
    }
}

pub(crate) struct ArbiterController {
    sys: Option<System>,
    rx: Receiver<ArbiterCommand>,
    stop: Option<oneshot::Sender<i32>>,
}

impl ArbiterController {
    pub(super) async fn run(mut self, sys: System) {
        self.sys = Some(sys);
        while let Ok(ArbiterCommand::Stop) = self.rx.recv().await {
            // the system arbiter has no `stop`, `System::stop()`
            // runs its callbacks
            if let Some(stop) = self.stop.take() {
                run_shutdown_callbacks();
                let _ = stop.send(0);
            }
        }
    }
}

/// Runs the callbacks registered with [`Arbiter::on_shutdown()`], including
/// ones registered while they run.
pub(crate) fn run_shutdown_callbacks() {
    loop {
        let callbacks = ON_SHUTDOWN
            .try_with(|cell| mem::take(&mut *cell.borrow_mut()))
            .unwrap_or_default();
        if callbacks.is_empty() {
            break;
        }
        for f in callbacks {
            f();
        }
    }
}

/// Inserts a value into the current arbiter's thread-local storage.
///
/// Values are keyed by type. Replaces and drops any previous value of type `T`;
/// the previous value is dropped after releasing the storage borrow.
///
/// If the storage has already been destroyed because the thread is
/// exiting, the value is dropped.
pub fn set_item<T: 'static>(item: T) {
    let item: Rc<dyn Any> = Rc::new(item);
    let old = STORAGE
        .try_with(move |cell| cell.borrow_mut().insert(TypeId::of::<T>(), item))
        .ok()
        .flatten();
    drop(old);
}

/// Returns a cloned value from the current arbiter's thread-local storage.
///
/// Returns `None` if no value of type `T` is stored, or if the storage has
/// already been destroyed because the thread is exiting.
pub fn get_item<T: Clone + 'static>() -> Option<T> {
    STORAGE
        .try_with(move |cell| {
            cell.borrow()
                .get(&TypeId::of::<T>())
                .and_then(|boxed| boxed.downcast_ref())
                .cloned()
        })
        .ok()
        .flatten()
}

/// Provides access to a value in the current arbiter's thread-local storage.
///
/// A default value is inserted if the requested type is not already present.
/// The callback holds a reference to that value independently of the storage,
/// so it may replace or remove stored values without invalidating its argument.
/// If the storage has already been destroyed because the thread is exiting,
/// `f` receives a temporary default value that is not stored.
pub fn with_item<T: Default + 'static, F, R>(f: F) -> R
where
    F: FnOnce(&T) -> R,
{
    // `f` holds its own reference, so the value stays alive even if `f`
    // replaces or removes it
    let val = STORAGE
        .try_with(|cell| {
            let existing = cell.borrow().get(&TypeId::of::<T>()).cloned();
            existing.unwrap_or_else(|| {
                let val: Rc<dyn Any> = Rc::new(T::default());
                cell.borrow_mut().insert(TypeId::of::<T>(), val.clone());
                val
            })
        })
        .ok();
    match val {
        Some(val) => f(val.downcast_ref::<T>().unwrap()),
        None => f(&T::default()),
    }
}

#[doc(hidden)]
/// Remove all items from storage.
///
/// # Safety
///
/// All outstanding calls to [`with_item`] must have completed.
pub unsafe fn remove_all_items() {
    clear_storage();
    System::remove_current();
}

/// Removes all items from the storage.
///
/// Each item is dropped outside of the storage borrow, so that its destructor
/// may access the storage. Items it inserts are removed too.
fn clear_storage() {
    let _ = STORAGE.try_with(|cell| {
        loop {
            let mut items = cell.borrow_mut();
            let Some(key) = items.keys().next().copied() else {
                break;
            };
            let item = items.remove(&key);
            drop(items);
            drop(item);
        }
    });
}

#[cfg(test)]
#[path = "../tests/unit/arbiter.rs"]
mod tests;