palladium_runtime/engine/

handle.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::fs::{FileSystem, TokioFileSystem};
use dashmap::DashMap;
use palladium_actor::{
    Actor, ActorPath, AddrHash, AskError, ChildSpec, EngineId, Envelope, Message, MessagePayload,
    PoolConfig, RemoteMessage, SendError, StopReason, WorkerPool,
};
use palladium_transport::network::{Network, TokioNetwork};
use palladium_transport::{
    InProcessTransport, MailboxMessage, MailboxReceiver, MailboxSender, QuicTransport,
    QuicTransportConfig, TcpTransport, TcpTransportConfig, TransportError, TransportRegistry,
    TypeRegistry,
};

use crate::addr::{make_ask_fn, make_remote_ask_fn, make_remote_send_fn, make_send_fn};
use crate::common::LifecycleSignal;
use crate::introspection::{ActorInfo, ActorQuery, EngineSnapshot};
use crate::reactor::{Reactor, TokioReactor};
use crate::registry::ActorRegistry;
use crate::responses::ResponseRegistry;

/// Handle for interacting with the engine from external (test) code.
#[derive(Clone)]
pub struct EngineHandle<
    R: Reactor = TokioReactor,
    N: Network = TokioNetwork,
    F: FileSystem = TokioFileSystem,
> {
    pub(crate) transport: Arc<InProcessTransport>,
    pub(crate) transport_registry: Arc<TransportRegistry>,
    pub(crate) type_registry: TypeRegistry,
    pub(crate) responses: Arc<ResponseRegistry>,
    pub(crate) registry: Arc<ActorRegistry<R>>,
    pub(crate) ask_timeout: Duration,
    pub(crate) reactor: R,
    pub(crate) network: N,
    pub(crate) _fs: F,
    pub(crate) start_time: Instant,
    pub(crate) source_addr: AddrHash,
    pub(crate) federated_routing: Option<Arc<crate::federation::FederatedRouting>>,
    /// Cache of direct `MailboxSender`s for local actors, keyed by `AddrHash`.
    ///
    /// Populated on the first `send_to` to a local actor; avoids the DashMap
    /// lookup in `TransportRegistry::try_route` on subsequent sends.
    /// Evicted when a send returns `ActorStopped` (receiver dropped).
    /// Shared across handle clones via `Arc` so all clones benefit.
    pub(crate) send_cache: Arc<DashMap<AddrHash, MailboxSender>>,
    /// Receiving end of the response pump mailbox.  Consumed by the first
    /// `remote_addr_for` ask, which spawns a long-lived pump task.
    pub(crate) pump_rx: Arc<std::sync::Mutex<Option<MailboxReceiver>>>,
}

impl<R: Reactor, N: Network, F: FileSystem> EngineHandle<R, N, F> {
    pub fn type_registry(&self) -> TypeRegistry {
        self.type_registry.clone()
    }

    /// Send a message to an actor by its `AddrHash`.
    ///
    /// **Fast path**: after the first send to a local actor, the `MailboxSender`
    /// is cached in `send_cache`, bypassing the `TransportRegistry` DashMap on
    /// subsequent sends.  Stale entries are evicted on `ActorStopped`.
    pub fn send_to<M: Message>(&self, target: AddrHash, msg: M) -> Result<(), SendError> {
        let envelope = Envelope::new(self.source_addr, target, M::TYPE_TAG, 0);
        let payload = MessagePayload::local(msg);

        // Fast path: use cached MailboxSender to avoid DashMap in TransportRegistry.
        if let Some(sender) = self.send_cache.get(&target) {
            return match sender.try_send(MailboxMessage { envelope, payload }) {
                Ok(()) => Ok(()),
                Err(TransportError::MailboxFull) => Err(SendError::MailboxFull),
                Err(_) => {
                    // Receiver dropped — evict stale cache entry.
                    drop(sender); // release DashMap shard lock before remove
                    self.send_cache.remove(&target);
                    Err(SendError::ActorStopped)
                }
            };
        }

        // Cache miss: check if this is a local actor, populate cache, and send.
        if let Some(sender) = self.transport_registry.get_local_sender(target) {
            let result = sender.try_send(MailboxMessage { envelope, payload });
            match result {
                Ok(()) => {
                    self.send_cache.insert(target, sender);
                    return Ok(());
                }
                Err(TransportError::MailboxFull) => {
                    // Actor is alive but full; cache for next send.
                    self.send_cache.insert(target, sender);
                    return Err(SendError::MailboxFull);
                }
                Err(_) => return Err(SendError::ActorStopped),
            }
        }

        // Non-local destination (remote, QUIC, TCP, etc.): full routing.
        self.transport_registry
            .try_route(envelope, payload)
            .map_err(|e| match e {
                TransportError::MailboxFull => SendError::MailboxFull,
                TransportError::ConnectionFailed => SendError::ConnectionFailed,
                TransportError::Unroutable(_) => SendError::Unroutable,
                TransportError::SerializationRequired | TransportError::SerializationError(_) => {
                    SendError::SerializationFailed
                }
                _ => SendError::ActorStopped,
            })
    }

    /// Send up to `n` messages to one target using a producer-side batch fast path.
    ///
    /// For local destinations this uses `MailboxSender::try_send_batch_owned`,
    /// reducing wake/schedule overhead vs repeated `send_to` calls.
    ///
    /// Returns the number of messages successfully enqueued. On local mailbox
    /// backpressure this may be less than `n` without returning an error.
    pub fn send_many<M: Message + Clone>(
        &self,
        target: AddrHash,
        msg: M,
        n: usize,
    ) -> Result<usize, SendError> {
        if n == 0 {
            return Ok(0);
        }

        let make_batch = || {
            let mut batch = Vec::with_capacity(n);
            for _ in 0..n {
                batch.push(MailboxMessage {
                    envelope: Envelope::new(self.source_addr, target, M::TYPE_TAG, 0),
                    payload: MessagePayload::local(msg.clone()),
                });
            }
            batch
        };

        // Fast path: use cached MailboxSender.
        if let Some(sender) = self.send_cache.get(&target) {
            let (sent, err, _remaining) = sender.try_send_batch_owned(make_batch());
            return match err {
                None => Ok(sent),
                Some(TransportError::MailboxFull) => Ok(sent),
                _ => {
                    drop(sender);
                    self.send_cache.remove(&target);
                    Err(SendError::ActorStopped)
                }
            };
        }

        // Cache miss: local actor lookup and batch send.
        if let Some(sender) = self.transport_registry.get_local_sender(target) {
            let (sent, err, _remaining) = sender.try_send_batch_owned(make_batch());
            match err {
                None => {
                    self.send_cache.insert(target, sender);
                    return Ok(sent);
                }
                Some(TransportError::MailboxFull) => {
                    self.send_cache.insert(target, sender);
                    return Ok(sent);
                }
                _ => return Err(SendError::ActorStopped),
            }
        }

        // Non-local fallback: preserve send_to semantics.
        let mut sent = 0usize;
        for _ in 0..n {
            match self.send_to::<M>(target, msg.clone()) {
                Ok(()) => sent += 1,
                Err(SendError::MailboxFull) => return Ok(sent),
                Err(e) => return Err(e),
            }
        }
        Ok(sent)
    }

    /// Create an `Addr<M>` that sends to `target` via this engine's transport.
    pub fn addr_for<M: Message>(&self, target: AddrHash) -> palladium_actor::Addr<M>
    where
        R: Send + Sync,
    {
        let send_fn = make_send_fn::<M>(target, self.source_addr, self.transport_registry.clone());
        let ask_fn = make_ask_fn::<M, R>(
            target,
            self.source_addr,
            self.transport.clone(),
            self.transport_registry.clone(),
            self.responses.clone(),
            self.ask_timeout,
            self.reactor.clone(),
        );
        palladium_actor::Addr::with_handlers(target, send_fn, ask_fn)
    }

    /// Create a remote `Addr<M>` that always serializes payloads before routing.
    ///
    /// This is intended for explicit remote sends; it rejects local destinations
    /// to avoid accidental local delivery of serialized payloads.
    pub fn remote_addr<M: RemoteMessage>(&self, target: AddrHash) -> palladium_actor::Addr<M> {
        let registry = self.transport_registry.clone();
        let source = self.source_addr;
        let send_fn = Arc::new(move |msg: M| -> Result<(), SendError> {
            if registry.is_local(target) {
                return Err(SendError::PolicyViolation);
            }
            let bytes = bincode::serialize(&msg).map_err(|_| SendError::SerializationFailed)?;
            let envelope = Envelope::new(source, target, M::TYPE_TAG, bytes.len() as u32);
            registry
                .try_route(envelope, MessagePayload::serialized(bytes))
                .map_err(|e| match e {
                    TransportError::MailboxFull => SendError::MailboxFull,
                    TransportError::ConnectionFailed => SendError::ConnectionFailed,
                    TransportError::Unroutable(_) => SendError::Unroutable,
                    TransportError::SerializationRequired
                    | TransportError::SerializationError(_) => SendError::SerializationFailed,
                    _ => SendError::ActorStopped,
                })
        });
        let ask_fn: palladium_actor::AskFn<M> =
            Arc::new(|_msg: M| -> palladium_actor::AskFuture<M> {
                Box::pin(async { Err(AskError::Unbound) })
            });
        palladium_actor::Addr::with_handlers(target, send_fn, ask_fn)
    }

    /// Create a typed remote `Addr<M>` that supports both `send()` and `ask()`.
    ///
    /// Both request and response payloads are serialized for transport delivery.
    /// Local destinations are rejected with `AskError::Send(PolicyViolation)`.
    ///
    /// Before calling, register the message type on both engines:
    /// ```ignore
    /// handle.type_registry().register_remote_ask::<MyMsg>();
    /// ```
    ///
    /// The first `ask()` on the returned address lazily starts the engine's
    /// response pump task (a lightweight background task that routes incoming
    /// serialized response frames into the local response registry).
    pub fn remote_addr_for<M: RemoteMessage>(&self, target: AddrHash) -> palladium_actor::Addr<M>
    where
        M::Response: serde::Serialize + for<'de> serde::Deserialize<'de> + Send + 'static,
        R: Clone + Send + Sync,
    {
        let send_fn =
            make_remote_send_fn::<M>(target, self.source_addr, self.transport_registry.clone());
        let ask_fn = make_remote_ask_fn::<M, R>(
            target,
            self.source_addr,
            self.transport_registry.clone(),
            self.responses.clone(),
            self.ask_timeout,
            self.reactor.clone(),
            Arc::clone(&self.pump_rx),
        );
        palladium_actor::Addr::with_handlers(target, send_fn, ask_fn)
    }

    pub fn transport(&self) -> &Arc<InProcessTransport> {
        &self.transport
    }

    pub fn transport_registry(&self) -> &Arc<TransportRegistry> {
        &self.transport_registry
    }

    pub async fn attach_quic_transport(
        &self,
        config: QuicTransportConfig,
        peers: HashMap<EngineId, SocketAddr>,
    ) -> Result<Arc<QuicTransport<R, N>>, TransportError>
    where
        R: Clone,
        N: Clone,
    {
        let transport = QuicTransport::bind(
            config,
            peers,
            self.transport_registry.clone(),
            self.type_registry.clone(),
            self.reactor.clone(),
            self.network.clone(),
        )
        .await?;
        self.transport_registry.add_transport(transport.clone())?;
        Ok(transport)
    }

    pub async fn attach_tcp_transport(
        &self,
        config: TcpTransportConfig,
        peers: HashMap<EngineId, SocketAddr>,
    ) -> Result<Arc<TcpTransport<R, N>>, TransportError>
    where
        R: Clone,
        N: Clone,
    {
        let transport = TcpTransport::bind(
            config,
            peers,
            self.transport_registry.clone(),
            self.type_registry.clone(),
            self.reactor.clone(),
            self.network.clone(),
        )
        .await?;
        self.transport_registry.add_transport(transport.clone())?;
        Ok(transport)
    }

    // ── Introspection (REQ-085) ───────────────────────────────────────────────

    /// List actors, optionally filtered by `query`.
    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> {
        self.registry.actor_info_by_path(path, 0)
    }

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

    /// Lookup an actor's address by its path.
    pub fn lookup_path(&self, path: &ActorPath) -> Option<AddrHash> {
        self.registry.get_by_path(path).map(|s| s.addr).or_else(|| {
            self.federated_routing
                .as_ref()
                .and_then(|routing| routing.resolve_optional(path))
        })
    }

    /// Resolve an actor's address by path with routing errors.
    pub fn resolve_path(&self, path: &ActorPath) -> Result<AddrHash, SendError> {
        if let Some(local) = self.registry.get_by_path(path).map(|s| s.addr) {
            return Ok(local);
        }
        if let Some(routing) = &self.federated_routing {
            return routing.resolve(path);
        }
        Err(SendError::Unroutable)
    }

    /// Send a hard-stop signal to an actor by path.
    ///
    /// Returns `true` if the actor was found and the signal was sent. The
    /// actor may still be running until the signal is processed by the
    /// reactor. Used by `FaultInjector::crash_actor`.
    pub fn stop_actor(&self, path: &ActorPath) -> bool {
        if let Some(slot) = self.registry.get_by_path(path) {
            slot.ctrl_tx
                .send(LifecycleSignal::Stop(StopReason::Supervisor))
                .ok();
            true
        } else {
            false
        }
    }

    /// Fill an actor's mailbox to capacity by sending dummy messages.
    ///
    /// Returns the number of messages that were successfully enqueued.
    /// Subsequent sends to the same actor will return `SendError::MailboxFull`
    /// until the actor drains messages. Used by `FaultInjector::fill_mailbox`.
    pub fn fill_mailbox(&self, path: &ActorPath) -> usize {
        let addr = self.registry.get_by_path(path).map(|s| s.addr);
        let Some(addr) = addr else { return 0 };
        let src = AddrHash::synthetic(b"fault-fill-mailbox");
        let mut count = 0;
        loop {
            let env = Envelope::new(src, addr, 0, 0);
            match self.transport.try_deliver(env, MessagePayload::local(())) {
                Ok(()) => count += 1,
                Err(_) => break,
            }
        }
        count
    }

    /// Dynamically spawn a user actor by sending a signal to the root /user supervisor.
    pub fn spawn_user_actor(&self, spec: ChildSpec<R>) {
        let user_path = ActorPath::parse("/user").unwrap();
        if let Some(slot) = self.registry.get_by_path(&user_path) {
            slot.ctrl_tx.send(LifecycleSignal::SpawnChild(spec)).ok();
        }
    }

    /// Spawn a pool of N identical worker actors under the `/user` supervisor
    /// and return a [`WorkerPool<M>`] routing handle.
    ///
    /// Each worker gets a [`StableAddr<M>`] backed by an [`AddrRefresher`] so
    /// that the pool transparently routes to new incarnations after restarts.
    ///
    /// # Ordering
    ///
    /// Spawn requests are sent asynchronously to the `/user` supervisor.  The
    /// returned pool's `StableAddr`s are initially `None`; they are populated
    /// once the supervisor task processes each `SpawnChild` signal.  Yield at
    /// least once (e.g. `tokio::task::yield_now().await`) before sending
    /// messages through the pool.
    ///
    /// # Panics
    ///
    /// Panics if `config.size` is 0 (no workers requested).
    pub fn spawn_worker_pool<M, G>(&self, config: &PoolConfig, factory: G) -> WorkerPool<M>
    where
        M: Message,
        R: Clone + Send + Sync + 'static,
        G: Fn() -> Box<dyn Actor<R>> + Send + Sync + Clone + 'static,
    {
        let user_path = ActorPath::parse("/user").unwrap();
        let handle = self.clone();
        crate::worker_pool::spawn_worker_pool(
            &user_path,
            config,
            factory,
            self.source_addr,
            &self.transport,
            &self.transport_registry,
            &self.responses,
            &self.reactor,
            move |spec| handle.spawn_user_actor(spec),
        )
    }
}