palladium_runtime/multi_core/

handle.rs

use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::oneshot;

use palladium_actor::{ActorPath, AddrHash};
use palladium_transport::{InProcessTransport, MailboxReceiver, TransportRegistry};

use crate::engine::EngineHandle;
use crate::introspection::{ActorInfo, ActorQuery, EngineSnapshot};
use crate::placement::{CoreStats, PlacementMap};
use crate::reactor::{Reactor, TokioReactor};
use crate::registry::ActorRegistry;
use crate::responses::ResponseRegistry;

use super::errors::ShutdownError;

use crate::fs::{FileSystem, TokioFileSystem};

pub(crate) struct CoreHandle {
    #[allow(dead_code)]
    pub(crate) core_id: usize,
    pub(crate) shutdown_tx: Option<oneshot::Sender<()>>,
    pub(crate) thread: Option<std::thread::JoinHandle<()>>,
    #[allow(dead_code)]
    pub(crate) stats: Arc<CoreStats>,
}

/// Handle for a running [`MultiCoreEngine`].
///
/// Provides access to the shared transport, registry, and placement map, plus
/// graceful shutdown.
pub struct MultiCoreHandle<F: FileSystem = TokioFileSystem, R: Reactor = TokioReactor> {
    pub(crate) handles: Vec<CoreHandle>,
    pub(crate) transport: Arc<InProcessTransport>,
    pub(crate) transport_registry: Arc<TransportRegistry>,
    pub(crate) registry: Arc<ActorRegistry<R>>,
    pub(crate) responses: Arc<ResponseRegistry>,
    pub(crate) placement: Arc<PlacementMap>,
    pub(crate) start_time: Instant,
    pub(crate) ask_timeout: Duration,
    pub(crate) num_cores: usize,
    pub(crate) per_core_stats: Vec<Arc<CoreStats>>,
    /// Fires when the control plane receives an `engine.stop` RPC.
    /// `None` when no control plane socket was configured.
    pub(crate) internal_shutdown_rx: Option<oneshot::Receiver<()>>,
    pub(crate) federated_routing: Option<Arc<crate::federation::FederatedRouting>>,
    pub(crate) fs: F,
    pub(crate) reactor: R,
}

impl<F: FileSystem, R: Reactor> MultiCoreHandle<F, R> {
    /// Return an [`EngineHandle`] backed by this engine's shared transport.
    ///
    /// Useful in tests to build `Addr<M>` objects and interact with actors.
    pub fn engine_handle(&self) -> EngineHandle<R, palladium_transport::network::TokioNetwork, F> {
        EngineHandle {
            transport: self.transport.clone(),
            transport_registry: self.transport_registry.clone(),
            type_registry: self.transport_registry.type_registry(),
            responses: self.responses.clone(),
            registry: self.registry.clone(),
            ask_timeout: self.ask_timeout,
            reactor: self.reactor.clone(),
            network: palladium_transport::network::TokioNetwork,
            _fs: self.fs.clone(),
            start_time: self.start_time,
            source_addr: AddrHash::synthetic(b"engine-handle"),
            federated_routing: self.federated_routing.clone(),
            send_cache: Arc::new(dashmap::DashMap::new()),
            // MultiCoreHandle does not register a pump mailbox; remote_addr_for
            // ask() will time out without a pump task running.
            pump_rx: Arc::new(std::sync::Mutex::new(None::<MailboxReceiver>)),
        }
    }

    /// Access the shared transport directly.
    pub fn transport(&self) -> &Arc<InProcessTransport> {
        &self.transport
    }

    /// Access the placement map.
    pub fn placement(&self) -> &Arc<PlacementMap> {
        &self.placement
    }

    /// List actors, optionally filtered.
    pub fn actor_list(&self, query: ActorQuery) -> Vec<ActorInfo> {
        self.registry.snapshot(&query, 0)
    }

    /// Get info for a single actor by path.
    pub fn actor_info(&self, path: &ActorPath) -> Option<ActorInfo> {
        let core_id = self
            .placement
            .get(&AddrHash::new(path, 0).path_hash()) // Generation 0 lookups are fine for discovery
            .map(|v| *v)
            .unwrap_or(0);
        self.registry.actor_info_by_path(path, core_id)
    }

    /// System-level snapshot aggregated across all cores.
    pub fn snapshot(&self) -> EngineSnapshot {
        let actors = self.actor_list(ActorQuery::default());
        EngineSnapshot {
            num_cores: self.num_cores,
            uptime_secs: self.start_time.elapsed().as_secs(),
            actors,
        }
    }

    /// Per-core statistics counters.
    pub fn core_stats(&self) -> &[Arc<CoreStats>] {
        &self.per_core_stats
    }

    /// Block until actor at `path` is running, or `timeout` elapses.
    ///
    /// Returns `true` if the actor became running within `timeout`.
    pub fn wait_for_actor(&self, path: &ActorPath, timeout: Duration) -> bool {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if self
                .registry
                .get_by_path(path)
                .is_some_and(|s| s.running.load(Ordering::Relaxed))
            {
                return true;
            }
            std::thread::sleep(Duration::from_millis(10));
        }
        false
    }

    /// Block until the engine shuts down, then gracefully stop all cores.
    ///
    /// Waits for the internal shutdown signal (fired by the control plane's
    /// `engine.stop` RPC) and then calls [`shutdown`](Self::shutdown).
    ///
    /// If no control plane socket was configured, this method returns
    /// immediately (the caller is responsible for calling `shutdown()`).
    pub fn run(self) -> Result<(), ShutdownError> {
        let reactor = self.reactor.clone();
        let Self {
            handles,
            transport,
            transport_registry,
            registry,
            responses,
            placement,
            start_time,
            ask_timeout,
            num_cores,
            per_core_stats,
            internal_shutdown_rx,
            federated_routing,
            fs,
            reactor: _,
        } = self;

        if let Some(rx) = internal_shutdown_rx {
            // Block the calling thread until the control plane fires engine.stop.
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("failed to build shutdown-wait runtime");
            let _ = rt.block_on(rx);
        }

        MultiCoreHandle {
            handles,
            transport,
            transport_registry,
            registry,
            responses,
            placement,
            start_time,
            ask_timeout,
            num_cores,
            per_core_stats,
            internal_shutdown_rx: None,
            federated_routing,
            fs,
            reactor,
        }
        .shutdown()
    }

    /// Gracefully shut down all cores.
    ///
    /// Sends shutdown signals to every core, waits for threads to join, then
    /// performs brutal cancellation of any remaining actor tasks.
    pub fn shutdown(mut self) -> Result<(), ShutdownError> {
        for h in &mut self.handles {
            if let Some(tx) = h.shutdown_tx.take() {
                let _ = tx.send(());
            }
        }
        let mut panicked = false;
        for h in &mut self.handles {
            if let Some(thread) = h.thread.take() {
                if thread.join().is_err() {
                    panicked = true;
                }
            }
        }
        // Brutal escalation: cancel any tasks that didn't drain in time.
        self.registry.cancel_all();
        self.responses.drain();

        if panicked {
            Err(ShutdownError::ThreadPanic)
        } else {
            Ok(())
        }
    }
}