palladium_runtime/multi_core/

engine.rs

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

use tokio::sync::oneshot;

use palladium_actor::{
    ActorPath, AddrHash, ChildSpec, NamespacePolicy, RestartPolicy, ShutdownPolicy,
};
use palladium_transport::{InProcessTransport, TransportRegistry};

use crate::control_plane::ControlPlaneActor;
use crate::engine::EngineConfig;
use crate::federation::FederatedRouting;
use crate::fs::{FileSystem, TokioFileSystem};
use crate::placement::{CoreStats, Placement, PlacementMap};
use crate::registry::ActorRegistry;
use crate::responses::ResponseRegistry;
use crate::ring_buffer::{CacheLineAtomicBool, InterCoreQueue};

use super::core_loop::core_run_async;
use super::errors::EngineError;
use super::handle::{CoreHandle, MultiCoreHandle};
use super::helpers::{default_num_cores, pin_thread_to_core, resolve_placement};

/// Builder for a multi-core actor engine.
///
/// Each core runs an independent `current_thread` Tokio runtime (optionally
/// pinned to a physical CPU). All cores share a single [`InProcessTransport`]
/// so cross-core sends work transparently: a message from core 0 to an actor
/// on core 1 is pushed directly into that actor's Tokio mpsc mailbox, which
/// the receiving core's task wakes up to process.
///
/// # Example
/// ```rust,no_run
/// use palladium_runtime::{MultiCoreEngine, EngineConfig};
/// use palladium_actor::ChildSpec;
/// use palladium_runtime::Placement;
///
/// let mut engine = MultiCoreEngine::new(2);
/// // engine.add_user_actor(spec, Placement::Core(0));
/// let handle = engine.start().unwrap();
/// handle.shutdown().unwrap();
/// ```
pub struct MultiCoreEngine<F: FileSystem = TokioFileSystem> {
    config: EngineConfig<crate::reactor::TokioReactor>,
    num_cores: usize,
    /// Actor specs indexed by core id.
    user_specs_per_core: Vec<Vec<ChildSpec<crate::reactor::TokioReactor>>>,
    placement: Arc<PlacementMap>,
    rings: Arc<Vec<Vec<InterCoreQueue>>>,
    fs: F,
}

impl MultiCoreEngine<TokioFileSystem> {
    /// Create a new engine with `num_cores` cores and default config.
    pub fn new(num_cores: usize) -> Self {
        Self::with_config(EngineConfig {
            num_cores: Some(num_cores),
            ..Default::default()
        })
    }

    /// Create a new engine from an [`EngineConfig`].
    ///
    /// `num_cores` is taken from `config.num_cores`, falling back to
    /// `std::thread::available_parallelism()`.
    pub fn with_config(config: EngineConfig<crate::reactor::TokioReactor>) -> Self {
        Self::with_fs(config, TokioFileSystem)
    }
}

impl<F: FileSystem> MultiCoreEngine<F> {
    /// Create a new engine with a specific filesystem.
    pub fn with_fs(config: EngineConfig<crate::reactor::TokioReactor>, fs: F) -> Self {
        let num_cores = config.num_cores.unwrap_or_else(default_num_cores).max(1);
        let capacity = config.inter_core_queue_capacity;
        let rings = (0..num_cores)
            .map(|_| {
                (0..num_cores)
                    .map(|_| InterCoreQueue::new(capacity))
                    .collect()
            })
            .collect();
        Self {
            placement: Arc::new(PlacementMap::new()),
            config,
            num_cores,
            user_specs_per_core: vec![Vec::new(); num_cores],
            rings: Arc::new(rings),
            fs,
        }
    }

    /// Register an actor to be spawned at startup on a specific core.
    ///
    /// Pre-populates the placement map so cross-core routing is immediately
    /// available once the engine starts.
    pub fn add_user_actor(
        &mut self,
        spec: ChildSpec<crate::reactor::TokioReactor>,
        placement: Placement,
    ) {
        let core_id = resolve_placement(&placement, &self.placement, self.num_cores, &spec.name);
        // Pre-compute the actor's path so the placement map is ready before start.
        let parent =
            ActorPath::parse(&format!("/user/core{core_id}")).expect("valid core supervisor path");
        if let Ok(child_path) = parent.child(&spec.name) {
            let addr = AddrHash::new(&child_path, 0);
            self.placement.insert(addr.path_hash(), core_id);
        }
        self.user_specs_per_core[core_id].push(spec);
    }

    /// Start the engine: spawn N threads, each running a Tokio runtime.
    ///
    /// Returns a [`MultiCoreHandle`] immediately — actors are started
    /// asynchronously inside the spawned threads. Use
    /// [`MultiCoreHandle::wait_for_actor`] to synchronise before sending.
    ///
    /// If [`EngineConfig::control_plane_socket`] is set, a
    /// [`ControlPlaneActor`] is added to core 0's user supervisor so the
    /// engine can be managed via `pd status`, `pd actor`, and `pd stop`.
    pub fn start(self) -> Result<MultiCoreHandle<F>, EngineError> {
        let registry = Arc::new(ActorRegistry::<crate::reactor::TokioReactor>::new());
        let responses = Arc::new(ResponseRegistry::new(self.config.response_capacity));
        let placement = self.placement;
        let rings = self.rings;
        let start_time = Instant::now();
        let num_cores = self.num_cores;
        let mut config = self.config;
        let fs = self.fs;

        // Shared state for all InProcessTransport instances.
        let global_locals = Arc::new(dashmap::DashMap::new());
        let death_listeners = Arc::new(parking_lot::Mutex::new(std::collections::BTreeMap::new()));

        // Per-core work-available flags (cache-line padded, one per core).
        // Sender cores set has_work[dst] after a ring push; the drain task on
        // core `dst` uses the flag to avoid unnecessary yield_now() / kernel
        // wakeups when more inter-core traffic is already in-flight.
        let has_work: Arc<Vec<CacheLineAtomicBool>> = Arc::new(
            (0..num_cores)
                .map(|_| CacheLineAtomicBool::new(false))
                .collect(),
        );
        let core_idle: Arc<Vec<CacheLineAtomicBool>> = Arc::new(
            (0..num_cores)
                .map(|_| CacheLineAtomicBool::new(true))
                .collect(),
        );

        let federated_routing = maybe_init_federation(&mut config);
        ensure_consensus_config(&mut config);

        // Internal shutdown channel
        let (internal_shutdown_tx, internal_shutdown_rx) = tokio::sync::oneshot::channel::<()>();
        let internal_shutdown_tx = Arc::new(Mutex::new(Some(internal_shutdown_tx)));
        let (tcp_addr, quic_addr, has_control_plane) = resolve_control_plane_addrs(&config);

        let mut user_specs_per_core = self.user_specs_per_core;

        // Pre-create Core 0's registry so the ControlPlaneActor can use it.
        // The CP runs on Core 0 and needs to talk to other cores via rings.
        let core0_local_t = Arc::new(InProcessTransport::with_shared_state(
            global_locals.clone(),
            death_listeners.clone(),
        ));
        let core0_tr = Arc::new(TransportRegistry::new_with_local(core0_local_t.local_map()));
        let core0_inter_core_t = Arc::new(super::transport::InterCoreTransport {
            my_core: 0,
            rings: rings.clone(),
            placement: placement.clone(),
            local: core0_local_t.clone(),
            has_work: has_work.clone(),
            core_idle: core_idle.clone(),
        });
        core0_tr.add_transport(core0_inter_core_t).ok();

        if has_control_plane {
            inject_control_plane_actor(
                &mut user_specs_per_core,
                &config,
                &registry,
                &responses,
                &internal_shutdown_tx,
                start_time,
                tcp_addr,
                quic_addr,
                federated_routing.as_ref().map(|r| r.policy()),
                self.num_cores,
                crate::reactor::TokioReactor,
                palladium_transport::network::TokioNetwork,
                fs.clone(),
            );
        }

        let mut handles = Vec::with_capacity(num_cores);
        let mut per_core_stats = Vec::with_capacity(num_cores);

        for (core_id, specs) in user_specs_per_core.into_iter().enumerate() {
            let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();

            let (local_t, tr) = if core_id == 0 {
                (core0_local_t.clone(), core0_tr.clone())
            } else {
                let local_t = Arc::new(InProcessTransport::with_shared_state(
                    global_locals.clone(),
                    death_listeners.clone(),
                ));
                let tr = Arc::new(TransportRegistry::new_with_local(local_t.local_map()));
                let inter_core_t = Arc::new(super::transport::InterCoreTransport {
                    my_core: core_id,
                    rings: rings.clone(),
                    placement: placement.clone(),
                    local: local_t.clone(),
                    has_work: has_work.clone(),
                    core_idle: core_idle.clone(),
                });
                tr.add_transport(inter_core_t).ok();
                (local_t, tr)
            };

            let reg = registry.clone();
            let res = responses.clone();
            let cfg = config.clone();
            let f = fs.clone();
            let stats = Arc::new(CoreStats::new());
            let stats_capture = stats.clone();
            let rings_capture = rings.clone();
            let has_work_capture = has_work.clone();
            let core_idle_capture = core_idle.clone();

            let thread = std::thread::Builder::new()
                .name(format!("pd-core-{core_id}"))
                .spawn(move || {
                    pin_thread_to_core(core_id);
                    let rt = tokio::runtime::Builder::new_current_thread()
                        .enable_all()
                        .build()
                        .expect("failed to build Tokio runtime");
                    let local = tokio::task::LocalSet::new();
                    local.block_on(
                        &rt,
                        core_run_async(
                            core_id,
                            num_cores,
                            local_t,
                            tr,
                            res,
                            reg,
                            rings_capture,
                            has_work_capture,
                            core_idle_capture,
                            shutdown_rx,
                            cfg,
                            specs,
                            f,
                            stats_capture,
                        ),
                    );
                })
                .map_err(EngineError::ThreadSpawnFailed)?;

            per_core_stats.push(stats.clone());
            handles.push(CoreHandle {
                core_id,
                shutdown_tx: Some(shutdown_tx),
                thread: Some(thread),
                stats,
            });
        }

        Ok(MultiCoreHandle {
            handles,
            transport: core0_local_t,
            transport_registry: core0_tr,
            registry,
            responses,
            placement,
            start_time,
            ask_timeout: config.ask_timeout,
            num_cores,
            per_core_stats,
            internal_shutdown_rx: if has_control_plane {
                Some(internal_shutdown_rx)
            } else {
                None
            },
            federated_routing,
            fs,
            reactor: crate::reactor::TokioReactor,
        })
    }
}

fn maybe_init_federation(
    config: &mut EngineConfig<crate::reactor::TokioReactor>,
) -> Option<Arc<FederatedRouting>> {
    if config.federation_enabled && config.cluster_membership.is_none() {
        if let Some(bind) = &config.federation_bind_addr {
            if let Ok(addr) = bind.parse::<std::net::SocketAddr>() {
                let mut membership_cfg = config.federation_membership.clone();
                if !config.federation_seed_nodes.is_empty() {
                    let mut seeds = Vec::new();
                    for seed in &config.federation_seed_nodes {
                        if let Ok(seed_addr) = seed.parse() {
                            seeds.push(seed_addr);
                        }
                    }
                    membership_cfg.seed_nodes = seeds;
                }
                let membership = palladium_federation::ClusterMembership::new(
                    config.engine_id.clone(),
                    addr,
                    membership_cfg,
                );
                let membership = Arc::new(parking_lot::RwLock::new(membership));
                let registry_arc = Arc::new(parking_lot::RwLock::new(
                    palladium_federation::FederatedRegistry::default(),
                ));
                config.cluster_membership = Some(Arc::clone(&membership));
                if config.federated_registry.is_none() {
                    config.federated_registry = Some(Arc::clone(&registry_arc));
                }
                return Some(Arc::new(FederatedRouting::new(
                    registry_arc,
                    config.federation_policy.clone(),
                )));
            }
            eprintln!("invalid federation_bind_addr: {bind}");
        } else {
            eprintln!("federation_enabled=true but federation_bind_addr is not set");
        }
    }
    config.federated_registry.clone().map(|registry_arc| {
        Arc::new(FederatedRouting::new(
            registry_arc,
            config.federation_policy.clone(),
        ))
    })
}

fn ensure_consensus_config(config: &mut EngineConfig<crate::reactor::TokioReactor>) {
    if config.consensus_engine.is_some() {
        if config.consensus_groups.is_none() {
            config.consensus_groups = Some(Arc::new(parking_lot::RwLock::new(
                std::collections::HashMap::new(),
            )));
        }
        if config.consensus_group_meta.is_none() {
            config.consensus_group_meta = Some(Arc::new(parking_lot::RwLock::new(
                std::collections::HashMap::new(),
            )));
        }
    }
}

fn resolve_control_plane_addrs<R: crate::reactor::Reactor>(
    config: &EngineConfig<R>,
) -> (
    Option<std::net::SocketAddr>,
    Option<std::net::SocketAddr>,
    bool,
) {
    let tcp_addr = config
        .control_plane_tcp_addr
        .as_deref()
        .and_then(|addr| crate::control_plane::parse_socket_addr(addr).ok());
    let quic_addr = config
        .control_plane_quic_addr
        .as_deref()
        .and_then(|addr| crate::control_plane::parse_socket_addr(addr).ok());
    let has_control_plane =
        config.control_plane_socket.is_some() || tcp_addr.is_some() || quic_addr.is_some();
    (tcp_addr, quic_addr, has_control_plane)
}

#[allow(clippy::too_many_arguments)]
fn inject_control_plane_actor<
    R: crate::reactor::Reactor,
    N: palladium_transport::network::Network,
    F: crate::fs::FileSystem,
>(
    user_specs_per_core: &mut [Vec<ChildSpec<R>>],
    config: &EngineConfig<R>,
    registry: &Arc<ActorRegistry<R>>,
    responses: &Arc<ResponseRegistry>,
    internal_shutdown_tx: &Arc<Mutex<Option<oneshot::Sender<()>>>>,
    start_time: Instant,
    tcp_addr: Option<std::net::SocketAddr>,
    quic_addr: Option<std::net::SocketAddr>,
    federation_policy: Option<Arc<parking_lot::RwLock<palladium_federation::FederationPolicy>>>,
    num_cores: usize,
    reactor: R,
    network: N,
    fs: F,
) {
    let cp_path = ActorPath::parse("/user/core0/control-plane").expect("valid control-plane path");
    let ns = NamespacePolicy::default_for(&cp_path).expect("valid namespace");
    let reg = Arc::clone(registry);
    let responses = Arc::clone(responses);
    let tx = Arc::clone(internal_shutdown_tx);
    let prpc = config.plugin_rpc.clone();
    let actor_spawn = config.actor_spawn.clone();
    let cluster_membership = config.cluster_membership.clone();
    let federated_registry = config.federated_registry.clone();
    let engine_id = config.engine_id.clone();
    let consensus_engine = config.consensus_engine.clone();
    let consensus_groups = config.consensus_groups.clone();
    let consensus_group_meta = config.consensus_group_meta.clone();
    let consensus_engine_info = config.consensus_engine_info.clone();
    let socket_path = config.control_plane_socket.clone();
    let tls = config.control_plane_tls.clone();
    let r = reactor.clone();
    let n = network.clone();
    let f = fs.clone();

    user_specs_per_core[0].push(ChildSpec::new(
        "control-plane",
        RestartPolicy::Permanent,
        ShutdownPolicy::Timeout(Duration::from_secs(5)),
        ns,
        move || {
            let tx_final = tx.lock().unwrap().take();
            Box::new(ControlPlaneActor::new(
                socket_path.clone(),
                tcp_addr,
                quic_addr,
                tls.clone(),
                Arc::clone(&reg),
                Arc::clone(&responses),
                start_time,
                tx_final,
                num_cores,
                prpc.clone(),
                actor_spawn.clone(),
                cluster_membership.clone(),
                federated_registry.clone(),
                federation_policy.clone(),
                engine_id.clone(),
                consensus_engine.clone(),
                consensus_groups.clone(),
                consensus_group_meta.clone(),
                consensus_engine_info.clone(),
                r.clone(),
                n.clone(),
                f.clone(),
            ))
        },
    ));
}