runa/crates/runa-image-api/src/lib.rs

//! Trusted, in-process application services for a running Runa image.

use std::collections::BTreeMap;
use std::num::NonZeroUsize;
use std::sync::Arc;

use runa_actor::{ActorRef, ActorSnapshot, Event, Message, SystemMessage};
use runa_code::{BehaviorSnapshot, CodeStore, SourceCommit, SourceFile};
use runa_core::{ImplementationId, Value};
use runa_runtime::{Runtime, RuntimeResult};

pub struct ImageApi {
    code: CodeStore,
    runtime: Runtime,
    actor_paths: BTreeMap<ActorRef, String>,
    pending_paths: BTreeMap<ActorRef, (String, ImplementationId)>,
    events: Vec<Event>,
}

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

impl ImageApi {
    pub fn new() -> Self {
        Self {
            code: CodeStore::default(),
            runtime: Runtime::default(),
            actor_paths: BTreeMap::new(),
            pending_paths: BTreeMap::new(),
            events: Vec::new(),
        }
    }

    pub fn publish_behavior(
        &mut self,
        name: &str,
        expected_revision: Option<u64>,
        source: &str,
    ) -> RuntimeResult<BehaviorSnapshot> {
        let artifact = self.code.compile(source)?;
        self.runtime.publish(Arc::clone(&artifact))?;
        let behavior = self
            .code
            .publish_behavior(name, expected_revision, source)?;
        if runa_code::validate_source_path(name).is_ok() {
            let expected_source_id = self.code.source(name).map(|file| file.source_id);
            self.code.write_source(name, expected_source_id, source)?;
        }
        Ok(behavior)
    }

    pub fn rollback_behavior(
        &mut self,
        name: &str,
        expected_revision: u64,
        target_revision: u64,
    ) -> RuntimeResult<BehaviorSnapshot> {
        self.code
            .rollback_behavior(name, expected_revision, target_revision)
    }

    pub fn behavior(&self, name: &str) -> Option<BehaviorSnapshot> {
        self.code.behavior(name)
    }

    pub fn behaviors(&self) -> Vec<BehaviorSnapshot> {
        self.code.behaviors()
    }

    pub fn behavior_source(&self, name: &str, revision: u64) -> RuntimeResult<&str> {
        let behavior = self.code.behavior(name).ok_or("unknown behavior")?;
        let version = behavior
            .history
            .iter()
            .find(|version| version.revision == revision)
            .ok_or("unknown behavior revision")?;
        self.code
            .artifact(version.implementation_id)
            .map(|artifact| artifact.original_source())
            .ok_or_else(|| "behavior artifact is unavailable".into())
    }

    pub fn write_source(
        &mut self,
        path: &str,
        expected_source_id: Option<runa_core::SourceId>,
        source: &str,
    ) -> RuntimeResult<SourceFile> {
        self.code.write_source(path, expected_source_id, source)
    }

    pub fn sources(&self) -> Vec<SourceFile> {
        self.code.sources()
    }

    pub fn source(&self, path: &str) -> Option<&SourceFile> {
        self.code.source(path)
    }

    pub fn commit_source(&mut self, path: &str) -> RuntimeResult<SourceCommit> {
        self.code.commit_source(path)
    }

    pub fn source_history(&self, path: &str) -> Option<&[SourceCommit]> {
        self.code.source_history(path)
    }

    pub fn source_commit(&self, path: &str, revision: &str) -> Option<&SourceCommit> {
        self.code.source_commit(path, revision)
    }

    pub fn source_commit_source(&self, path: &str, revision: &str) -> Option<&str> {
        self.code.source_commit_source(path, revision)
    }

    pub fn rollback_source(&mut self, path: &str, revision: &str) -> RuntimeResult<SourceFile> {
        self.code.rollback_source(path, revision)
    }

    pub fn spawn(
        &mut self,
        path: &str,
        state: Value,
        capacity: NonZeroUsize,
    ) -> RuntimeResult<ActorRef> {
        let source = self
            .code
            .source(path)
            .ok_or("unknown source path")?
            .source
            .clone();
        let artifact = self.code.compile(&source)?;
        self.runtime.publish(Arc::clone(&artifact))?;
        let actor = self
            .runtime
            .spawn(artifact.implementation_id, state, capacity)?;
        self.actor_paths.insert(actor, path.to_owned());
        Ok(actor)
    }

    pub fn reload(&mut self, actor: ActorRef) -> RuntimeResult<ImplementationId> {
        let path = self
            .actor_paths
            .get(&actor)
            .cloned()
            .ok_or("actor has no source path")?;
        self.schedule_source_install(actor, &path)
    }

    pub fn r#become(&mut self, actor: ActorRef, path: &str) -> RuntimeResult<ImplementationId> {
        self.schedule_source_install(actor, path)
    }

    fn schedule_source_install(
        &mut self,
        actor: ActorRef,
        path: &str,
    ) -> RuntimeResult<ImplementationId> {
        let source = self
            .code
            .source(path)
            .ok_or("unknown source path")?
            .source
            .clone();
        let artifact = self.code.compile(&source)?;
        self.runtime.publish(Arc::clone(&artifact))?;
        self.runtime
            .schedule_install(actor, artifact.implementation_id)?;
        self.pending_paths
            .insert(actor, (path.to_owned(), artifact.implementation_id));
        Ok(artifact.implementation_id)
    }

    pub fn send(&mut self, actor: ActorRef, message: Message) -> RuntimeResult<()> {
        self.runtime.send(actor, message)
    }

    pub fn tick(&mut self, ticks: usize, budget: NonZeroUsize) {
        for _ in 0..ticks {
            self.runtime.tick(budget);
            self.collect_runtime_events();
        }
    }

    pub fn inspect(&self, actor: ActorRef) -> RuntimeResult<ActorSnapshot> {
        self.runtime.inspect(actor)
    }

    pub fn edit_actor(
        &mut self,
        actor: ActorRef,
        state: Option<Value>,
        messages: Option<(Option<Value>, Vec<Value>)>,
    ) -> RuntimeResult<()> {
        self.runtime.edit_actor(actor, state, messages)
    }

    pub fn set_log_level(
        &mut self,
        actor: ActorRef,
        level: runa_core::LogLevel,
    ) -> RuntimeResult<()> {
        self.runtime.set_log_level(actor, level)
    }

    pub fn log_level(&self, actor: ActorRef) -> RuntimeResult<runa_core::LogLevel> {
        self.runtime.log_level(actor)
    }

    pub fn logs(&self, actors: &[ActorRef], after: u64) -> RuntimeResult<Vec<runa_core::LogEntry>> {
        self.runtime.logs(actors, after)
    }

    pub fn set_log_sink(&mut self, sink: impl Fn(runa_core::LogEntry) + 'static) {
        self.runtime.set_log_sink(sink);
    }

    pub fn actors(&self) -> Vec<ActorSnapshot> {
        self.runtime.actors()
    }

    pub fn events(&mut self) -> Vec<Event> {
        self.collect_runtime_events();
        std::mem::take(&mut self.events)
    }

    fn collect_runtime_events(&mut self) {
        for event in self.runtime.drain_events() {
            if let Event::BehaviorInstalled(actor, implementation) = &event
                && self
                    .pending_paths
                    .get(actor)
                    .is_some_and(|(_, pending)| pending == implementation)
            {
                let (path, _) = self
                    .pending_paths
                    .remove(actor)
                    .expect("pending path exists");
                self.actor_paths.insert(*actor, path);
            }
            self.events.push(event);
        }
    }

    pub fn system(&mut self, actor: ActorRef, message: SystemMessage) -> RuntimeResult<()> {
        self.collect_runtime_events();
        self.runtime.system(actor, message)?;
        if matches!(message, SystemMessage::Kill | SystemMessage::Restart) {
            self.pending_paths.remove(&actor);
        }
        Ok(())
    }

    pub fn stop(&mut self, actor: ActorRef) -> RuntimeResult<()> {
        self.system(actor, SystemMessage::Stop)
    }

    pub fn restart(&mut self, actor: ActorRef) -> RuntimeResult<()> {
        self.system(actor, SystemMessage::Restart)
    }

    pub fn pause(&mut self, actor: ActorRef) -> RuntimeResult<()> {
        self.system(actor, SystemMessage::Pause)
    }

    pub fn resume(&mut self, actor: ActorRef) -> RuntimeResult<()> {
        self.system(actor, SystemMessage::Resume)
    }

    pub fn cancel(&mut self, actor: ActorRef) -> RuntimeResult<()> {
        self.runtime.cancel(actor)
    }

    pub fn install(
        &mut self,
        actor: ActorRef,
        behavior_name: &str,
    ) -> RuntimeResult<ImplementationId> {
        self.install_with_migration(actor, behavior_name, Ok)
    }

    pub fn install_with_migration(
        &mut self,
        actor: ActorRef,
        behavior_name: &str,
        migrate: impl FnOnce(Value) -> RuntimeResult<Value>,
    ) -> RuntimeResult<ImplementationId> {
        let behavior = self
            .code
            .behavior(behavior_name)
            .ok_or("unknown behavior")?;
        self.runtime
            .install_with_migration(actor, behavior.current, migrate)?;
        Ok(behavior.current)
    }
}

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

    fn test_message(kind: i64, value: i64, target: Option<runa_core::ActorRef>) -> Value {
        Value::Record(std::collections::BTreeMap::from([
            ("kind".to_owned(), Value::Integer(kind)),
            ("value".to_owned(), Value::Integer(value)),
            (
                "target".to_owned(),
                target.map_or(Value::Null, Value::ActorRef),
            ),
        ]))
    }

    fn capacity(n: usize) -> NonZeroUsize {
        NonZeroUsize::new(n).expect("capacity must be positive")
    }

    #[test]
    fn image_api_publishes_spawns_inspects_and_exposes_events() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        let behavior = image.publish_behavior(
            "counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message.value }",
        )?;
        assert_eq!(
            image.behavior_source("counter", behavior.revision)?,
            "pub fn handle(ctx, state, sender, message) { state + message.value }"
        );
        let actor = image.spawn("counter", Value::Integer(5), capacity(2))?;
        image.send(actor, test_message(0, 7, None))?;
        image.tick(100, capacity(10));

        assert_eq!(image.inspect(actor)?.state, Value::Integer(12));
        assert_eq!(image.inspect(actor)?.implementation, behavior.current);
        assert!(
            image
                .events()
                .iter()
                .any(|event| matches!(event, Event::Committed(id) if id == &actor))
        );
        Ok(())
    }

    #[test]
    fn image_api_rejects_stale_behavior_updates_without_changing_current() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        let first = image.publish_behavior(
            "counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message.value }",
        )?;
        assert!(
            image
                .publish_behavior(
                    "counter",
                    None,
                    "pub fn handle(ctx, state, sender, message) { state - message.value }"
                )
                .is_err()
        );
        assert_eq!(image.behavior("counter").unwrap().current, first.current);
        Ok(())
    }

    #[test]
    fn image_api_installs_behavior_with_atomic_state_migration() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        let first = image.publish_behavior(
            "counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message.value }",
        )?;
        image.publish_behavior(
            "counter",
            Some(first.revision),
            "pub fn handle(ctx, state, sender, message) { state * message.value }",
        )?;
        let actor = image.spawn("counter", Value::Integer(5), capacity(2))?;
        let implementation = image.install_with_migration(actor, "counter", |state| {
            let Value::Integer(state) = state else {
                return Err("expected integer".into());
            };
            Ok(Value::Integer(
                state.checked_mul(10).ok_or("migration overflow")?,
            ))
        })?;

        assert_eq!(image.inspect(actor)?.state, Value::Integer(50));
        assert_eq!(image.inspect(actor)?.implementation, implementation);
        Ok(())
    }

    #[test]
    fn spawn_compiles_current_uncommitted_source() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        let file = image.write_source(
            "actors/counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message.value }",
        )?;
        let actor = image.spawn("actors/counter", Value::Integer(3), capacity(2))?;
        assert_eq!(
            image.inspect(actor)?.implementation,
            image.code.compile(&file.source)?.implementation_id
        );
        Ok(())
    }

    #[test]
    fn invalid_uncommitted_source_fails_spawn_without_creating_actor() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.write_source("actors/broken", None, "pub fn handle( {")?;
        assert!(image.commit_source("actors/broken").is_ok());
        assert!(
            image
                .spawn("actors/broken", Value::Null, capacity(2))
                .is_err()
        );
        assert!(image.actors().is_empty());
        Ok(())
    }

    #[test]
    fn compatibility_behavior_seeds_only_its_exact_working_path() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.publish_behavior(
            "counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message.value }",
        )?;
        assert!(image.source("counter").is_some());
        assert!(image.source("actors/counter").is_none());
        Ok(())
    }

    #[test]
    fn killed_pending_become_preserves_the_installed_path() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.write_source(
            "actors/first",
            None,
            "pub fn handle(ctx, state, sender, message) { state + 1 }",
        )?;
        image.write_source(
            "actors/second",
            None,
            "pub fn handle(ctx, state, sender, message) { state + 10 }",
        )?;
        let actor = image.spawn("actors/first", Value::Integer(0), capacity(2))?;
        let original = image.inspect(actor)?.implementation;
        image.r#become(actor, "actors/second")?;
        assert_eq!(image.actor_paths[&actor], "actors/first");
        image.system(actor, SystemMessage::Kill)?;
        image.restart(actor)?;
        assert_eq!(image.reload(actor)?, original);
        image.tick(1, capacity(100));
        assert_eq!(image.inspect(actor)?.implementation, original);
        assert_eq!(image.actor_paths[&actor], "actors/first");

        image.r#become(actor, "actors/second")?;
        image.tick(1, capacity(100));
        assert_eq!(image.actor_paths[&actor], "actors/second");
        let current = image.inspect(actor)?.implementation;
        assert_eq!(image.reload(actor)?, current);
        Ok(())
    }

    #[test]
    fn restart_forces_a_new_incarnation_during_active_execution() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.write_source(
            "actors/loop",
            None,
            "pub fn handle(ctx, state, sender, message) { loop {} }",
        )?;
        let actor = image.spawn("actors/loop", Value::Null, capacity(2))?;
        image.send(actor, Value::Null)?;
        image.tick(1, capacity(100));
        let before = image.inspect(actor)?;
        assert!(before.execution.is_some());
        image.restart(actor)?;
        let after = image.inspect(actor)?;
        assert_eq!(after.incarnation, before.incarnation + 1);
        assert_eq!(after.state, before.state);
        assert!(after.execution.is_none());
        Ok(())
    }

    #[test]
    fn reload_compiles_latest_working_source_and_preserves_state_and_mailbox() -> RuntimeResult<()>
    {
        let mut image = ImageApi::new();
        let file = image.write_source(
            "actors/counter",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message }",
        )?;
        image.commit_source("actors/counter")?;
        let actor = image.spawn("actors/counter", Value::Integer(3), capacity(2))?;
        image.send(actor, Value::Integer(2))?;
        let before = image.inspect(actor)?;
        image.write_source(
            "actors/counter",
            Some(file.source_id),
            "pub fn handle(ctx, state, sender, message) { state + message * 10 }",
        )?;
        let implementation = image.reload(actor)?;
        let scheduled = image.inspect(actor)?;
        assert_eq!(scheduled.implementation, before.implementation);
        assert_eq!(scheduled.state, before.state);
        assert_eq!(scheduled.queued, before.queued);
        image.tick(100, capacity(100));
        let after = image.inspect(actor)?;
        assert_eq!(after.implementation, implementation);
        assert_eq!(after.state, Value::Integer(23));
        assert_eq!(after.incarnation, before.incarnation);
        assert_eq!(after.mailbox, before.mailbox);
        assert_eq!(image.actor_paths[&actor], "actors/counter");
        Ok(())
    }

    #[test]
    fn become_binds_only_after_install_and_reload_follows_the_new_path() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.write_source(
            "actors/first",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message }",
        )?;
        let second = image.write_source(
            "actors/second",
            None,
            "pub fn handle(ctx, state, sender, message) { state * message }",
        )?;
        let actor = image.spawn("actors/first", Value::Integer(3), capacity(2))?;
        image.send(actor, Value::Integer(2))?;
        let before = image.inspect(actor)?;
        let implementation = image.r#become(actor, "actors/second")?;
        assert_eq!(image.actor_paths[&actor], "actors/first");
        assert_eq!(image.inspect(actor)?.queued, before.queued);
        image.tick(100, capacity(100));
        assert_eq!(image.actor_paths[&actor], "actors/second");
        assert_eq!(image.inspect(actor)?.implementation, implementation);
        assert_eq!(image.inspect(actor)?.state, Value::Integer(6));
        image.write_source(
            "actors/second",
            Some(second.source_id),
            "pub fn handle(ctx, state, sender, message) { state * message * 10 }",
        )?;
        let updated = image.reload(actor)?;
        assert_ne!(updated, implementation);
        image.send(actor, Value::Integer(2))?;
        image.tick(100, capacity(100));
        assert_eq!(image.inspect(actor)?.implementation, updated);
        assert_eq!(image.inspect(actor)?.state, Value::Integer(120));
        assert_eq!(image.inspect(actor)?.mailbox, before.mailbox);
        Ok(())
    }

    #[test]
    fn restart_and_superseding_reload_cancel_pending_become_path_changes() -> RuntimeResult<()> {
        let mut image = ImageApi::new();
        image.write_source(
            "actors/first",
            None,
            "pub fn handle(ctx, state, sender, message) { state + message }",
        )?;
        image.write_source(
            "actors/second",
            None,
            "pub fn handle(ctx, state, sender, message) { state * message }",
        )?;
        let actor = image.spawn("actors/first", Value::Integer(3), capacity(2))?;
        let before = image.inspect(actor)?;
        image.r#become(actor, "actors/second")?;
        image.restart(actor)?;
        image.tick(1, capacity(100));
        assert_eq!(image.actor_paths[&actor], "actors/first");
        assert_eq!(image.inspect(actor)?.implementation, before.implementation);
        assert_eq!(image.inspect(actor)?.incarnation, before.incarnation + 1);
        assert!(!image.pending_paths.contains_key(&actor));

        image.r#become(actor, "actors/second")?;
        assert_eq!(image.reload(actor)?, before.implementation);
        image.tick(1, capacity(100));
        assert_eq!(image.actor_paths[&actor], "actors/first");
        assert_eq!(image.inspect(actor)?.implementation, before.implementation);
        image.write_source("actors/broken", None, "pub fn broken( {")?;
        assert!(image.r#become(actor, "actors/broken").is_err());
        assert_eq!(image.actor_paths[&actor], "actors/first");
        assert!(image.r#become(actor, "actors/missing").is_err());
        assert_eq!(image.actor_paths[&actor], "actors/first");
        Ok(())
    }
}