sprite_core/

engine.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use parking_lot::RwLock;
use crossbeam_channel::{unbounded, Sender};

use crate::actor::{Context, StateStore};
use crate::arena::Arena;
use crate::message::Message;
use crate::registry::Registry;
use crate::error::SpriteError;

#[derive(Clone)]
pub struct Handle {
    pub id: u64,
    pub name: String,
    pub(crate) tx: Sender<Message>,
}

impl Handle {
    pub fn send(&self, msg: Message) {
        let _ = self.tx.send(msg);
    }
    pub fn send_msg<T: crate::util::IntoMessage>(&self, msg: T) {
        self.send(msg.into_message());
    }
    pub fn request(&self, msg: Message, timeout: Duration) -> Result<crate::request::Response, SpriteError> {
        let (req, rx) = crate::request::Request::new(msg);
        self.send(req.payload);
        rx.recv_timeout(timeout)
            .map_err(|_| SpriteError::RequestTimeout)
    }
}

pub struct Engine {
    pub(crate) inner: Arc<EngineInner>,
}

pub(crate) struct EngineInner {
    pub(crate) next_id: AtomicU64,
    pub(crate) registry: Arc<Registry>,
    pub(crate) channels: RwLock<HashMap<u64, Sender<Message>>>,
    pub(crate) running: AtomicBool,
}

impl EngineInner {
    pub(crate) fn new() -> Self {
        Self {
            next_id: AtomicU64::new(1),
            registry: Arc::new(Registry::new()),
            channels: RwLock::new(HashMap::new()),
            running: AtomicBool::new(true),
        }
    }

    pub(crate) fn send_to(&self, id: u64, msg: Message) {
        let channels = self.channels.read();
        if let Some(tx) = channels.get(&id) {
            let _ = tx.send(msg);
        }
    }

    pub(crate) fn request(&self, id: u64, msg: Message, timeout: Duration) -> Option<Message> {
        let channels = self.channels.read();
        if let Some(tx) = channels.get(&id) {
            let (req, rx) = crate::request::Request::new(msg);
            let _ = tx.send(req.payload);
            rx.recv_timeout(timeout).ok().map(|r| r.into_message())
        } else {
            None
        }
    }

    /// Spawn with defaults (used by Context::spawn)
    pub(crate) fn spawn_simple<F>(&self, name: &str, setup: F) -> Handle
    where F: Fn(&mut Context) + Send + Sync + 'static,
    {
        self.spawn(name, setup, 1024 * 64, 10, Duration::from_secs(5), Arc::new(self.clone_shallow()))
    }

    pub(crate) fn spawn<F>(
        &self,
        name: &str,
        setup: F,
        arena_size: usize,
        max_recoveries: u32,
        recovery_window: Duration,
        engine_arc: Arc<EngineInner>,
    ) -> Handle
    where
        F: Fn(&mut Context) + Send + Sync + 'static,
    {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let (tx, rx) = unbounded();

        {
            let mut channels = self.channels.write();
            channels.insert(id, tx.clone());
        }
        self.registry.register(name, id);

        let state_store: StateStore = Arc::new(RwLock::new(HashMap::new()));
        let setup = Arc::new(setup);
        let name_owned = name.to_string();
        let tx_for_handle = tx.clone();
        let registry = self.registry.clone();
        let engine_weak = Arc::downgrade(&engine_arc);

        std::thread::spawn(move || {
            let mut arena = Arena::with_capacity(arena_size);
            let mut recovery_count = 0u32;
            let mut last_recovery = Instant::now();
            let mut is_first_mount = true;

            loop {
                let engine_ref = match engine_weak.upgrade() {
                    Some(arc) => arc,
                    None => break,
                };

                let mut ctx = Context::new(
                    id, name_owned.clone(), state_store.clone(),
                    rx.clone(), tx.clone(), engine_ref,
                );
                ctx.is_first_mount = is_first_mount;
                let setup_clone = setup.clone();

                let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    setup_clone(&mut ctx);
                    if is_first_mount {
                        if let Some(ref h) = ctx.mount_handler {
                            h();
                        }
                    }
                    is_first_mount = false;

                    loop {
                        match ctx.rx.recv_timeout(Duration::from_millis(100)) {
                            Ok(msg) => {
                                ctx.metrics.inc_received();
                                if let Some(ref handler) = ctx.message_handler {
                                    handler(msg);
                                }
                            }
                            Err(_) => {}
                        }
                        if !ctx.engine.running.load(Ordering::SeqCst) {
                            break;
                        }
                    }
                    if let Some(ref h) = ctx.unmount_handler {
                        h();
                    }
                }));

                match result {
                    Ok(()) => break,
                    Err(_) => {
                        ctx.metrics.inc_panic();
                        recovery_count += 1;
                        if recovery_count > max_recoveries && last_recovery.elapsed() < recovery_window {
                            tracing::error!(
                                "[Actor {}] CIRCUIT BREAKER TRIPPED after {} recoveries — halting.",
                                id, recovery_count
                            );
                            break;
                        }
                        last_recovery = Instant::now();
                        let start = Instant::now();
                        arena.reset();
                        let elapsed = start.elapsed();
                        ctx.metrics.inc_recovery();
                        tracing::debug!("[Actor {}] recovered in {:?}", id, elapsed);
                        if let Some(ref h) = ctx.panic_handler {
                            h();
                        }
                    }
                }
            }
            registry.unregister(&name_owned);
        });

        Handle { id, name: name.to_string(), tx: tx_for_handle }
    }

    fn clone_shallow(&self) -> Self {
        Self {
            next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
            registry: self.registry.clone(),
            channels: RwLock::new(self.channels.read().clone()),
            running: AtomicBool::new(self.running.load(Ordering::SeqCst)),
        }
    }
}

impl Engine {
    pub fn new() -> Self {
        Self { inner: Arc::new(EngineInner::new()) }
    }

    pub fn spawn<F>(&self, name: &str, setup: F) -> Handle
    where F: Fn(&mut Context) + Send + Sync + 'static,
    {
        self.inner.spawn(name, setup, 1024 * 64, 10, Duration::from_secs(5), self.inner.clone())
    }

    pub fn spawn_with_config<F>(
        &self, name: &str, setup: F,
        arena_size: usize, max_recoveries: u32, recovery_window: Duration,
    ) -> Handle
    where F: Fn(&mut Context) + Send + Sync + 'static,
    {
        self.inner.spawn(name, setup, arena_size, max_recoveries, recovery_window, self.inner.clone())
    }

    pub fn send_to(&self, id: u64, msg: Message) {
        self.inner.send_to(id, msg);
    }

    pub fn send_named(&self, name: &str, msg: Message) {
        if let Some(id) = self.inner.registry.lookup(name) {
            self.inner.send_to(id, msg);
        }
    }

    pub fn lookup(&self, name: &str) -> Option<u64> {
        self.inner.registry.lookup(name)
    }

    pub fn broadcast(&self, msg: Message) -> usize {
        let channels = self.inner.channels.read();
        let mut sent = 0;
        for (_, tx) in channels.iter() {
            if tx.send(msg.clone()).is_ok() { sent += 1; }
        }
        sent
    }

    pub fn shutdown(&self) {
        self.inner.running.store(false, Ordering::SeqCst);
    }

    pub fn is_running(&self) -> bool {
        self.inner.running.load(Ordering::SeqCst)
    }

    pub fn actor_count(&self) -> usize {
        self.inner.channels.read().len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn spawn_and_send() {
        let engine = Engine::new();
        let handle = engine.spawn("test", |ctx| {
            ctx.on_message(|msg| { println!("got: {:?}", msg); });
        });
        std::thread::sleep(Duration::from_millis(20));
        handle.send(Message::text("hello"));
        std::thread::sleep(Duration::from_millis(50));
    }

    #[test]
    fn state_persists_across_panics() {
        let engine = Engine::new();
        let handle = engine.spawn("fragile", |ctx| {
            let count = ctx.use_state("count", 0i64);
            ctx.on_message(move |msg| {
                if msg == "set" { count.set(42); }
                if msg == "panic" { panic!("boom"); }
                if msg == "check" { assert_eq!(count.get(), 42); }
            });
        });
        std::thread::sleep(Duration::from_millis(20));
        handle.send(Message::text("set"));
        std::thread::sleep(Duration::from_millis(20));
        handle.send(Message::text("panic"));
        std::thread::sleep(Duration::from_millis(50));
        handle.send(Message::text("check"));
        std::thread::sleep(Duration::from_millis(50));
    }

    #[test]
    fn named_lookup() {
        let engine = Engine::new();
        let h = engine.spawn("logger", |ctx| {
            ctx.on_message(|msg| println!("{:?}", msg));
        });
        std::thread::sleep(Duration::from_millis(10));
        assert_eq!(engine.lookup("logger"), Some(h.id));
        engine.send_named("logger", Message::text("hi"));
        std::thread::sleep(Duration::from_millis(50));
    }

    #[test]
    fn broadcast_works() {
        let engine = Engine::new();
        let _ = engine.spawn("a", |ctx| {
            ctx.on_message(|msg| println!("a: {:?}", msg));
        });
        let _ = engine.spawn("b", |ctx| {
            ctx.on_message(|msg| println!("b: {:?}", msg));
        });
        std::thread::sleep(Duration::from_millis(20));
        let sent = engine.broadcast(Message::text("all"));
        assert_eq!(sent, 2);
        std::thread::sleep(Duration::from_millis(50));
    }

    #[test]
    fn mount_and_unmount() {
        let engine = Engine::new();
        let handle = engine.spawn("lifecycle", |ctx| {
            ctx.on_mount(|| println!("mounted"));
            ctx.on_unmount(|| println!("unmounted"));
            ctx.on_message(|_| {});
        });
        std::thread::sleep(Duration::from_millis(20));
        handle.send(Message::text("hi"));
        std::thread::sleep(Duration::from_millis(50));
    }
}