1use std::collections::BTreeMap;
64use std::sync::{Arc, OnceLock};
65use std::time::{Duration, Instant};
66
67use aion::{ActivityDispatch, ActivityDispatcher};
68use aion_core::{ActivityErrorKind, ActivityId, ContentType, Payload, RunId, WorkflowId};
69use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
70use dashmap::DashMap;
71use dashmap::mapref::entry::Entry;
72
73use super::dispatch::{ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink};
74use super::envelope::{CompletionFences, CompletionToken, idempotency_key};
75use super::heartbeat::{HeartbeatTracker, InFlightActivity};
76use super::queue_service::{
77 DeliveryRefusal, ExpiredClock, PARK_POLL_INTERVAL, PoolCensus, QueueDeclarationSource,
78 QueueServiceConfig, QueueServiceReason, QueueServiceState, SelectionRefusal, ServiceAddress,
79 ServiceWait, WorkerUnavailable, deliver_within_schedule_to_start, select_worker_or_refuse,
80};
81use super::registry::{
82 ConnectedWorkerRegistry, WorkerArrival, WorkerDelivery, WorkerHandle, WorkerId, WorkerMessage,
83};
84use crate::error::ServerError;
85use crate::shutdown::DrainState;
86use tracing::info_span;
87
88type SyncSender = std::sync::mpsc::SyncSender<Result<String, String>>;
89type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
90
91type PendingActivityKey = (WorkflowId, ActivityId);
109
110pub trait OutboxDeliveryCallback: Send + Sync {
120 fn deliver_completion(
130 &self,
131 workflow_id: &WorkflowId,
132 activity_id: &ActivityId,
133 run_id: Option<&RunId>,
134 result: String,
135 ) -> Result<bool, ServerError>;
136
137 fn deliver_failure(
144 &self,
145 workflow_id: &WorkflowId,
146 activity_id: &ActivityId,
147 run_id: Option<&RunId>,
148 reason: String,
149 ) -> Result<bool, ServerError>;
150}
151
152#[derive(Clone)]
164pub struct PendingActivities {
165 pending: Arc<DashMap<PendingActivityKey, SyncSender>>,
166 completion_fences: CompletionFences,
167 outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
168 transport_losses: super::transport_loss::TransportLossLedger,
174}
175
176impl std::fmt::Debug for PendingActivities {
177 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178 formatter
179 .debug_struct("PendingActivities")
180 .field("pending", &self.pending.len())
181 .field("completion_fences", &self.completion_fences)
182 .field(
183 "outbox_delivery_installed",
184 &self.outbox_delivery.get().is_some(),
185 )
186 .field("transport_losses", &self.transport_losses)
187 .finish()
188 }
189}
190
191impl PendingActivities {
192 fn insert(
193 &self,
194 workflow_id: WorkflowId,
195 run_id: &RunId,
196 activity_id: ActivityId,
197 attempt: u32,
198 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
199 let completion_token =
200 self.completion_fences
201 .issue(&workflow_id, run_id, &activity_id, attempt)?;
202 let (tx, rx) = std::sync::mpsc::sync_channel(1);
203 let key = (workflow_id, activity_id);
204 match self.pending.entry(key) {
205 Entry::Vacant(entry) => {
206 entry.insert(tx);
207 Ok((completion_token, rx))
208 }
209 Entry::Occupied(entry) => {
210 let (workflow_id, activity_id) = entry.key();
211 Err(ServerError::PendingActivityCollision {
212 workflow_id: workflow_id.clone(),
213 activity_id: activity_id.clone(),
214 })
215 }
216 }
217 }
218
219 #[must_use]
221 pub fn completion_fences(&self) -> CompletionFences {
222 self.completion_fences.clone()
223 }
224
225 #[cfg(test)]
229 pub(crate) fn insert_for_test(
230 &self,
231 workflow_id: WorkflowId,
232 run_id: &RunId,
233 activity_id: ActivityId,
234 attempt: u32,
235 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
236 self.insert(workflow_id, run_id, activity_id, attempt)
237 }
238
239 pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
245 if self.outbox_delivery.set(callback).is_err() {
246 tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
247 }
248 }
249
250 fn complete(
258 &self,
259 workflow_id: &WorkflowId,
260 activity_id: &ActivityId,
261 run_id: Option<&RunId>,
262 result: Result<String, String>,
263 ) -> bool {
264 let matched = self
267 .pending
268 .remove(&(workflow_id.clone(), activity_id.clone()));
269 if let Some((_, sender)) = matched {
270 return sender.send(result).is_ok();
271 }
272 let Some(callback) = self.outbox_delivery.get() else {
273 return false;
275 };
276 let outcome = match result {
277 Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
278 Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
279 };
280 match outcome {
281 Ok(true) => true,
282 Ok(false) => {
283 tracing::debug!(
285 workflow_id = %workflow_id,
286 activity_id = %activity_id,
287 "unmatched outbox completion for a workflow that is not currently live; \
288 recovery will re-arm it"
289 );
290 false
291 }
292 Err(error) => {
293 tracing::warn!(
294 workflow_id = %workflow_id,
295 activity_id = %activity_id,
296 %error,
297 "failed to deliver unmatched outbox completion to the live workflow"
298 );
299 false
300 }
301 }
302 }
303
304 fn complete_fenced_after_accept(
305 &self,
306 workflow_id: &WorkflowId,
307 activity_id: &ActivityId,
308 run_id: Option<&RunId>,
309 completion_token: &CompletionToken,
310 result: Result<String, String>,
311 after_accept: impl FnOnce() -> Result<bool, ServerError>,
312 ) -> Result<bool, ServerError> {
313 let accepted = self
318 .completion_fences
319 .accept(workflow_id, activity_id, completion_token)
320 .inspect_err(|error| {
321 tracing::warn!(
322 workflow_id = %workflow_id,
323 activity_id = %activity_id,
324 %error,
325 "activity completion rejected by execution-generation fence"
326 );
327 })?;
328 let should_publish = match after_accept() {
329 Ok(should_publish) => should_publish,
330 Err(error) => {
331 if let Err(restore_error) =
332 self.completion_fences
333 .restore_if_absent(workflow_id, activity_id, &accepted)
334 {
335 tracing::error!(%restore_error, "failed to restore completion generation after settlement failure");
336 }
337 return Err(error);
338 }
339 };
340 if !should_publish {
341 if let Err(restore_error) =
348 self.completion_fences
349 .restore_if_absent(workflow_id, activity_id, &accepted)
350 {
351 tracing::error!(%restore_error, "failed to restore completion generation after a non-publishing settlement");
352 }
353 return Ok(false);
354 }
355 let transport_domain = result
361 .as_ref()
362 .err()
363 .is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
364 if !transport_domain
365 && let Err(error) = self.transport_losses.clear(workflow_id, activity_id)
366 {
367 tracing::warn!(
368 workflow_id = %workflow_id,
369 activity_id = %activity_id,
370 %error,
371 "failed to retire the transport-loss budget for a resolved activity"
372 );
373 }
374 Ok(self.complete(workflow_id, activity_id, run_id, result))
375 }
376
377 #[must_use]
389 pub fn new(heartbeat_window: std::time::Duration) -> Self {
390 Self {
391 pending: Arc::default(),
392 completion_fences: CompletionFences::default(),
393 outbox_delivery: Arc::default(),
394 transport_losses: super::transport_loss::TransportLossLedger::new(heartbeat_window),
395 }
396 }
397
398 #[must_use]
400 pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
401 &self.transport_losses
402 }
403
404 fn classify_worker_loss(
412 &self,
413 workflow_id: &WorkflowId,
414 activity_id: &ActivityId,
415 worker_id: crate::worker::registry::WorkerId,
416 ) -> String {
417 let detail = super::transport_loss::worker_lost_detail(worker_id);
418 match self
419 .transport_losses
420 .record_loss(workflow_id, activity_id, &detail)
421 {
422 Ok(verdict) => {
423 if verdict.exhausted {
424 tracing::error!(
425 operation = "activity_complete",
426 workflow_id = %workflow_id,
427 activity_id = %activity_id,
428 worker_id = ?worker_id,
429 error_type = "TransportExhausted",
430 losses = verdict.losses,
431 budget_ms = self.transport_losses.budget().as_millis(),
432 "activity abandoned: the transport kept losing its worker past the \
433 transport-loss budget"
434 );
435 } else {
436 tracing::warn!(
437 operation = "activity_complete",
438 workflow_id = %workflow_id,
439 activity_id = %activity_id,
440 worker_id = ?worker_id,
441 error_type = "WorkerLost",
442 losses = verdict.losses,
443 budget_ms = self.transport_losses.budget().as_millis(),
444 "worker lost before reporting an activity result; the activity never ran \
445 and will be re-dispatched attempt-neutrally"
446 );
447 }
448 verdict.reason
449 }
450 Err(error) => {
451 tracing::error!(
452 workflow_id = %workflow_id,
453 activity_id = %activity_id,
454 %error,
455 "transport-loss ledger is unreadable; abandoning the activity rather than \
456 re-dispatching it without a budget"
457 );
458 format!(
459 "{}{detail} (transport-loss budget unreadable: {error})",
460 super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
461 )
462 }
463 }
464 }
465
466 pub(crate) fn complete_activity_after_accept(
467 &self,
468 completion: ActivityCompletion,
469 after_accept: impl FnOnce() -> Result<(), ServerError>,
470 ) -> Result<(), ServerError> {
471 let result = match completion.outcome {
472 ActivityCompletionOutcome::Succeeded(payload) => {
473 payload_to_string(&payload).map_err(|reason| {
474 tracing::error!(
475 operation = "activity_complete",
476 workflow_id = %completion.workflow_id,
477 activity_id = %completion.activity_id,
478 error_type = "ActivityResultDecode",
479 %reason,
480 "activity completion failed"
481 );
482 ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
483 })?
484 }
485 ActivityCompletionOutcome::Failed(error) => {
486 let prefix = match error.kind {
487 ActivityErrorKind::Retryable => "retryable",
488 ActivityErrorKind::PolicyRefused => "policy_refused",
489 ActivityErrorKind::Terminal => "terminal",
490 };
491 tracing::error!(
492 operation = "activity_complete",
493 workflow_id = %completion.workflow_id,
494 activity_id = %completion.activity_id,
495 error_type = "ActivityFailed",
496 error_kind = prefix,
497 reason = %error.message,
498 "activity completion failed"
499 );
500 Err(format!("{prefix}:{}", error.message))
501 }
502 ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
503 &completion.workflow_id,
504 &completion.activity_id,
505 worker_id,
506 )),
507 };
508 let accepted_settlement = || after_accept().map(|()| true);
509 self.complete_fenced_after_accept(
510 &completion.workflow_id,
511 &completion.activity_id,
512 completion.run_id.as_ref(),
513 &completion.completion_token,
514 result,
515 accepted_settlement,
516 )?;
517 Ok(())
518 }
519}
520
521impl ActivityCompletionSink for PendingActivities {
522 fn park_activity(
533 &self,
534 workflow_id: &WorkflowId,
535 activity_id: &ActivityId,
536 ) -> Result<(), ServerError> {
537 self.completion_fences
541 .revoke_current(workflow_id, activity_id)?;
542 let matched = self
543 .pending
544 .remove(&(workflow_id.clone(), activity_id.clone()));
545 if let Some((_, sender)) = matched {
546 let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
550 }
551 Ok(())
552 }
553
554 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
555 self.complete_activity_after_accept(completion, || Ok(()))
556 }
557}
558
559fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
560 match payload.content_type() {
561 ContentType::Json => String::from_utf8(payload.bytes().to_vec())
562 .map(Ok)
563 .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
564 }
565}
566
567pub struct WorkerActivityDispatcher {
575 registry: ConnectedWorkerRegistry,
576 namespace: String,
577 pending: PendingActivities,
578 heartbeat_tracker: HeartbeatTracker,
579 drain_state: DrainState,
580 tokio_handle: Option<tokio::runtime::Handle>,
581 attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
587 queue_service: QueueServiceConfig,
591 queue_declarations: QueueDeclarationSource,
594 queue_state: QueueServiceState,
596 cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
600}
601
602impl std::fmt::Debug for WorkerActivityDispatcher {
603 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
604 f.debug_struct("WorkerActivityDispatcher")
605 .field("namespace", &self.namespace)
606 .finish_non_exhaustive()
607 }
608}
609
610impl WorkerActivityDispatcher {
611 #[must_use]
619 pub fn new(
620 registry: ConnectedWorkerRegistry,
621 namespace: impl Into<String>,
622 heartbeat_tracker: HeartbeatTracker,
623 ) -> Self {
624 Self {
625 registry,
626 namespace: namespace.into(),
627 pending: PendingActivities::new(heartbeat_tracker.heartbeat_window()),
631 heartbeat_tracker,
632 drain_state: DrainState::default(),
633 tokio_handle: None,
634 attempt_owners: None,
635 queue_service: QueueServiceConfig::default(),
636 queue_declarations: QueueDeclarationSource::default(),
637 queue_state: QueueServiceState::default(),
638 cluster_publisher: None,
639 }
640 }
641
642 #[must_use]
646 pub fn with_cluster_publisher(
647 mut self,
648 cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
649 ) -> Self {
650 self.cluster_publisher = Some(cluster_publisher);
651 self
652 }
653
654 #[must_use]
657 pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
658 self.queue_service = queue_service;
659 self
660 }
661
662 #[must_use]
666 pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
667 self.queue_declarations = queue_declarations;
668 self
669 }
670
671 #[must_use]
674 pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
675 self.queue_state = queue_state;
676 self
677 }
678
679 #[must_use]
686 pub fn with_attempt_owners(
687 mut self,
688 attempt_owners: super::intervention::AttemptOwnerIndex,
689 ) -> Self {
690 self.attempt_owners = Some(attempt_owners);
691 self
692 }
693
694 #[must_use]
699 pub const fn pending(&self) -> &PendingActivities {
700 &self.pending
701 }
702
703 #[must_use]
705 pub fn with_pending(mut self, pending: PendingActivities) -> Self {
706 self.pending = pending;
707 self
708 }
709
710 #[must_use]
712 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
713 self.drain_state = drain_state;
714 self
715 }
716
717 #[must_use]
719 pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
720 self.tokio_handle = Some(tokio_handle);
721 self
722 }
723}
724
725impl WorkerActivityDispatcher {
726 fn ensure_accepting(
750 &self,
751 namespace: &str,
752 activity_type: &str,
753 workflow_id: &WorkflowId,
754 activity_id: &ActivityId,
755 worker_id: Option<WorkerId>,
756 ) -> Result<(), String> {
757 self.drain_state
758 .ensure_accepting(namespace, activity_type)
759 .map_err(|error| {
760 log_worker_error(
761 "WorkerDispatch",
762 namespace,
763 activity_type,
764 workflow_id,
765 activity_id,
766 worker_id,
767 &error.to_string(),
768 );
769 aion::PARKED_ACTIVITY_REASON.to_owned()
770 })
771 }
772
773 fn select_worker_or_wait(
781 &self,
782 address: &ServiceAddress,
783 workflow_id: &WorkflowId,
784 activity_id: &ActivityId,
785 ) -> Result<WorkerHandle, String> {
786 let wait = ServiceWait {
787 registry: &self.registry,
788 declarations: &self.queue_declarations,
789 config: &self.queue_service,
790 state: &self.queue_state,
791 address,
792 workflow_id,
793 activity_id,
794 publisher: self.cluster_publisher.as_ref(),
795 };
796 let mut accepting = || {
797 self.ensure_accepting(
798 &address.namespace,
799 &address.activity_type,
800 workflow_id,
801 activity_id,
802 None,
803 )
804 };
805 let mut park = |arrival: WorkerArrival, budget: Option<Duration>| {
806 self.park_for_worker(arrival, budget);
807 };
808 select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
809 let reason = refusal.reason_string();
810 if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
813 let error_type = match refusal {
814 SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
815 _ => "WorkerRegistry",
816 };
817 log_worker_error(
818 error_type,
819 &address.namespace,
820 &address.activity_type,
821 workflow_id,
822 activity_id,
823 None,
824 &reason,
825 );
826 }
827 reason
828 })
829 }
830
831 fn park_for_worker(&self, arrival: WorkerArrival, budget: Option<Duration>) {
866 let handle = self
867 .tokio_handle
868 .clone()
869 .or_else(|| tokio::runtime::Handle::try_current().ok());
870 let Some(handle) = handle else {
871 std::thread::sleep(
883 budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
884 );
885 return;
886 };
887 handle.block_on(async {
888 let raced = async {
889 tokio::select! {
890 () = arrival => {}
891 () = self.drain_state.wait_for_drain() => {}
892 }
893 };
894 match budget {
895 None => raced.await,
896 Some(budget) => {
897 drop(tokio::time::timeout(budget, raced).await);
900 }
901 }
902 });
903 }
904
905 fn track_worker_task(
906 &self,
907 worker_id: WorkerId,
908 activity_type: &str,
909 workflow_id: &WorkflowId,
910 activity_id: &ActivityId,
911 attempt: u32,
912 completion_token: CompletionToken,
913 ) -> Result<(), String> {
914 self.heartbeat_tracker
915 .track_task(
916 worker_id,
917 InFlightActivity {
918 workflow_id: workflow_id.clone(),
919 activity_id: activity_id.clone(),
920 attempt,
925 completion_token,
926 },
927 Instant::now(),
928 )
929 .map_err(|error| {
930 let reason = error.to_string();
931 log_worker_error(
932 "WorkerHeartbeatTracker",
933 &self.namespace,
934 activity_type,
935 workflow_id,
936 activity_id,
937 Some(worker_id),
938 &reason,
939 );
940 reason
941 })
942 }
943
944 fn cleanup_activity(
945 &self,
946 worker_id: WorkerId,
947 workflow_id: &WorkflowId,
948 activity_id: &ActivityId,
949 completion_token: &CompletionToken,
950 ) {
951 self.pending
952 .pending
953 .remove(&(workflow_id.clone(), activity_id.clone()));
954 if let Err(error) =
958 self.pending
959 .completion_fences
960 .revoke(workflow_id, activity_id, completion_token)
961 {
962 tracing::error!(
963 workflow_id = %workflow_id,
964 activity_id = %activity_id,
965 %error,
966 "failed to revoke undelivered activity generation"
967 );
968 }
969 let _ = self
970 .heartbeat_tracker
971 .complete_task(worker_id, workflow_id, activity_id);
972 self.drain_state.notify_activity_drained();
973 }
974
975 fn send_activity_task(
985 &self,
986 worker: &WorkerHandle,
987 task: ProtoActivityTask,
988 address: &ServiceAddress,
989 workflow_id: &WorkflowId,
990 activity_id: &ActivityId,
991 completion_token: &CompletionToken,
992 ) -> Result<(), String> {
993 match worker.delivery() {
994 WorkerDelivery::Grpc(sender) => {
995 let worker_id = worker.id();
996 let mut accepting = || {
997 self.ensure_accepting(
998 &address.namespace,
999 &address.activity_type,
1000 workflow_id,
1001 activity_id,
1002 Some(worker_id),
1003 )
1004 };
1005 let handed_over = deliver_within_schedule_to_start(
1006 sender,
1007 WorkerMessage::ActivityTask(Box::new(task)),
1008 self.queue_service.schedule_to_start_timeout,
1009 &mut accepting,
1010 );
1011 let Err(refusal) = handed_over else {
1012 return Ok(());
1013 };
1014 self.cleanup_activity(worker_id, workflow_id, activity_id, completion_token);
1015 let (error_type, reason) = self.hand_off_failure(&refusal, address);
1016 if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
1017 log_worker_error(
1018 error_type,
1019 &address.namespace,
1020 &address.activity_type,
1021 workflow_id,
1022 activity_id,
1023 Some(worker_id),
1024 &reason,
1025 );
1026 }
1027 Err(reason)
1028 }
1029 #[cfg(feature = "liminal-transport")]
1036 WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
1037 worker.id(),
1038 delivery,
1039 task,
1040 &address.activity_type,
1041 workflow_id,
1042 activity_id,
1043 ),
1044 }
1045 }
1046
1047 fn hand_off_failure(
1054 &self,
1055 refusal: &DeliveryRefusal,
1056 address: &ServiceAddress,
1057 ) -> (&'static str, String) {
1058 match refusal {
1059 DeliveryRefusal::Saturated { waited } => {
1060 let census = self
1063 .registry
1064 .pool_census(
1065 &address.namespace,
1066 &address.task_queue,
1067 &address.activity_type,
1068 address.node.as_deref(),
1069 )
1070 .unwrap_or_else(|error| {
1071 tracing::error!(
1072 namespace = %address.namespace,
1073 task_queue = %address.task_queue,
1074 activity_type = %address.activity_type,
1075 %error,
1076 "poller census failed while reporting a saturated queue; \
1077 the refusal carries an empty census"
1078 );
1079 PoolCensus::default()
1080 });
1081 let unavailable = WorkerUnavailable {
1082 reason: QueueServiceReason::Saturated,
1083 clock: Some(ExpiredClock::ScheduleToStart),
1084 waited: *waited,
1085 address: address.clone(),
1086 census,
1087 };
1088 ("WorkerUnavailable", unavailable.reason_string())
1089 }
1090 DeliveryRefusal::Full => (
1091 "WorkerChannelClosed",
1092 "worker task channel full or closed: no available capacity".to_owned(),
1093 ),
1094 DeliveryRefusal::Closed => (
1095 "WorkerChannelClosed",
1096 "worker task channel full or closed: channel closed".to_owned(),
1097 ),
1098 DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
1099 }
1100 }
1101
1102 #[cfg(feature = "liminal-transport")]
1126 fn send_liminal_activity_task(
1127 &self,
1128 worker_id: WorkerId,
1129 delivery: &super::liminal_transport::LiminalWorkerDelivery,
1130 task: ProtoActivityTask,
1131 activity_type: &str,
1132 workflow_id: &WorkflowId,
1133 activity_id: &ActivityId,
1134 ) -> Result<(), String> {
1135 let completion_token =
1136 CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1137 .map_err(|error| error.to_string())?;
1138 let heartbeat_window_ms =
1139 u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1140 .unwrap_or(u64::MAX);
1141 let attempt = task.attempt;
1142 let run_id = task
1149 .run_id
1150 .map(RunId::try_from)
1151 .transpose()
1152 .map_err(|error| error.to_string())?
1153 .ok_or_else(|| {
1154 "activity task run id is missing; refusing to dispatch an unidentified run"
1155 .to_owned()
1156 })?;
1157 let request = super::liminal_transport::DispatchRequest {
1158 activity_type: activity_type.to_owned(),
1159 workflow_id: workflow_id.clone(),
1160 ordinal: activity_id.sequence_position(),
1161 run_id: Some(run_id.clone()),
1162 attempt,
1163 completion_token: task.completion_token,
1164 idempotency_key: task.idempotency_key,
1165 labels: task.labels.into_iter().collect(),
1166 heartbeat_window_ms,
1167 input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1168 };
1169 let awaiter = match delivery.push_dispatch(&request) {
1173 Ok(awaiter) => awaiter,
1174 Err(error) => {
1175 let reason = format!("worker liminal push failed: {error}");
1176 self.cleanup_activity(worker_id, workflow_id, activity_id, &completion_token);
1177 log_worker_error(
1178 "WorkerChannelClosed",
1179 &self.namespace,
1180 activity_type,
1181 workflow_id,
1182 activity_id,
1183 Some(worker_id),
1184 &reason,
1185 );
1186 return Err(reason);
1187 }
1188 };
1189 let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1193 super::liminal_transport::AttemptOwnerGuard::bind(
1194 owners.clone(),
1195 super::intervention::AttemptKey::new(
1196 workflow_id.clone(),
1197 run_id.clone(),
1198 activity_id.clone(),
1199 attempt,
1200 ),
1201 worker_id,
1202 )
1203 });
1204 self.spawn_liminal_reply_router(
1205 worker_id,
1206 awaiter,
1207 workflow_id,
1208 activity_id,
1209 &completion_token,
1210 owner_binding,
1211 );
1212 Ok(())
1213 }
1214
1215 #[cfg(feature = "liminal-transport")]
1255 fn spawn_liminal_reply_router(
1256 &self,
1257 worker_id: WorkerId,
1258 awaiter: liminal_server::server::connection::PushReplyAwaiter,
1259 workflow_id: &WorkflowId,
1260 activity_id: &ActivityId,
1261 completion_token: &CompletionToken,
1262 owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1263 ) {
1264 let pending = self.pending.clone();
1265 let heartbeat_tracker = self.heartbeat_tracker.clone();
1266 let drain_state = self.drain_state.clone();
1267 let workflow_id = workflow_id.clone();
1268 let activity_id = activity_id.clone();
1269 let completion_token = completion_token.clone();
1270 std::thread::spawn(move || {
1271 let _owner_binding = owner_binding;
1275 route_liminal_reply(
1276 &pending,
1277 &heartbeat_tracker,
1278 &drain_state,
1279 &awaiter,
1280 (worker_id, &workflow_id, &activity_id, &completion_token),
1281 );
1282 });
1283 }
1284
1285 fn await_activity_result(
1289 &self,
1290 context: &ActivityDispatchContext<'_>,
1291 rx: &SyncReceiver,
1292 ) -> Result<String, String> {
1293 match self.registry.is_registered(context.worker_id) {
1301 Ok(true) => {}
1302 Ok(false) => {
1303 if let Ok(result) = rx.try_recv() {
1307 return self.deliver_result(context, result);
1308 }
1309 self.cleanup_activity(
1310 context.worker_id,
1311 context.workflow_id,
1312 context.activity_id,
1313 &context.completion_token,
1314 );
1315 let reason = self.pending.classify_worker_loss(
1320 context.workflow_id,
1321 context.activity_id,
1322 context.worker_id,
1323 );
1324 log_worker_error(
1325 "WorkerLost",
1326 &self.namespace,
1327 context.activity_type,
1328 context.workflow_id,
1329 context.activity_id,
1330 Some(context.worker_id),
1331 &reason,
1332 );
1333 return Err(reason);
1334 }
1335 Err(error) => {
1336 self.cleanup_activity(
1337 context.worker_id,
1338 context.workflow_id,
1339 context.activity_id,
1340 &context.completion_token,
1341 );
1342 let reason = format!("worker registry inspection failed: {error}");
1343 log_worker_error(
1344 "WorkerRegistry",
1345 &self.namespace,
1346 context.activity_type,
1347 context.workflow_id,
1348 context.activity_id,
1349 Some(context.worker_id),
1350 &reason,
1351 );
1352 return Err(reason);
1353 }
1354 }
1355 if let Ok(result) = rx.recv() {
1356 return self.deliver_result(context, result);
1357 }
1358 self.cleanup_activity(
1361 context.worker_id,
1362 context.workflow_id,
1363 context.activity_id,
1364 &context.completion_token,
1365 );
1366 let reason = "activity response channel dropped".to_owned();
1367 log_worker_error(
1368 "WorkerChannelClosed",
1369 &self.namespace,
1370 context.activity_type,
1371 context.workflow_id,
1372 context.activity_id,
1373 Some(context.worker_id),
1374 &reason,
1375 );
1376 Err(reason)
1377 }
1378
1379 fn deliver_result(
1380 &self,
1381 context: &ActivityDispatchContext<'_>,
1382 result: Result<String, String>,
1383 ) -> Result<String, String> {
1384 self.pending
1385 .pending
1386 .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
1387 if let Err(reason) = &result
1392 && aion::is_parked_reason(reason)
1393 {
1394 tracing::info!(
1395 operation = "activity_dispatch",
1396 namespace = %self.namespace,
1397 workflow_id = %context.workflow_id,
1398 activity_id = %context.activity_id,
1399 activity_type = context.activity_type,
1400 worker_id = ?context.worker_id,
1401 "activity parked for restart recovery"
1402 );
1403 return result;
1404 }
1405 log_activity_completion(context, result.is_ok());
1406 result.inspect_err(|reason| {
1407 log_worker_error(
1408 "ActivityFailed",
1409 &self.namespace,
1410 context.activity_type,
1411 context.workflow_id,
1412 context.activity_id,
1413 Some(context.worker_id),
1414 reason,
1415 );
1416 })
1417 }
1418}
1419
1420impl ActivityDispatcher for WorkerActivityDispatcher {
1421 fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1422 match tokio::runtime::Handle::try_current() {
1423 Ok(handle) => match handle.runtime_flavor() {
1424 tokio::runtime::RuntimeFlavor::MultiThread => {
1425 tokio::task::block_in_place(|| self.dispatch_blocking(request))
1432 }
1433 flavor => Err(format!(
1434 "activity dispatch blocks the calling thread until the worker responds; \
1435 a {flavor:?} tokio runtime cannot host that wait because the worker \
1436 stream forwarder shares its only executor thread and the task could \
1437 never be delivered — run the engine on a multi-thread tokio runtime"
1438 )),
1439 },
1440 Err(_) => self.dispatch_blocking(request),
1444 }
1445 }
1446}
1447
1448impl WorkerActivityDispatcher {
1449 fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1466 let ActivityDispatch {
1467 namespace,
1468 task_queue,
1469 node,
1473 workflow_id,
1474 run_id,
1475 activity_id,
1476 name,
1477 input,
1478 config: _,
1479 attempt,
1480 labels,
1481 advisory: _,
1486 } = request;
1487 let started_at = Instant::now();
1488 self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1489 let address = ServiceAddress {
1490 namespace: namespace.clone(),
1491 task_queue: task_queue.clone(),
1492 activity_type: name.clone(),
1493 node: node.clone(),
1494 };
1495 let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1496 let worker_id = worker.id();
1497 let span = info_span!(
1498 "activity_dispatch",
1499 operation = "activity_dispatch",
1500 namespace = %namespace,
1501 task_queue = %task_queue,
1502 node = node.as_deref(),
1503 workflow_id = %workflow_id,
1504 activity_id = %activity_id,
1505 activity_type = %name,
1506 worker_id = ?worker_id,
1507 );
1508 let _span_guard = span.enter();
1509 self.ensure_accepting(
1510 &namespace,
1511 &name,
1512 &workflow_id,
1513 &activity_id,
1514 Some(worker_id),
1515 )?;
1516
1517 let (completion_token, rx) = self
1518 .pending
1519 .insert(workflow_id.clone(), &run_id, activity_id.clone(), attempt)
1520 .map_err(|error| error.to_string())?;
1521 let task = activity_task(
1522 &name,
1523 &input,
1524 (&workflow_id, &run_id, &activity_id),
1525 attempt,
1526 labels,
1527 &completion_token,
1528 );
1529 if let Err(error) = self.track_worker_task(
1530 worker_id,
1531 &name,
1532 &workflow_id,
1533 &activity_id,
1534 attempt,
1535 completion_token.clone(),
1536 ) {
1537 self.cleanup_activity(worker_id, &workflow_id, &activity_id, &completion_token);
1538 return Err(error);
1539 }
1540 self.send_activity_task(
1541 &worker,
1542 task,
1543 &address,
1544 &workflow_id,
1545 &activity_id,
1546 &completion_token,
1547 )?;
1548 let context = ActivityDispatchContext {
1549 namespace: &namespace,
1550 activity_type: &name,
1551 worker_id,
1552 workflow_id: &workflow_id,
1553 activity_id: &activity_id,
1554 completion_token,
1555 started_at,
1556 };
1557 self.await_activity_result(&context, &rx)
1558 }
1559}
1560
1561#[cfg(feature = "liminal-transport")]
1568fn route_liminal_reply(
1569 pending: &PendingActivities,
1570 heartbeat_tracker: &HeartbeatTracker,
1571 drain_state: &DrainState,
1572 awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1573 execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1574) {
1575 let (worker_id, workflow_id, activity_id, current_token) = execution;
1576 let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1581 heartbeat_tracker
1582 .is_tracked(worker_id, workflow_id, activity_id)
1583 .unwrap_or(false)
1584 });
1585 let (run_id, submitted_token, outcome, synthesized) = match waited {
1589 Ok(Some(response)) => {
1590 let submitted_token = match CompletionToken::from_wire(
1591 workflow_id,
1592 activity_id,
1593 response.completion_token,
1594 ) {
1595 Ok(token) => token,
1596 Err(error) => {
1597 tracing::warn!(
1598 worker_id = ?worker_id,
1599 workflow_id = %workflow_id,
1600 activity_id = %activity_id,
1601 %error,
1602 "liminal activity completion omitted its generation proof"
1603 );
1604 return;
1605 }
1606 };
1607 (response.run_id, submitted_token, response.outcome, false)
1608 }
1609 Ok(None) => {
1610 tracing::debug!(
1611 worker_id = ?worker_id,
1612 workflow_id = %workflow_id,
1613 activity_id = %activity_id,
1614 "liminal dispatch resolved by another path; abandoning reply wait"
1615 );
1616 return;
1617 }
1618 Err(error) if error.is_worker_connection_lost() => (
1619 None,
1620 current_token.clone(),
1621 Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1625 true,
1626 ),
1627 Err(error) => (
1628 None,
1629 current_token.clone(),
1630 Err(format!("retryable:worker liminal reply failed: {error}")),
1631 true,
1632 ),
1633 };
1634 let after_accept = || {
1639 if synthesized {
1640 let was_tracked = clear_completed_task_tracking(
1641 heartbeat_tracker,
1642 worker_id,
1643 workflow_id,
1644 activity_id,
1645 );
1646 if !was_tracked {
1647 tracing::debug!(
1650 worker_id = ?worker_id,
1651 workflow_id = %workflow_id,
1652 activity_id = %activity_id,
1653 "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1654 );
1655 return Ok(false);
1656 }
1657 } else {
1658 let _ = clear_completed_task_tracking(
1664 heartbeat_tracker,
1665 worker_id,
1666 workflow_id,
1667 activity_id,
1668 );
1669 }
1670 drain_state.notify_activity_drained();
1671 Ok(true)
1672 };
1673 if let Err(error) = pending.complete_fenced_after_accept(
1674 workflow_id,
1675 activity_id,
1676 run_id.as_ref(),
1677 &submitted_token,
1678 outcome,
1679 after_accept,
1680 ) {
1681 tracing::warn!(
1682 worker_id = ?worker_id,
1683 workflow_id = %workflow_id,
1684 activity_id = %activity_id,
1685 %error,
1686 "liminal activity completion handoff rejected"
1687 );
1688 }
1689}
1690
1691pub(crate) fn clear_completed_task_tracking(
1698 heartbeat_tracker: &HeartbeatTracker,
1699 worker_id: WorkerId,
1700 workflow_id: &WorkflowId,
1701 activity_id: &ActivityId,
1702) -> bool {
1703 heartbeat_tracker
1704 .complete_task(worker_id, workflow_id, activity_id)
1705 .unwrap_or_else(|error| {
1706 tracing::error!(
1707 worker_id = ?worker_id,
1708 workflow_id = %workflow_id,
1709 activity_id = %activity_id,
1710 %error,
1711 "failed to clear in-flight tracking for a completed activity; delivering anyway — a poisoned tracker must not withhold a result"
1712 );
1713 true
1714 })
1715}
1716
1717struct ActivityDispatchContext<'a> {
1718 namespace: &'a str,
1719 activity_type: &'a str,
1720 worker_id: WorkerId,
1721 workflow_id: &'a WorkflowId,
1722 activity_id: &'a ActivityId,
1723 completion_token: CompletionToken,
1724 started_at: Instant,
1725}
1726
1727fn activity_task(
1728 activity_type: &str,
1729 input: &str,
1730 execution: (&WorkflowId, &RunId, &ActivityId),
1731 attempt: u32,
1732 labels: BTreeMap<String, String>,
1733 completion_token: &CompletionToken,
1734) -> ProtoActivityTask {
1735 let (workflow_id, run_id, activity_id) = execution;
1736 ProtoActivityTask {
1737 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1738 activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1739 activity_type: activity_type.to_owned(),
1740 input: Some(ProtoPayload {
1741 content_type: String::from("application/json"),
1742 bytes: input.as_bytes().to_vec(),
1743 }),
1744 attempt,
1745 labels: labels.into_iter().collect(),
1746 run_id: Some(run_id.clone().into()),
1747 completion_token: completion_token.as_str().to_owned(),
1748 idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
1749 }
1750}
1751
1752fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1753 let duration_ms = duration_ms(context.started_at.elapsed());
1754 tracing::info!(
1755 operation = "activity_complete",
1756 namespace = context.namespace,
1757 workflow_id = %context.workflow_id,
1758 activity_id = %context.activity_id,
1759 activity_type = context.activity_type,
1760 worker_id = ?context.worker_id,
1761 duration_ms,
1762 outcome = if succeeded { "succeeded" } else { "failed" },
1763 "activity completed"
1764 );
1765}
1766
1767fn duration_ms(duration: Duration) -> u64 {
1768 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1769}
1770
1771fn log_worker_error(
1772 error_type: &'static str,
1773 namespace: &str,
1774 activity_type: &str,
1775 workflow_id: &WorkflowId,
1776 activity_id: &ActivityId,
1777 worker_id: Option<super::registry::WorkerId>,
1778 reason: &str,
1779) {
1780 tracing::error!(
1781 operation = "activity_dispatch",
1782 namespace,
1783 workflow_id = %workflow_id,
1784 activity_id = %activity_id,
1785 activity_type,
1786 worker_id = ?worker_id,
1787 error_type,
1788 reason,
1789 "worker interaction failed"
1790 );
1791}
1792
1793#[cfg(test)]
1794mod tests {
1795 use std::sync::{
1796 Mutex,
1797 atomic::{AtomicBool, Ordering},
1798 };
1799
1800 use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1801
1802 use super::*;
1803
1804 fn activity_id(pos: u64) -> ActivityId {
1805 ActivityId::from_sequence_position(pos)
1806 }
1807
1808 fn insert_attempt_one(
1814 pending: &PendingActivities,
1815 workflow_id: &WorkflowId,
1816 id: &ActivityId,
1817 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
1818 pending.insert(workflow_id.clone(), &RunId::new_v4(), id.clone(), 1)
1819 }
1820
1821 #[test]
1822 fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
1823 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1824 let workflow_id = WorkflowId::new_v4();
1825 let id = activity_id(1);
1826 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
1827
1828 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1829 assert_eq!(
1830 rx.recv_timeout(Duration::from_millis(50)),
1831 Ok(Ok("done".to_owned()))
1832 );
1833 Ok(())
1834 }
1835
1836 #[test]
1841 fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
1842 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1843 let workflow_id = WorkflowId::new_v4();
1844 let id = activity_id(2);
1845 let run_id = RunId::new_v4();
1849 let _first = pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
1850
1851 let error = pending
1852 .insert(workflow_id, &run_id, id, 1)
1853 .err()
1854 .ok_or_else(|| {
1855 ServerError::worker_dispatch(
1856 "test",
1857 "same-worker-collision",
1858 "a live responder collision was silently replaced",
1859 )
1860 })?;
1861
1862 assert!(
1863 error.to_string().contains("pending activity collision"),
1864 "{error}"
1865 );
1866 Ok(())
1867 }
1868
1869 #[test]
1870 fn pending_complete_unknown_returns_false() {
1871 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1872 assert!(!pending.complete(
1873 &WorkflowId::new_v4(),
1874 &activity_id(99),
1875 None,
1876 Ok("orphan".to_owned())
1877 ));
1878 }
1879
1880 #[derive(Default)]
1881 struct RecordingOutboxCallback {
1882 completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1883 failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1884 live: bool,
1885 }
1886
1887 impl OutboxDeliveryCallback for RecordingOutboxCallback {
1888 fn deliver_completion(
1889 &self,
1890 workflow_id: &WorkflowId,
1891 activity_id: &ActivityId,
1892 run_id: Option<&RunId>,
1893 result: String,
1894 ) -> Result<bool, ServerError> {
1895 let _ = run_id;
1896 self.completions
1897 .lock()
1898 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1899 .push((workflow_id.clone(), activity_id.clone(), result));
1900 Ok(self.live)
1901 }
1902
1903 fn deliver_failure(
1904 &self,
1905 workflow_id: &WorkflowId,
1906 activity_id: &ActivityId,
1907 run_id: Option<&RunId>,
1908 reason: String,
1909 ) -> Result<bool, ServerError> {
1910 let _ = run_id;
1911 self.failures
1912 .lock()
1913 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1914 .push((workflow_id.clone(), activity_id.clone(), reason));
1915 Ok(self.live)
1916 }
1917 }
1918
1919 #[test]
1920 fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1921 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1922 let callback = Arc::new(RecordingOutboxCallback {
1923 live: true,
1924 ..RecordingOutboxCallback::default()
1925 });
1926 pending.clone().set_outbox_delivery(callback.clone());
1928
1929 let workflow_id = WorkflowId::new_v4();
1930 let id = activity_id(7);
1931
1932 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1935 let completions = callback
1936 .completions
1937 .lock()
1938 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1939 assert_eq!(completions.len(), 1);
1940 assert_eq!(completions[0].0, workflow_id);
1941 assert_eq!(completions[0].1, id);
1942 assert_eq!(completions[0].2, "done");
1943 Ok(())
1944 }
1945
1946 #[test]
1947 fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1948 -> Result<(), ServerError> {
1949 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1950 let callback = Arc::new(RecordingOutboxCallback::default());
1952 pending.set_outbox_delivery(callback.clone());
1953
1954 let workflow_id = WorkflowId::new_v4();
1955 let id = activity_id(8);
1956
1957 assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1958 let failures = callback
1959 .failures
1960 .lock()
1961 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1962 assert_eq!(failures.len(), 1);
1963 assert_eq!(failures[0].2, "retryable:boom");
1964 Ok(())
1965 }
1966
1967 #[test]
1968 fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1969 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1972 assert!(!pending.complete(
1973 &WorkflowId::new_v4(),
1974 &activity_id(9),
1975 None,
1976 Ok("x".to_owned())
1977 ));
1978 }
1979
1980 #[test]
1981 fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1982 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1983 let callback = Arc::new(RecordingOutboxCallback {
1984 live: true,
1985 ..RecordingOutboxCallback::default()
1986 });
1987 pending.set_outbox_delivery(callback.clone());
1988
1989 let workflow_id = WorkflowId::new_v4();
1990 let id = activity_id(10);
1991 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
1992
1993 assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1994 assert_eq!(
1995 rx.recv_timeout(Duration::from_millis(50)),
1996 Ok(Ok("matched".to_owned()))
1997 );
1998 assert!(
1999 callback
2000 .completions
2001 .lock()
2002 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2003 .is_empty(),
2004 "a matched completion must deliver to its waiter, not the outbox callback"
2005 );
2006 Ok(())
2007 }
2008
2009 #[test]
2013 fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
2014 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2015 let workflow_id = WorkflowId::new_v4();
2016 let id = activity_id(11);
2017 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
2018
2019 pending.park_activity(&workflow_id, &id)?;
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, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
2024 Ok(())
2025 }
2026
2027 #[test]
2031 fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
2032 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2033 let callback = Arc::new(RecordingOutboxCallback {
2034 live: true,
2035 ..RecordingOutboxCallback::default()
2036 });
2037 pending.set_outbox_delivery(callback.clone());
2038
2039 pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
2040
2041 assert!(
2042 callback
2043 .failures
2044 .lock()
2045 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2046 .is_empty(),
2047 "a park must never be delivered as an outbox failure"
2048 );
2049 assert!(
2050 callback
2051 .completions
2052 .lock()
2053 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2054 .is_empty(),
2055 "a park must never be delivered as an outbox completion"
2056 );
2057 Ok(())
2058 }
2059
2060 #[test]
2061 fn completion_sink_routes_success() -> Result<(), ServerError> {
2062 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2063 let workflow_id = WorkflowId::new_v4();
2064 let id = activity_id(2);
2065 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2066 let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
2067
2068 pending.complete_activity(ActivityCompletion {
2069 workflow_id,
2070 activity_id: id,
2071 run_id: None,
2072 completion_token,
2073 outcome: ActivityCompletionOutcome::Succeeded(payload),
2074 })?;
2075
2076 let result = rx
2077 .recv_timeout(Duration::from_millis(50))
2078 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2079 assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
2080 Ok(())
2081 }
2082
2083 #[test]
2084 fn accepted_hook_runs_before_the_result_is_observable() -> Result<(), ServerError> {
2085 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2086 let workflow_id = WorkflowId::new_v4();
2087 let id = activity_id(13);
2088 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2089 let accepted = AtomicBool::new(false);
2090
2091 pending.complete_activity_after_accept(
2092 ActivityCompletion {
2093 workflow_id,
2094 activity_id: id,
2095 run_id: None,
2096 completion_token,
2097 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2098 ContentType::Json,
2099 br#""ordered""#.to_vec(),
2100 )),
2101 },
2102 || {
2103 accepted.store(true, Ordering::SeqCst);
2104 Ok(())
2105 },
2106 )?;
2107
2108 let result = rx
2109 .recv_timeout(Duration::from_millis(50))
2110 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2111 assert_eq!(result, Ok(r#""ordered""#.to_owned()));
2112 assert!(
2113 accepted.load(Ordering::SeqCst),
2114 "the accepted-path cleanup hook must finish before the waiter can observe the result"
2115 );
2116 Ok(())
2117 }
2118
2119 #[test]
2120 fn a_non_publishing_settlement_restores_the_generation_for_the_true_resolver()
2121 -> Result<(), ServerError> {
2122 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2123 let workflow_id = WorkflowId::new_v4();
2124 let id = activity_id(16);
2125 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2126
2127 let settled = pending.complete_fenced_after_accept(
2132 &workflow_id,
2133 &id,
2134 None,
2135 &completion_token,
2136 Err("retryable:worker lost".to_owned()),
2137 || Ok(false),
2138 )?;
2139 assert!(!settled, "a non-publishing settlement must report false");
2140 assert!(
2141 rx.try_recv().is_err(),
2142 "a non-publishing settlement must not publish"
2143 );
2144
2145 let published = pending.complete_fenced_after_accept(
2148 &workflow_id,
2149 &id,
2150 None,
2151 &completion_token,
2152 Err("retryable:worker lost".to_owned()),
2153 || Ok(true),
2154 )?;
2155 assert!(
2156 published,
2157 "the true resolver must be accepted after a non-publishing settlement"
2158 );
2159 assert_eq!(
2160 rx.recv_timeout(Duration::from_millis(50))
2161 .map_err(|error| ServerError::worker_dispatch(
2162 "",
2163 "",
2164 format!("channel: {error}")
2165 ))?,
2166 Err("retryable:worker lost".to_owned())
2167 );
2168 Ok(())
2169 }
2170
2171 #[test]
2172 fn a_poisoned_tracker_fails_open_and_never_withholds_a_completion() {
2173 let tracker = HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW);
2174 tracker.poison_for_tests();
2175 let workflow_id = WorkflowId::new_v4();
2176 let id = activity_id(17);
2177 let was_tracked =
2178 clear_completed_task_tracking(&tracker, WorkerId::from_value(1), &workflow_id, &id);
2179 assert!(
2180 was_tracked,
2181 "a poisoned tracker must fail open (deliver), never withhold"
2182 );
2183 }
2184
2185 #[test]
2186 fn accepted_hook_failure_blocks_publication_and_restores_generation() -> Result<(), ServerError>
2187 {
2188 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2189 let workflow_id = WorkflowId::new_v4();
2190 let id = activity_id(15);
2191 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2192
2193 let rejected = pending.complete_activity_after_accept(
2194 ActivityCompletion {
2195 workflow_id: workflow_id.clone(),
2196 activity_id: id.clone(),
2197 run_id: None,
2198 completion_token: completion_token.clone(),
2199 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2200 ContentType::Json,
2201 br#""ordered""#.to_vec(),
2202 )),
2203 },
2204 || Err(ServerError::lock_poisoned("accepted-path settlement")),
2205 );
2206
2207 assert!(matches!(rejected, Err(ServerError::LockPoisoned { .. })));
2208 assert!(
2209 rx.try_recv().is_err(),
2210 "failed accepted-path settlement must not publish the completion"
2211 );
2212
2213 pending.complete_activity_after_accept(
2214 ActivityCompletion {
2215 workflow_id,
2216 activity_id: id,
2217 run_id: None,
2218 completion_token,
2219 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2220 ContentType::Json,
2221 br#""ordered""#.to_vec(),
2222 )),
2223 },
2224 || Ok(()),
2225 )?;
2226 assert_eq!(
2227 rx.recv_timeout(Duration::from_millis(50))
2228 .map_err(|error| ServerError::worker_dispatch(
2229 "",
2230 "",
2231 format!("channel: {error}")
2232 ))?,
2233 Ok(r#""ordered""#.to_owned())
2234 );
2235 Ok(())
2236 }
2237
2238 #[test]
2239 fn refused_generation_never_runs_the_accepted_hook() -> Result<(), ServerError> {
2240 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2241 let workflow_id = WorkflowId::new_v4();
2242 let id = activity_id(14);
2243 let (_completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2244 let accepted = AtomicBool::new(false);
2245
2246 let rejected = pending.complete_activity_after_accept(
2247 ActivityCompletion {
2248 workflow_id,
2249 activity_id: id,
2250 run_id: None,
2251 completion_token: CompletionToken::for_test(),
2252 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2253 ContentType::Json,
2254 br#""stale""#.to_vec(),
2255 )),
2256 },
2257 || {
2258 accepted.store(true, Ordering::SeqCst);
2259 Ok(())
2260 },
2261 );
2262
2263 assert!(matches!(
2264 rejected,
2265 Err(ServerError::ActivityCompletionRejected { .. })
2266 ));
2267 assert!(
2268 !accepted.load(Ordering::SeqCst),
2269 "a refused generation must preserve liveness by skipping accepted-path cleanup"
2270 );
2271 assert!(
2272 rx.try_recv().is_err(),
2273 "a refused generation must leave the current waiter unresolved"
2274 );
2275 Ok(())
2276 }
2277
2278 #[test]
2279 fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
2280 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2281 let workflow_id = WorkflowId::new_v4();
2282 let id = activity_id(12);
2283 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2284
2285 let malformed = pending.complete_activity(ActivityCompletion {
2286 workflow_id: workflow_id.clone(),
2287 activity_id: id.clone(),
2288 run_id: None,
2289 completion_token: completion_token.clone(),
2290 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2291 ContentType::Json,
2292 vec![0xff],
2293 )),
2294 });
2295 assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
2296 assert!(
2297 rx.try_recv().is_err(),
2298 "an invalid result must leave the waiter unresolved"
2299 );
2300
2301 pending.complete_activity(ActivityCompletion {
2302 workflow_id,
2303 activity_id: id,
2304 run_id: None,
2305 completion_token,
2306 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2307 ContentType::Json,
2308 br#""valid""#.to_vec(),
2309 )),
2310 })?;
2311 let result = rx
2312 .recv_timeout(Duration::from_millis(50))
2313 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2314 assert_eq!(result, Ok(r#""valid""#.to_owned()));
2315 Ok(())
2316 }
2317
2318 #[test]
2319 fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
2320 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2321 let workflow_id = WorkflowId::new_v4();
2322 let id = activity_id(3);
2323 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2324
2325 pending.complete_activity(ActivityCompletion {
2326 workflow_id,
2327 activity_id: id,
2328 run_id: None,
2329 completion_token,
2330 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2331 kind: ActivityErrorKind::Retryable,
2332 message: "temporary".to_owned(),
2333 details: None,
2334 }),
2335 })?;
2336
2337 let result = rx
2338 .recv_timeout(Duration::from_millis(50))
2339 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2340 assert_eq!(result, Err("retryable:temporary".to_owned()));
2341 Ok(())
2342 }
2343
2344 #[test]
2345 fn completion_sink_routes_policy_refusal_with_its_own_prefix() -> Result<(), ServerError> {
2346 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2347 let workflow_id = WorkflowId::new_v4();
2348 let id = activity_id(4);
2349 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2350
2351 pending.complete_activity(ActivityCompletion {
2352 workflow_id,
2353 activity_id: id,
2354 run_id: None,
2355 completion_token,
2356 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2357 kind: ActivityErrorKind::PolicyRefused,
2358 message: "provider safety policy".to_owned(),
2359 details: None,
2360 }),
2361 })?;
2362
2363 let result = rx
2364 .recv_timeout(Duration::from_millis(50))
2365 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2366 assert_eq!(
2367 result,
2368 Err("policy_refused:provider safety policy".to_owned())
2369 );
2370 Ok(())
2371 }
2372
2373 #[test]
2382 fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
2383 -> Result<(), ServerError> {
2384 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2385 let post_restart_workflow = WorkflowId::new_v4();
2386 let pre_restart_workflow = WorkflowId::new_v4();
2387 let id = activity_id(1);
2389 let (completion_token, rx) = insert_attempt_one(&pending, &post_restart_workflow, &id)?;
2390
2391 let rejected = pending.complete_activity(ActivityCompletion {
2393 workflow_id: pre_restart_workflow,
2394 activity_id: id.clone(),
2395 run_id: None,
2396 completion_token: CompletionToken::for_test(),
2397 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2398 ContentType::Json,
2399 br#""stale""#.to_vec(),
2400 )),
2401 });
2402 assert!(matches!(
2403 rejected,
2404 Err(ServerError::ActivityCompletionRejected { .. })
2405 ));
2406 assert!(
2407 rx.try_recv().is_err(),
2408 "stale result for a different workflow must not complete this dispatch"
2409 );
2410
2411 pending.complete_activity(ActivityCompletion {
2413 workflow_id: post_restart_workflow,
2414 activity_id: id,
2415 run_id: None,
2416 completion_token,
2417 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2418 ContentType::Json,
2419 br#""fresh""#.to_vec(),
2420 )),
2421 })?;
2422 let result = rx
2423 .recv_timeout(Duration::from_millis(50))
2424 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2425 assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2426 Ok(())
2427 }
2428
2429 const TEST_HEARTBEAT_WINDOW: Duration = Duration::from_secs(5);
2436
2437 #[test]
2446 fn a_sink_derives_its_transport_budget_from_the_operator_window() {
2447 let window = Duration::from_secs(7);
2448 let sink = PendingActivities::new(window);
2449
2450 assert_eq!(
2451 sink.transport_losses().budget(),
2452 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2453 "the transport-loss budget must be the operator's heartbeat window times \
2454 TRANSPORT_LOSS_BUDGET_WINDOWS; a sink carrying any other budget is one no \
2455 operator declared"
2456 );
2457 }
2458
2459 #[test]
2467 fn a_dispatcher_carries_the_transport_budget_of_the_tracker_it_was_given() {
2468 let window = Duration::from_secs(11);
2469 let registry = ConnectedWorkerRegistry::default();
2470 let dispatcher =
2471 WorkerActivityDispatcher::new(registry, "default", HeartbeatTracker::new(window));
2472
2473 assert_eq!(
2474 dispatcher.pending().transport_losses().budget(),
2475 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2476 "the dispatcher's transport budget must be derived from the heartbeat window it \
2477 was constructed with; a zero or mismatched budget here means a worker loss is \
2478 declared transport-exhausted after ONE re-dispatch instead of after the \
2479 operator's window"
2480 );
2481 assert_ne!(
2482 dispatcher.pending().transport_losses().budget(),
2483 Duration::ZERO,
2484 "a ZERO budget is the specific regression this pin exists for: it grants exactly \
2485 one re-dispatchable loss and then declares the infrastructure flapping"
2486 );
2487 }
2488
2489 fn test_tracker() -> HeartbeatTracker {
2490 HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW)
2491 }
2492
2493 fn greet_request() -> ActivityDispatch {
2496 ActivityDispatch {
2497 namespace: "default".to_owned(),
2498 task_queue: "default".to_owned(),
2499 node: None,
2500 workflow_id: WorkflowId::new_v4(),
2501 run_id: RunId::new_v4(),
2502 activity_id: ActivityId::from_sequence_position(0),
2503 name: "greet".to_owned(),
2504 input: "{}".to_owned(),
2505 config: "{}".to_owned(),
2506 attempt: 1,
2507 labels: std::collections::BTreeMap::new(),
2508 advisory: false,
2509 }
2510 }
2511
2512 #[test]
2513 fn dispatcher_fails_immediately_when_draining_without_workers() {
2514 let registry = ConnectedWorkerRegistry::default();
2515 let drain = DrainState::default();
2516 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2517 .with_drain_state(drain.clone());
2518
2519 let _ = drain.begin();
2520
2521 let result = dispatcher.dispatch(greet_request());
2522
2523 assert!(result.is_err());
2524 let err = result.err().unwrap_or_default();
2525 assert!(
2526 err.contains("drain"),
2527 "expected drain rejection, got: {err}"
2528 );
2529 }
2530
2531 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2547 async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
2548 -> Result<(), Box<dyn std::error::Error>> {
2549 let registry = ConnectedWorkerRegistry::default();
2550 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2551 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2552 let activity_types = [String::from("greet")];
2553 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2554
2555 let sink = pending.clone();
2556 let echo_worker = tokio::spawn(async move {
2557 let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
2558 return Err("expected an activity task on the worker channel".to_owned());
2559 };
2560 let workflow_id = task
2561 .workflow_id
2562 .ok_or("task missing workflow id")
2563 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2564 let activity_id = task
2565 .activity_id
2566 .map(ActivityId::from)
2567 .ok_or("task missing activity id")?;
2568 let completion_token =
2569 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2570 .map_err(|error| error.to_string())?;
2571 sink.complete_activity(ActivityCompletion {
2572 workflow_id,
2573 activity_id,
2574 run_id: None,
2575 completion_token,
2576 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2577 ContentType::Json,
2578 br#"{"greeting":"hello"}"#.to_vec(),
2579 )),
2580 })
2581 .map_err(|error| error.to_string())
2582 });
2583
2584 let dispatcher = Arc::new(
2585 WorkerActivityDispatcher::new(registry, "default", test_tracker())
2586 .with_pending(pending),
2587 );
2588 let started = Instant::now();
2589 let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
2592 dispatcher.dispatch(greet_request())
2593 }));
2594 let result = dispatch_task.await.map_err(|error| error.to_string())?;
2595 let elapsed = started.elapsed();
2596
2597 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2598 assert!(
2599 elapsed < Duration::from_secs(5),
2600 "dispatch round trip took {elapsed:?}; task delivery must not \
2601 depend on the blocked dispatch thread"
2602 );
2603 echo_worker.await.map_err(|error| error.to_string())??;
2604 registration.deregister()?;
2605 Ok(())
2606 }
2607
2608 #[tokio::test]
2612 async fn dispatch_on_current_thread_runtime_fails_fast()
2613 -> Result<(), Box<dyn std::error::Error>> {
2614 let registry = ConnectedWorkerRegistry::default();
2615 let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
2616 let activity_types = [String::from("greet")];
2617 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2618 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
2619
2620 let started = Instant::now();
2621 let result = dispatcher.dispatch(greet_request());
2622 let elapsed = started.elapsed();
2623
2624 let err = result.err().ok_or("expected dispatch to fail")?;
2625 assert!(
2626 err.contains("multi-thread tokio runtime"),
2627 "unexpected error: {err}"
2628 );
2629 assert!(
2630 elapsed < Duration::from_secs(5),
2631 "fail-fast path took {elapsed:?}"
2632 );
2633 registration.deregister()?;
2634 Ok(())
2635 }
2636
2637 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2643 async fn dispatch_pinned_to_node_reaches_only_that_node()
2644 -> Result<(), Box<dyn std::error::Error>> {
2645 let registry = ConnectedWorkerRegistry::default();
2646 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2647 let activity_types = [String::from("greet")];
2648 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2649 let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
2650 let on_n2 = registry.register_namespaces(
2659 [String::from("default")],
2660 "default",
2661 Some(String::from("n2")),
2662 activity_types.iter(),
2663 n2_tx,
2664 )?;
2665 let on_n1 = registry.register_namespaces(
2666 [String::from("default")],
2667 "default",
2668 Some(String::from("n1")),
2669 activity_types.iter(),
2670 n1_tx,
2671 )?;
2672
2673 let sink = pending.clone();
2677 let echo_n1 = tokio::spawn(async move {
2678 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2679 return Err("expected an activity task on the n1 worker channel".to_owned());
2680 };
2681 let workflow_id = task
2682 .workflow_id
2683 .ok_or("task missing workflow id")
2684 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2685 let activity_id = task
2686 .activity_id
2687 .map(ActivityId::from)
2688 .ok_or("task missing activity id")?;
2689 let completion_token =
2690 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2691 .map_err(|error| error.to_string())?;
2692 sink.complete_activity(ActivityCompletion {
2693 workflow_id,
2694 activity_id,
2695 run_id: None,
2696 completion_token,
2697 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2698 ContentType::Json,
2699 br#"{"greeting":"hello"}"#.to_vec(),
2700 )),
2701 })
2702 .map_err(|error| error.to_string())
2703 });
2704
2705 let dispatcher = Arc::new(
2706 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2707 .with_pending(pending),
2708 );
2709
2710 let pinned = ActivityDispatch {
2711 node: Some(String::from("n1")),
2712 ..greet_request()
2713 };
2714 let started = Instant::now();
2715 let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
2716 .await
2717 .map_err(|error| error.to_string())?;
2718 let elapsed = started.elapsed();
2719
2720 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2721 assert!(
2722 elapsed < Duration::from_secs(5),
2723 "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
2724 );
2725 echo_n1.await.map_err(|error| error.to_string())??;
2726
2727 assert!(
2729 n2_rx.try_recv().is_err(),
2730 "node=Some(\"n1\") dispatch must not reach the n2 worker"
2731 );
2732
2733 on_n1.deregister()?;
2734 on_n2.deregister()?;
2735 Ok(())
2736 }
2737
2738 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2741 async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
2742 -> Result<(), Box<dyn std::error::Error>> {
2743 let registry = ConnectedWorkerRegistry::default();
2744 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2745 let activity_types = [String::from("greet")];
2746 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2747 let on_n1 = registry.register_namespaces(
2748 [String::from("default")],
2749 "default",
2750 Some(String::from("n1")),
2751 activity_types.iter(),
2752 n1_tx,
2753 )?;
2754
2755 let sink = pending.clone();
2756 let echo = tokio::spawn(async move {
2757 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2758 return Err("expected an activity task on the worker channel".to_owned());
2759 };
2760 let workflow_id = task
2761 .workflow_id
2762 .ok_or("task missing workflow id")
2763 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2764 let activity_id = task
2765 .activity_id
2766 .map(ActivityId::from)
2767 .ok_or("task missing activity id")?;
2768 let completion_token =
2769 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2770 .map_err(|error| error.to_string())?;
2771 sink.complete_activity(ActivityCompletion {
2772 workflow_id,
2773 activity_id,
2774 run_id: None,
2775 completion_token,
2776 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2777 ContentType::Json,
2778 br#"{"greeting":"hello"}"#.to_vec(),
2779 )),
2780 })
2781 .map_err(|error| error.to_string())
2782 });
2783
2784 let dispatcher = Arc::new(
2785 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2786 .with_pending(pending),
2787 );
2788
2789 let result = tokio::spawn(futures::future::lazy(move |_| {
2791 dispatcher.dispatch(greet_request())
2792 }))
2793 .await
2794 .map_err(|error| error.to_string())?;
2795
2796 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2797 echo.await.map_err(|error| error.to_string())??;
2798 on_n1.deregister()?;
2799 Ok(())
2800 }
2801}