aion-rs 0.29.0

//! Mailbox-delivery engine handle for live workflow queries.
//!
//! Delivery never executes the handler itself: it queues a pending-query
//! record for the target pid, parks the reply sender keyed by query id, and
//! enqueues an `aion_query` wake marker. The workflow's next suspending-await
//! invocation drains the queue through the query pump and replies over the
//! parked sender — nothing on this path touches the recorder or resolver.
//!
//! The caller's query arguments ride the pending-query record as the JSON
//! text the workflow handler receives. This module is the boundary that
//! validates that text (see [`arguments_text`]): an arguments payload that is
//! not well-formed UTF-8 JSON is refused as [`QueryError::InvalidArguments`]
//! without disturbing the workflow at all.

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

use aion_core::{Payload, WorkflowId, WorkflowStatus};
use uuid::Uuid;

use crate::Pid;
use crate::engine_seam::{
    ChildWorkflowSpawnRequest, ChildWorkflowSpawnResult, EngineHandle, EngineSeamError,
    TimerWheelEntry, WorkflowMailboxMessage, WorkflowProcessHandle, WorkflowResidency,
};
use crate::query::QueryError;
use crate::registry::{HandleResidency, Registry};
use crate::runtime::RuntimeHandle;

use super::nif_query::{
    insert_pending_reply, is_query_registered, prune_closed_pending_replies, take_pending_reply,
};
use super::nif_state::{EngineNifState, PendingQuery};

pub(super) struct QueryMailboxEngine {
    registry: Arc<Registry>,
    // Weak: the engine state owns this engine through its query bridge slot.
    nif_state: Weak<EngineNifState>,
    // Weak: the runtime owns the engine state that owns this engine.
    runtime: Weak<RuntimeHandle>,
}

impl QueryMailboxEngine {
    pub(super) fn new(
        registry: Arc<Registry>,
        nif_state: Weak<EngineNifState>,
        runtime: Weak<RuntimeHandle>,
    ) -> Self {
        Self {
            registry,
            nif_state,
            runtime,
        }
    }

    /// Park the reply, queue the query, and wake the workflow process.
    ///
    /// On marker-delivery failure every inserted entry is removed, so a
    /// failed delivery leaves no stale state. A failure after exit cleanup
    /// starts, or after scheduler retirement, is the query-racing-completion
    /// window — the caller observes typed `ReplyDropped` through its dropped
    /// sender. Only a failure while neither terminal condition holds is
    /// reported as an engine fault.
    fn enqueue_query(
        &self,
        state: &EngineNifState,
        pid: u64,
        name: String,
        arguments: String,
        reply_to: crate::engine_seam::QueryReplySender,
    ) -> Result<(), QueryError> {
        let runtime = self
            .runtime
            .upgrade()
            .ok_or_else(|| QueryError::Engine(delivery_error("engine runtime has shut down")))?;
        // Hygiene: drop senders whose caller already timed out, so a
        // never-woken workflow does not accumulate stale reply channels.
        prune_closed_pending_replies(state)
            .map_err(|error| QueryError::Engine(delivery_error(error)))?;
        let query_id = Uuid::new_v4().to_string();
        insert_pending_reply(state, query_id.clone(), pid, reply_to)
            .map_err(|error| QueryError::Engine(delivery_error(error)))?;
        state
            .pending_queries
            .entry(pid)
            .or_default()
            .push_back(PendingQuery {
                query_id: query_id.clone(),
                name,
                arguments,
            });
        if let Err(error) = runtime.deliver_query_request(pid) {
            // Roll back both entries before classifying the failure.
            if let Some(mut queue) = state.pending_queries.get_mut(&pid) {
                queue.retain(|pending| pending.query_id != query_id);
            }
            let removed = take_pending_reply(state, &query_id)
                .map_err(|reason| QueryError::Engine(delivery_error(reason)))?;
            drop(removed);
            // A workflow that exited between the caller's residency/terminal
            // checks and the wake-marker enqueue can never answer — this is
            // the query-racing-completion window, not an engine fault.
            // Cleanup stamps its tombstone before removing handlers, and the
            // handlers mutex orders readers after that stamp; scheduler
            // retirement may independently win first, hence the disjunction.
            // Dropping the parked sender above resolves the waiting caller
            // with the same typed `ReplyDropped` as exit-time cleanup
            // ("workflow ended before answering").
            if query_reply_was_dropped_by_completion(&runtime, pid) {
                return Ok(());
            }
            return Err(QueryError::Engine(delivery_error(format!(
                "query wake marker delivery failed: {error}"
            ))));
        }
        Ok(())
    }
}

impl EngineHandle for QueryMailboxEngine {
    fn resolve_workflow(
        &self,
        workflow_id: &WorkflowId,
    ) -> Result<WorkflowResidency, EngineSeamError> {
        // 🔴 THE INDEX ANSWERS "WHICH RUN IS CURRENT"; THE HANDLE MAP ANSWERS
        // ONLY "DOES ANY HANDLE EXIST" — `Slots::any_handle_for` draws exactly
        // that distinction, and this call site wants the first question. It
        // previously asked the second: a `find` over `Registry::list`, which is
        // `slots.handles.values().cloned().collect()`, a bare `HashMap::values`
        // in arbitrary order. After a continue-as-new TWO handles share one
        // workflow id — on the nif path the predecessor's handle is never
        // removed at all, and on the API path `lifecycle::continue_as_new`
        // records the terminal, starts the successor, and only then removes it,
        // so both are registered and resolvable inside that window. The scan
        // could therefore return the predecessor, and which one it returned was
        // determined by nothing the caller controls.
        //
        // `live_run_pid` is the run-aware lookup and its own documentation names
        // this case (OBX-011): the index is upserted newest-run-wins on insert,
        // and `forget_live_index_entry` drops an entry only while it still
        // points at the run being removed — so through the whole window the
        // index names the SUCCESSOR and nothing else.
        //
        // ⚠️ THE PID COMES FROM THE HANDLE, NOT FROM THE INDEX TUPLE, which is
        // why it is discarded here. The index is read under its own lock and the
        // handle under another; pairing a handle with a pid sampled from the
        // other lock is how the two come apart. The handle is the single
        // authority for both residency and pid, so both are taken from it.
        let Some((run, _pid)) = self
            .registry
            .live_run_pid(workflow_id)
            .map_err(|error| delivery_error(error.to_string()))?
        else {
            return Ok(WorkflowResidency::Unknown);
        };
        // A live handle whose index entry is gone is a workflow with no current
        // run, and `Unknown` is the honest answer for it. That is a deliberate
        // change from the scan, which reported such a workflow as `Terminal`:
        // both refuse the query — `QueryError::Unknown` against
        // `QueryError::NotRunning` — and `Unknown` is the accurate one when
        // nothing is current. Pinned by test rather than left to be discovered.
        let handle = self
            .registry
            .get(workflow_id, &run)
            .map_err(|error| delivery_error(error.to_string()))?;
        match handle {
            Some(handle) if handle.cached_status() != WorkflowStatus::Running => {
                Ok(WorkflowResidency::Terminal)
            }
            Some(handle) => match handle.residency() {
                HandleResidency::Resident => Ok(WorkflowResidency::Resident(
                    WorkflowProcessHandle::new(handle.pid()),
                )),
                // A suspended workflow has no live process to answer from;
                // AT-007 forbids resuming solely to answer a query.
                HandleResidency::Suspended => Ok(WorkflowResidency::NonResident),
            },
            None => Ok(WorkflowResidency::Unknown),
        }
    }

    fn deliver_workflow_message(
        &self,
        process: WorkflowProcessHandle,
        message: WorkflowMailboxMessage,
    ) -> Result<(), EngineSeamError> {
        let WorkflowMailboxMessage::Query {
            name,
            reply_to,
            payload,
        } = message
        else {
            return Err(delivery_error(
                "query mailbox engine only accepts query messages",
            ));
        };
        let Some(state) = self.nif_state.upgrade() else {
            return reply_to
                .send(Err(QueryError::Engine(delivery_error(
                    "engine NIF state has been dropped",
                ))))
                .map_err(|_| delivery_error("query caller dropped reply receiver"));
        };
        // Arguments become a JSON string value inside the pump sentinel, so
        // they must be well-formed UTF-8 JSON before anything is queued. A
        // malformed payload is the caller's fault and is refused here rather
        // than waking the workflow to fail inside a handler.
        let arguments = match arguments_text(&payload) {
            Ok(arguments) => arguments,
            Err(reason) => {
                return reply_to
                    .send(Err(QueryError::InvalidArguments { reason }))
                    .map_err(|_| delivery_error("query caller dropped reply receiver"));
            }
        };
        match is_query_registered(&state, process.pid(), &name) {
            Ok(true) => {
                match self.enqueue_query(&state, process.pid(), name, arguments, reply_to) {
                    Ok(()) => Ok(()),
                    // The reply sender was consumed by the rollback inside
                    // `enqueue_query`; surface the failure to the seam caller.
                    Err(error) => Err(delivery_error(format!("query enqueue failed: {error}"))),
                }
            }
            // An unregistered name never disturbs the workflow process. Exit
            // cleanup stamps its tombstone before removing handlers, and the
            // handlers mutex orders this absent-registration read after that
            // stamp. Scheduler retirement can independently happen first, so
            // either terminal observation means completion, not author error.
            Ok(false) => {
                let reply = match self.runtime.upgrade() {
                    Some(runtime)
                        if !query_reply_was_dropped_by_completion(&runtime, process.pid()) =>
                    {
                        Err(QueryError::UnknownQuery(name))
                    }
                    // A shut-down runtime cannot answer either; name this
                    // branch explicitly rather than relying on Weak semantics.
                    Some(_) | None => Err(QueryError::ReplyDropped),
                };
                reply_to
                    .send(reply)
                    .map_err(|_| delivery_error("query caller dropped reply receiver"))
            }
            Err(error) => reply_to
                .send(Err(QueryError::Engine(delivery_error(error))))
                .map_err(|_| delivery_error("query caller dropped reply receiver")),
        }
    }

    fn spawn_child_workflow(
        &self,
        request: ChildWorkflowSpawnRequest,
    ) -> Result<ChildWorkflowSpawnResult, EngineSeamError> {
        Err(EngineSeamError::ChildSpawn {
            reason: format!(
                "query mailbox engine does not spawn child workflow {}",
                request.workflow_type
            ),
        })
    }

    fn terminate_linked_child_workflow(
        &self,
        parent_workflow_id: &WorkflowId,
        child_process: WorkflowProcessHandle,
        correlation: u64,
    ) -> Result<(), EngineSeamError> {
        Err(EngineSeamError::ChildTermination {
            reason: format!(
                "query mailbox engine does not terminate child {child_process:?} for {parent_workflow_id} with correlation {correlation}"
            ),
        })
    }

    fn terminate_linked_activity(
        &self,
        parent_workflow_id: &WorkflowId,
        activity_process: Pid,
        correlation: u64,
    ) -> Result<(), EngineSeamError> {
        Err(EngineSeamError::ChildTermination {
            reason: format!(
                "query mailbox engine does not terminate activity {activity_process} for {parent_workflow_id} with correlation {correlation}"
            ),
        })
    }

    fn arm_timer(&self, entry: TimerWheelEntry) -> Result<(), EngineSeamError> {
        Err(EngineSeamError::TimerWheel {
            reason: format!("query mailbox engine does not arm timer {}", entry.timer_id),
        })
    }

    fn disarm_timer(
        &self,
        process: WorkflowProcessHandle,
        timer_id: &aion_core::TimerId,
    ) -> Result<(), EngineSeamError> {
        Err(EngineSeamError::TimerWheel {
            reason: format!(
                "query mailbox engine does not disarm timer {timer_id} for {process:?}"
            ),
        })
    }

    fn record_workflow_event(
        &self,
        workflow_id: &WorkflowId,
        event: aion_core::Event,
    ) -> Result<crate::engine_seam::RecordOutcome, EngineSeamError> {
        Err(EngineSeamError::Recorder {
            reason: format!(
                "queries must not record event {} for workflow {workflow_id}",
                event.seq()
            ),
        })
    }

    fn record_redelivered_timer_fire(
        &self,
        workflow_id: &WorkflowId,
        timer_id: &aion_core::TimerId,
    ) -> Result<crate::engine_seam::RedeliveredFire, EngineSeamError> {
        Err(EngineSeamError::Recorder {
            reason: format!(
                "queries must not answer timer redelivery of `{timer_id}` for {workflow_id}"
            ),
        })
    }
}

fn query_reply_was_dropped_by_completion(runtime: &RuntimeHandle, pid: Pid) -> bool {
    runtime.process_cleanup_started(pid) || !runtime.is_live(pid)
}

fn delivery_error(reason: impl Into<String>) -> EngineSeamError {
    EngineSeamError::Delivery {
        reason: reason.into(),
    }
}

/// Render a caller's query arguments payload as the JSON text a handler reads.
///
/// The pump sentinel embeds this text as a JSON string value, and the SDK
/// hands it to the registered handler's arguments codec. Both steps require a
/// well-formed UTF-8 JSON document, so both properties are checked here — the
/// last point at which the failure is still attributable to the caller. The
/// caller's exact bytes are forwarded on success: re-serializing would rewrite
/// key order and number formatting the author's codec may depend on.
fn arguments_text(payload: &Payload) -> Result<String, String> {
    let text = std::str::from_utf8(payload.bytes())
        .map_err(|error| format!("arguments payload is not valid UTF-8: {error}"))?;
    serde_json::from_str::<serde_json::Value>(text).map_err(|error| {
        format!("arguments payload is not a well-formed JSON document: {error}")
    })?;
    Ok(text.to_owned())
}

#[cfg(test)]
#[path = "nif_query_mailbox_tests.rs"]
mod tests;