palladium_plugin/

wasm_actor.rs

//! [`WasmActor`] — an `Actor` implementation backed by a WASM module instance.
//!
//! # Lifecycle
//!
//! 1. A `WasmActor` is constructed by [`PluginRegistry::create_actor`] (or
//!    directly in tests) with a compiled [`WasmModule`] and a [`WasmHost`].
//! 2. `on_start` instantiates the module with real [`WasmImports`] wired to
//!    the outbox, sets the initial fuel budget, and calls the module's
//!    `pd_actor_on_start` export (if present).
//! 3. `on_message` writes the incoming envelope + payload into WASM linear
//!    memory, calls the module's `pd_actor_on_message` export, then drains the
//!    outbox by forwarding each queued message through `ctx.send_raw()`.
//! 4. `on_stop` calls the module's `pd_actor_on_stop` export (if present).
//!
//! # Outbox pattern
//!
//! Because `Actor` is not `Send`, the `pd_send` host import can use an
//! `Arc<Mutex<…>>` outbox that is `Send + Sync + 'static` (required by the
//! wasmtime linker) while remaining accessible to the single-threaded actor.
//! After each WASM call the outbox is drained synchronously.
//!
//! # Memory layout
//!
//! The host reserves two slots in WASM linear memory for passing messages:
//!
//! | Offset | Size | Contents |
//! |--------|------|----------|
//! | 0      | 80   | Serialised `Envelope` (`Envelope::to_bytes()`) |
//! | 80     | ≤ payload | Raw payload bytes |
//!
//! These offsets are passed as `i32` arguments to the WASM exports.

use std::sync::{Arc, Mutex};

use palladium_actor::{Actor, ActorContext, ActorError, Envelope, MessagePayload, StopReason};

use crate::wasm::{WasmError, WasmHost, WasmImports, WasmInstance, WasmModule, WasmVal};

// ── Memory layout constants ───────────────────────────────────────────────────

/// Offset in WASM linear memory where the host writes the incoming envelope.
const ENVELOPE_OFFSET: u32 = 0;

/// Offset in WASM linear memory where the host writes the incoming payload.
const PAYLOAD_OFFSET: u32 = 80; // immediately after the 80-byte envelope

// ── WasmActor ────────────────────────────────────────────────────────────────

/// An actor whose lifecycle hooks are implemented in a WASM module.
///
/// Construct via [`WasmActor::new`] and spawn with the normal `ChildSpec`
/// factory API.  The actor creates its WASM instance lazily in `on_start`.
pub struct WasmActor<R: palladium_actor::Reactor> {
    /// Shared compiled module (may be reused across multiple actor instances).
    module: Arc<Box<dyn WasmModule>>,
    /// WASM runtime host — used to instantiate the module in `on_start`.
    host: Arc<dyn WasmHost>,
    /// Live WASM instance; `None` until `on_start` has run.
    instance: Option<Box<dyn WasmInstance>>,
    /// Opaque state handle passed to every WASM lifecycle call.
    /// Returned by `pd_actor_create` if that export exists; otherwise 0.
    state_handle: i32,
    /// Buffer populated by the `pd_send` import during a WASM call.
    /// Drained by `on_message` after the call returns.
    outbox: Arc<Mutex<Vec<(Envelope, MessagePayload)>>>,
    /// Fuel units granted per `on_message` call.  0 = use the instance default.
    pub fuel_per_message: u64,
    _phantom: std::marker::PhantomData<R>,
}

// SAFETY: WasmActor is managed by the pd-runtime which ensures it only ever
// runs on a single core's local executor. Box<dyn Actor> requires Send so
// we can move it into the spawn_local task.
unsafe impl<R: palladium_actor::Reactor> Send for WasmActor<R> {}

impl<R: palladium_actor::Reactor> WasmActor<R> {
    /// Create a new actor backed by `module`.
    ///
    /// * `module` — a compiled (but not yet instantiated) WASM module.
    /// * `host` — the WASM runtime that will instantiate the module in
    ///   `on_start`.
    /// * `fuel_per_message` — fuel budget granted before each `on_message`
    ///   call. Use `0` to rely on the instance's built-in default.
    pub fn new(
        module: Arc<Box<dyn WasmModule>>,
        host: Arc<dyn WasmHost>,
        fuel_per_message: u64,
    ) -> Self {
        Self {
            module,
            host,
            instance: None,
            state_handle: 0,
            outbox: Arc::new(Mutex::new(Vec::new())),
            fuel_per_message,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Build a [`WasmImports`] struct with the `pd_send` closure wired to this
    /// actor's outbox.
    fn make_imports(&self) -> WasmImports {
        let outbox = Arc::clone(&self.outbox);
        WasmImports {
            pd_send: Box::new(move |env_ptr, env_len, payload_ptr, payload_len| {
                if env_ptr.is_null() || env_len as usize != Envelope::SIZE {
                    return -1;
                }
                // Safety: env_ptr is valid for `env_len` bytes — guaranteed by
                // the WasmtimeInstance func_wrap wrapper, which copies the bytes
                // from WASM linear memory immediately before this call.
                let env_bytes: [u8; Envelope::SIZE] =
                    unsafe { *(env_ptr as *const [u8; Envelope::SIZE]) };
                let envelope = Envelope::from_bytes(&env_bytes);

                let payload = if payload_len > 0 && !payload_ptr.is_null() {
                    // Safety: same guarantee — wasmtime validated bounds.
                    let slice =
                        unsafe { std::slice::from_raw_parts(payload_ptr, payload_len as usize) };
                    MessagePayload::serialized(slice.to_vec())
                } else {
                    MessagePayload::serialized(Vec::new())
                };

                outbox.lock().unwrap().push((envelope, payload));
                0
            }),
            pd_now_micros: Box::new(move || {
                // In simulation, this should ideally be virtual wall-clock time.
                // For now, we return a value that is deterministic if the reactor is deterministic.
                // TODO: Add wall-clock support to Reactor trait.
                0
            }),
            pd_log: Box::new(|_level, _msg_ptr, _msg_len| {}),
        }
    }

    /// Return the exports of the compiled module.
    fn exports(&self) -> Vec<String> {
        self.module.exports()
    }

    /// Drain the outbox and forward each message through `ctx.send_raw()`.
    /// Returns Ok(()) if all messages were sent, or Err(ActorError::Handler)
    /// if any send failed (e.g. mailbox full).
    fn drain_outbox(&self, ctx: &mut ActorContext<R>) -> Result<(), ActorError> {
        let messages = {
            let mut lock = self.outbox.lock().unwrap();
            std::mem::take(&mut *lock)
        };
        for (env, payload) in messages.into_iter() {
            ctx.send_raw(env, payload)
                .map_err(|_| ActorError::Handler)?;
        }
        Ok(())
    }

    /// Map a [`WasmError`] from a lifecycle call to an [`ActorError`].
    fn map_wasm_err(e: WasmError) -> ActorError {
        match e {
            WasmError::FuelExhausted => ActorError::ResourceExhausted,
            WasmError::Trap(msg) => ActorError::WasmTrap(msg),
            _ => ActorError::Handler,
        }
    }
}

impl<R: palladium_actor::Reactor> Actor<R> for WasmActor<R> {
    fn on_start(&mut self, _ctx: &mut ActorContext<R>) -> Result<(), ActorError> {
        // Pre-check exports before borrowing self.instance mutably.
        let has_on_start = self.exports().contains(&"pd_actor_on_start".to_string());

        // Instantiate with wired-up imports.
        let imports = self.make_imports();
        let instance = self
            .host
            .instantiate(self.module.as_ref().as_ref(), imports)
            .map_err(|_| ActorError::Init)?;
        self.instance = Some(instance);

        let inst = self.instance.as_mut().unwrap();

        // Set initial fuel budget if requested.
        if self.fuel_per_message > 0 {
            inst.set_fuel(self.fuel_per_message)
                .map_err(|_| ActorError::Init)?;
        }

        // Call pd_actor_on_start if the module exports it.
        if has_on_start {
            let rets = inst
                .call(
                    "pd_actor_on_start",
                    &[WasmVal::I32(self.state_handle), WasmVal::I32(0)],
                )
                .map_err(Self::map_wasm_err)?;
            let code = first_i32(&rets);
            if code != 0 {
                return Err(ActorError::Init);
            }
        }

        Ok(())
    }

    fn on_message(
        &mut self,
        ctx: &mut ActorContext<R>,
        envelope: &Envelope,
        payload: MessagePayload,
    ) -> Result<(), ActorError> {
        let inst = match self.instance.as_mut() {
            Some(i) => i,
            None => return Err(ActorError::Init),
        };

        // Refresh fuel for this message.
        if self.fuel_per_message > 0 {
            inst.set_fuel(self.fuel_per_message)
                .map_err(|_| ActorError::ResourceExhausted)?;
        }

        // Write envelope to WASM memory at ENVELOPE_OFFSET.
        let env_bytes = envelope.to_bytes();
        inst.memory_write(ENVELOPE_OFFSET, &env_bytes)
            .map_err(|_| ActorError::Handler)?;

        // Write payload bytes to WASM memory at PAYLOAD_OFFSET.
        let payload_slice: &[u8] = match &payload {
            MessagePayload::Serialized(b) => b.as_ref(),
            MessagePayload::Local(_) => &[],
        };
        let payload_len = payload_slice.len() as i32;
        if payload_len > 0 {
            inst.memory_write(PAYLOAD_OFFSET, payload_slice)
                .map_err(|_| ActorError::Handler)?;
        }

        // Call pd_actor_on_message.
        let payload_ptr = if payload_len > 0 {
            PAYLOAD_OFFSET as i32
        } else {
            0
        };
        let rets = inst
            .call(
                "pd_actor_on_message",
                &[
                    WasmVal::I32(self.state_handle),
                    WasmVal::I32(0), // ctx: unused in WASM
                    WasmVal::I32(ENVELOPE_OFFSET as i32),
                    WasmVal::I32(payload_ptr),
                    WasmVal::I32(payload_len),
                ],
            )
            .map_err(Self::map_wasm_err)?;

        if first_i32(&rets) != 0 {
            return Err(ActorError::Handler);
        }

        // Drain outbox: deliver queued sends through the runtime bridge.
        self.drain_outbox(ctx)
    }

    fn on_stop(&mut self, _ctx: &mut ActorContext<R>, reason: StopReason) {
        // Pre-check before borrowing self.instance mutably.
        let has_on_stop = self.exports().contains(&"pd_actor_on_stop".to_string());

        let inst = match self.instance.as_mut() {
            Some(i) => i,
            None => return,
        };

        if !has_on_stop {
            return;
        }

        let reason_code = match reason {
            StopReason::Normal => 0i32,
            StopReason::Requested => 1,
            StopReason::Supervisor => 2,
            StopReason::Hierarchical => 2,
            StopReason::Shutdown => 3,
            StopReason::Killed => 4,
            StopReason::Error(_) => -1,
        };

        // Best-effort: ignore errors in on_stop.
        let _ = inst.call(
            "pd_actor_on_stop",
            &[
                WasmVal::I32(self.state_handle),
                WasmVal::I32(0),
                WasmVal::I32(reason_code),
            ],
        );
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Extract the first `i32` return value, or 0 if absent.
fn first_i32(rets: &[WasmVal]) -> i32 {
    match rets.first() {
        Some(WasmVal::I32(n)) => *n,
        _ => 0,
    }
}