wasmeld/crates/wasmeld-runtime/tests/runtime_components.rs

use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
    process::Command,
    sync::{Arc, Mutex, OnceLock},
    thread,
    time::{Duration, Instant},
};

use wasmeld_runtime::{
    KV_MAX_VALUE_BYTES, KvBackend, KvBackendError, ResourceLimits, Runtime, RuntimeConfig,
    RuntimeError, ServiceManifest,
};

static COMPONENTS_BUILT: OnceLock<()> = OnceLock::new();

#[test]
fn echo_component_round_trips_bytes() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let (key, actor) = start_component(&runtime, "echo", "echo_component.wasm");

    assert!(runtime.capabilities(&key).unwrap().is_empty());
    assert_eq!(actor.invoke(b"hello wasm".to_vec()).unwrap(), b"hello wasm");

    runtime.stop(&key).unwrap();
}

#[test]
fn counter_memory_is_resident_and_resets_after_restart() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let (key, actor) = start_component(&runtime, "counter", "counter_component.wasm");

    assert_eq!(decode_counter(actor.invoke(Vec::new()).unwrap()), 1);
    assert_eq!(decode_counter(actor.invoke(Vec::new()).unwrap()), 2);

    runtime.stop(&key).unwrap();
    let restarted = runtime.start(&key, Vec::new()).unwrap();
    assert_eq!(decode_counter(restarted.invoke(Vec::new()).unwrap()), 1);

    runtime.stop(&key).unwrap();
}

#[test]
fn actor_serializes_concurrent_counter_calls() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let (key, actor) = start_component(&runtime, "counter-concurrent", "counter_component.wasm");

    let mut workers = Vec::new();
    for _ in 0..8 {
        let actor = actor.clone();
        workers.push(thread::spawn(move || {
            decode_counter(actor.invoke(Vec::new()).unwrap())
        }));
    }

    let mut values = workers
        .into_iter()
        .map(|worker| worker.join().expect("worker should not panic"))
        .collect::<Vec<_>>();
    values.sort_unstable();
    assert_eq!(values, (1..=8).collect::<Vec<_>>());

    runtime.stop(&key).unwrap();
}

#[test]
fn runtime_reports_stats_and_stops_all_actors() {
    let runtime = Runtime::new(RuntimeConfig::default()).unwrap();
    let (key, actor) = start_component(&runtime, "runtime-control", "echo_component.wasm");

    assert_eq!(runtime.stats().unwrap().registered_services, 1);
    assert_eq!(runtime.stats().unwrap().running_actors, 1);

    runtime.stop_all().unwrap();
    assert_eq!(runtime.stats().unwrap().registered_services, 1);
    assert_eq!(runtime.stats().unwrap().running_actors, 0);
    assert!(matches!(
        actor.invoke(b"stopped".to_vec()),
        Err(RuntimeError::ActorUnavailable(unavailable)) if unavailable == key
    ));

    runtime.start(&key, Vec::new()).unwrap();
    assert_eq!(
        runtime.invoke(&key, b"started-again".to_vec()).unwrap(),
        b"started-again"
    );
}

#[test]
fn faulted_actor_can_be_started_again() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let (key, actor) = start_component(&runtime, "fault", "fault_component.wasm");

    assert!(matches!(
        actor.invoke(b"fault".to_vec()),
        Err(RuntimeError::ActorFault { .. })
    ));

    let restarted = runtime.start(&key, Vec::new()).unwrap();
    assert_eq!(restarted.invoke(b"ready".to_vec()).unwrap(), b"ready");

    runtime.stop(&key).unwrap();
}

#[test]
fn fuel_limit_stops_cpu_bound_component() {
    let runtime = Runtime::new(RuntimeConfig {
        epoch_tick: Duration::from_millis(10),
        ..RuntimeConfig::default()
    })
    .expect("runtime should start");
    let key = register_component(
        &runtime,
        "spin-fuel",
        "spin_component.wasm",
        ResourceLimits {
            fuel_per_call: 10_000,
            deadline_ms: 50,
            ..test_limits()
        },
    );
    let actor = runtime.start(&key, Vec::new()).unwrap();

    let error = actor.invoke(u64::MAX.to_le_bytes().to_vec()).unwrap_err();
    assert!(
        matches!(&error, RuntimeError::ActorFault { .. }),
        "expected a Wasm trap after fuel exhaustion, got {error:?}"
    );
    assert!(matches!(
        actor.invoke(Vec::new()),
        Err(RuntimeError::ActorUnavailable(_))
    ));
}

#[test]
fn epoch_interrupts_cpu_bound_initialization() {
    let runtime = Runtime::new(RuntimeConfig {
        epoch_tick: Duration::from_millis(2),
        ..RuntimeConfig::default()
    })
    .expect("runtime should start");
    let key = register_component(
        &runtime,
        "spin-epoch",
        "spin_component.wasm",
        ResourceLimits {
            fuel_per_call: 10_000_000_000,
            deadline_ms: 100,
            ..test_limits()
        },
    );

    let started_at = Instant::now();
    let error = match runtime.start(&key, b"spin".to_vec()) {
        Ok(_) => panic!("epoch budget should interrupt the spin component"),
        Err(error) => error,
    };
    assert!(
        matches!(&error, RuntimeError::ActorInitialization(_)),
        "expected initialization to be interrupted, got {error:?}"
    );
    let elapsed = started_at.elapsed();
    assert!(
        elapsed >= Duration::from_millis(20),
        "initialization failed before the epoch budget could expire"
    );
    assert!(
        elapsed < Duration::from_secs(2),
        "epoch interruption took too long"
    );
}

#[test]
fn explicit_clock_capability_is_linked() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let (key, actor) = start_component(&runtime, "clock-probe", "clock_probe_component.wasm");

    let capabilities = runtime.capabilities(&key).unwrap();
    assert_eq!(capabilities.len(), 1);
    assert_eq!(
        capabilities[0].interface(),
        "wasmeld:clock/monotonic-clock@0.1.0"
    );

    let first = actor.invoke(b"first".to_vec()).unwrap();
    let second = actor.invoke(b"second".to_vec()).unwrap();
    assert_eq!(&first[8..], b"first");
    assert_eq!(&second[8..], b"second");
    assert!(decode_clock(&second) >= decode_clock(&first));

    runtime.stop(&key).unwrap();
}

#[test]
fn kv_capability_is_service_scoped_and_shared_across_revisions() {
    let backend = Arc::new(MemoryKvBackend::default());
    let runtime = Runtime::new(RuntimeConfig {
        kv_backend: Some(backend),
        ..RuntimeConfig::default()
    })
    .expect("runtime should start");
    let first_key = register_component_revision(
        &runtime,
        "shared-kv",
        "0.1.0",
        "kv_probe_component.wasm",
        test_limits(),
    );
    let second_key = register_component_revision(
        &runtime,
        "shared-kv",
        "0.2.0",
        "kv_probe_component.wasm",
        test_limits(),
    );
    let isolated_key = register_component_revision(
        &runtime,
        "isolated-kv",
        "0.1.0",
        "kv_probe_component.wasm",
        test_limits(),
    );

    let first = runtime.start(&first_key, Vec::new()).unwrap();
    let second = runtime.start(&second_key, Vec::new()).unwrap();
    let isolated = runtime.start(&isolated_key, Vec::new()).unwrap();

    assert_eq!(first.invoke(b"set:answer:forty-two".to_vec()).unwrap(), b"");
    assert_eq!(second.invoke(b"get:answer".to_vec()).unwrap(), b"forty-two");
    assert_eq!(isolated.invoke(b"get:answer".to_vec()).unwrap(), b"");
    assert_eq!(
        runtime.capabilities(&first_key).unwrap()[0].interface(),
        "wasmeld:kv/store@0.1.0"
    );

    runtime.stop_all().unwrap();
}

#[test]
fn kv_capability_requires_a_backend_and_enforces_limits() {
    let artifact = component_artifact("kv_probe_component.wasm");
    let manifest = ServiceManifest {
        id: "kv-without-backend".to_owned(),
        revision: "0.1.0".to_owned(),
        component: artifact.display().to_string(),
        world: component_world("kv_probe_component.wasm").to_owned(),
        limits: test_limits(),
    };
    let bytes = fs::read(&artifact).unwrap();
    let runtime = Runtime::new(RuntimeConfig::default()).unwrap();
    assert!(matches!(
        runtime.register(manifest, &bytes),
        Err(RuntimeError::CapabilityUnavailable(interface))
            if interface == "wasmeld:kv/store@0.1.0"
    ));

    let runtime = Runtime::new(RuntimeConfig {
        kv_backend: Some(Arc::new(MemoryKvBackend::default())),
        ..RuntimeConfig::default()
    })
    .unwrap();
    let mut limits = test_limits();
    limits.max_input_bytes = 128 * 1024;
    let key = register_component(&runtime, "kv-limits", "kv_probe_component.wasm", limits);
    let actor = runtime.start(&key, Vec::new()).unwrap();

    let long_key = format!("get:{}", "k".repeat(257));
    assert!(matches!(
        actor.invoke(long_key.into_bytes()),
        Err(RuntimeError::ComponentError { message, .. })
            if message.contains("InvalidKey")
    ));

    let mut oversized = b"set:key:".to_vec();
    oversized.resize(oversized.len() + KV_MAX_VALUE_BYTES + 1, b'x');
    assert!(matches!(
        actor.invoke(oversized),
        Err(RuntimeError::ComponentError { message, .. })
            if message.contains("ValueTooLarge")
    ));

    runtime.stop(&key).unwrap();
}

#[test]
fn unregister_stops_actor_and_releases_compiled_revision() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let key = register_component(&runtime, "unregister", "echo_component.wasm", test_limits());
    runtime.start(&key, Vec::new()).unwrap();

    runtime.unregister(&key).unwrap();

    let stats = runtime.stats().unwrap();
    assert_eq!(stats.registered_services, 0);
    assert_eq!(stats.running_actors, 0);
    assert!(matches!(
        runtime.start(&key, Vec::new()),
        Err(RuntimeError::ServiceNotRegistered(missing)) if missing == key
    ));
}

#[test]
fn registration_rejects_non_whitelisted_wasi_imports() {
    let runtime = Runtime::new(RuntimeConfig::default()).expect("runtime should start");
    let artifact_name = "wasi_clock_probe_component.wasm";
    let artifact = component_artifact(artifact_name);
    let manifest = ServiceManifest {
        id: "wasi-clock-probe".to_owned(),
        revision: "0.1.0".to_owned(),
        component: artifact.display().to_string(),
        world: component_world(artifact_name).to_owned(),
        limits: test_limits(),
    };

    let error = runtime
        .register(manifest, fs::read(artifact).unwrap())
        .expect_err("monotonic clock must not be a V1 capability");
    assert!(matches!(
        error,
        RuntimeError::UnsupportedImport(name) if name.starts_with("wasi:clocks/monotonic-clock@")
    ));
}

fn start_component(
    runtime: &Runtime,
    service_id: &str,
    artifact_name: &str,
) -> (wasmeld_runtime::ServiceKey, wasmeld_runtime::ActorHandle) {
    let key = register_component(runtime, service_id, artifact_name, test_limits());
    let actor = runtime.start(&key, Vec::new()).unwrap();
    (key, actor)
}

fn register_component(
    runtime: &Runtime,
    service_id: &str,
    artifact_name: &str,
    limits: ResourceLimits,
) -> wasmeld_runtime::ServiceKey {
    register_component_revision(runtime, service_id, "0.1.0", artifact_name, limits)
}

fn register_component_revision(
    runtime: &Runtime,
    service_id: &str,
    revision: &str,
    artifact_name: &str,
    limits: ResourceLimits,
) -> wasmeld_runtime::ServiceKey {
    let artifact = component_artifact(artifact_name);
    let manifest = ServiceManifest {
        id: service_id.to_owned(),
        revision: revision.to_owned(),
        component: artifact.display().to_string(),
        world: component_world(artifact_name).to_owned(),
        limits,
    };

    let bytes = fs::read(&artifact).expect("component artifact should be readable");
    runtime.register(manifest, bytes).unwrap()
}

fn component_artifact(name: &str) -> PathBuf {
    build_components_once();
    workspace_root()
        .join("target/wasm32-wasip2/release")
        .join(name)
}

fn build_components_once() {
    COMPONENTS_BUILT.get_or_init(|| {
        let toolchain =
            std::env::var("WASMELD_COMPONENT_TOOLCHAIN").unwrap_or_else(|_| "1.90.0".to_owned());
        for manifest in [
            "components/echo/Cargo.toml",
            "components/counter/Cargo.toml",
            "components/fault/Cargo.toml",
            "components/spin/Cargo.toml",
            "components/clock-probe/Cargo.toml",
            "components/kv-probe/Cargo.toml",
            "components/wasi-clock-probe/Cargo.toml",
        ] {
            let status = Command::new("rustup")
                .args([
                    "run",
                    &toolchain,
                    "cargo",
                    "build",
                    "--manifest-path",
                    manifest,
                    "--target",
                    "wasm32-wasip2",
                    "--release",
                ])
                .current_dir(workspace_root())
                .status()
                .expect("component build command should start");
            assert!(status.success(), "component build should succeed");
        }
    });
}

fn workspace_root() -> &'static Path {
    static ROOT: OnceLock<PathBuf> = OnceLock::new();
    ROOT.get_or_init(|| {
        PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .and_then(Path::parent)
            .expect("runtime crate must live in the workspace")
            .to_path_buf()
    })
    .as_path()
}

fn test_limits() -> ResourceLimits {
    ResourceLimits {
        memory_bytes: 64 * 1024 * 1024,
        fuel_per_call: 1_000_000,
        deadline_ms: 500,
        mailbox_capacity: 16,
        max_input_bytes: 16 * 1024,
        max_output_bytes: 16 * 1024,
    }
}

fn decode_counter(bytes: Vec<u8>) -> u64 {
    let array: [u8; 8] = bytes
        .try_into()
        .expect("counter response must be a u64 payload");
    u64::from_le_bytes(array)
}

fn decode_clock(bytes: &[u8]) -> u64 {
    u64::from_le_bytes(
        bytes[..8]
            .try_into()
            .expect("clock response must start with a u64 payload"),
    )
}

fn component_world(artifact_name: &str) -> &'static str {
    match artifact_name {
        "echo_component.wasm" => "component:echo/echo-component@0.1.0",
        "counter_component.wasm" => "component:counter/counter-component@0.1.0",
        "fault_component.wasm" => "component:fault/fault-component@0.1.0",
        "spin_component.wasm" => "component:spin/spin-component@0.1.0",
        "clock_probe_component.wasm" => "component:clock-probe/clock-probe-component@0.1.0",
        "kv_probe_component.wasm" => "component:kv-probe/kv-probe-component@0.1.0",
        "wasi_clock_probe_component.wasm" => {
            "component:wasi-clock-probe/wasi-clock-probe-component@0.1.0"
        }
        other => panic!("no WIT world registered for {other}"),
    }
}

#[derive(Debug, Default)]
struct MemoryKvBackend {
    entries: Mutex<HashMap<(String, String), Vec<u8>>>,
}

impl KvBackend for MemoryKvBackend {
    fn get(
        &self,
        service_id: &str,
        key: &str,
        _timeout: Duration,
    ) -> Result<Option<Vec<u8>>, KvBackendError> {
        Ok(self
            .entries
            .lock()
            .unwrap()
            .get(&(service_id.to_owned(), key.to_owned()))
            .cloned())
    }

    fn set(
        &self,
        service_id: &str,
        key: &str,
        value: &[u8],
        _timeout: Duration,
    ) -> Result<(), KvBackendError> {
        self.entries
            .lock()
            .unwrap()
            .insert((service_id.to_owned(), key.to_owned()), value.to_vec());
        Ok(())
    }

    fn delete(
        &self,
        service_id: &str,
        key: &str,
        _timeout: Duration,
    ) -> Result<(), KvBackendError> {
        self.entries
            .lock()
            .unwrap()
            .remove(&(service_id.to_owned(), key.to_owned()));
        Ok(())
    }
}