Skip to main content

aion_server/worker/
bridge.rs

1//! NIF bridge dispatcher that routes `run_activity` calls to connected workers.
2//!
3//! `WorkerActivityDispatcher` implements `aion::ActivityDispatcher` so the
4//! engine's activity NIFs can synchronously dispatch to a remote worker and
5//! block until the result comes back.
6//!
7//! # Threading contract
8//!
9//! The engine invokes [`aion::ActivityDispatcher::dispatch`] from two kinds of
10//! threads: beamr scheduler threads (concurrency combinators) and spawned
11//! tokio tasks (the two-phase `dispatch_activity` completion task). The task
12//! send uses `try_send()` (non-blocking channel push) and the response wait
13//! blocks on `std::sync::mpsc::Receiver::recv`.
14//!
15//! Blocking is harmless on a beamr thread, but on a tokio runtime worker it
16//! must be wrapped in `tokio::task::block_in_place`: the `try_send` wakes the
17//! per-worker gRPC stream forwarder task, and tokio schedules a task woken
18//! from task context into the *current* worker's LIFO slot, which no other
19//! runtime worker can steal. Without the `block_in_place` core handoff the
20//! forwarder sits trapped in that slot while this thread blocks, so the queued
21//! `ActivityTask` is never flushed to the worker even though the worker is
22//! healthy. `block_in_place` moves the worker's scheduler core (LIFO slot
23//! included) to another thread before the wait begins, so dispatch-to-delivery
24//! stays in the millisecond range and the runtime keeps full parallelism.
25//!
26//! # Wait termination
27//!
28//! The engine imposes no activity timeout of its own: agent-style activities
29//! legitimately run for over an hour, so the completion wait is unbounded.
30//! The blocking `recv` terminates on exactly one of:
31//!
32//! - **Completion** — the worker reports a result and the stream handler
33//!   delivers it through [`ActivityCompletionSink::complete_activity`].
34//! - **Worker loss** — the worker's gRPC stream ends (process death,
35//!   disconnect, expired token); the stream teardown sweeps the worker's
36//!   in-flight tasks through the same sink as retryable lost-worker
37//!   failures ([`HeartbeatTracker::fail_disconnected_worker`]). A
38//!   liminal-delivered worker's loss is observed by its reply router thread
39//!   instead (the correlated-reply awaiter wakes the moment the connection
40//!   closes) and resolves the dispatch with the same retryable lost-worker
41//!   failure.
42//! - **Graceful-drain park (#207)** — during a drain, a worker stream ending
43//!   (or the drain-timeout backstop) PARKS the worker's in-flight tasks
44//!   through [`ActivityCompletionSink::park_activity`]: the waiter resolves
45//!   with the ephemeral parked sentinel
46//!   ([`aion::PARKED_ACTIVITY_REASON`]), nothing is recorded or delivered,
47//!   and restart recovery re-dispatches the dangling ordinal — kill -9
48//!   convergence.
49//! - **Drain timeout at shutdown** — the shutdown coordinator parks all
50//!   remaining in-flight tasks through the sink
51//!   (`HeartbeatTracker::park_all_in_flight_workers`).
52//! - **Channel teardown** — every sender for the pending entry is dropped
53//!   (a cleanup path removed the entry without completing it); surfaced as
54//!   a channel-closed dispatch error, never a hang.
55//!
56//! An activity's duration is bounded only by the workflow's own
57//! `timeout_seconds` and by worker liveness — never by an engine constant.
58
59use std::collections::BTreeMap;
60use std::sync::{Arc, OnceLock};
61use std::time::{Duration, Instant};
62
63use aion::{ActivityDispatch, ActivityDispatcher};
64use aion_core::{ActivityId, ContentType, Payload, RunId, WorkflowId};
65use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
66use dashmap::DashMap;
67
68use super::dispatch::{ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink};
69use super::heartbeat::{HeartbeatTracker, InFlightActivity};
70use super::registry::{
71    ConnectedWorkerRegistry, WorkerDelivery, WorkerHandle, WorkerId, WorkerMessage,
72};
73use crate::error::ServerError;
74use crate::shutdown::DrainState;
75use tracing::info_span;
76
77type SyncSender = std::sync::mpsc::SyncSender<Result<String, String>>;
78type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
79
80/// Execution-scoped key for an in-flight activity dispatch.
81///
82/// The engine seam ([`ActivityDispatch`]) carries the *real* workflow id and
83/// the *real* per-workflow activity ordinal recorded in history, so this pair
84/// uniquely and stably identifies one execution. Keying by bare [`ActivityId`]
85/// would be unsafe across server restarts — a stale result re-reported from a
86/// worker's previous session could complete a *different* post-restart
87/// dispatch reusing the same ordinal — but pairing it with the real workflow
88/// id closes that race: two different workflow executions never share a
89/// workflow id, so a stale `(workflow_id, activity_id)` from a previous server
90/// life can only ever match the exact execution it belongs to.
91///
92/// The wire (`ActivityResult`) carries both ids, plus an attempt discriminator
93/// (`ActivityTask.attempt`). The pending key stays attempt-free for now: a
94/// retry re-dispatches under the same `(workflow_id, activity_id)` and the
95/// outstanding entry is the one awaiting completion. Redelivery bookkeeping
96/// can widen this key with the wire attempt later — no protocol change needed.
97type PendingActivityKey = (WorkflowId, ActivityId);
98
99/// Routes an unmatched durable-outbox completion into the live workflow.
100///
101/// When the outbox is ON a worker completion can arrive at the sink with no
102/// pending oneshot (the dispatch was non-blocking fan-out, or the original
103/// waiter was lost). Rather than dropping it, [`PendingActivities::complete`]
104/// hands it to this callback, which resolves the workflow to its live engine
105/// process and delivers the terminal into its mailbox. The callback is only
106/// installed when the outbox is enabled, so flag-off the unmatched branch
107/// stays a silent drop.
108pub trait OutboxDeliveryCallback: Send + Sync {
109    /// Deliver a successful completion to the live workflow.
110    ///
111    /// Returns `Ok(true)` when delivered to a live workflow and `Ok(false)`
112    /// when no run is currently live (the expected stale-completion case that
113    /// recovery re-arms).
114    ///
115    /// # Errors
116    ///
117    /// Returns [`ServerError`] when the engine rejects the delivery.
118    fn deliver_completion(
119        &self,
120        workflow_id: &WorkflowId,
121        activity_id: &ActivityId,
122        run_id: Option<&RunId>,
123        result: String,
124    ) -> Result<bool, ServerError>;
125
126    /// Deliver a failure to the live workflow. Same `bool`/error contract as
127    /// [`Self::deliver_completion`].
128    ///
129    /// # Errors
130    ///
131    /// Returns [`ServerError`] when the engine rejects the delivery.
132    fn deliver_failure(
133        &self,
134        workflow_id: &WorkflowId,
135        activity_id: &ActivityId,
136        run_id: Option<&RunId>,
137        reason: String,
138    ) -> Result<bool, ServerError>;
139}
140
141/// Tracks in-flight activity dispatches waiting for worker results.
142///
143/// When the server's worker stream handler receives an `ActivityResult`, it
144/// calls [`complete_activity`](ActivityCompletionSink::complete_activity) to
145/// deliver the result to the blocked NIF thread. Entries are keyed by
146/// [`PendingActivityKey`] so a stale result from a previous server life can
147/// never be matched to a different execution (#59).
148///
149/// Clones share both the pending map and the outbox-delivery callback through
150/// `Arc`, so [`set_outbox_delivery`](Self::set_outbox_delivery) called once on
151/// any clone after construction is visible to the clone the dispatcher holds.
152#[derive(Clone, Default)]
153pub struct PendingActivities {
154    pending: Arc<DashMap<PendingActivityKey, SyncSender>>,
155    outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
156}
157
158impl std::fmt::Debug for PendingActivities {
159    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160        formatter
161            .debug_struct("PendingActivities")
162            .field("pending", &self.pending.len())
163            .field(
164                "outbox_delivery_installed",
165                &self.outbox_delivery.get().is_some(),
166            )
167            .finish()
168    }
169}
170
171impl PendingActivities {
172    fn insert(&self, workflow_id: WorkflowId, activity_id: ActivityId) -> SyncReceiver {
173        let (tx, rx) = std::sync::mpsc::sync_channel(1);
174        self.pending.insert((workflow_id, activity_id), tx);
175        rx
176    }
177
178    /// Test seam: register a pending waiter exactly as a live dispatch does,
179    /// so crate-internal tests outside this module (the stream-teardown
180    /// park-vs-fail pins) can observe how a sweep resolves it.
181    #[cfg(test)]
182    pub(crate) fn insert_for_test(
183        &self,
184        workflow_id: WorkflowId,
185        activity_id: ActivityId,
186    ) -> SyncReceiver {
187        self.insert(workflow_id, activity_id)
188    }
189
190    /// Install the unmatched-completion delivery callback (idempotent).
191    ///
192    /// Set once, after construction, when the durable outbox is enabled. A
193    /// second set is ignored and logged: the callback is process-wide and must
194    /// not silently change identity.
195    pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
196        if self.outbox_delivery.set(callback).is_err() {
197            tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
198        }
199    }
200
201    /// Complete a pending dispatch, or route an unmatched completion to the
202    /// outbox delivery callback when one is installed.
203    ///
204    /// A matched entry delivers to its waiting oneshot exactly as before. An
205    /// unmatched completion is dropped silently when no callback is installed
206    /// (outbox OFF — byte-identical to the prior behaviour); with a callback
207    /// installed (outbox ON) it is routed into the live workflow's mailbox.
208    fn complete(
209        &self,
210        workflow_id: &WorkflowId,
211        activity_id: &ActivityId,
212        run_id: Option<&RunId>,
213        result: Result<String, String>,
214    ) -> bool {
215        // Take and drop the DashMap guard before any callback runs: the engine
216        // delivery the callback invokes must never execute under a shard lock.
217        let matched = self
218            .pending
219            .remove(&(workflow_id.clone(), activity_id.clone()));
220        if let Some((_, sender)) = matched {
221            return sender.send(result).is_ok();
222        }
223        let Some(callback) = self.outbox_delivery.get() else {
224            // Outbox OFF: silent drop, byte-identical to the prior behaviour.
225            return false;
226        };
227        let outcome = match result {
228            Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
229            Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
230        };
231        match outcome {
232            Ok(true) => true,
233            Ok(false) => {
234                // Not live: the expected stale-completion case recovery re-arms.
235                tracing::debug!(
236                    workflow_id = %workflow_id,
237                    activity_id = %activity_id,
238                    "unmatched outbox completion for a workflow that is not currently live; \
239                     recovery will re-arm it"
240                );
241                false
242            }
243            Err(error) => {
244                tracing::warn!(
245                    workflow_id = %workflow_id,
246                    activity_id = %activity_id,
247                    %error,
248                    "failed to deliver unmatched outbox completion to the live workflow"
249                );
250                false
251            }
252        }
253    }
254}
255
256impl ActivityCompletionSink for PendingActivities {
257    /// Park one in-flight dispatch for restart recovery (#207): resolve the
258    /// matched waiter with the ephemeral parked sentinel, and nothing else.
259    ///
260    /// This resolution is MANDATORY, not an optimization: the default
261    /// `ActivityDispatcher::dispatch_async` runs the dispatcher's blocking
262    /// `std::sync::mpsc::recv()` on tokio's blocking pool, and tokio `Runtime`
263    /// drop joins blocking threads — an unresolved waiter would wedge process
264    /// exit indefinitely. An unmatched dispatch (already resolved by another
265    /// path) is a no-op: a park is NEVER routed to the outbox delivery
266    /// callback, because it is not a failure and must never reach a workflow.
267    fn park_activity(
268        &self,
269        workflow_id: &WorkflowId,
270        activity_id: &ActivityId,
271    ) -> Result<(), ServerError> {
272        let matched = self
273            .pending
274            .remove(&(workflow_id.clone(), activity_id.clone()));
275        if let Some((_, sender)) = matched {
276            // A send failure means the waiter side already dropped its
277            // receiver (the dispatch is being cleaned up concurrently) —
278            // benign: there is no thread left to unblock.
279            let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
280        }
281        Ok(())
282    }
283
284    fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
285        let result = match completion.outcome {
286            ActivityCompletionOutcome::Succeeded(payload) => {
287                payload_to_string(&payload).map_err(|reason| {
288                    tracing::error!(
289                        operation = "activity_complete",
290                        workflow_id = %completion.workflow_id,
291                        activity_id = %completion.activity_id,
292                        error_type = "ActivityResultDecode",
293                        %reason,
294                        "activity completion failed"
295                    );
296                    ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
297                })?
298            }
299            ActivityCompletionOutcome::Failed(error) => {
300                let prefix = if error.is_retryable() {
301                    "retryable"
302                } else {
303                    "terminal"
304                };
305                tracing::error!(
306                    operation = "activity_complete",
307                    workflow_id = %completion.workflow_id,
308                    activity_id = %completion.activity_id,
309                    error_type = "ActivityFailed",
310                    error_kind = prefix,
311                    reason = %error.message,
312                    "activity completion failed"
313                );
314                Err(format!("{prefix}:{}", error.message))
315            }
316        };
317        self.complete(
318            &completion.workflow_id,
319            &completion.activity_id,
320            completion.run_id.as_ref(),
321            result,
322        );
323        Ok(())
324    }
325}
326
327fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
328    match payload.content_type() {
329        ContentType::Json => String::from_utf8(payload.bytes().to_vec())
330            .map(Ok)
331            .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
332    }
333}
334
335/// Dispatcher that routes `run_activity` NIF calls to connected workers.
336///
337/// Synchronous interface — uses `try_send` for the task channel and
338/// `std::sync::mpsc::Receiver::recv` for the response. Callers on a
339/// multi-thread tokio runtime are detected and moved into
340/// `tokio::task::block_in_place` so the blocking wait never starves the
341/// runtime tasks that flush the worker stream (see the module docs).
342pub struct WorkerActivityDispatcher {
343    registry: ConnectedWorkerRegistry,
344    namespace: String,
345    pending: PendingActivities,
346    heartbeat_tracker: HeartbeatTracker,
347    drain_state: DrainState,
348    tokio_handle: Option<tokio::runtime::Handle>,
349    /// NOI-6 attempt→owner back-index. When installed, each liminal-delivered
350    /// dispatch binds its `(workflow, activity, attempt)` to the owning worker
351    /// for the dispatch's lifetime, so the intervention router (and the ops
352    /// console's live-attempts enumeration) can see and target it. `None`
353    /// (isolated tests) binds nothing.
354    attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
355}
356
357impl std::fmt::Debug for WorkerActivityDispatcher {
358    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
359        f.debug_struct("WorkerActivityDispatcher")
360            .field("namespace", &self.namespace)
361            .finish_non_exhaustive()
362    }
363}
364
365impl WorkerActivityDispatcher {
366    /// Build a dispatcher for the given namespace, worker registry, and
367    /// liveness tracker.
368    ///
369    /// The tracker must be the same instance the worker stream handler and
370    /// shutdown coordinator share: the unbounded completion wait relies on
371    /// stream teardown sweeping this tracker's in-flight entries to fail
372    /// dispatches whose worker was lost.
373    #[must_use]
374    pub fn new(
375        registry: ConnectedWorkerRegistry,
376        namespace: impl Into<String>,
377        heartbeat_tracker: HeartbeatTracker,
378    ) -> Self {
379        Self {
380            registry,
381            namespace: namespace.into(),
382            pending: PendingActivities::default(),
383            heartbeat_tracker,
384            drain_state: DrainState::default(),
385            tokio_handle: None,
386            attempt_owners: None,
387        }
388    }
389
390    /// Share the server's NOI-6 attempt→owner back-index so liminal-delivered
391    /// dispatches are visible (and targetable) to the intervention router for
392    /// exactly as long as they are in flight. The production boot passes
393    /// `ServerState`'s index — the SAME instance `intervenable_attempts` and
394    /// `intervene` read — or the console's live-attempts list stays empty for
395    /// every bridge-dispatched agent step.
396    #[must_use]
397    pub fn with_attempt_owners(
398        mut self,
399        attempt_owners: super::intervention::AttemptOwnerIndex,
400    ) -> Self {
401        self.attempt_owners = Some(attempt_owners);
402        self
403    }
404
405    /// Share a caller-supplied pending-activities tracker.
406    #[must_use]
407    pub fn with_pending(mut self, pending: PendingActivities) -> Self {
408        self.pending = pending;
409        self
410    }
411
412    /// Share the server drain gate.
413    #[must_use]
414    pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
415        self.drain_state = drain_state;
416        self
417    }
418
419    /// Share the server runtime handle for sync history writes from dirty NIF threads.
420    #[must_use]
421    pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
422        self.tokio_handle = Some(tokio_handle);
423        self
424    }
425}
426
427impl WorkerActivityDispatcher {
428    fn ensure_accepting(
429        &self,
430        namespace: &str,
431        activity_type: &str,
432        workflow_id: &WorkflowId,
433        activity_id: &ActivityId,
434        worker_id: Option<WorkerId>,
435    ) -> Result<(), String> {
436        self.drain_state
437            .ensure_accepting(namespace, activity_type)
438            .map_err(|error| {
439                let reason = error.to_string();
440                log_worker_error(
441                    "WorkerDispatch",
442                    namespace,
443                    activity_type,
444                    workflow_id,
445                    activity_id,
446                    worker_id,
447                    &reason,
448                );
449                reason
450            })
451    }
452
453    /// Select a worker for the namespace and activity type, waiting if none is
454    /// currently available. Blocks until a matching worker registers or the
455    /// server begins draining.
456    fn select_worker_or_wait(
457        &self,
458        namespace: &str,
459        task_queue: &str,
460        activity_type: &str,
461        node: Option<&str>,
462        workflow_id: &WorkflowId,
463        activity_id: &ActivityId,
464    ) -> Result<WorkerHandle, String> {
465        loop {
466            // `node` is the OPTIONAL within-pool affinity carried on the
467            // dispatch: `Some(n)` pins selection to workers advertising node
468            // `n` (require semantics — it waits, via the no-worker path below,
469            // if none are present); `None` is unpinned and reaches any worker
470            // in the (namespace, task_queue) pool — byte-identical to the
471            // pre-NODE behaviour.
472            match self
473                .registry
474                .select_worker(namespace, task_queue, activity_type, node)
475            {
476                Ok(Some(worker)) => return Ok(worker),
477                Ok(None) => {
478                    self.ensure_accepting(
479                        namespace,
480                        activity_type,
481                        workflow_id,
482                        activity_id,
483                        None,
484                    )?;
485                    tracing::info!(
486                        namespace,
487                        activity_type,
488                        node,
489                        workflow_id = %workflow_id,
490                        activity_id = %activity_id,
491                        "no connected worker; waiting for a matching worker to register"
492                    );
493                    match &self.tokio_handle {
494                        Some(handle) => {
495                            handle.block_on(self.registry.wait_for_worker());
496                        }
497                        None => match tokio::runtime::Handle::try_current() {
498                            Ok(handle) => {
499                                handle.block_on(self.registry.wait_for_worker());
500                            }
501                            Err(_) => {
502                                std::thread::sleep(Duration::from_millis(500));
503                            }
504                        },
505                    }
506                }
507                Err(error) => {
508                    let reason = format!("registry error: {error}");
509                    log_worker_error(
510                        "WorkerRegistry",
511                        namespace,
512                        activity_type,
513                        workflow_id,
514                        activity_id,
515                        None,
516                        &reason,
517                    );
518                    return Err(reason);
519                }
520            }
521        }
522    }
523
524    fn track_worker_task(
525        &self,
526        worker_id: WorkerId,
527        activity_type: &str,
528        workflow_id: &WorkflowId,
529        activity_id: &ActivityId,
530    ) -> Result<(), String> {
531        self.heartbeat_tracker
532            .track_task(
533                worker_id,
534                InFlightActivity {
535                    workflow_id: workflow_id.clone(),
536                    activity_id: activity_id.clone(),
537                },
538                Instant::now(),
539            )
540            .map_err(|error| {
541                let reason = error.to_string();
542                log_worker_error(
543                    "WorkerHeartbeatTracker",
544                    &self.namespace,
545                    activity_type,
546                    workflow_id,
547                    activity_id,
548                    Some(worker_id),
549                    &reason,
550                );
551                reason
552            })
553    }
554
555    fn cleanup_activity(
556        &self,
557        worker_id: WorkerId,
558        workflow_id: &WorkflowId,
559        activity_id: &ActivityId,
560    ) {
561        self.pending
562            .pending
563            .remove(&(workflow_id.clone(), activity_id.clone()));
564        let _ = self
565            .heartbeat_tracker
566            .complete_task(worker_id, workflow_id, activity_id);
567        self.drain_state.notify_activity_drained();
568    }
569
570    /// Deliver one dispatched task to the selected worker over ITS transport.
571    ///
572    /// The bridge is transport-agnostic at this seam: selection
573    /// ([`Self::select_worker_or_wait`]) already treats every registry member
574    /// identically, and this match delivers on whichever [`WorkerDelivery`] leg
575    /// the worker registered with — the gRPC stream `mpsc` push, or the liminal
576    /// server-push on the worker's existing connection. Both legs resolve
577    /// through the SAME pending map, so `await_activity_result` is oblivious to
578    /// the transport.
579    fn send_activity_task(
580        &self,
581        worker: &WorkerHandle,
582        task: ProtoActivityTask,
583        activity_type: &str,
584        workflow_id: &WorkflowId,
585        activity_id: &ActivityId,
586    ) -> Result<(), String> {
587        match worker.delivery() {
588            WorkerDelivery::Grpc(sender) => {
589                match sender.try_send(WorkerMessage::ActivityTask(task)) {
590                    Ok(()) => Ok(()),
591                    Err(error) => {
592                        let worker_id = worker.id();
593                        let reason = format!("worker task channel full or closed: {error}");
594                        self.cleanup_activity(worker_id, workflow_id, activity_id);
595                        log_worker_error(
596                            "WorkerChannelClosed",
597                            &self.namespace,
598                            activity_type,
599                            workflow_id,
600                            activity_id,
601                            Some(worker_id),
602                            &reason,
603                        );
604                        Err(reason)
605                    }
606                }
607            }
608            #[cfg(feature = "liminal-transport")]
609            WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
610                worker.id(),
611                delivery,
612                task,
613                activity_type,
614                workflow_id,
615                activity_id,
616            ),
617        }
618    }
619
620    /// Deliver one dispatched task to a liminal-connected worker: push the SAME
621    /// wire frame the outbox liminal path pushes (a
622    /// [`DispatchRequest`](super::liminal_transport::DispatchRequest) — the
623    /// worker's serve loop cannot tell a bridge dispatch from an outbox row) and
624    /// hand the correlated-reply awaiter to a dedicated router thread that
625    /// resolves this dispatch's pending entry exactly like a gRPC completion.
626    ///
627    /// The wire carries the SAME engine-provided `attempt` and `labels` the gRPC
628    /// arm's `ActivityTask` carries (a retry over liminal executes with the real
629    /// attempt, not a re-stamped first delivery), plus the server's heartbeat
630    /// window so the worker's automatic liveness pump keeps this TRACKED
631    /// dispatch alive under the #176 expiry sweeper. It carries no run context
632    /// (`run_id: None`), byte-identical to the gRPC bridge task's `run_id: None`
633    /// (OBX-011).
634    ///
635    /// A successful push also binds the attempt into the NOI-6 attempt→owner
636    /// back-index (when installed) with the SAME `(workflow, activity, attempt)`
637    /// key the worker stamps its intervention session with, exactly as the
638    /// outbox liminal arm binds each row dispatch — so the ops console can
639    /// enumerate this live attempt and route interventions to its worker. The
640    /// binding is released when the reply router exits (reply, abandonment, or
641    /// disconnect — every path). The gRPC arm carries no bind because the agent
642    /// harness seam exists only on the liminal worker transport.
643    #[cfg(feature = "liminal-transport")]
644    fn send_liminal_activity_task(
645        &self,
646        worker_id: WorkerId,
647        delivery: &super::liminal_transport::LiminalWorkerDelivery,
648        task: ProtoActivityTask,
649        activity_type: &str,
650        workflow_id: &WorkflowId,
651        activity_id: &ActivityId,
652    ) -> Result<(), String> {
653        let heartbeat_window_ms =
654            u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
655                .unwrap_or(u64::MAX);
656        let attempt = task.attempt;
657        let request = super::liminal_transport::DispatchRequest {
658            activity_type: activity_type.to_owned(),
659            workflow_id: workflow_id.clone(),
660            ordinal: activity_id.sequence_position(),
661            run_id: None,
662            attempt,
663            labels: task.labels.into_iter().collect(),
664            heartbeat_window_ms,
665            input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
666        };
667        // A push-enqueue failure means the worker's connection was already gone
668        // at push time — the same synchronous-failure contract as a closed gRPC
669        // stream channel above.
670        let awaiter = match delivery.push_dispatch(&request) {
671            Ok(awaiter) => awaiter,
672            Err(error) => {
673                let reason = format!("worker liminal push failed: {error}");
674                self.cleanup_activity(worker_id, workflow_id, activity_id);
675                log_worker_error(
676                    "WorkerChannelClosed",
677                    &self.namespace,
678                    activity_type,
679                    workflow_id,
680                    activity_id,
681                    Some(worker_id),
682                    &reason,
683                );
684                return Err(reason);
685            }
686        };
687        // NOI-6: the attempt is live on `worker_id` from this push until the
688        // router resolves it — bind it for exactly that window (the guard is
689        // dropped when the router thread exits).
690        let owner_binding = self.attempt_owners.as_ref().map(|owners| {
691            super::liminal_transport::AttemptOwnerGuard::bind(
692                owners.clone(),
693                super::intervention::AttemptKey::new(
694                    workflow_id.clone(),
695                    activity_id.clone(),
696                    attempt,
697                ),
698                worker_id,
699            )
700        });
701        self.spawn_liminal_reply_router(
702            worker_id,
703            awaiter,
704            workflow_id,
705            activity_id,
706            owner_binding,
707        );
708        Ok(())
709    }
710
711    /// Waits (on a dedicated router thread, bounded by the dispatch's own
712    /// lifetime) for the worker's correlated
713    /// [`DispatchResponse`](super::liminal_transport::DispatchResponse) and
714    /// re-enters it through the SAME completion bookkeeping the gRPC inbound
715    /// stream applies (`process_inbound` in `worker_grpc.rs`): clear the
716    /// in-flight liveness entry, wake any drain waiter, then resolve the
717    /// bridge's pending map — result, failure, and retryable classification
718    /// identical (the worker encodes the `retryable:`/`terminal:` reason
719    /// vocabulary on the wire). An unmatched (already-resolved) REAL reply
720    /// routes through the outbox delivery callback exactly like a late gRPC
721    /// result.
722    ///
723    /// The #176 heartbeat sweeper covers this dispatch exactly as it covers a
724    /// gRPC one: the dispatch is tracked in the shared [`HeartbeatTracker`] and
725    /// the worker's runtime pumps automatic liveness beats over the reserved
726    /// liminal channel (`WORKER_LIVENESS_CHANNEL`), so a healthy worker running
727    /// an over-window activity is never falsely expired while a wedged one
728    /// still is. Prompt worker-DEATH detection additionally rides the
729    /// connection itself — the awaiter wakes with the typed Disconnected error
730    /// the moment the connection closes, resolving the SAME retryable
731    /// lost-worker failure the gRPC stream-teardown sweep reports.
732    ///
733    /// Two structural guards mirror the gRPC arm's tracker gating:
734    ///
735    /// - A SYNTHESIZED failure (disconnect / receive fault) is delivered only
736    ///   when this router's own `complete_task` actually retired the tracked
737    ///   entry — the same "fail only still-tracked tasks" gate
738    ///   `remove_worker_tasks` gives the gRPC sweeps — so a dispatch already
739    ///   resolved elsewhere (expiry sweep, shutdown drain, deregistered
740    ///   fast path) never has a spurious failure injected for an ordinal whose
741    ///   retry may be live on another worker.
742    /// - The wait itself ends one reply-poll after the tracked entry
743    ///   disappears, so an abandoned dispatch never parks this thread for the
744    ///   remaining life of the worker's connection. A real reply arriving
745    ///   AFTER that exit is dropped (the resolving path owns the ordinal — its
746    ///   retry re-executes); this is the one deliberate divergence from the
747    ///   gRPC arm, whose shared stream task routes any late result to the
748    ///   outbox callback, and it is the safer half of the trade because a
749    ///   stale attempt's result can never resolve a newer attempt's entry.
750    #[cfg(feature = "liminal-transport")]
751    fn spawn_liminal_reply_router(
752        &self,
753        worker_id: WorkerId,
754        awaiter: liminal_server::server::connection::PushReplyAwaiter,
755        workflow_id: &WorkflowId,
756        activity_id: &ActivityId,
757        owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
758    ) {
759        let pending = self.pending.clone();
760        let heartbeat_tracker = self.heartbeat_tracker.clone();
761        let drain_state = self.drain_state.clone();
762        let workflow_id = workflow_id.clone();
763        let activity_id = activity_id.clone();
764        std::thread::spawn(move || {
765            // Owns the NOI-6 attempt binding for the dispatch's lifetime: it
766            // drops (releasing the back-index entry) when this router exits,
767            // on every path — reply, abandonment, disconnect, or panic.
768            let _owner_binding = owner_binding;
769            route_liminal_reply(
770                &pending,
771                &heartbeat_tracker,
772                &drain_state,
773                worker_id,
774                &awaiter,
775                &workflow_id,
776                &activity_id,
777            );
778        });
779    }
780
781    /// Block until the dispatch terminates (see the module docs for the
782    /// exhaustive termination list). The wait is deliberately unbounded:
783    /// the engine imposes no activity timeout of its own.
784    fn await_activity_result(
785        &self,
786        context: &ActivityDispatchContext<'_>,
787        rx: &SyncReceiver,
788    ) -> Result<String, String> {
789        // Close the dispatch/disconnect race before blocking. A worker whose
790        // stream tore down *before* this dispatch tracked its task was swept
791        // without this entry, so nothing would ever deliver through `rx`.
792        // `fail_lost_worker` deregisters before it collects tasks, and this
793        // dispatch tracked its task before sending, so: if the worker is
794        // still registered here, any later sweep is guaranteed to include
795        // this task and unblock the `recv` below.
796        match self.registry.is_registered(context.worker_id) {
797            Ok(true) => {}
798            Ok(false) => {
799                // A sweep that did include this task may have delivered
800                // already; prefer its verdict (or a genuine result that
801                // raced the disconnect) over fabricating one.
802                if let Ok(result) = rx.try_recv() {
803                    return self.deliver_result(context, result);
804                }
805                self.cleanup_activity(context.worker_id, context.workflow_id, context.activity_id);
806                let reason = format!(
807                    "retryable:{}",
808                    super::dispatch::lost_worker_error(context.worker_id).message
809                );
810                log_worker_error(
811                    "WorkerLost",
812                    &self.namespace,
813                    context.activity_type,
814                    context.workflow_id,
815                    context.activity_id,
816                    Some(context.worker_id),
817                    &reason,
818                );
819                return Err(reason);
820            }
821            Err(error) => {
822                self.cleanup_activity(context.worker_id, context.workflow_id, context.activity_id);
823                let reason = format!("worker registry inspection failed: {error}");
824                log_worker_error(
825                    "WorkerRegistry",
826                    &self.namespace,
827                    context.activity_type,
828                    context.workflow_id,
829                    context.activity_id,
830                    Some(context.worker_id),
831                    &reason,
832                );
833                return Err(reason);
834            }
835        }
836        if let Ok(result) = rx.recv() {
837            return self.deliver_result(context, result);
838        }
839        // Every sender was dropped without completing: a cleanup path
840        // removed the pending entry. Surface it instead of hanging.
841        self.cleanup_activity(context.worker_id, context.workflow_id, context.activity_id);
842        let reason = "activity response channel dropped".to_owned();
843        log_worker_error(
844            "WorkerChannelClosed",
845            &self.namespace,
846            context.activity_type,
847            context.workflow_id,
848            context.activity_id,
849            Some(context.worker_id),
850            &reason,
851        );
852        Err(reason)
853    }
854
855    fn deliver_result(
856        &self,
857        context: &ActivityDispatchContext<'_>,
858        result: Result<String, String>,
859    ) -> Result<String, String> {
860        self.pending
861            .pending
862            .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
863        // A parked dispatch (#207) is not a failure: the server is draining and
864        // restart recovery re-dispatches the ordinal. Info, never error — a
865        // routine deploy must not emit an ActivityFailed log per in-flight
866        // dispatch (the incident's alarm noise).
867        if let Err(reason) = &result
868            && aion::is_parked_reason(reason)
869        {
870            tracing::info!(
871                operation = "activity_dispatch",
872                namespace = %self.namespace,
873                workflow_id = %context.workflow_id,
874                activity_id = %context.activity_id,
875                activity_type = context.activity_type,
876                worker_id = ?context.worker_id,
877                "activity parked for restart recovery"
878            );
879            return result;
880        }
881        log_activity_completion(context, result.is_ok());
882        result.inspect_err(|reason| {
883            log_worker_error(
884                "ActivityFailed",
885                &self.namespace,
886                context.activity_type,
887                context.workflow_id,
888                context.activity_id,
889                Some(context.worker_id),
890                reason,
891            );
892        })
893    }
894}
895
896impl ActivityDispatcher for WorkerActivityDispatcher {
897    fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
898        match tokio::runtime::Handle::try_current() {
899            Ok(handle) => match handle.runtime_flavor() {
900                tokio::runtime::RuntimeFlavor::MultiThread => {
901                    // We are inside a tokio runtime (the engine spawns the
902                    // sync dispatch onto its handle). Hand this worker's
903                    // scheduler core to another thread before blocking so the
904                    // stream forwarder woken by our `try_send` can actually
905                    // run — otherwise it is trapped in this worker's
906                    // non-stealable LIFO slot for as long as we block.
907                    tokio::task::block_in_place(|| self.dispatch_blocking(request))
908                }
909                flavor => Err(format!(
910                    "activity dispatch blocks the calling thread until the worker responds; \
911                     a {flavor:?} tokio runtime cannot host that wait because the worker \
912                     stream forwarder shares its only executor thread and the task could \
913                     never be delivered — run the engine on a multi-thread tokio runtime"
914                )),
915            },
916            // No tokio context: a beamr scheduler thread or other plain OS
917            // thread. Blocking here is the designed contract and cannot starve
918            // the server runtime.
919            Err(_) => self.dispatch_blocking(request),
920        }
921    }
922}
923
924impl WorkerActivityDispatcher {
925    /// Dispatch the activity and block the calling thread until the worker
926    /// responds, the worker is declared lost, or the server drains (see the
927    /// module docs for the exhaustive termination list).
928    ///
929    /// The request carries the *real* workflow and activity ids the engine
930    /// recorded in history, so the worker logs, the pending-completion key,
931    /// and the heartbeat tracker all correlate directly against the event
932    /// store. `config` is forwarded by the engine seam but not yet consumed
933    /// here (the retry executor that reads it is unbuilt).
934    ///
935    /// Must never run while the calling thread still owns a tokio scheduler
936    /// core: the response can only arrive after the runtime's stream
937    /// forwarder flushes the queued [`WorkerMessage::ActivityTask`] to the
938    /// worker, so the thread blocking here must not be the one responsible
939    /// for polling that forwarder. [`ActivityDispatcher::dispatch`] enforces
940    /// this with `tokio::task::block_in_place`.
941    fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
942        let ActivityDispatch {
943            namespace,
944            task_queue,
945            // OPTIONAL within-pool node affinity (NODE-4): `Some(n)` pins this
946            // dispatch to workers advertising node `n` (require semantics);
947            // `None` is unpinned and reaches any worker in the pool.
948            node,
949            workflow_id,
950            activity_id,
951            name,
952            input,
953            config: _,
954            attempt,
955            labels,
956        } = request;
957        let started_at = Instant::now();
958        self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
959        let worker = self.select_worker_or_wait(
960            &namespace,
961            &task_queue,
962            &name,
963            node.as_deref(),
964            &workflow_id,
965            &activity_id,
966        )?;
967        let worker_id = worker.id();
968        let span = info_span!(
969            "activity_dispatch",
970            operation = "activity_dispatch",
971            namespace = %namespace,
972            task_queue = %task_queue,
973            node = node.as_deref(),
974            workflow_id = %workflow_id,
975            activity_id = %activity_id,
976            activity_type = %name,
977            worker_id = ?worker_id,
978        );
979        let _span_guard = span.enter();
980        self.ensure_accepting(
981            &namespace,
982            &name,
983            &workflow_id,
984            &activity_id,
985            Some(worker_id),
986        )?;
987
988        let task = activity_task(&name, &input, &workflow_id, &activity_id, attempt, labels);
989        let rx = self
990            .pending
991            .insert(workflow_id.clone(), activity_id.clone());
992        self.track_worker_task(worker_id, &name, &workflow_id, &activity_id)?;
993        self.send_activity_task(&worker, task, &name, &workflow_id, &activity_id)?;
994        let context = ActivityDispatchContext {
995            namespace: &namespace,
996            activity_type: &name,
997            worker_id,
998            workflow_id: &workflow_id,
999            activity_id: &activity_id,
1000            started_at,
1001        };
1002        self.await_activity_result(&context, &rx)
1003    }
1004}
1005
1006/// Body of one liminal reply-router thread (see
1007/// [`WorkerActivityDispatcher::spawn_liminal_reply_router`] for the contract).
1008///
1009/// Resolves by the key THIS push dispatched (the awaiter is already
1010/// correlation-scoped to it), never by the reply's echoed ids: a buggy echo
1011/// must not cross executions.
1012#[cfg(feature = "liminal-transport")]
1013fn route_liminal_reply(
1014    pending: &PendingActivities,
1015    heartbeat_tracker: &HeartbeatTracker,
1016    drain_state: &DrainState,
1017    worker_id: WorkerId,
1018    awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1019    workflow_id: &WorkflowId,
1020    activity_id: &ActivityId,
1021) {
1022    // The wait re-arms only while this dispatch is still tracked in-flight, so
1023    // a dispatch resolved elsewhere (expiry sweep, shutdown drain, cleanup)
1024    // releases this thread within one reply poll instead of parking it for the
1025    // remaining life of the worker's connection.
1026    let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1027        heartbeat_tracker
1028            .is_tracked(worker_id, workflow_id, activity_id)
1029            .unwrap_or(false)
1030    });
1031    // `synthesized` marks a failure this router FABRICATED (disconnect or
1032    // receive fault) as opposed to a real worker reply: only fabricated
1033    // failures are gated on the tracker below.
1034    let (run_id, outcome, synthesized) = match waited {
1035        Ok(Some(response)) => (response.run_id, response.outcome, false),
1036        Ok(None) => {
1037            tracing::debug!(
1038                worker_id = ?worker_id,
1039                workflow_id = %workflow_id,
1040                activity_id = %activity_id,
1041                "liminal dispatch resolved by another path; abandoning reply wait"
1042            );
1043            return;
1044        }
1045        Err(error) if error.is_worker_connection_lost() => (
1046            None,
1047            Err(format!(
1048                "retryable:{}",
1049                super::dispatch::lost_worker_error(worker_id).message
1050            )),
1051            true,
1052        ),
1053        Err(error) => (
1054            None,
1055            Err(format!("retryable:worker liminal reply failed: {error}")),
1056            true,
1057        ),
1058    };
1059    // The gRPC sweeps fail only still-tracked tasks (`remove_worker_tasks`);
1060    // this is the same structural gate: `complete_task` reports whether THIS
1061    // call retired the tracked entry. A poisoned tracker fails open (deliver)
1062    // so the blocked dispatch thread is never left hanging on a broken lock.
1063    let was_tracked = heartbeat_tracker
1064        .complete_task(worker_id, workflow_id, activity_id)
1065        .unwrap_or_else(|error| {
1066            tracing::error!(
1067                worker_id = ?worker_id,
1068                workflow_id = %workflow_id,
1069                activity_id = %activity_id,
1070                %error,
1071                "failed to clear in-flight tracking for completed liminal activity"
1072            );
1073            true
1074        });
1075    if synthesized && !was_tracked {
1076        // Another path already resolved this dispatch (and notified drain):
1077        // injecting the fabricated lost-worker failure now could reach a retry
1078        // attempt's entry or the outbox callback for an ordinal that is no
1079        // longer this router's to fail.
1080        tracing::debug!(
1081            worker_id = ?worker_id,
1082            workflow_id = %workflow_id,
1083            activity_id = %activity_id,
1084            "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1085        );
1086        return;
1087    }
1088    drain_state.notify_activity_drained();
1089    pending.complete(workflow_id, activity_id, run_id.as_ref(), outcome);
1090}
1091
1092struct ActivityDispatchContext<'a> {
1093    namespace: &'a str,
1094    activity_type: &'a str,
1095    worker_id: WorkerId,
1096    workflow_id: &'a WorkflowId,
1097    activity_id: &'a ActivityId,
1098    started_at: Instant,
1099}
1100
1101fn activity_task(
1102    activity_type: &str,
1103    input: &str,
1104    workflow_id: &WorkflowId,
1105    activity_id: &ActivityId,
1106    attempt: u32,
1107    labels: BTreeMap<String, String>,
1108) -> ProtoActivityTask {
1109    ProtoActivityTask {
1110        workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1111        activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1112        activity_type: activity_type.to_owned(),
1113        input: Some(ProtoPayload {
1114            content_type: String::from("application/json"),
1115            bytes: input.as_bytes().to_vec(),
1116        }),
1117        attempt,
1118        labels: labels.into_iter().collect(),
1119        // The synchronous `ActivityDispatch` bridge path carries no run context
1120        // (run scoping is threaded through the durable-outbox path; OBX-011).
1121        run_id: None,
1122    }
1123}
1124
1125fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1126    let duration_ms = duration_ms(context.started_at.elapsed());
1127    tracing::info!(
1128        operation = "activity_complete",
1129        namespace = context.namespace,
1130        workflow_id = %context.workflow_id,
1131        activity_id = %context.activity_id,
1132        activity_type = context.activity_type,
1133        worker_id = ?context.worker_id,
1134        duration_ms,
1135        outcome = if succeeded { "succeeded" } else { "failed" },
1136        "activity completed"
1137    );
1138}
1139
1140fn duration_ms(duration: Duration) -> u64 {
1141    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1142}
1143
1144fn log_worker_error(
1145    error_type: &'static str,
1146    namespace: &str,
1147    activity_type: &str,
1148    workflow_id: &WorkflowId,
1149    activity_id: &ActivityId,
1150    worker_id: Option<super::registry::WorkerId>,
1151    reason: &str,
1152) {
1153    tracing::error!(
1154        operation = "activity_dispatch",
1155        namespace,
1156        workflow_id = %workflow_id,
1157        activity_id = %activity_id,
1158        activity_type,
1159        worker_id = ?worker_id,
1160        error_type,
1161        reason,
1162        "worker interaction failed"
1163    );
1164}
1165
1166#[cfg(test)]
1167mod tests {
1168    use std::sync::Mutex;
1169
1170    use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1171
1172    use super::*;
1173
1174    fn activity_id(pos: u64) -> ActivityId {
1175        ActivityId::from_sequence_position(pos)
1176    }
1177
1178    #[test]
1179    fn pending_insert_and_complete_delivers_result() {
1180        let pending = PendingActivities::default();
1181        let workflow_id = WorkflowId::new_v4();
1182        let id = activity_id(1);
1183        let rx = pending.insert(workflow_id.clone(), id.clone());
1184
1185        assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1186        assert_eq!(
1187            rx.recv_timeout(Duration::from_millis(50)),
1188            Ok(Ok("done".to_owned()))
1189        );
1190    }
1191
1192    #[test]
1193    fn pending_complete_unknown_returns_false() {
1194        let pending = PendingActivities::default();
1195        assert!(!pending.complete(
1196            &WorkflowId::new_v4(),
1197            &activity_id(99),
1198            None,
1199            Ok("orphan".to_owned())
1200        ));
1201    }
1202
1203    #[derive(Default)]
1204    struct RecordingOutboxCallback {
1205        completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1206        failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1207        live: bool,
1208    }
1209
1210    impl OutboxDeliveryCallback for RecordingOutboxCallback {
1211        fn deliver_completion(
1212            &self,
1213            workflow_id: &WorkflowId,
1214            activity_id: &ActivityId,
1215            run_id: Option<&RunId>,
1216            result: String,
1217        ) -> Result<bool, ServerError> {
1218            let _ = run_id;
1219            self.completions
1220                .lock()
1221                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1222                .push((workflow_id.clone(), activity_id.clone(), result));
1223            Ok(self.live)
1224        }
1225
1226        fn deliver_failure(
1227            &self,
1228            workflow_id: &WorkflowId,
1229            activity_id: &ActivityId,
1230            run_id: Option<&RunId>,
1231            reason: String,
1232        ) -> Result<bool, ServerError> {
1233            let _ = run_id;
1234            self.failures
1235                .lock()
1236                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1237                .push((workflow_id.clone(), activity_id.clone(), reason));
1238            Ok(self.live)
1239        }
1240    }
1241
1242    #[test]
1243    fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1244        let pending = PendingActivities::default();
1245        let callback = Arc::new(RecordingOutboxCallback {
1246            live: true,
1247            ..RecordingOutboxCallback::default()
1248        });
1249        // Install on one clone; the wiring must be visible to every clone.
1250        pending.clone().set_outbox_delivery(callback.clone());
1251
1252        let workflow_id = WorkflowId::new_v4();
1253        let id = activity_id(7);
1254
1255        // No pending entry: the completion is unmatched and must route to the
1256        // callback rather than being dropped. A live workflow reports true.
1257        assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1258        let completions = callback
1259            .completions
1260            .lock()
1261            .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1262        assert_eq!(completions.len(), 1);
1263        assert_eq!(completions[0].0, workflow_id);
1264        assert_eq!(completions[0].1, id);
1265        assert_eq!(completions[0].2, "done");
1266        Ok(())
1267    }
1268
1269    #[test]
1270    fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1271    -> Result<(), ServerError> {
1272        let pending = PendingActivities::default();
1273        // live = false models the expected stale-completion case.
1274        let callback = Arc::new(RecordingOutboxCallback::default());
1275        pending.set_outbox_delivery(callback.clone());
1276
1277        let workflow_id = WorkflowId::new_v4();
1278        let id = activity_id(8);
1279
1280        assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1281        let failures = callback
1282            .failures
1283            .lock()
1284            .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1285        assert_eq!(failures.len(), 1);
1286        assert_eq!(failures[0].2, "retryable:boom");
1287        Ok(())
1288    }
1289
1290    #[test]
1291    fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1292        // Flag-off byte-identical behaviour: no callback, unmatched returns
1293        // false (silent drop) exactly as before.
1294        let pending = PendingActivities::default();
1295        assert!(!pending.complete(
1296            &WorkflowId::new_v4(),
1297            &activity_id(9),
1298            None,
1299            Ok("x".to_owned())
1300        ));
1301    }
1302
1303    #[test]
1304    fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1305        let pending = PendingActivities::default();
1306        let callback = Arc::new(RecordingOutboxCallback {
1307            live: true,
1308            ..RecordingOutboxCallback::default()
1309        });
1310        pending.set_outbox_delivery(callback.clone());
1311
1312        let workflow_id = WorkflowId::new_v4();
1313        let id = activity_id(10);
1314        let rx = pending.insert(workflow_id.clone(), id.clone());
1315
1316        assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1317        assert_eq!(
1318            rx.recv_timeout(Duration::from_millis(50)),
1319            Ok(Ok("matched".to_owned()))
1320        );
1321        assert!(
1322            callback
1323                .completions
1324                .lock()
1325                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1326                .is_empty(),
1327            "a matched completion must deliver to its waiter, not the outbox callback"
1328        );
1329        Ok(())
1330    }
1331
1332    /// #207: parking resolves the matched waiter with the ephemeral parked
1333    /// sentinel — the exact string the engine's retry loop classifies as
1334    /// `Parked` — and nothing else.
1335    #[test]
1336    fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
1337        let pending = PendingActivities::default();
1338        let workflow_id = WorkflowId::new_v4();
1339        let id = activity_id(11);
1340        let rx = pending.insert(workflow_id.clone(), id.clone());
1341
1342        pending.park_activity(&workflow_id, &id)?;
1343        let result = rx
1344            .recv_timeout(Duration::from_millis(50))
1345            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1346        assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
1347        Ok(())
1348    }
1349
1350    /// #207: an unmatched park is a no-op and is NEVER routed to the outbox
1351    /// delivery callback — a park is not a failure and must never reach a
1352    /// workflow.
1353    #[test]
1354    fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
1355        let pending = PendingActivities::default();
1356        let callback = Arc::new(RecordingOutboxCallback {
1357            live: true,
1358            ..RecordingOutboxCallback::default()
1359        });
1360        pending.set_outbox_delivery(callback.clone());
1361
1362        pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
1363
1364        assert!(
1365            callback
1366                .failures
1367                .lock()
1368                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1369                .is_empty(),
1370            "a park must never be delivered as an outbox failure"
1371        );
1372        assert!(
1373            callback
1374                .completions
1375                .lock()
1376                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1377                .is_empty(),
1378            "a park must never be delivered as an outbox completion"
1379        );
1380        Ok(())
1381    }
1382
1383    #[test]
1384    fn completion_sink_routes_success() -> Result<(), ServerError> {
1385        let pending = PendingActivities::default();
1386        let workflow_id = WorkflowId::new_v4();
1387        let id = activity_id(2);
1388        let rx = pending.insert(workflow_id.clone(), id.clone());
1389        let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
1390
1391        pending.complete_activity(ActivityCompletion {
1392            workflow_id,
1393            activity_id: id,
1394            run_id: None,
1395            outcome: ActivityCompletionOutcome::Succeeded(payload),
1396        })?;
1397
1398        let result = rx
1399            .recv_timeout(Duration::from_millis(50))
1400            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1401        assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
1402        Ok(())
1403    }
1404
1405    #[test]
1406    fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
1407        let pending = PendingActivities::default();
1408        let workflow_id = WorkflowId::new_v4();
1409        let id = activity_id(3);
1410        let rx = pending.insert(workflow_id.clone(), id.clone());
1411
1412        pending.complete_activity(ActivityCompletion {
1413            workflow_id,
1414            activity_id: id,
1415            run_id: None,
1416            outcome: ActivityCompletionOutcome::Failed(ActivityError {
1417                kind: ActivityErrorKind::Retryable,
1418                message: "temporary".to_owned(),
1419                details: None,
1420            }),
1421        })?;
1422
1423        let result = rx
1424            .recv_timeout(Duration::from_millis(50))
1425            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1426        assert_eq!(result, Err("retryable:temporary".to_owned()));
1427        Ok(())
1428    }
1429
1430    /// Regression test (#59, brief D12): pending tracking must be keyed by
1431    /// the full `(WorkflowId, ActivityId)` pair. The dispatcher fabricates
1432    /// activity ids from a process-local counter that resets on server
1433    /// restart, so a stale result re-reported from a worker's previous
1434    /// session carries the same bare `ActivityId` as a fresh post-restart
1435    /// dispatch. Under bare-`ActivityId` keying the stale result completed
1436    /// the wrong execution; with pair keying it is dropped and the genuine
1437    /// result still completes.
1438    #[test]
1439    fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
1440    -> Result<(), ServerError> {
1441        let pending = PendingActivities::default();
1442        let post_restart_workflow = WorkflowId::new_v4();
1443        let pre_restart_workflow = WorkflowId::new_v4();
1444        // Counter resets to the same sequence position after restart.
1445        let id = activity_id(1);
1446        let rx = pending.insert(post_restart_workflow.clone(), id.clone());
1447
1448        // Stale pre-restart result: same activity id, different workflow.
1449        pending.complete_activity(ActivityCompletion {
1450            workflow_id: pre_restart_workflow,
1451            activity_id: id.clone(),
1452            run_id: None,
1453            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1454                ContentType::Json,
1455                br#""stale""#.to_vec(),
1456            )),
1457        })?;
1458        assert!(
1459            rx.try_recv().is_err(),
1460            "stale result for a different workflow must not complete this dispatch"
1461        );
1462
1463        // The genuine result for the pending execution still completes.
1464        pending.complete_activity(ActivityCompletion {
1465            workflow_id: post_restart_workflow,
1466            activity_id: id,
1467            run_id: None,
1468            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1469                ContentType::Json,
1470                br#""fresh""#.to_vec(),
1471            )),
1472        })?;
1473        let result = rx
1474            .recv_timeout(Duration::from_millis(50))
1475            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1476        assert_eq!(result, Ok(r#""fresh""#.to_owned()));
1477        Ok(())
1478    }
1479
1480    /// Liveness tracker for dispatcher unit tests; the window only matters
1481    /// to expiry checks, which nothing in these tests drives.
1482    fn test_tracker() -> HeartbeatTracker {
1483        HeartbeatTracker::new(Duration::from_secs(5))
1484    }
1485
1486    /// A `greet` dispatch request carrying real (test-synthesized) ids, the
1487    /// engine-seam shape `WorkerActivityDispatcher::dispatch` now consumes.
1488    fn greet_request() -> ActivityDispatch {
1489        ActivityDispatch {
1490            namespace: "default".to_owned(),
1491            task_queue: "default".to_owned(),
1492            node: None,
1493            workflow_id: WorkflowId::new_v4(),
1494            activity_id: ActivityId::from_sequence_position(0),
1495            name: "greet".to_owned(),
1496            input: "{}".to_owned(),
1497            config: "{}".to_owned(),
1498            attempt: 1,
1499            labels: std::collections::BTreeMap::new(),
1500        }
1501    }
1502
1503    #[test]
1504    fn dispatcher_fails_immediately_when_draining_without_workers() {
1505        let registry = ConnectedWorkerRegistry::default();
1506        let drain = DrainState::default();
1507        let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
1508            .with_drain_state(drain.clone());
1509
1510        let _ = drain.begin();
1511
1512        let result = dispatcher.dispatch(greet_request());
1513
1514        assert!(result.is_err());
1515        let err = result.err().unwrap_or_default();
1516        assert!(
1517            err.contains("drain"),
1518            "expected drain rejection, got: {err}"
1519        );
1520    }
1521
1522    /// Regression test for the production stall where every remote activity
1523    /// timed out: the engine invoked the sync `dispatch` from inside a
1524    /// spawned tokio task (`futures::future::lazy` polled on a runtime
1525    /// worker), and the woken stream-consumer task landed in that blocked
1526    /// worker's non-stealable LIFO slot, so the queued `ActivityTask` was
1527    /// only delivered when the then-extant 30s dispatch timeout fired (the
1528    /// dispatch wait is unbounded today; the stall would now be a hang).
1529    ///
1530    /// Mirrors the real wiring minus tonic: the real registry channel that
1531    /// the gRPC stream forwarder drains, a worker task awaiting that channel
1532    /// on the same runtime, completion through the production
1533    /// `ActivityCompletionSink`, and the sync dispatch invoked from a
1534    /// runtime worker task — the worst case the `block_in_place` guard in
1535    /// `dispatch` defends against (the engine itself now routes through
1536    /// `dispatch_async`, off the async workers).
1537    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1538    async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
1539    -> Result<(), Box<dyn std::error::Error>> {
1540        let registry = ConnectedWorkerRegistry::default();
1541        let pending = PendingActivities::default();
1542        let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
1543        let activity_types = [String::from("greet")];
1544        let registration = registry.register("default", activity_types.iter(), worker_tx)?;
1545
1546        let sink = pending.clone();
1547        let echo_worker = tokio::spawn(async move {
1548            let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
1549                return Err("expected an activity task on the worker channel".to_owned());
1550            };
1551            let workflow_id = task
1552                .workflow_id
1553                .ok_or("task missing workflow id")
1554                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
1555            let activity_id = task
1556                .activity_id
1557                .map(ActivityId::from)
1558                .ok_or("task missing activity id")?;
1559            sink.complete_activity(ActivityCompletion {
1560                workflow_id,
1561                activity_id,
1562                run_id: None,
1563                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1564                    ContentType::Json,
1565                    br#"{"greeting":"hello"}"#.to_vec(),
1566                )),
1567            })
1568            .map_err(|error| error.to_string())
1569        });
1570
1571        let dispatcher = Arc::new(
1572            WorkerActivityDispatcher::new(registry, "default", test_tracker())
1573                .with_pending(pending),
1574        );
1575        let started = Instant::now();
1576        // Invoke the sync dispatch inside the first poll of a spawned task:
1577        // the worst-case calling context for the `block_in_place` guard.
1578        let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
1579            dispatcher.dispatch(greet_request())
1580        }));
1581        let result = dispatch_task.await.map_err(|error| error.to_string())?;
1582        let elapsed = started.elapsed();
1583
1584        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
1585        assert!(
1586            elapsed < Duration::from_secs(5),
1587            "dispatch round trip took {elapsed:?}; task delivery must not \
1588             depend on the blocked dispatch thread"
1589        );
1590        echo_worker.await.map_err(|error| error.to_string())??;
1591        registration.deregister()?;
1592        Ok(())
1593    }
1594
1595    /// A current-thread runtime cannot host the blocking wait (the stream
1596    /// forwarder would share its only executor thread), so dispatch must
1597    /// fail fast with a precise error instead of blocking forever.
1598    #[tokio::test]
1599    async fn dispatch_on_current_thread_runtime_fails_fast()
1600    -> Result<(), Box<dyn std::error::Error>> {
1601        let registry = ConnectedWorkerRegistry::default();
1602        let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
1603        let activity_types = [String::from("greet")];
1604        let registration = registry.register("default", activity_types.iter(), worker_tx)?;
1605        let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
1606
1607        let started = Instant::now();
1608        let result = dispatcher.dispatch(greet_request());
1609        let elapsed = started.elapsed();
1610
1611        let err = result.err().ok_or("expected dispatch to fail")?;
1612        assert!(
1613            err.contains("multi-thread tokio runtime"),
1614            "unexpected error: {err}"
1615        );
1616        assert!(
1617            elapsed < Duration::from_secs(5),
1618            "fail-fast path took {elapsed:?}"
1619        );
1620        registration.deregister()?;
1621        Ok(())
1622    }
1623
1624    /// Bridge-level mirror of the e2e node-pin proof: two workers share the
1625    /// `(namespace, task_queue)` pool but advertise different nodes; an
1626    /// `ActivityDispatch` pinned to one node must reach ONLY the worker on that
1627    /// node through the live engine-seam `WorkerActivityDispatcher`. This is the
1628    /// regression guard for the bridge discarding the dispatch's `node`.
1629    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1630    async fn dispatch_pinned_to_node_reaches_only_that_node()
1631    -> Result<(), Box<dyn std::error::Error>> {
1632        let registry = ConnectedWorkerRegistry::default();
1633        let pending = PendingActivities::default();
1634        let activity_types = [String::from("greet")];
1635        let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
1636        let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
1637        // Register the DECOY (n2) FIRST so it owns the lowest worker id. The
1638        // bridge's `select_worker` picks the lowest-id matching worker, so a
1639        // bridge that DISCARDED the node would route to n2 (the decoy) here —
1640        // the n1 echo would never fire and the round trip would time out. With
1641        // the node threaded through, selection is filtered to n1.
1642        let on_n2 = registry.register_namespaces(
1643            [String::from("default")],
1644            "default",
1645            Some(String::from("n2")),
1646            activity_types.iter(),
1647            n2_tx,
1648        )?;
1649        let on_n1 = registry.register_namespaces(
1650            [String::from("default")],
1651            "default",
1652            Some(String::from("n1")),
1653            activity_types.iter(),
1654            n1_tx,
1655        )?;
1656
1657        // Echo only on the n1 channel: the dispatch can only complete if the
1658        // task was routed to n1. If it leaked to n2, the n1 wait would stall and
1659        // the round trip below would time out instead.
1660        let sink = pending.clone();
1661        let echo_n1 = tokio::spawn(async move {
1662            let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
1663                return Err("expected an activity task on the n1 worker channel".to_owned());
1664            };
1665            let workflow_id = task
1666                .workflow_id
1667                .ok_or("task missing workflow id")
1668                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
1669            let activity_id = task
1670                .activity_id
1671                .map(ActivityId::from)
1672                .ok_or("task missing activity id")?;
1673            sink.complete_activity(ActivityCompletion {
1674                workflow_id,
1675                activity_id,
1676                run_id: None,
1677                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1678                    ContentType::Json,
1679                    br#"{"greeting":"hello"}"#.to_vec(),
1680                )),
1681            })
1682            .map_err(|error| error.to_string())
1683        });
1684
1685        let dispatcher = Arc::new(
1686            WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
1687                .with_pending(pending),
1688        );
1689
1690        let pinned = ActivityDispatch {
1691            node: Some(String::from("n1")),
1692            ..greet_request()
1693        };
1694        let started = Instant::now();
1695        let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
1696            .await
1697            .map_err(|error| error.to_string())?;
1698        let elapsed = started.elapsed();
1699
1700        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
1701        assert!(
1702            elapsed < Duration::from_secs(5),
1703            "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
1704        );
1705        echo_n1.await.map_err(|error| error.to_string())??;
1706
1707        // The n2 worker (wrong node) must never have been handed the task.
1708        assert!(
1709            n2_rx.try_recv().is_err(),
1710            "node=Some(\"n1\") dispatch must not reach the n2 worker"
1711        );
1712
1713        on_n1.deregister()?;
1714        on_n2.deregister()?;
1715        Ok(())
1716    }
1717
1718    /// An unpinned (`node = None`) dispatch is byte-identical to today: it
1719    /// reaches a worker in the pool regardless of the worker's advertised node.
1720    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1721    async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
1722    -> Result<(), Box<dyn std::error::Error>> {
1723        let registry = ConnectedWorkerRegistry::default();
1724        let pending = PendingActivities::default();
1725        let activity_types = [String::from("greet")];
1726        let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
1727        let on_n1 = registry.register_namespaces(
1728            [String::from("default")],
1729            "default",
1730            Some(String::from("n1")),
1731            activity_types.iter(),
1732            n1_tx,
1733        )?;
1734
1735        let sink = pending.clone();
1736        let echo = tokio::spawn(async move {
1737            let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
1738                return Err("expected an activity task on the worker channel".to_owned());
1739            };
1740            let workflow_id = task
1741                .workflow_id
1742                .ok_or("task missing workflow id")
1743                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
1744            let activity_id = task
1745                .activity_id
1746                .map(ActivityId::from)
1747                .ok_or("task missing activity id")?;
1748            sink.complete_activity(ActivityCompletion {
1749                workflow_id,
1750                activity_id,
1751                run_id: None,
1752                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1753                    ContentType::Json,
1754                    br#"{"greeting":"hello"}"#.to_vec(),
1755                )),
1756            })
1757            .map_err(|error| error.to_string())
1758        });
1759
1760        let dispatcher = Arc::new(
1761            WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
1762                .with_pending(pending),
1763        );
1764
1765        // greet_request() carries node: None — the unpinned path.
1766        let result = tokio::spawn(futures::future::lazy(move |_| {
1767            dispatcher.dispatch(greet_request())
1768        }))
1769        .await
1770        .map_err(|error| error.to_string())?;
1771
1772        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
1773        echo.await.map_err(|error| error.to_string())??;
1774        on_n1.deregister()?;
1775        Ok(())
1776    }
1777}