aion-rs 0.31.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, UndeliveredWake};

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 — unless the reply table itself
    /// is unreadable, and then the caller is told so with both causes. A
    /// failure after exit cleanup starts, or after the exit registry holds
    /// the process' terminal, is the query-racing-completion window: the
    /// caller observes typed `ReplyDropped` through its dropped sender. A pid
    /// that is merely absent from the process table is never read as an
    /// ending; an exit in flight is waited out by the classifier, and only a
    /// refusal it cannot excuse is reported as the engine fault it is,
    /// carrying the delivery refusal.
    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(format!(
                    "the parked reply could not be rolled back after wake marker delivery \
                     failed ({error}): {reason}"
                )))
            })?;
            drop(removed);
            // A workflow that exited between the caller's residency/terminal
            // checks and the wake-marker enqueue can never answer. That is
            // the query-racing-completion window, not an engine fault.
            // Dropping the parked sender above resolves the waiting caller
            // with the same typed `ReplyDropped` as exit-time cleanup
            // ("workflow ended before answering"). The classifier waits out
            // an exit that is still in flight, and this seam is synchronous,
            // so the park lands on the caller's tokio worker. It is bounded by
            // `ready_timeout` — the same window `deliver_query_request` blocks
            // this worker on in its own ready gate, on the branch where the
            // pid is still in the table — and the thread that publishes the
            // terminal is the exit drainer's own, never an executor worker, so
            // the park cannot starve its own release.
            match runtime.classify_undelivered_wake(pid) {
                Ok(UndeliveredWake::ProcessEnded) => return Ok(()),
                Ok(UndeliveredWake::ProcessDidNotEnd { .. }) => {
                    return Err(QueryError::Engine(delivery_error(format!(
                        "query wake marker delivery failed: {error}"
                    ))));
                }
                Ok(UndeliveredWake::ExitInFlight) => {
                    return Err(QueryError::Engine(delivery_error(format!(
                        "query wake marker delivery failed ({error}): process {pid} left the \
                         scheduler process table with its exit still in flight, and no terminal \
                         was published on its exit record within the readiness window"
                    ))));
                }
                Err(reason) => {
                    return Err(QueryError::Engine(delivery_error(format!(
                        "query completion classification failed after wake marker delivery \
                         failed ({error}): {reason}"
                    ))));
                }
            }
        }
        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, and
            // three different facts hide behind one absent registration;
            // `unregistered_query_error` separates them.
            Ok(false) => {
                let reply = match self.runtime.upgrade() {
                    Some(runtime) => Err(unregistered_query_error(&runtime, process.pid(), name)),
                    // The runtime is gone, so no ending can be classified;
                    // name the engine's own shutdown rather than attributing
                    // it to the workflow.
                    None => Err(QueryError::Engine(delivery_error(
                        "engine runtime has shut down",
                    ))),
                };
                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}"
            ),
        })
    }
}

/// The typed answer for a query name the workflow process holds no handler for.
///
/// Three distinct facts hide behind one absent registration, and only one of
/// them is the workflow author's to fix:
///
/// * the engine holds a record that the process ended — cleanup stamps its
///   tombstone before removing handlers, and the handlers mutex orders this
///   absent-registration read after that stamp, so the handler set was torn
///   down with the process. That is `ReplyDropped`, the same answer exit-time
///   cleanup gives a parked reply;
/// * the process was in the scheduler's table when the classifier looked and
///   declares no such name — the author's error, `UnknownQuery`, which the
///   wire reports as a bad request;
/// * the process left the scheduler's table with its exit still in flight and
///   published no terminal inside the readiness window the classifier waits
///   out. Calling that the author's error would hand a caller a non-retryable
///   verdict about a query name over a process that no longer exists, so it
///   is reported as the engine fault it is: a server-side failure the caller
///   can re-issue once the exit lands, never a request the caller is told to
///   correct.
fn unregistered_query_error(runtime: &RuntimeHandle, pid: Pid, name: String) -> QueryError {
    match runtime.classify_undelivered_wake(pid) {
        Ok(UndeliveredWake::ProcessEnded) => QueryError::ReplyDropped,
        Ok(UndeliveredWake::ProcessDidNotEnd { in_table: true }) => QueryError::UnknownQuery(name),
        Ok(UndeliveredWake::ProcessDidNotEnd { in_table: false }) => {
            QueryError::Engine(delivery_error(format!(
                "query `{name}` cannot be answered: process {pid} is absent from the scheduler \
                 process table and the engine holds no record of its ending"
            )))
        }
        Ok(UndeliveredWake::ExitInFlight) => QueryError::Engine(delivery_error(format!(
            "query `{name}` cannot be answered: process {pid} left the scheduler process table \
             with its exit still in flight, and no terminal was published on its exit record \
             within the readiness window"
        ))),
        Err(reason) => QueryError::Engine(delivery_error(format!(
            "query completion classification failed: {reason}"
        ))),
    }
}

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;