1use std::collections::BTreeMap;
64use std::sync::{Arc, OnceLock};
65use std::time::{Duration, Instant};
66
67use aion::{ActivityDispatch, ActivityDispatcher};
68use aion_core::{ActivityId, ContentType, Payload, RunId, WorkflowId};
69use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
70use dashmap::DashMap;
71
72use super::dispatch::{ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink};
73use super::envelope::{CompletionFences, CompletionToken, idempotency_key};
74use super::heartbeat::{HeartbeatTracker, InFlightActivity};
75use super::queue_service::{
76 DeliveryRefusal, ExpiredClock, PARK_POLL_INTERVAL, PoolCensus, QueueDeclarationSource,
77 QueueServiceConfig, QueueServiceReason, QueueServiceState, SelectionRefusal, ServiceAddress,
78 ServiceWait, WorkerUnavailable, deliver_within_schedule_to_start, select_worker_or_refuse,
79};
80use super::registry::{
81 ConnectedWorkerRegistry, WorkerDelivery, WorkerHandle, WorkerId, WorkerMessage,
82};
83use crate::error::ServerError;
84use crate::shutdown::DrainState;
85use tracing::info_span;
86
87type SyncSender = std::sync::mpsc::SyncSender<Result<String, String>>;
88type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
89
90type PendingActivityKey = (WorkflowId, ActivityId);
108
109pub trait OutboxDeliveryCallback: Send + Sync {
119 fn deliver_completion(
129 &self,
130 workflow_id: &WorkflowId,
131 activity_id: &ActivityId,
132 run_id: Option<&RunId>,
133 result: String,
134 ) -> Result<bool, ServerError>;
135
136 fn deliver_failure(
143 &self,
144 workflow_id: &WorkflowId,
145 activity_id: &ActivityId,
146 run_id: Option<&RunId>,
147 reason: String,
148 ) -> Result<bool, ServerError>;
149}
150
151#[derive(Clone, Default)]
163pub struct PendingActivities {
164 pending: Arc<DashMap<PendingActivityKey, SyncSender>>,
165 completion_fences: CompletionFences,
166 outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
167 transport_losses: super::transport_loss::TransportLossLedger,
173}
174
175impl std::fmt::Debug for PendingActivities {
176 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177 formatter
178 .debug_struct("PendingActivities")
179 .field("pending", &self.pending.len())
180 .field("completion_fences", &self.completion_fences)
181 .field(
182 "outbox_delivery_installed",
183 &self.outbox_delivery.get().is_some(),
184 )
185 .field("transport_losses", &self.transport_losses)
186 .finish()
187 }
188}
189
190impl PendingActivities {
191 fn insert(
192 &self,
193 workflow_id: WorkflowId,
194 activity_id: ActivityId,
195 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
196 let completion_token = self.completion_fences.issue(&workflow_id, &activity_id)?;
197 let (tx, rx) = std::sync::mpsc::sync_channel(1);
198 self.pending.insert((workflow_id, activity_id), tx);
199 Ok((completion_token, rx))
200 }
201
202 #[must_use]
204 pub fn completion_fences(&self) -> CompletionFences {
205 self.completion_fences.clone()
206 }
207
208 #[cfg(test)]
212 pub(crate) fn insert_for_test(
213 &self,
214 workflow_id: WorkflowId,
215 activity_id: ActivityId,
216 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
217 self.insert(workflow_id, activity_id)
218 }
219
220 pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
226 if self.outbox_delivery.set(callback).is_err() {
227 tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
228 }
229 }
230
231 fn complete(
239 &self,
240 workflow_id: &WorkflowId,
241 activity_id: &ActivityId,
242 run_id: Option<&RunId>,
243 result: Result<String, String>,
244 ) -> bool {
245 let matched = self
248 .pending
249 .remove(&(workflow_id.clone(), activity_id.clone()));
250 if let Some((_, sender)) = matched {
251 return sender.send(result).is_ok();
252 }
253 let Some(callback) = self.outbox_delivery.get() else {
254 return false;
256 };
257 let outcome = match result {
258 Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
259 Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
260 };
261 match outcome {
262 Ok(true) => true,
263 Ok(false) => {
264 tracing::debug!(
266 workflow_id = %workflow_id,
267 activity_id = %activity_id,
268 "unmatched outbox completion for a workflow that is not currently live; \
269 recovery will re-arm it"
270 );
271 false
272 }
273 Err(error) => {
274 tracing::warn!(
275 workflow_id = %workflow_id,
276 activity_id = %activity_id,
277 %error,
278 "failed to deliver unmatched outbox completion to the live workflow"
279 );
280 false
281 }
282 }
283 }
284
285 fn complete_fenced(
288 &self,
289 workflow_id: &WorkflowId,
290 activity_id: &ActivityId,
291 run_id: Option<&RunId>,
292 completion_token: &CompletionToken,
293 result: Result<String, String>,
294 ) -> Result<bool, ServerError> {
295 self.completion_fences
296 .accept(workflow_id, activity_id, completion_token)
297 .inspect_err(|error| {
298 tracing::warn!(
299 workflow_id = %workflow_id,
300 activity_id = %activity_id,
301 %error,
302 "activity completion rejected by execution-generation fence"
303 );
304 })?;
305 let transport_domain = result
311 .as_ref()
312 .err()
313 .is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
314 if !transport_domain {
315 if let Err(error) = self.transport_losses.clear(workflow_id, activity_id) {
316 tracing::warn!(
317 workflow_id = %workflow_id,
318 activity_id = %activity_id,
319 %error,
320 "failed to retire the transport-loss budget for a resolved activity"
321 );
322 }
323 }
324 Ok(self.complete(workflow_id, activity_id, run_id, result))
325 }
326
327 #[must_use]
335 pub fn with_heartbeat_window(mut self, heartbeat_window: std::time::Duration) -> Self {
336 self.transport_losses = super::transport_loss::TransportLossLedger::new(heartbeat_window);
337 self
338 }
339
340 #[must_use]
342 pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
343 &self.transport_losses
344 }
345
346 fn classify_worker_loss(
354 &self,
355 workflow_id: &WorkflowId,
356 activity_id: &ActivityId,
357 worker_id: crate::worker::registry::WorkerId,
358 ) -> String {
359 let detail = super::transport_loss::worker_lost_detail(worker_id);
360 match self
361 .transport_losses
362 .record_loss(workflow_id, activity_id, &detail)
363 {
364 Ok(verdict) => {
365 if verdict.exhausted {
366 tracing::error!(
367 operation = "activity_complete",
368 workflow_id = %workflow_id,
369 activity_id = %activity_id,
370 worker_id = ?worker_id,
371 error_type = "TransportExhausted",
372 losses = verdict.losses,
373 budget_ms = self.transport_losses.budget().as_millis(),
374 "activity abandoned: the transport kept losing its worker past the \
375 transport-loss budget"
376 );
377 } else {
378 tracing::warn!(
379 operation = "activity_complete",
380 workflow_id = %workflow_id,
381 activity_id = %activity_id,
382 worker_id = ?worker_id,
383 error_type = "WorkerLost",
384 losses = verdict.losses,
385 budget_ms = self.transport_losses.budget().as_millis(),
386 "worker lost before reporting an activity result; the activity never ran \
387 and will be re-dispatched attempt-neutrally"
388 );
389 }
390 verdict.reason
391 }
392 Err(error) => {
393 tracing::error!(
394 workflow_id = %workflow_id,
395 activity_id = %activity_id,
396 %error,
397 "transport-loss ledger is unreadable; abandoning the activity rather than \
398 re-dispatching it without a budget"
399 );
400 format!(
401 "{}{detail} (transport-loss budget unreadable: {error})",
402 super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
403 )
404 }
405 }
406 }
407}
408
409impl ActivityCompletionSink for PendingActivities {
410 fn park_activity(
421 &self,
422 workflow_id: &WorkflowId,
423 activity_id: &ActivityId,
424 ) -> Result<(), ServerError> {
425 self.completion_fences
426 .revoke_current(workflow_id, activity_id)?;
427 let matched = self
428 .pending
429 .remove(&(workflow_id.clone(), activity_id.clone()));
430 if let Some((_, sender)) = matched {
431 let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
435 }
436 Ok(())
437 }
438
439 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
440 let result = match completion.outcome {
441 ActivityCompletionOutcome::Succeeded(payload) => {
442 payload_to_string(&payload).map_err(|reason| {
443 tracing::error!(
444 operation = "activity_complete",
445 workflow_id = %completion.workflow_id,
446 activity_id = %completion.activity_id,
447 error_type = "ActivityResultDecode",
448 %reason,
449 "activity completion failed"
450 );
451 ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
452 })?
453 }
454 ActivityCompletionOutcome::Failed(error) => {
455 let prefix = if error.is_retryable() {
456 "retryable"
457 } else {
458 "terminal"
459 };
460 tracing::error!(
461 operation = "activity_complete",
462 workflow_id = %completion.workflow_id,
463 activity_id = %completion.activity_id,
464 error_type = "ActivityFailed",
465 error_kind = prefix,
466 reason = %error.message,
467 "activity completion failed"
468 );
469 Err(format!("{prefix}:{}", error.message))
470 }
471 ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
476 &completion.workflow_id,
477 &completion.activity_id,
478 worker_id,
479 )),
480 };
481 self.complete_fenced(
482 &completion.workflow_id,
483 &completion.activity_id,
484 completion.run_id.as_ref(),
485 &completion.completion_token,
486 result,
487 )?;
488 Ok(())
489 }
490}
491
492fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
493 match payload.content_type() {
494 ContentType::Json => String::from_utf8(payload.bytes().to_vec())
495 .map(Ok)
496 .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
497 }
498}
499
500pub struct WorkerActivityDispatcher {
508 registry: ConnectedWorkerRegistry,
509 namespace: String,
510 pending: PendingActivities,
511 heartbeat_tracker: HeartbeatTracker,
512 drain_state: DrainState,
513 tokio_handle: Option<tokio::runtime::Handle>,
514 attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
520 queue_service: QueueServiceConfig,
524 queue_declarations: QueueDeclarationSource,
527 queue_state: QueueServiceState,
529}
530
531impl std::fmt::Debug for WorkerActivityDispatcher {
532 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
533 f.debug_struct("WorkerActivityDispatcher")
534 .field("namespace", &self.namespace)
535 .finish_non_exhaustive()
536 }
537}
538
539impl WorkerActivityDispatcher {
540 #[must_use]
548 pub fn new(
549 registry: ConnectedWorkerRegistry,
550 namespace: impl Into<String>,
551 heartbeat_tracker: HeartbeatTracker,
552 ) -> Self {
553 Self {
554 registry,
555 namespace: namespace.into(),
556 pending: PendingActivities::default(),
557 heartbeat_tracker,
558 drain_state: DrainState::default(),
559 tokio_handle: None,
560 attempt_owners: None,
561 queue_service: QueueServiceConfig::default(),
562 queue_declarations: QueueDeclarationSource::default(),
563 queue_state: QueueServiceState::default(),
564 }
565 }
566
567 #[must_use]
570 pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
571 self.queue_service = queue_service;
572 self
573 }
574
575 #[must_use]
579 pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
580 self.queue_declarations = queue_declarations;
581 self
582 }
583
584 #[must_use]
587 pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
588 self.queue_state = queue_state;
589 self
590 }
591
592 #[must_use]
599 pub fn with_attempt_owners(
600 mut self,
601 attempt_owners: super::intervention::AttemptOwnerIndex,
602 ) -> Self {
603 self.attempt_owners = Some(attempt_owners);
604 self
605 }
606
607 #[must_use]
609 pub fn with_pending(mut self, pending: PendingActivities) -> Self {
610 self.pending = pending;
611 self
612 }
613
614 #[must_use]
616 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
617 self.drain_state = drain_state;
618 self
619 }
620
621 #[must_use]
623 pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
624 self.tokio_handle = Some(tokio_handle);
625 self
626 }
627}
628
629impl WorkerActivityDispatcher {
630 fn ensure_accepting(
654 &self,
655 namespace: &str,
656 activity_type: &str,
657 workflow_id: &WorkflowId,
658 activity_id: &ActivityId,
659 worker_id: Option<WorkerId>,
660 ) -> Result<(), String> {
661 self.drain_state
662 .ensure_accepting(namespace, activity_type)
663 .map_err(|error| {
664 log_worker_error(
665 "WorkerDispatch",
666 namespace,
667 activity_type,
668 workflow_id,
669 activity_id,
670 worker_id,
671 &error.to_string(),
672 );
673 aion::PARKED_ACTIVITY_REASON.to_owned()
674 })
675 }
676
677 fn select_worker_or_wait(
685 &self,
686 address: &ServiceAddress,
687 workflow_id: &WorkflowId,
688 activity_id: &ActivityId,
689 ) -> Result<WorkerHandle, String> {
690 let wait = ServiceWait {
691 registry: &self.registry,
692 declarations: &self.queue_declarations,
693 config: &self.queue_service,
694 state: &self.queue_state,
695 address,
696 workflow_id,
697 activity_id,
698 };
699 let mut accepting = || {
700 self.ensure_accepting(
701 &address.namespace,
702 &address.activity_type,
703 workflow_id,
704 activity_id,
705 None,
706 )
707 };
708 let mut park = |budget: Option<Duration>| self.park_for_worker(budget);
709 select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
710 let reason = refusal.reason_string();
711 if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
714 let error_type = match refusal {
715 SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
716 _ => "WorkerRegistry",
717 };
718 log_worker_error(
719 error_type,
720 &address.namespace,
721 &address.activity_type,
722 workflow_id,
723 activity_id,
724 None,
725 &reason,
726 );
727 }
728 reason
729 })
730 }
731
732 fn park_for_worker(&self, budget: Option<Duration>) {
760 let handle = self
761 .tokio_handle
762 .clone()
763 .or_else(|| tokio::runtime::Handle::try_current().ok());
764 let Some(handle) = handle else {
765 std::thread::sleep(
770 budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
771 );
772 return;
773 };
774 handle.block_on(async {
775 let arrival = async {
776 tokio::select! {
777 () = self.registry.wait_for_worker() => {}
778 () = self.drain_state.wait_for_drain() => {}
779 }
780 };
781 match budget {
782 None => arrival.await,
783 Some(budget) => {
784 drop(tokio::time::timeout(budget, arrival).await);
787 }
788 }
789 });
790 }
791
792 fn track_worker_task(
793 &self,
794 worker_id: WorkerId,
795 activity_type: &str,
796 workflow_id: &WorkflowId,
797 activity_id: &ActivityId,
798 attempt: u32,
799 completion_token: CompletionToken,
800 ) -> Result<(), String> {
801 self.heartbeat_tracker
802 .track_task(
803 worker_id,
804 InFlightActivity {
805 workflow_id: workflow_id.clone(),
806 activity_id: activity_id.clone(),
807 attempt,
812 completion_token,
813 },
814 Instant::now(),
815 )
816 .map_err(|error| {
817 let reason = error.to_string();
818 log_worker_error(
819 "WorkerHeartbeatTracker",
820 &self.namespace,
821 activity_type,
822 workflow_id,
823 activity_id,
824 Some(worker_id),
825 &reason,
826 );
827 reason
828 })
829 }
830
831 fn cleanup_activity(
832 &self,
833 worker_id: WorkerId,
834 workflow_id: &WorkflowId,
835 activity_id: &ActivityId,
836 completion_token: &CompletionToken,
837 ) {
838 self.pending
839 .pending
840 .remove(&(workflow_id.clone(), activity_id.clone()));
841 if let Err(error) =
842 self.pending
843 .completion_fences
844 .revoke(workflow_id, activity_id, completion_token)
845 {
846 tracing::error!(
847 workflow_id = %workflow_id,
848 activity_id = %activity_id,
849 %error,
850 "failed to revoke undelivered activity generation"
851 );
852 }
853 let _ = self
854 .heartbeat_tracker
855 .complete_task(worker_id, workflow_id, activity_id);
856 self.drain_state.notify_activity_drained();
857 }
858
859 fn send_activity_task(
869 &self,
870 worker: &WorkerHandle,
871 task: ProtoActivityTask,
872 address: &ServiceAddress,
873 workflow_id: &WorkflowId,
874 activity_id: &ActivityId,
875 completion_token: &CompletionToken,
876 ) -> Result<(), String> {
877 match worker.delivery() {
878 WorkerDelivery::Grpc(sender) => {
879 let worker_id = worker.id();
880 let mut accepting = || {
881 self.ensure_accepting(
882 &address.namespace,
883 &address.activity_type,
884 workflow_id,
885 activity_id,
886 Some(worker_id),
887 )
888 };
889 let handed_over = deliver_within_schedule_to_start(
890 sender,
891 WorkerMessage::ActivityTask(Box::new(task)),
892 self.queue_service.schedule_to_start_timeout,
893 &mut accepting,
894 );
895 let Err(refusal) = handed_over else {
896 return Ok(());
897 };
898 self.cleanup_activity(worker_id, workflow_id, activity_id, completion_token);
899 let (error_type, reason) = self.hand_off_failure(&refusal, address);
900 if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
901 log_worker_error(
902 error_type,
903 &address.namespace,
904 &address.activity_type,
905 workflow_id,
906 activity_id,
907 Some(worker_id),
908 &reason,
909 );
910 }
911 Err(reason)
912 }
913 #[cfg(feature = "liminal-transport")]
920 WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
921 worker.id(),
922 delivery,
923 task,
924 &address.activity_type,
925 workflow_id,
926 activity_id,
927 ),
928 }
929 }
930
931 fn hand_off_failure(
938 &self,
939 refusal: &DeliveryRefusal,
940 address: &ServiceAddress,
941 ) -> (&'static str, String) {
942 match refusal {
943 DeliveryRefusal::Saturated { waited } => {
944 let census = self
947 .registry
948 .pool_census(
949 &address.namespace,
950 &address.task_queue,
951 &address.activity_type,
952 address.node.as_deref(),
953 )
954 .unwrap_or_else(|error| {
955 tracing::error!(
956 namespace = %address.namespace,
957 task_queue = %address.task_queue,
958 activity_type = %address.activity_type,
959 %error,
960 "poller census failed while reporting a saturated queue; \
961 the refusal carries an empty census"
962 );
963 PoolCensus::default()
964 });
965 let unavailable = WorkerUnavailable {
966 reason: QueueServiceReason::Saturated,
967 clock: Some(ExpiredClock::ScheduleToStart),
968 waited: *waited,
969 address: address.clone(),
970 census,
971 };
972 ("WorkerUnavailable", unavailable.reason_string())
973 }
974 DeliveryRefusal::Full => (
975 "WorkerChannelClosed",
976 "worker task channel full or closed: no available capacity".to_owned(),
977 ),
978 DeliveryRefusal::Closed => (
979 "WorkerChannelClosed",
980 "worker task channel full or closed: channel closed".to_owned(),
981 ),
982 DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
983 }
984 }
985
986 #[cfg(feature = "liminal-transport")]
1010 fn send_liminal_activity_task(
1011 &self,
1012 worker_id: WorkerId,
1013 delivery: &super::liminal_transport::LiminalWorkerDelivery,
1014 task: ProtoActivityTask,
1015 activity_type: &str,
1016 workflow_id: &WorkflowId,
1017 activity_id: &ActivityId,
1018 ) -> Result<(), String> {
1019 let completion_token =
1020 CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1021 .map_err(|error| error.to_string())?;
1022 let heartbeat_window_ms =
1023 u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1024 .unwrap_or(u64::MAX);
1025 let attempt = task.attempt;
1026 let run_id = task
1033 .run_id
1034 .map(RunId::try_from)
1035 .transpose()
1036 .map_err(|error| error.to_string())?
1037 .ok_or_else(|| {
1038 "activity task run id is missing; refusing to dispatch an unidentified run"
1039 .to_owned()
1040 })?;
1041 let request = super::liminal_transport::DispatchRequest {
1042 activity_type: activity_type.to_owned(),
1043 workflow_id: workflow_id.clone(),
1044 ordinal: activity_id.sequence_position(),
1045 run_id: Some(run_id.clone()),
1046 attempt,
1047 completion_token: task.completion_token,
1048 idempotency_key: task.idempotency_key,
1049 labels: task.labels.into_iter().collect(),
1050 heartbeat_window_ms,
1051 input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1052 };
1053 let awaiter = match delivery.push_dispatch(&request) {
1057 Ok(awaiter) => awaiter,
1058 Err(error) => {
1059 let reason = format!("worker liminal push failed: {error}");
1060 self.cleanup_activity(worker_id, workflow_id, activity_id, &completion_token);
1061 log_worker_error(
1062 "WorkerChannelClosed",
1063 &self.namespace,
1064 activity_type,
1065 workflow_id,
1066 activity_id,
1067 Some(worker_id),
1068 &reason,
1069 );
1070 return Err(reason);
1071 }
1072 };
1073 let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1077 super::liminal_transport::AttemptOwnerGuard::bind(
1078 owners.clone(),
1079 super::intervention::AttemptKey::new(
1080 workflow_id.clone(),
1081 run_id.clone(),
1082 activity_id.clone(),
1083 attempt,
1084 ),
1085 worker_id,
1086 )
1087 });
1088 self.spawn_liminal_reply_router(
1089 worker_id,
1090 awaiter,
1091 workflow_id,
1092 activity_id,
1093 &completion_token,
1094 owner_binding,
1095 );
1096 Ok(())
1097 }
1098
1099 #[cfg(feature = "liminal-transport")]
1139 fn spawn_liminal_reply_router(
1140 &self,
1141 worker_id: WorkerId,
1142 awaiter: liminal_server::server::connection::PushReplyAwaiter,
1143 workflow_id: &WorkflowId,
1144 activity_id: &ActivityId,
1145 completion_token: &CompletionToken,
1146 owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1147 ) {
1148 let pending = self.pending.clone();
1149 let heartbeat_tracker = self.heartbeat_tracker.clone();
1150 let drain_state = self.drain_state.clone();
1151 let workflow_id = workflow_id.clone();
1152 let activity_id = activity_id.clone();
1153 let completion_token = completion_token.clone();
1154 std::thread::spawn(move || {
1155 let _owner_binding = owner_binding;
1159 route_liminal_reply(
1160 &pending,
1161 &heartbeat_tracker,
1162 &drain_state,
1163 &awaiter,
1164 (worker_id, &workflow_id, &activity_id, &completion_token),
1165 );
1166 });
1167 }
1168
1169 fn await_activity_result(
1173 &self,
1174 context: &ActivityDispatchContext<'_>,
1175 rx: &SyncReceiver,
1176 ) -> Result<String, String> {
1177 match self.registry.is_registered(context.worker_id) {
1185 Ok(true) => {}
1186 Ok(false) => {
1187 if let Ok(result) = rx.try_recv() {
1191 return self.deliver_result(context, result);
1192 }
1193 self.cleanup_activity(
1194 context.worker_id,
1195 context.workflow_id,
1196 context.activity_id,
1197 &context.completion_token,
1198 );
1199 let reason = self.pending.classify_worker_loss(
1204 context.workflow_id,
1205 context.activity_id,
1206 context.worker_id,
1207 );
1208 log_worker_error(
1209 "WorkerLost",
1210 &self.namespace,
1211 context.activity_type,
1212 context.workflow_id,
1213 context.activity_id,
1214 Some(context.worker_id),
1215 &reason,
1216 );
1217 return Err(reason);
1218 }
1219 Err(error) => {
1220 self.cleanup_activity(
1221 context.worker_id,
1222 context.workflow_id,
1223 context.activity_id,
1224 &context.completion_token,
1225 );
1226 let reason = format!("worker registry inspection failed: {error}");
1227 log_worker_error(
1228 "WorkerRegistry",
1229 &self.namespace,
1230 context.activity_type,
1231 context.workflow_id,
1232 context.activity_id,
1233 Some(context.worker_id),
1234 &reason,
1235 );
1236 return Err(reason);
1237 }
1238 }
1239 if let Ok(result) = rx.recv() {
1240 return self.deliver_result(context, result);
1241 }
1242 self.cleanup_activity(
1245 context.worker_id,
1246 context.workflow_id,
1247 context.activity_id,
1248 &context.completion_token,
1249 );
1250 let reason = "activity response channel dropped".to_owned();
1251 log_worker_error(
1252 "WorkerChannelClosed",
1253 &self.namespace,
1254 context.activity_type,
1255 context.workflow_id,
1256 context.activity_id,
1257 Some(context.worker_id),
1258 &reason,
1259 );
1260 Err(reason)
1261 }
1262
1263 fn deliver_result(
1264 &self,
1265 context: &ActivityDispatchContext<'_>,
1266 result: Result<String, String>,
1267 ) -> Result<String, String> {
1268 self.pending
1269 .pending
1270 .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
1271 if let Err(reason) = &result
1276 && aion::is_parked_reason(reason)
1277 {
1278 tracing::info!(
1279 operation = "activity_dispatch",
1280 namespace = %self.namespace,
1281 workflow_id = %context.workflow_id,
1282 activity_id = %context.activity_id,
1283 activity_type = context.activity_type,
1284 worker_id = ?context.worker_id,
1285 "activity parked for restart recovery"
1286 );
1287 return result;
1288 }
1289 log_activity_completion(context, result.is_ok());
1290 result.inspect_err(|reason| {
1291 log_worker_error(
1292 "ActivityFailed",
1293 &self.namespace,
1294 context.activity_type,
1295 context.workflow_id,
1296 context.activity_id,
1297 Some(context.worker_id),
1298 reason,
1299 );
1300 })
1301 }
1302}
1303
1304impl ActivityDispatcher for WorkerActivityDispatcher {
1305 fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1306 match tokio::runtime::Handle::try_current() {
1307 Ok(handle) => match handle.runtime_flavor() {
1308 tokio::runtime::RuntimeFlavor::MultiThread => {
1309 tokio::task::block_in_place(|| self.dispatch_blocking(request))
1316 }
1317 flavor => Err(format!(
1318 "activity dispatch blocks the calling thread until the worker responds; \
1319 a {flavor:?} tokio runtime cannot host that wait because the worker \
1320 stream forwarder shares its only executor thread and the task could \
1321 never be delivered — run the engine on a multi-thread tokio runtime"
1322 )),
1323 },
1324 Err(_) => self.dispatch_blocking(request),
1328 }
1329 }
1330}
1331
1332impl WorkerActivityDispatcher {
1333 fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1350 let ActivityDispatch {
1351 namespace,
1352 task_queue,
1353 node,
1357 workflow_id,
1358 run_id,
1359 activity_id,
1360 name,
1361 input,
1362 config: _,
1363 attempt,
1364 labels,
1365 advisory: _,
1370 } = request;
1371 let started_at = Instant::now();
1372 self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1373 let address = ServiceAddress {
1374 namespace: namespace.clone(),
1375 task_queue: task_queue.clone(),
1376 activity_type: name.clone(),
1377 node: node.clone(),
1378 };
1379 let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1380 let worker_id = worker.id();
1381 let span = info_span!(
1382 "activity_dispatch",
1383 operation = "activity_dispatch",
1384 namespace = %namespace,
1385 task_queue = %task_queue,
1386 node = node.as_deref(),
1387 workflow_id = %workflow_id,
1388 activity_id = %activity_id,
1389 activity_type = %name,
1390 worker_id = ?worker_id,
1391 );
1392 let _span_guard = span.enter();
1393 self.ensure_accepting(
1394 &namespace,
1395 &name,
1396 &workflow_id,
1397 &activity_id,
1398 Some(worker_id),
1399 )?;
1400
1401 let (completion_token, rx) = self
1402 .pending
1403 .insert(workflow_id.clone(), activity_id.clone())
1404 .map_err(|error| error.to_string())?;
1405 let task = activity_task(
1406 &name,
1407 &input,
1408 (&workflow_id, &run_id, &activity_id),
1409 attempt,
1410 labels,
1411 &completion_token,
1412 );
1413 if let Err(error) = self.track_worker_task(
1414 worker_id,
1415 &name,
1416 &workflow_id,
1417 &activity_id,
1418 attempt,
1419 completion_token.clone(),
1420 ) {
1421 self.cleanup_activity(worker_id, &workflow_id, &activity_id, &completion_token);
1422 return Err(error);
1423 }
1424 self.send_activity_task(
1425 &worker,
1426 task,
1427 &address,
1428 &workflow_id,
1429 &activity_id,
1430 &completion_token,
1431 )?;
1432 let context = ActivityDispatchContext {
1433 namespace: &namespace,
1434 activity_type: &name,
1435 worker_id,
1436 workflow_id: &workflow_id,
1437 activity_id: &activity_id,
1438 completion_token,
1439 started_at,
1440 };
1441 self.await_activity_result(&context, &rx)
1442 }
1443}
1444
1445#[cfg(feature = "liminal-transport")]
1452fn route_liminal_reply(
1453 pending: &PendingActivities,
1454 heartbeat_tracker: &HeartbeatTracker,
1455 drain_state: &DrainState,
1456 awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1457 execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1458) {
1459 let (worker_id, workflow_id, activity_id, current_token) = execution;
1460 let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1465 heartbeat_tracker
1466 .is_tracked(worker_id, workflow_id, activity_id)
1467 .unwrap_or(false)
1468 });
1469 let (run_id, submitted_token, outcome, synthesized) = match waited {
1473 Ok(Some(response)) => {
1474 let submitted_token = match CompletionToken::from_wire(
1475 workflow_id,
1476 activity_id,
1477 response.completion_token,
1478 ) {
1479 Ok(token) => token,
1480 Err(error) => {
1481 tracing::warn!(
1482 worker_id = ?worker_id,
1483 workflow_id = %workflow_id,
1484 activity_id = %activity_id,
1485 %error,
1486 "liminal activity completion omitted its generation proof"
1487 );
1488 return;
1489 }
1490 };
1491 (response.run_id, submitted_token, response.outcome, false)
1492 }
1493 Ok(None) => {
1494 tracing::debug!(
1495 worker_id = ?worker_id,
1496 workflow_id = %workflow_id,
1497 activity_id = %activity_id,
1498 "liminal dispatch resolved by another path; abandoning reply wait"
1499 );
1500 return;
1501 }
1502 Err(error) if error.is_worker_connection_lost() => (
1503 None,
1504 current_token.clone(),
1505 Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1509 true,
1510 ),
1511 Err(error) => (
1512 None,
1513 current_token.clone(),
1514 Err(format!("retryable:worker liminal reply failed: {error}")),
1515 true,
1516 ),
1517 };
1518 if synthesized {
1523 let was_tracked =
1524 complete_liminal_tracking(heartbeat_tracker, worker_id, workflow_id, activity_id);
1525 if !was_tracked {
1526 tracing::debug!(
1529 worker_id = ?worker_id,
1530 workflow_id = %workflow_id,
1531 activity_id = %activity_id,
1532 "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1533 );
1534 return;
1535 }
1536 }
1537 if let Err(error) = pending.complete_fenced(
1538 workflow_id,
1539 activity_id,
1540 run_id.as_ref(),
1541 &submitted_token,
1542 outcome,
1543 ) {
1544 tracing::warn!(
1545 worker_id = ?worker_id,
1546 workflow_id = %workflow_id,
1547 activity_id = %activity_id,
1548 %error,
1549 "liminal activity completion handoff rejected"
1550 );
1551 return;
1552 }
1553 if !synthesized
1554 && let Err(error) = heartbeat_tracker.complete_task(worker_id, workflow_id, activity_id)
1555 {
1556 tracing::error!(
1557 worker_id = ?worker_id,
1558 workflow_id = %workflow_id,
1559 activity_id = %activity_id,
1560 %error,
1561 "failed to clear in-flight tracking for completed liminal activity"
1562 );
1563 }
1564 drain_state.notify_activity_drained();
1565}
1566
1567#[cfg(feature = "liminal-transport")]
1568fn complete_liminal_tracking(
1569 heartbeat_tracker: &HeartbeatTracker,
1570 worker_id: WorkerId,
1571 workflow_id: &WorkflowId,
1572 activity_id: &ActivityId,
1573) -> bool {
1574 heartbeat_tracker
1575 .complete_task(worker_id, workflow_id, activity_id)
1576 .unwrap_or_else(|error| {
1577 tracing::error!(
1578 worker_id = ?worker_id,
1579 workflow_id = %workflow_id,
1580 activity_id = %activity_id,
1581 %error,
1582 "failed to clear in-flight tracking for completed liminal activity"
1583 );
1584 true
1585 })
1586}
1587
1588struct ActivityDispatchContext<'a> {
1589 namespace: &'a str,
1590 activity_type: &'a str,
1591 worker_id: WorkerId,
1592 workflow_id: &'a WorkflowId,
1593 activity_id: &'a ActivityId,
1594 completion_token: CompletionToken,
1595 started_at: Instant,
1596}
1597
1598fn activity_task(
1599 activity_type: &str,
1600 input: &str,
1601 execution: (&WorkflowId, &RunId, &ActivityId),
1602 attempt: u32,
1603 labels: BTreeMap<String, String>,
1604 completion_token: &CompletionToken,
1605) -> ProtoActivityTask {
1606 let (workflow_id, run_id, activity_id) = execution;
1607 ProtoActivityTask {
1608 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1609 activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1610 activity_type: activity_type.to_owned(),
1611 input: Some(ProtoPayload {
1612 content_type: String::from("application/json"),
1613 bytes: input.as_bytes().to_vec(),
1614 }),
1615 attempt,
1616 labels: labels.into_iter().collect(),
1617 run_id: Some(run_id.clone().into()),
1618 completion_token: completion_token.as_str().to_owned(),
1619 idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
1620 }
1621}
1622
1623fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1624 let duration_ms = duration_ms(context.started_at.elapsed());
1625 tracing::info!(
1626 operation = "activity_complete",
1627 namespace = context.namespace,
1628 workflow_id = %context.workflow_id,
1629 activity_id = %context.activity_id,
1630 activity_type = context.activity_type,
1631 worker_id = ?context.worker_id,
1632 duration_ms,
1633 outcome = if succeeded { "succeeded" } else { "failed" },
1634 "activity completed"
1635 );
1636}
1637
1638fn duration_ms(duration: Duration) -> u64 {
1639 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1640}
1641
1642fn log_worker_error(
1643 error_type: &'static str,
1644 namespace: &str,
1645 activity_type: &str,
1646 workflow_id: &WorkflowId,
1647 activity_id: &ActivityId,
1648 worker_id: Option<super::registry::WorkerId>,
1649 reason: &str,
1650) {
1651 tracing::error!(
1652 operation = "activity_dispatch",
1653 namespace,
1654 workflow_id = %workflow_id,
1655 activity_id = %activity_id,
1656 activity_type,
1657 worker_id = ?worker_id,
1658 error_type,
1659 reason,
1660 "worker interaction failed"
1661 );
1662}
1663
1664#[cfg(test)]
1665mod tests {
1666 use std::sync::Mutex;
1667
1668 use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1669
1670 use super::*;
1671
1672 fn activity_id(pos: u64) -> ActivityId {
1673 ActivityId::from_sequence_position(pos)
1674 }
1675
1676 #[test]
1677 fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
1678 let pending = PendingActivities::default();
1679 let workflow_id = WorkflowId::new_v4();
1680 let id = activity_id(1);
1681 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1682
1683 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1684 assert_eq!(
1685 rx.recv_timeout(Duration::from_millis(50)),
1686 Ok(Ok("done".to_owned()))
1687 );
1688 Ok(())
1689 }
1690
1691 #[test]
1692 fn pending_complete_unknown_returns_false() {
1693 let pending = PendingActivities::default();
1694 assert!(!pending.complete(
1695 &WorkflowId::new_v4(),
1696 &activity_id(99),
1697 None,
1698 Ok("orphan".to_owned())
1699 ));
1700 }
1701
1702 #[derive(Default)]
1703 struct RecordingOutboxCallback {
1704 completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1705 failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1706 live: bool,
1707 }
1708
1709 impl OutboxDeliveryCallback for RecordingOutboxCallback {
1710 fn deliver_completion(
1711 &self,
1712 workflow_id: &WorkflowId,
1713 activity_id: &ActivityId,
1714 run_id: Option<&RunId>,
1715 result: String,
1716 ) -> Result<bool, ServerError> {
1717 let _ = run_id;
1718 self.completions
1719 .lock()
1720 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1721 .push((workflow_id.clone(), activity_id.clone(), result));
1722 Ok(self.live)
1723 }
1724
1725 fn deliver_failure(
1726 &self,
1727 workflow_id: &WorkflowId,
1728 activity_id: &ActivityId,
1729 run_id: Option<&RunId>,
1730 reason: String,
1731 ) -> Result<bool, ServerError> {
1732 let _ = run_id;
1733 self.failures
1734 .lock()
1735 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1736 .push((workflow_id.clone(), activity_id.clone(), reason));
1737 Ok(self.live)
1738 }
1739 }
1740
1741 #[test]
1742 fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1743 let pending = PendingActivities::default();
1744 let callback = Arc::new(RecordingOutboxCallback {
1745 live: true,
1746 ..RecordingOutboxCallback::default()
1747 });
1748 pending.clone().set_outbox_delivery(callback.clone());
1750
1751 let workflow_id = WorkflowId::new_v4();
1752 let id = activity_id(7);
1753
1754 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1757 let completions = callback
1758 .completions
1759 .lock()
1760 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1761 assert_eq!(completions.len(), 1);
1762 assert_eq!(completions[0].0, workflow_id);
1763 assert_eq!(completions[0].1, id);
1764 assert_eq!(completions[0].2, "done");
1765 Ok(())
1766 }
1767
1768 #[test]
1769 fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1770 -> Result<(), ServerError> {
1771 let pending = PendingActivities::default();
1772 let callback = Arc::new(RecordingOutboxCallback::default());
1774 pending.set_outbox_delivery(callback.clone());
1775
1776 let workflow_id = WorkflowId::new_v4();
1777 let id = activity_id(8);
1778
1779 assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1780 let failures = callback
1781 .failures
1782 .lock()
1783 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1784 assert_eq!(failures.len(), 1);
1785 assert_eq!(failures[0].2, "retryable:boom");
1786 Ok(())
1787 }
1788
1789 #[test]
1790 fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1791 let pending = PendingActivities::default();
1794 assert!(!pending.complete(
1795 &WorkflowId::new_v4(),
1796 &activity_id(9),
1797 None,
1798 Ok("x".to_owned())
1799 ));
1800 }
1801
1802 #[test]
1803 fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1804 let pending = PendingActivities::default();
1805 let callback = Arc::new(RecordingOutboxCallback {
1806 live: true,
1807 ..RecordingOutboxCallback::default()
1808 });
1809 pending.set_outbox_delivery(callback.clone());
1810
1811 let workflow_id = WorkflowId::new_v4();
1812 let id = activity_id(10);
1813 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1814
1815 assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1816 assert_eq!(
1817 rx.recv_timeout(Duration::from_millis(50)),
1818 Ok(Ok("matched".to_owned()))
1819 );
1820 assert!(
1821 callback
1822 .completions
1823 .lock()
1824 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1825 .is_empty(),
1826 "a matched completion must deliver to its waiter, not the outbox callback"
1827 );
1828 Ok(())
1829 }
1830
1831 #[test]
1835 fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
1836 let pending = PendingActivities::default();
1837 let workflow_id = WorkflowId::new_v4();
1838 let id = activity_id(11);
1839 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1840
1841 pending.park_activity(&workflow_id, &id)?;
1842 let result = rx
1843 .recv_timeout(Duration::from_millis(50))
1844 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1845 assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
1846 Ok(())
1847 }
1848
1849 #[test]
1853 fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
1854 let pending = PendingActivities::default();
1855 let callback = Arc::new(RecordingOutboxCallback {
1856 live: true,
1857 ..RecordingOutboxCallback::default()
1858 });
1859 pending.set_outbox_delivery(callback.clone());
1860
1861 pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
1862
1863 assert!(
1864 callback
1865 .failures
1866 .lock()
1867 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1868 .is_empty(),
1869 "a park must never be delivered as an outbox failure"
1870 );
1871 assert!(
1872 callback
1873 .completions
1874 .lock()
1875 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1876 .is_empty(),
1877 "a park must never be delivered as an outbox completion"
1878 );
1879 Ok(())
1880 }
1881
1882 #[test]
1883 fn completion_sink_routes_success() -> Result<(), ServerError> {
1884 let pending = PendingActivities::default();
1885 let workflow_id = WorkflowId::new_v4();
1886 let id = activity_id(2);
1887 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1888 let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
1889
1890 pending.complete_activity(ActivityCompletion {
1891 workflow_id,
1892 activity_id: id,
1893 run_id: None,
1894 completion_token,
1895 outcome: ActivityCompletionOutcome::Succeeded(payload),
1896 })?;
1897
1898 let result = rx
1899 .recv_timeout(Duration::from_millis(50))
1900 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1901 assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
1902 Ok(())
1903 }
1904
1905 #[test]
1906 fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
1907 let pending = PendingActivities::default();
1908 let workflow_id = WorkflowId::new_v4();
1909 let id = activity_id(12);
1910 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1911
1912 let malformed = pending.complete_activity(ActivityCompletion {
1913 workflow_id: workflow_id.clone(),
1914 activity_id: id.clone(),
1915 run_id: None,
1916 completion_token: completion_token.clone(),
1917 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1918 ContentType::Json,
1919 vec![0xff],
1920 )),
1921 });
1922 assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
1923 assert!(
1924 rx.try_recv().is_err(),
1925 "an invalid result must leave the waiter unresolved"
1926 );
1927
1928 pending.complete_activity(ActivityCompletion {
1929 workflow_id,
1930 activity_id: id,
1931 run_id: None,
1932 completion_token,
1933 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1934 ContentType::Json,
1935 br#""valid""#.to_vec(),
1936 )),
1937 })?;
1938 let result = rx
1939 .recv_timeout(Duration::from_millis(50))
1940 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
1941 assert_eq!(result, Ok(r#""valid""#.to_owned()));
1942 Ok(())
1943 }
1944
1945 #[test]
1946 fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
1947 let pending = PendingActivities::default();
1948 let workflow_id = WorkflowId::new_v4();
1949 let id = activity_id(3);
1950 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1951
1952 pending.complete_activity(ActivityCompletion {
1953 workflow_id,
1954 activity_id: id,
1955 run_id: None,
1956 completion_token,
1957 outcome: ActivityCompletionOutcome::Failed(ActivityError {
1958 kind: ActivityErrorKind::Retryable,
1959 message: "temporary".to_owned(),
1960 details: None,
1961 }),
1962 })?;
1963
1964 let result = rx
1965 .recv_timeout(Duration::from_millis(50))
1966 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1967 assert_eq!(result, Err("retryable:temporary".to_owned()));
1968 Ok(())
1969 }
1970
1971 #[test]
1980 fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
1981 -> Result<(), ServerError> {
1982 let pending = PendingActivities::default();
1983 let post_restart_workflow = WorkflowId::new_v4();
1984 let pre_restart_workflow = WorkflowId::new_v4();
1985 let id = activity_id(1);
1987 let (completion_token, rx) = pending.insert(post_restart_workflow.clone(), id.clone())?;
1988
1989 let rejected = pending.complete_activity(ActivityCompletion {
1991 workflow_id: pre_restart_workflow,
1992 activity_id: id.clone(),
1993 run_id: None,
1994 completion_token: CompletionToken::for_test(),
1995 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1996 ContentType::Json,
1997 br#""stale""#.to_vec(),
1998 )),
1999 });
2000 assert!(matches!(
2001 rejected,
2002 Err(ServerError::ActivityCompletionRejected { .. })
2003 ));
2004 assert!(
2005 rx.try_recv().is_err(),
2006 "stale result for a different workflow must not complete this dispatch"
2007 );
2008
2009 pending.complete_activity(ActivityCompletion {
2011 workflow_id: post_restart_workflow,
2012 activity_id: id,
2013 run_id: None,
2014 completion_token,
2015 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2016 ContentType::Json,
2017 br#""fresh""#.to_vec(),
2018 )),
2019 })?;
2020 let result = rx
2021 .recv_timeout(Duration::from_millis(50))
2022 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2023 assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2024 Ok(())
2025 }
2026
2027 fn test_tracker() -> HeartbeatTracker {
2030 HeartbeatTracker::new(Duration::from_secs(5))
2031 }
2032
2033 fn greet_request() -> ActivityDispatch {
2036 ActivityDispatch {
2037 namespace: "default".to_owned(),
2038 task_queue: "default".to_owned(),
2039 node: None,
2040 workflow_id: WorkflowId::new_v4(),
2041 run_id: RunId::new_v4(),
2042 activity_id: ActivityId::from_sequence_position(0),
2043 name: "greet".to_owned(),
2044 input: "{}".to_owned(),
2045 config: "{}".to_owned(),
2046 attempt: 1,
2047 labels: std::collections::BTreeMap::new(),
2048 advisory: false,
2049 }
2050 }
2051
2052 #[test]
2053 fn dispatcher_fails_immediately_when_draining_without_workers() {
2054 let registry = ConnectedWorkerRegistry::default();
2055 let drain = DrainState::default();
2056 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2057 .with_drain_state(drain.clone());
2058
2059 let _ = drain.begin();
2060
2061 let result = dispatcher.dispatch(greet_request());
2062
2063 assert!(result.is_err());
2064 let err = result.err().unwrap_or_default();
2065 assert!(
2066 err.contains("drain"),
2067 "expected drain rejection, got: {err}"
2068 );
2069 }
2070
2071 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2087 async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
2088 -> Result<(), Box<dyn std::error::Error>> {
2089 let registry = ConnectedWorkerRegistry::default();
2090 let pending = PendingActivities::default();
2091 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2092 let activity_types = [String::from("greet")];
2093 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2094
2095 let sink = pending.clone();
2096 let echo_worker = tokio::spawn(async move {
2097 let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
2098 return Err("expected an activity task on the worker channel".to_owned());
2099 };
2100 let workflow_id = task
2101 .workflow_id
2102 .ok_or("task missing workflow id")
2103 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2104 let activity_id = task
2105 .activity_id
2106 .map(ActivityId::from)
2107 .ok_or("task missing activity id")?;
2108 let completion_token =
2109 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2110 .map_err(|error| error.to_string())?;
2111 sink.complete_activity(ActivityCompletion {
2112 workflow_id,
2113 activity_id,
2114 run_id: None,
2115 completion_token,
2116 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2117 ContentType::Json,
2118 br#"{"greeting":"hello"}"#.to_vec(),
2119 )),
2120 })
2121 .map_err(|error| error.to_string())
2122 });
2123
2124 let dispatcher = Arc::new(
2125 WorkerActivityDispatcher::new(registry, "default", test_tracker())
2126 .with_pending(pending),
2127 );
2128 let started = Instant::now();
2129 let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
2132 dispatcher.dispatch(greet_request())
2133 }));
2134 let result = dispatch_task.await.map_err(|error| error.to_string())?;
2135 let elapsed = started.elapsed();
2136
2137 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2138 assert!(
2139 elapsed < Duration::from_secs(5),
2140 "dispatch round trip took {elapsed:?}; task delivery must not \
2141 depend on the blocked dispatch thread"
2142 );
2143 echo_worker.await.map_err(|error| error.to_string())??;
2144 registration.deregister()?;
2145 Ok(())
2146 }
2147
2148 #[tokio::test]
2152 async fn dispatch_on_current_thread_runtime_fails_fast()
2153 -> Result<(), Box<dyn std::error::Error>> {
2154 let registry = ConnectedWorkerRegistry::default();
2155 let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
2156 let activity_types = [String::from("greet")];
2157 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2158 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
2159
2160 let started = Instant::now();
2161 let result = dispatcher.dispatch(greet_request());
2162 let elapsed = started.elapsed();
2163
2164 let err = result.err().ok_or("expected dispatch to fail")?;
2165 assert!(
2166 err.contains("multi-thread tokio runtime"),
2167 "unexpected error: {err}"
2168 );
2169 assert!(
2170 elapsed < Duration::from_secs(5),
2171 "fail-fast path took {elapsed:?}"
2172 );
2173 registration.deregister()?;
2174 Ok(())
2175 }
2176
2177 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2183 async fn dispatch_pinned_to_node_reaches_only_that_node()
2184 -> Result<(), Box<dyn std::error::Error>> {
2185 let registry = ConnectedWorkerRegistry::default();
2186 let pending = PendingActivities::default();
2187 let activity_types = [String::from("greet")];
2188 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2189 let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
2190 let on_n2 = registry.register_namespaces(
2196 [String::from("default")],
2197 "default",
2198 Some(String::from("n2")),
2199 activity_types.iter(),
2200 n2_tx,
2201 )?;
2202 let on_n1 = registry.register_namespaces(
2203 [String::from("default")],
2204 "default",
2205 Some(String::from("n1")),
2206 activity_types.iter(),
2207 n1_tx,
2208 )?;
2209
2210 let sink = pending.clone();
2214 let echo_n1 = tokio::spawn(async move {
2215 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2216 return Err("expected an activity task on the n1 worker channel".to_owned());
2217 };
2218 let workflow_id = task
2219 .workflow_id
2220 .ok_or("task missing workflow id")
2221 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2222 let activity_id = task
2223 .activity_id
2224 .map(ActivityId::from)
2225 .ok_or("task missing activity id")?;
2226 let completion_token =
2227 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2228 .map_err(|error| error.to_string())?;
2229 sink.complete_activity(ActivityCompletion {
2230 workflow_id,
2231 activity_id,
2232 run_id: None,
2233 completion_token,
2234 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2235 ContentType::Json,
2236 br#"{"greeting":"hello"}"#.to_vec(),
2237 )),
2238 })
2239 .map_err(|error| error.to_string())
2240 });
2241
2242 let dispatcher = Arc::new(
2243 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2244 .with_pending(pending),
2245 );
2246
2247 let pinned = ActivityDispatch {
2248 node: Some(String::from("n1")),
2249 ..greet_request()
2250 };
2251 let started = Instant::now();
2252 let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
2253 .await
2254 .map_err(|error| error.to_string())?;
2255 let elapsed = started.elapsed();
2256
2257 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2258 assert!(
2259 elapsed < Duration::from_secs(5),
2260 "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
2261 );
2262 echo_n1.await.map_err(|error| error.to_string())??;
2263
2264 assert!(
2266 n2_rx.try_recv().is_err(),
2267 "node=Some(\"n1\") dispatch must not reach the n2 worker"
2268 );
2269
2270 on_n1.deregister()?;
2271 on_n2.deregister()?;
2272 Ok(())
2273 }
2274
2275 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2278 async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
2279 -> Result<(), Box<dyn std::error::Error>> {
2280 let registry = ConnectedWorkerRegistry::default();
2281 let pending = PendingActivities::default();
2282 let activity_types = [String::from("greet")];
2283 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2284 let on_n1 = registry.register_namespaces(
2285 [String::from("default")],
2286 "default",
2287 Some(String::from("n1")),
2288 activity_types.iter(),
2289 n1_tx,
2290 )?;
2291
2292 let sink = pending.clone();
2293 let echo = tokio::spawn(async move {
2294 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2295 return Err("expected an activity task on the worker channel".to_owned());
2296 };
2297 let workflow_id = task
2298 .workflow_id
2299 .ok_or("task missing workflow id")
2300 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2301 let activity_id = task
2302 .activity_id
2303 .map(ActivityId::from)
2304 .ok_or("task missing activity id")?;
2305 let completion_token =
2306 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2307 .map_err(|error| error.to_string())?;
2308 sink.complete_activity(ActivityCompletion {
2309 workflow_id,
2310 activity_id,
2311 run_id: None,
2312 completion_token,
2313 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2314 ContentType::Json,
2315 br#"{"greeting":"hello"}"#.to_vec(),
2316 )),
2317 })
2318 .map_err(|error| error.to_string())
2319 });
2320
2321 let dispatcher = Arc::new(
2322 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2323 .with_pending(pending),
2324 );
2325
2326 let result = tokio::spawn(futures::future::lazy(move |_| {
2328 dispatcher.dispatch(greet_request())
2329 }))
2330 .await
2331 .map_err(|error| error.to_string())?;
2332
2333 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2334 echo.await.map_err(|error| error.to_string())??;
2335 on_n1.deregister()?;
2336 Ok(())
2337 }
2338}