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>>;
89
90struct PendingWaiter {
101 sender: SyncSender,
102 attempt: u32,
103}
104
105fn superseded_site_reason(new_attempt: u32) -> String {
114 format!(
115 "retryable: superseded by attempt {new_attempt} at this execution site — the prior \
116 attempt reached its bound, its completion token is fence-refused, and its worker \
117 has been asked to stop"
118 )
119}
120type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
121
122type PendingActivityKey = (WorkflowId, ActivityId);
140
141pub trait OutboxDeliveryCallback: Send + Sync {
151 fn deliver_completion(
161 &self,
162 workflow_id: &WorkflowId,
163 activity_id: &ActivityId,
164 run_id: Option<&RunId>,
165 result: String,
166 ) -> Result<bool, ServerError>;
167
168 fn deliver_failure(
175 &self,
176 workflow_id: &WorkflowId,
177 activity_id: &ActivityId,
178 run_id: Option<&RunId>,
179 reason: String,
180 ) -> Result<bool, ServerError>;
181}
182
183#[derive(Clone)]
195pub struct PendingActivities {
196 pending: Arc<DashMap<PendingActivityKey, PendingWaiter>>,
197 completion_fences: CompletionFences,
198 outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
199 lease_recorder: super::lease_record::LeaseRecorderSeam,
203 transport_losses: super::transport_loss::TransportLossLedger,
209}
210
211impl std::fmt::Debug for PendingActivities {
212 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213 formatter
214 .debug_struct("PendingActivities")
215 .field("pending", &self.pending.len())
216 .field("completion_fences", &self.completion_fences)
217 .field(
218 "outbox_delivery_installed",
219 &self.outbox_delivery.get().is_some(),
220 )
221 .field("lease_recorder", &self.lease_recorder)
222 .field("transport_losses", &self.transport_losses)
223 .finish()
224 }
225}
226
227impl PendingActivities {
228 fn insert(
237 &self,
238 workflow_id: WorkflowId,
239 run_id: &RunId,
240 activity_id: ActivityId,
241 attempt: u32,
242 ) -> Result<(CompletionToken, SyncReceiver, Option<u32>), ServerError> {
243 let completion_token =
244 self.completion_fences
245 .issue(&workflow_id, run_id, &activity_id, attempt)?;
246 let (tx, rx) = std::sync::mpsc::sync_channel(1);
247 let key = (workflow_id, activity_id);
248 match self.pending.entry(key) {
249 Entry::Vacant(entry) => {
250 entry.insert(PendingWaiter {
251 sender: tx,
252 attempt,
253 });
254 Ok((completion_token, rx, None))
255 }
256 Entry::Occupied(mut entry) => {
257 let held_attempt = entry.get().attempt;
258 if attempt > held_attempt {
259 let superseded = entry.insert(PendingWaiter {
265 sender: tx,
266 attempt,
267 });
268 let _ = superseded.sender.send(Err(superseded_site_reason(attempt)));
269 Ok((completion_token, rx, Some(held_attempt)))
270 } else {
271 let (workflow_id, activity_id) = entry.key();
272 Err(ServerError::PendingActivityCollision {
273 workflow_id: workflow_id.clone(),
274 activity_id: activity_id.clone(),
275 held_attempt,
276 incoming_attempt: attempt,
277 })
278 }
279 }
280 }
281 }
282
283 #[must_use]
285 pub fn completion_fences(&self) -> CompletionFences {
286 self.completion_fences.clone()
287 }
288
289 #[cfg(test)]
293 pub(crate) fn insert_for_test(
294 &self,
295 workflow_id: WorkflowId,
296 run_id: &RunId,
297 activity_id: ActivityId,
298 attempt: u32,
299 ) -> Result<(CompletionToken, SyncReceiver, Option<u32>), ServerError> {
300 self.insert(workflow_id, run_id, activity_id, attempt)
301 }
302
303 pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
309 if self.outbox_delivery.set(callback).is_err() {
310 tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
311 }
312 }
313
314 fn complete(
322 &self,
323 workflow_id: &WorkflowId,
324 activity_id: &ActivityId,
325 run_id: Option<&RunId>,
326 result: Result<String, String>,
327 ) -> bool {
328 let matched = self
331 .pending
332 .remove(&(workflow_id.clone(), activity_id.clone()));
333 if let Some((_, waiter)) = matched {
334 return waiter.sender.send(result).is_ok();
335 }
336 let Some(callback) = self.outbox_delivery.get() else {
337 return false;
339 };
340 let outcome = match result {
341 Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
342 Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
343 };
344 match outcome {
345 Ok(true) => true,
346 Ok(false) => {
347 tracing::debug!(
349 workflow_id = %workflow_id,
350 activity_id = %activity_id,
351 "unmatched outbox completion for a workflow that is not currently live; \
352 recovery will re-arm it"
353 );
354 false
355 }
356 Err(error) => {
357 tracing::warn!(
358 workflow_id = %workflow_id,
359 activity_id = %activity_id,
360 %error,
361 "failed to deliver unmatched outbox completion to the live workflow"
362 );
363 false
364 }
365 }
366 }
367
368 fn complete_fenced_after_accept(
369 &self,
370 workflow_id: &WorkflowId,
371 activity_id: &ActivityId,
372 run_id: Option<&RunId>,
373 completion_token: &CompletionToken,
374 result: Result<String, String>,
375 after_accept: impl FnOnce() -> Result<bool, ServerError>,
376 ) -> Result<bool, ServerError> {
377 let accepted = self
382 .completion_fences
383 .accept(workflow_id, activity_id, completion_token)
384 .inspect_err(|error| {
385 tracing::warn!(
386 workflow_id = %workflow_id,
387 activity_id = %activity_id,
388 %error,
389 "activity completion rejected by execution-generation fence"
390 );
391 })?;
392 let should_publish = match after_accept() {
393 Ok(should_publish) => should_publish,
394 Err(error) => {
395 if let Err(restore_error) =
396 self.completion_fences
397 .restore_if_absent(workflow_id, activity_id, &accepted)
398 {
399 tracing::error!(%restore_error, "failed to restore completion generation after settlement failure");
400 }
401 return Err(error);
402 }
403 };
404 if !should_publish {
405 if let Err(restore_error) =
412 self.completion_fences
413 .restore_if_absent(workflow_id, activity_id, &accepted)
414 {
415 tracing::error!(%restore_error, "failed to restore completion generation after a non-publishing settlement");
416 }
417 return Ok(false);
418 }
419 let transport_domain = result
425 .as_ref()
426 .err()
427 .is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
428 if !transport_domain
429 && let Err(error) = self.transport_losses.clear(workflow_id, activity_id)
430 {
431 tracing::warn!(
432 workflow_id = %workflow_id,
433 activity_id = %activity_id,
434 %error,
435 "failed to retire the transport-loss budget for a resolved activity"
436 );
437 }
438 Ok(self.complete(workflow_id, activity_id, run_id, result))
439 }
440
441 #[must_use]
453 pub fn new(heartbeat_window: std::time::Duration) -> Self {
454 Self {
455 pending: Arc::default(),
456 completion_fences: CompletionFences::default(),
457 outbox_delivery: Arc::default(),
458 lease_recorder: super::lease_record::LeaseRecorderSeam::default(),
459 transport_losses: super::transport_loss::TransportLossLedger::new(heartbeat_window),
460 }
461 }
462
463 #[must_use]
466 pub fn lease_recorder(&self) -> super::lease_record::LeaseRecorderSeam {
467 self.lease_recorder.clone()
468 }
469
470 pub fn set_lease_recorder(
473 &self,
474 recorder: Arc<dyn super::lease_record::ActivityLeaseRecorder>,
475 metrics: Option<crate::observability::Metrics>,
476 ) {
477 self.lease_recorder.install(recorder, metrics);
478 }
479
480 #[must_use]
482 pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
483 &self.transport_losses
484 }
485
486 fn classify_worker_loss(
494 &self,
495 workflow_id: &WorkflowId,
496 activity_id: &ActivityId,
497 worker_id: crate::worker::registry::WorkerId,
498 ) -> String {
499 let detail = super::transport_loss::worker_lost_detail(worker_id);
500 match self
501 .transport_losses
502 .record_loss(workflow_id, activity_id, &detail)
503 {
504 Ok(verdict) => {
505 if verdict.exhausted {
506 tracing::error!(
507 operation = "activity_complete",
508 workflow_id = %workflow_id,
509 activity_id = %activity_id,
510 worker_id = ?worker_id,
511 error_type = "TransportExhausted",
512 losses = verdict.losses,
513 budget_ms = self.transport_losses.budget().as_millis(),
514 "activity abandoned: the transport kept losing its worker past the \
515 transport-loss budget"
516 );
517 } else {
518 tracing::warn!(
519 operation = "activity_complete",
520 workflow_id = %workflow_id,
521 activity_id = %activity_id,
522 worker_id = ?worker_id,
523 error_type = "WorkerLost",
524 losses = verdict.losses,
525 budget_ms = self.transport_losses.budget().as_millis(),
526 "worker lost before reporting an activity result; the activity never ran \
527 and will be re-dispatched attempt-neutrally"
528 );
529 }
530 verdict.reason
531 }
532 Err(error) => {
533 tracing::error!(
534 workflow_id = %workflow_id,
535 activity_id = %activity_id,
536 %error,
537 "transport-loss ledger is unreadable; abandoning the activity rather than \
538 re-dispatching it without a budget"
539 );
540 format!(
541 "{}{detail} (transport-loss budget unreadable: {error})",
542 super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
543 )
544 }
545 }
546 }
547
548 pub(crate) fn complete_activity_after_accept(
549 &self,
550 completion: ActivityCompletion,
551 after_accept: impl FnOnce() -> Result<(), ServerError>,
552 ) -> Result<(), ServerError> {
553 let result = match completion.outcome {
554 ActivityCompletionOutcome::Succeeded(payload) => {
555 payload_to_string(&payload).map_err(|reason| {
556 tracing::error!(
557 operation = "activity_complete",
558 workflow_id = %completion.workflow_id,
559 activity_id = %completion.activity_id,
560 error_type = "ActivityResultDecode",
561 %reason,
562 "activity completion failed"
563 );
564 ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
565 })?
566 }
567 ActivityCompletionOutcome::Failed(error) => {
568 let prefix = match error.kind {
569 ActivityErrorKind::Retryable => "retryable",
570 ActivityErrorKind::PolicyRefused => "policy_refused",
571 ActivityErrorKind::Terminal => "terminal",
572 };
573 tracing::error!(
574 operation = "activity_complete",
575 workflow_id = %completion.workflow_id,
576 activity_id = %completion.activity_id,
577 error_type = "ActivityFailed",
578 error_kind = prefix,
579 reason = %error.message,
580 "activity completion failed"
581 );
582 Err(format!("{prefix}:{}", error.message))
583 }
584 ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
585 &completion.workflow_id,
586 &completion.activity_id,
587 worker_id,
588 )),
589 };
590 let accepted_settlement = || after_accept().map(|()| true);
591 self.complete_fenced_after_accept(
592 &completion.workflow_id,
593 &completion.activity_id,
594 completion.run_id.as_ref(),
595 &completion.completion_token,
596 result,
597 accepted_settlement,
598 )?;
599 Ok(())
600 }
601}
602
603impl ActivityCompletionSink for PendingActivities {
604 fn park_activity(
615 &self,
616 workflow_id: &WorkflowId,
617 activity_id: &ActivityId,
618 ) -> Result<(), ServerError> {
619 self.completion_fences
623 .revoke_current(workflow_id, activity_id)?;
624 let matched = self
625 .pending
626 .remove(&(workflow_id.clone(), activity_id.clone()));
627 if let Some((_, waiter)) = matched {
628 let _ = waiter
632 .sender
633 .send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
634 }
635 Ok(())
636 }
637
638 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
639 self.complete_activity_after_accept(completion, || Ok(()))
640 }
641}
642
643fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
644 match payload.content_type() {
645 ContentType::Json => String::from_utf8(payload.bytes().to_vec())
646 .map(Ok)
647 .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
648 }
649}
650
651pub struct WorkerActivityDispatcher {
659 registry: ConnectedWorkerRegistry,
660 namespace: String,
661 pending: PendingActivities,
662 heartbeat_tracker: HeartbeatTracker,
663 drain_state: DrainState,
664 tokio_handle: Option<tokio::runtime::Handle>,
665 attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
671 queue_service: QueueServiceConfig,
675 queue_declarations: QueueDeclarationSource,
678 queue_state: QueueServiceState,
680 cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
684}
685
686impl std::fmt::Debug for WorkerActivityDispatcher {
687 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
688 f.debug_struct("WorkerActivityDispatcher")
689 .field("namespace", &self.namespace)
690 .finish_non_exhaustive()
691 }
692}
693
694impl WorkerActivityDispatcher {
695 #[must_use]
703 pub fn new(
704 registry: ConnectedWorkerRegistry,
705 namespace: impl Into<String>,
706 heartbeat_tracker: HeartbeatTracker,
707 ) -> Self {
708 Self {
709 registry,
710 namespace: namespace.into(),
711 pending: PendingActivities::new(heartbeat_tracker.heartbeat_window()),
715 heartbeat_tracker,
716 drain_state: DrainState::default(),
717 tokio_handle: None,
718 attempt_owners: None,
719 queue_service: QueueServiceConfig::default(),
720 queue_declarations: QueueDeclarationSource::default(),
721 queue_state: QueueServiceState::default(),
722 cluster_publisher: None,
723 }
724 }
725
726 #[must_use]
730 pub fn with_cluster_publisher(
731 mut self,
732 cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
733 ) -> Self {
734 self.cluster_publisher = Some(cluster_publisher);
735 self
736 }
737
738 #[must_use]
741 pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
742 self.queue_service = queue_service;
743 self
744 }
745
746 #[must_use]
750 pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
751 self.queue_declarations = queue_declarations;
752 self
753 }
754
755 #[must_use]
758 pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
759 self.queue_state = queue_state;
760 self
761 }
762
763 #[must_use]
770 pub fn with_attempt_owners(
771 mut self,
772 attempt_owners: super::intervention::AttemptOwnerIndex,
773 ) -> Self {
774 self.attempt_owners = Some(attempt_owners);
775 self
776 }
777
778 #[must_use]
783 pub const fn pending(&self) -> &PendingActivities {
784 &self.pending
785 }
786
787 #[must_use]
789 pub fn with_pending(mut self, pending: PendingActivities) -> Self {
790 self.pending = pending;
791 self
792 }
793
794 #[must_use]
796 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
797 self.drain_state = drain_state;
798 self
799 }
800
801 #[must_use]
803 pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
804 self.tokio_handle = Some(tokio_handle);
805 self
806 }
807}
808
809impl WorkerActivityDispatcher {
810 fn ensure_accepting(
834 &self,
835 namespace: &str,
836 activity_type: &str,
837 workflow_id: &WorkflowId,
838 activity_id: &ActivityId,
839 worker_id: Option<WorkerId>,
840 ) -> Result<(), String> {
841 self.drain_state
842 .ensure_accepting(namespace, activity_type)
843 .map_err(|error| {
844 log_worker_error(
845 "WorkerDispatch",
846 namespace,
847 activity_type,
848 workflow_id,
849 activity_id,
850 worker_id,
851 &error.to_string(),
852 );
853 aion::PARKED_ACTIVITY_REASON.to_owned()
854 })
855 }
856
857 fn select_worker_or_wait(
865 &self,
866 address: &ServiceAddress,
867 workflow_id: &WorkflowId,
868 activity_id: &ActivityId,
869 ) -> Result<WorkerHandle, String> {
870 let wait = ServiceWait {
871 registry: &self.registry,
872 declarations: &self.queue_declarations,
873 config: &self.queue_service,
874 state: &self.queue_state,
875 address,
876 workflow_id,
877 activity_id,
878 publisher: self.cluster_publisher.as_ref(),
879 };
880 let mut accepting = || {
881 self.ensure_accepting(
882 &address.namespace,
883 &address.activity_type,
884 workflow_id,
885 activity_id,
886 None,
887 )
888 };
889 let mut park = |arrival: WorkerArrival, budget: Option<Duration>| {
890 self.park_for_worker(arrival, budget);
891 };
892 select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
893 let reason = refusal.reason_string();
894 if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
897 let error_type = match refusal {
898 SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
899 _ => "WorkerRegistry",
900 };
901 log_worker_error(
902 error_type,
903 &address.namespace,
904 &address.activity_type,
905 workflow_id,
906 activity_id,
907 None,
908 &reason,
909 );
910 }
911 reason
912 })
913 }
914
915 fn park_for_worker(&self, arrival: WorkerArrival, budget: Option<Duration>) {
950 let handle = self
951 .tokio_handle
952 .clone()
953 .or_else(|| tokio::runtime::Handle::try_current().ok());
954 let Some(handle) = handle else {
955 std::thread::sleep(
967 budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
968 );
969 return;
970 };
971 handle.block_on(async {
972 let raced = async {
973 tokio::select! {
974 () = arrival => {}
975 () = self.drain_state.wait_for_drain() => {}
976 }
977 };
978 match budget {
979 None => raced.await,
980 Some(budget) => {
981 drop(tokio::time::timeout(budget, raced).await);
984 }
985 }
986 });
987 }
988
989 fn track_worker_task(
990 &self,
991 worker_id: WorkerId,
992 activity_type: &str,
993 workflow_id: &WorkflowId,
994 activity_id: &ActivityId,
995 attempt: u32,
996 completion_token: CompletionToken,
997 ) -> Result<(), String> {
998 self.heartbeat_tracker
999 .track_task(
1000 worker_id,
1001 InFlightActivity {
1002 workflow_id: workflow_id.clone(),
1003 activity_id: activity_id.clone(),
1004 attempt,
1009 completion_token,
1010 },
1011 Instant::now(),
1012 )
1013 .map_err(|error| {
1014 let reason = error.to_string();
1015 log_worker_error(
1016 "WorkerHeartbeatTracker",
1017 &self.namespace,
1018 activity_type,
1019 workflow_id,
1020 activity_id,
1021 Some(worker_id),
1022 &reason,
1023 );
1024 reason
1025 })
1026 }
1027
1028 fn cleanup_activity(
1029 &self,
1030 worker_id: WorkerId,
1031 workflow_id: &WorkflowId,
1032 activity_id: &ActivityId,
1033 attempt: u32,
1034 completion_token: &CompletionToken,
1035 ) {
1036 self.pending
1041 .pending
1042 .remove_if(&(workflow_id.clone(), activity_id.clone()), |_, waiter| {
1043 waiter.attempt == attempt
1044 });
1045 if let Err(error) =
1049 self.pending
1050 .completion_fences
1051 .revoke(workflow_id, activity_id, completion_token)
1052 {
1053 tracing::error!(
1054 workflow_id = %workflow_id,
1055 activity_id = %activity_id,
1056 %error,
1057 "failed to revoke undelivered activity generation"
1058 );
1059 }
1060 let _ = self
1061 .heartbeat_tracker
1062 .complete_task(worker_id, workflow_id, activity_id);
1063 self.drain_state.notify_activity_drained();
1064 }
1065
1066 fn send_activity_task(
1076 &self,
1077 worker: &WorkerHandle,
1078 task: ProtoActivityTask,
1079 address: &ServiceAddress,
1080 handoff: &super::lease_record::LeaseHandoff,
1081 ) -> Result<(), String> {
1082 let workflow_id = &handoff.key().workflow_id;
1083 let activity_id = &handoff.key().activity_id;
1084 let completion_token = handoff.token();
1085 match worker.delivery() {
1086 WorkerDelivery::Grpc(sender) => {
1087 let worker_id = worker.id();
1088 let mut accepting = || {
1089 self.ensure_accepting(
1090 &address.namespace,
1091 &address.activity_type,
1092 workflow_id,
1093 activity_id,
1094 Some(worker_id),
1095 )
1096 };
1097 let handed_over = deliver_within_schedule_to_start(
1098 sender,
1099 WorkerMessage::ActivityTask(Box::new(task)),
1100 self.queue_service.schedule_to_start_timeout,
1101 &mut accepting,
1102 );
1103 let Err(refusal) = handed_over else {
1104 handoff.accepted_blocking(self.tokio_handle.as_ref());
1108 return Ok(());
1109 };
1110 self.cleanup_activity(
1111 worker_id,
1112 workflow_id,
1113 activity_id,
1114 handoff.key().attempt,
1115 completion_token,
1116 );
1117 let (error_type, reason) = self.hand_off_failure(&refusal, address);
1118 if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
1119 log_worker_error(
1120 error_type,
1121 &address.namespace,
1122 &address.activity_type,
1123 workflow_id,
1124 activity_id,
1125 Some(worker_id),
1126 &reason,
1127 );
1128 }
1129 Err(reason)
1130 }
1131 #[cfg(feature = "liminal-transport")]
1138 WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
1139 worker.id(),
1140 delivery,
1141 task,
1142 &address.activity_type,
1143 handoff,
1144 ),
1145 }
1146 }
1147
1148 fn hand_off_failure(
1155 &self,
1156 refusal: &DeliveryRefusal,
1157 address: &ServiceAddress,
1158 ) -> (&'static str, String) {
1159 match refusal {
1160 DeliveryRefusal::Saturated { waited } => {
1161 let census = self
1164 .registry
1165 .pool_census(
1166 &address.namespace,
1167 &address.task_queue,
1168 &address.activity_type,
1169 address.node.as_deref(),
1170 )
1171 .unwrap_or_else(|error| {
1172 tracing::error!(
1173 namespace = %address.namespace,
1174 task_queue = %address.task_queue,
1175 activity_type = %address.activity_type,
1176 %error,
1177 "poller census failed while reporting a saturated queue; \
1178 the refusal carries an empty census"
1179 );
1180 PoolCensus::default()
1181 });
1182 let unavailable = WorkerUnavailable {
1183 reason: QueueServiceReason::Saturated,
1184 clock: Some(ExpiredClock::ScheduleToStart),
1185 waited: *waited,
1186 address: address.clone(),
1187 census,
1188 };
1189 ("WorkerUnavailable", unavailable.reason_string())
1190 }
1191 DeliveryRefusal::Full => (
1192 "WorkerChannelClosed",
1193 "worker task channel full or closed: no available capacity".to_owned(),
1194 ),
1195 DeliveryRefusal::Closed => (
1196 "WorkerChannelClosed",
1197 "worker task channel full or closed: channel closed".to_owned(),
1198 ),
1199 DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
1200 }
1201 }
1202
1203 #[cfg(feature = "liminal-transport")]
1227 fn send_liminal_activity_task(
1228 &self,
1229 worker_id: WorkerId,
1230 delivery: &super::liminal_transport::LiminalWorkerDelivery,
1231 task: ProtoActivityTask,
1232 activity_type: &str,
1233 handoff: &super::lease_record::LeaseHandoff,
1234 ) -> Result<(), String> {
1235 let workflow_id = &handoff.key().workflow_id;
1236 let activity_id = &handoff.key().activity_id;
1237 let completion_token =
1238 CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1239 .map_err(|error| error.to_string())?;
1240 let heartbeat_window_ms =
1241 u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1242 .unwrap_or(u64::MAX);
1243 let attempt = task.attempt;
1244 let run_id = task
1251 .run_id
1252 .map(RunId::try_from)
1253 .transpose()
1254 .map_err(|error| error.to_string())?
1255 .ok_or_else(|| {
1256 "activity task run id is missing; refusing to dispatch an unidentified run"
1257 .to_owned()
1258 })?;
1259 let request = super::liminal_transport::DispatchRequest {
1260 activity_type: activity_type.to_owned(),
1261 workflow_id: workflow_id.clone(),
1262 ordinal: activity_id.sequence_position(),
1263 run_id: Some(run_id.clone()),
1264 attempt,
1265 completion_token: task.completion_token,
1266 idempotency_key: task.idempotency_key,
1267 labels: task.labels.into_iter().collect(),
1268 heartbeat_window_ms,
1269 input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1270 };
1271 let awaiter = match delivery.push_dispatch(&request) {
1275 Ok(awaiter) => awaiter,
1276 Err(error) => {
1277 let reason = format!("worker liminal push failed: {error}");
1278 self.cleanup_activity(
1279 worker_id,
1280 workflow_id,
1281 activity_id,
1282 attempt,
1283 &completion_token,
1284 );
1285 log_worker_error(
1286 "WorkerChannelClosed",
1287 &self.namespace,
1288 activity_type,
1289 workflow_id,
1290 activity_id,
1291 Some(worker_id),
1292 &reason,
1293 );
1294 return Err(reason);
1295 }
1296 };
1297 handoff.accepted_blocking(self.tokio_handle.as_ref());
1304 let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1305 super::liminal_transport::AttemptOwnerGuard::bind(
1306 owners.clone(),
1307 super::intervention::AttemptKey::new(
1308 workflow_id.clone(),
1309 run_id.clone(),
1310 activity_id.clone(),
1311 attempt,
1312 ),
1313 worker_id,
1314 )
1315 });
1316 self.spawn_liminal_reply_router(
1317 worker_id,
1318 awaiter,
1319 workflow_id,
1320 activity_id,
1321 &completion_token,
1322 owner_binding,
1323 );
1324 Ok(())
1325 }
1326
1327 #[cfg(feature = "liminal-transport")]
1367 fn spawn_liminal_reply_router(
1368 &self,
1369 worker_id: WorkerId,
1370 awaiter: liminal_server::server::connection::PushReplyAwaiter,
1371 workflow_id: &WorkflowId,
1372 activity_id: &ActivityId,
1373 completion_token: &CompletionToken,
1374 owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1375 ) {
1376 let pending = self.pending.clone();
1377 let heartbeat_tracker = self.heartbeat_tracker.clone();
1378 let drain_state = self.drain_state.clone();
1379 let workflow_id = workflow_id.clone();
1380 let activity_id = activity_id.clone();
1381 let completion_token = completion_token.clone();
1382 std::thread::spawn(move || {
1383 let _owner_binding = owner_binding;
1387 route_liminal_reply(
1388 &pending,
1389 &heartbeat_tracker,
1390 &drain_state,
1391 &awaiter,
1392 (worker_id, &workflow_id, &activity_id, &completion_token),
1393 );
1394 });
1395 }
1396
1397 fn await_activity_result(
1401 &self,
1402 context: &ActivityDispatchContext<'_>,
1403 rx: &SyncReceiver,
1404 ) -> Result<String, String> {
1405 match self.registry.is_registered(context.worker_id) {
1413 Ok(true) => {}
1414 Ok(false) => {
1415 if let Ok(result) = rx.try_recv() {
1419 return self.deliver_result(context, result);
1420 }
1421 self.cleanup_activity(
1422 context.worker_id,
1423 context.workflow_id,
1424 context.activity_id,
1425 context.attempt,
1426 &context.completion_token,
1427 );
1428 let reason = self.pending.classify_worker_loss(
1433 context.workflow_id,
1434 context.activity_id,
1435 context.worker_id,
1436 );
1437 log_worker_error(
1438 "WorkerLost",
1439 &self.namespace,
1440 context.activity_type,
1441 context.workflow_id,
1442 context.activity_id,
1443 Some(context.worker_id),
1444 &reason,
1445 );
1446 return Err(reason);
1447 }
1448 Err(error) => {
1449 self.cleanup_activity(
1450 context.worker_id,
1451 context.workflow_id,
1452 context.activity_id,
1453 context.attempt,
1454 &context.completion_token,
1455 );
1456 let reason = format!("worker registry inspection failed: {error}");
1457 log_worker_error(
1458 "WorkerRegistry",
1459 &self.namespace,
1460 context.activity_type,
1461 context.workflow_id,
1462 context.activity_id,
1463 Some(context.worker_id),
1464 &reason,
1465 );
1466 return Err(reason);
1467 }
1468 }
1469 if let Ok(result) = rx.recv() {
1470 return self.deliver_result(context, result);
1471 }
1472 self.cleanup_activity(
1475 context.worker_id,
1476 context.workflow_id,
1477 context.activity_id,
1478 context.attempt,
1479 &context.completion_token,
1480 );
1481 let reason = "activity response channel dropped".to_owned();
1482 log_worker_error(
1483 "WorkerChannelClosed",
1484 &self.namespace,
1485 context.activity_type,
1486 context.workflow_id,
1487 context.activity_id,
1488 Some(context.worker_id),
1489 &reason,
1490 );
1491 Err(reason)
1492 }
1493
1494 fn deliver_result(
1495 &self,
1496 context: &ActivityDispatchContext<'_>,
1497 result: Result<String, String>,
1498 ) -> Result<String, String> {
1499 self.pending.pending.remove_if(
1503 &(context.workflow_id.clone(), context.activity_id.clone()),
1504 |_, waiter| waiter.attempt == context.attempt,
1505 );
1506 if let Err(reason) = &result
1511 && aion::is_parked_reason(reason)
1512 {
1513 tracing::info!(
1514 operation = "activity_dispatch",
1515 namespace = %self.namespace,
1516 workflow_id = %context.workflow_id,
1517 activity_id = %context.activity_id,
1518 activity_type = context.activity_type,
1519 worker_id = ?context.worker_id,
1520 "activity parked for restart recovery"
1521 );
1522 return result;
1523 }
1524 log_activity_completion(context, result.is_ok());
1525 result.inspect_err(|reason| {
1526 log_worker_error(
1527 "ActivityFailed",
1528 &self.namespace,
1529 context.activity_type,
1530 context.workflow_id,
1531 context.activity_id,
1532 Some(context.worker_id),
1533 reason,
1534 );
1535 })
1536 }
1537}
1538
1539impl ActivityDispatcher for WorkerActivityDispatcher {
1540 fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1541 match tokio::runtime::Handle::try_current() {
1542 Ok(handle) => match handle.runtime_flavor() {
1543 tokio::runtime::RuntimeFlavor::MultiThread => {
1544 tokio::task::block_in_place(|| self.dispatch_blocking(request))
1551 }
1552 flavor => Err(format!(
1553 "activity dispatch blocks the calling thread until the worker responds; \
1554 a {flavor:?} tokio runtime cannot host that wait because the worker \
1555 stream forwarder shares its only executor thread and the task could \
1556 never be delivered — run the engine on a multi-thread tokio runtime"
1557 )),
1558 },
1559 Err(_) => self.dispatch_blocking(request),
1563 }
1564 }
1565}
1566
1567impl WorkerActivityDispatcher {
1568 fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1585 let ActivityDispatch {
1586 namespace,
1587 task_queue,
1588 node,
1592 workflow_id,
1593 run_id,
1594 activity_id,
1595 name,
1596 input,
1597 config: _,
1598 attempt,
1599 labels,
1600 advisory: _,
1605 } = request;
1606 let started_at = Instant::now();
1607 self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1608 let address = ServiceAddress {
1609 namespace: namespace.clone(),
1610 task_queue: task_queue.clone(),
1611 activity_type: name.clone(),
1612 node: node.clone(),
1613 };
1614 let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1615 let worker_id = worker.id();
1616 let span = info_span!(
1617 "activity_dispatch",
1618 operation = "activity_dispatch",
1619 namespace = %namespace,
1620 task_queue = %task_queue,
1621 node = node.as_deref(),
1622 workflow_id = %workflow_id,
1623 activity_id = %activity_id,
1624 activity_type = %name,
1625 worker_id = ?worker_id,
1626 );
1627 let _span_guard = span.enter();
1628 self.ensure_accepting(
1629 &namespace,
1630 &name,
1631 &workflow_id,
1632 &activity_id,
1633 Some(worker_id),
1634 )?;
1635
1636 let (completion_token, rx, superseded_attempt) = self
1637 .pending
1638 .insert(workflow_id.clone(), &run_id, activity_id.clone(), attempt)
1639 .map_err(|error| error.to_string())?;
1640 if let Some(held_attempt) = superseded_attempt {
1641 self.release_superseded_attempt(&workflow_id, &activity_id, held_attempt, attempt);
1647 }
1648 let handoff = super::lease_record::LeaseHandoff::arm(
1653 self.pending.lease_recorder(),
1654 super::lease_record::LeaseKey {
1655 workflow_id: workflow_id.clone(),
1656 run_id: run_id.clone(),
1657 activity_id: activity_id.clone(),
1658 attempt,
1659 },
1660 super::lease_record::attribution_for(&worker),
1661 self.pending.completion_fences(),
1662 completion_token.clone(),
1663 )
1664 .map_err(|error| error.to_string())?;
1665 let task = activity_task(
1666 &name,
1667 &input,
1668 (&workflow_id, &run_id, &activity_id),
1669 attempt,
1670 labels,
1671 &completion_token,
1672 );
1673 if let Err(error) = self.track_worker_task(
1674 worker_id,
1675 &name,
1676 &workflow_id,
1677 &activity_id,
1678 attempt,
1679 completion_token.clone(),
1680 ) {
1681 self.cleanup_activity(
1682 worker_id,
1683 &workflow_id,
1684 &activity_id,
1685 attempt,
1686 &completion_token,
1687 );
1688 return Err(error);
1689 }
1690 self.send_activity_task(&worker, task, &address, &handoff)?;
1691 let context = ActivityDispatchContext {
1692 namespace: &namespace,
1693 activity_type: &name,
1694 worker_id,
1695 workflow_id: &workflow_id,
1696 activity_id: &activity_id,
1697 attempt,
1698 completion_token,
1699 started_at,
1700 };
1701 self.await_activity_result(&context, &rx)
1702 }
1703
1704 fn release_superseded_attempt(
1713 &self,
1714 workflow_id: &WorkflowId,
1715 activity_id: &ActivityId,
1716 held_attempt: u32,
1717 new_attempt: u32,
1718 ) {
1719 let in_flight = match self.heartbeat_tracker.in_flight_for_workflow(workflow_id) {
1720 Ok(in_flight) => in_flight,
1721 Err(error) => {
1722 tracing::error!(
1723 workflow_id = %workflow_id,
1724 activity_id = %activity_id,
1725 held_attempt,
1726 new_attempt,
1727 %error,
1728 "superseded-attempt release could not read the in-flight tracker; \
1729 the corpse's worker was not asked to stop"
1730 );
1731 return;
1732 }
1733 };
1734 let holders = in_flight
1735 .into_iter()
1736 .filter(|task| &task.activity_id == activity_id && task.attempt == held_attempt);
1737 let mut asked_anyone = false;
1738 for task in holders {
1739 asked_anyone = true;
1740 let delivery = match super::activity_cancel::ask_worker_to_stop(
1741 &self.registry,
1742 task.worker_id,
1743 workflow_id,
1744 activity_id,
1745 ) {
1746 Ok(delivery) => delivery,
1747 Err(error) => {
1748 tracing::error!(
1749 workflow_id = %workflow_id,
1750 activity_id = %activity_id,
1751 held_attempt,
1752 worker_id = task.worker_id.value(),
1753 %error,
1754 "superseded-attempt release could not ask the holding worker to stop"
1755 );
1756 continue;
1757 }
1758 };
1759 if let Err(error) =
1760 self.heartbeat_tracker
1761 .complete_task(task.worker_id, workflow_id, activity_id)
1762 {
1763 tracing::warn!(
1764 workflow_id = %workflow_id,
1765 activity_id = %activity_id,
1766 held_attempt,
1767 worker_id = task.worker_id.value(),
1768 %error,
1769 "superseded-attempt release could not untrack the corpse's heartbeat task"
1770 );
1771 }
1772 if delivery.was_requested() {
1773 tracing::info!(
1774 workflow_id = %workflow_id,
1775 activity_id = %activity_id,
1776 held_attempt,
1777 new_attempt,
1778 worker_id = task.worker_id.value(),
1779 "attempt takeover: asked the worker still executing the superseded \
1780 attempt to stop"
1781 );
1782 } else {
1783 tracing::warn!(
1784 workflow_id = %workflow_id,
1785 activity_id = %activity_id,
1786 held_attempt,
1787 new_attempt,
1788 worker_id = task.worker_id.value(),
1789 outcome = ?delivery,
1790 "attempt takeover: the superseded attempt's worker could NOT be asked \
1791 to stop — its work runs on unwitnessed until it next reports, and its \
1792 completion will be refused as a stale generation"
1793 );
1794 }
1795 }
1796 if !asked_anyone {
1797 tracing::info!(
1798 workflow_id = %workflow_id,
1799 activity_id = %activity_id,
1800 held_attempt,
1801 new_attempt,
1802 "attempt takeover: no tracked worker holds the superseded attempt \
1803 (already untracked, or its transport records no heartbeat task); \
1804 its completion token is fence-refused regardless"
1805 );
1806 }
1807 }
1808}
1809
1810#[cfg(feature = "liminal-transport")]
1817fn route_liminal_reply(
1818 pending: &PendingActivities,
1819 heartbeat_tracker: &HeartbeatTracker,
1820 drain_state: &DrainState,
1821 awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1822 execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1823) {
1824 let (worker_id, workflow_id, activity_id, current_token) = execution;
1825 let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1830 heartbeat_tracker
1831 .is_tracked(worker_id, workflow_id, activity_id)
1832 .unwrap_or(false)
1833 });
1834 let (run_id, submitted_token, outcome, synthesized) = match waited {
1838 Ok(Some(response)) => {
1839 let submitted_token = match CompletionToken::from_wire(
1840 workflow_id,
1841 activity_id,
1842 response.completion_token,
1843 ) {
1844 Ok(token) => token,
1845 Err(error) => {
1846 tracing::warn!(
1847 worker_id = ?worker_id,
1848 workflow_id = %workflow_id,
1849 activity_id = %activity_id,
1850 %error,
1851 "liminal activity completion omitted its generation proof"
1852 );
1853 return;
1854 }
1855 };
1856 (response.run_id, submitted_token, response.outcome, false)
1857 }
1858 Ok(None) => {
1859 tracing::debug!(
1860 worker_id = ?worker_id,
1861 workflow_id = %workflow_id,
1862 activity_id = %activity_id,
1863 "liminal dispatch resolved by another path; abandoning reply wait"
1864 );
1865 return;
1866 }
1867 Err(error) if error.is_worker_connection_lost() => (
1868 None,
1869 current_token.clone(),
1870 Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1874 true,
1875 ),
1876 Err(error) => (
1877 None,
1878 current_token.clone(),
1879 Err(format!("retryable:worker liminal reply failed: {error}")),
1880 true,
1881 ),
1882 };
1883 let after_accept = || {
1888 if synthesized {
1889 let was_tracked = clear_completed_task_tracking(
1890 heartbeat_tracker,
1891 worker_id,
1892 workflow_id,
1893 activity_id,
1894 );
1895 if !was_tracked {
1896 tracing::debug!(
1899 worker_id = ?worker_id,
1900 workflow_id = %workflow_id,
1901 activity_id = %activity_id,
1902 "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1903 );
1904 return Ok(false);
1905 }
1906 } else {
1907 let _ = clear_completed_task_tracking(
1913 heartbeat_tracker,
1914 worker_id,
1915 workflow_id,
1916 activity_id,
1917 );
1918 }
1919 drain_state.notify_activity_drained();
1920 Ok(true)
1921 };
1922 if let Err(error) = pending.complete_fenced_after_accept(
1923 workflow_id,
1924 activity_id,
1925 run_id.as_ref(),
1926 &submitted_token,
1927 outcome,
1928 after_accept,
1929 ) {
1930 tracing::warn!(
1931 worker_id = ?worker_id,
1932 workflow_id = %workflow_id,
1933 activity_id = %activity_id,
1934 %error,
1935 "liminal activity completion handoff rejected"
1936 );
1937 }
1938}
1939
1940pub(crate) fn clear_completed_task_tracking(
1947 heartbeat_tracker: &HeartbeatTracker,
1948 worker_id: WorkerId,
1949 workflow_id: &WorkflowId,
1950 activity_id: &ActivityId,
1951) -> bool {
1952 heartbeat_tracker
1953 .complete_task(worker_id, workflow_id, activity_id)
1954 .unwrap_or_else(|error| {
1955 tracing::error!(
1956 worker_id = ?worker_id,
1957 workflow_id = %workflow_id,
1958 activity_id = %activity_id,
1959 %error,
1960 "failed to clear in-flight tracking for a completed activity; delivering anyway — a poisoned tracker must not withhold a result"
1961 );
1962 true
1963 })
1964}
1965
1966struct ActivityDispatchContext<'a> {
1967 namespace: &'a str,
1968 activity_type: &'a str,
1969 worker_id: WorkerId,
1970 workflow_id: &'a WorkflowId,
1971 activity_id: &'a ActivityId,
1972 attempt: u32,
1975 completion_token: CompletionToken,
1976 started_at: Instant,
1977}
1978
1979fn activity_task(
1980 activity_type: &str,
1981 input: &str,
1982 execution: (&WorkflowId, &RunId, &ActivityId),
1983 attempt: u32,
1984 labels: BTreeMap<String, String>,
1985 completion_token: &CompletionToken,
1986) -> ProtoActivityTask {
1987 let (workflow_id, run_id, activity_id) = execution;
1988 ProtoActivityTask {
1989 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1990 activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1991 activity_type: activity_type.to_owned(),
1992 input: Some(ProtoPayload {
1993 content_type: String::from("application/json"),
1994 bytes: input.as_bytes().to_vec(),
1995 }),
1996 attempt,
1997 labels: labels.into_iter().collect(),
1998 run_id: Some(run_id.clone().into()),
1999 completion_token: completion_token.as_str().to_owned(),
2000 idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
2001 }
2002}
2003
2004fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
2005 let duration_ms = duration_ms(context.started_at.elapsed());
2006 tracing::info!(
2007 operation = "activity_complete",
2008 namespace = context.namespace,
2009 workflow_id = %context.workflow_id,
2010 activity_id = %context.activity_id,
2011 activity_type = context.activity_type,
2012 worker_id = ?context.worker_id,
2013 duration_ms,
2014 outcome = if succeeded { "succeeded" } else { "failed" },
2015 "activity completed"
2016 );
2017}
2018
2019fn duration_ms(duration: Duration) -> u64 {
2020 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
2021}
2022
2023fn log_worker_error(
2024 error_type: &'static str,
2025 namespace: &str,
2026 activity_type: &str,
2027 workflow_id: &WorkflowId,
2028 activity_id: &ActivityId,
2029 worker_id: Option<super::registry::WorkerId>,
2030 reason: &str,
2031) {
2032 tracing::error!(
2033 operation = "activity_dispatch",
2034 namespace,
2035 workflow_id = %workflow_id,
2036 activity_id = %activity_id,
2037 activity_type,
2038 worker_id = ?worker_id,
2039 error_type,
2040 reason,
2041 "worker interaction failed"
2042 );
2043}
2044
2045#[cfg(test)]
2046mod tests {
2047 use std::sync::{
2048 Mutex,
2049 atomic::{AtomicBool, Ordering},
2050 };
2051
2052 use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
2053
2054 use super::*;
2055
2056 fn activity_id(pos: u64) -> ActivityId {
2057 ActivityId::from_sequence_position(pos)
2058 }
2059
2060 fn insert_attempt_one(
2066 pending: &PendingActivities,
2067 workflow_id: &WorkflowId,
2068 id: &ActivityId,
2069 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
2070 let (token, rx, superseded) =
2071 pending.insert(workflow_id.clone(), &RunId::new_v4(), id.clone(), 1)?;
2072 assert!(
2073 superseded.is_none(),
2074 "a first insert must never report a superseded holder"
2075 );
2076 Ok((token, rx))
2077 }
2078
2079 #[test]
2080 fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
2081 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2082 let workflow_id = WorkflowId::new_v4();
2083 let id = activity_id(1);
2084 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
2085
2086 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
2087 assert_eq!(
2088 rx.recv_timeout(Duration::from_millis(50)),
2089 Ok(Ok("done".to_owned()))
2090 );
2091 Ok(())
2092 }
2093
2094 #[test]
2099 fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
2100 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2101 let workflow_id = WorkflowId::new_v4();
2102 let id = activity_id(2);
2103 let run_id = RunId::new_v4();
2107 let _first = pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
2108
2109 let error = pending
2110 .insert(workflow_id, &run_id, id, 1)
2111 .err()
2112 .ok_or_else(|| {
2113 ServerError::worker_dispatch(
2114 "test",
2115 "same-worker-collision",
2116 "a live responder collision was silently replaced",
2117 )
2118 })?;
2119
2120 assert!(
2121 error.to_string().contains("pending activity collision"),
2122 "{error}"
2123 );
2124 Ok(())
2125 }
2126
2127 #[test]
2134 fn a_higher_attempt_takes_over_a_stale_execution_site() -> Result<(), ServerError> {
2135 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2136 let workflow_id = WorkflowId::new_v4();
2137 let id = activity_id(3);
2138 let run_id = RunId::new_v4();
2139 let (_stale_token, stale_rx, none) =
2140 pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
2141 assert!(none.is_none());
2142
2143 let (fresh_token, fresh_rx, superseded) =
2144 pending.insert(workflow_id.clone(), &run_id, id.clone(), 2)?;
2145 assert_eq!(
2146 superseded,
2147 Some(1),
2148 "the takeover must name the attempt it displaced"
2149 );
2150
2151 let stale_resolution = stale_rx
2154 .recv_timeout(Duration::from_millis(50))
2155 .map_err(|_| {
2156 ServerError::worker_dispatch(
2157 "test",
2158 "takeover",
2159 "the superseded waiter was never woken",
2160 )
2161 })?;
2162 let reason = match stale_resolution {
2163 Ok(delivered) => {
2164 return Err(ServerError::worker_dispatch(
2165 "test",
2166 "takeover",
2167 format!("the superseded waiter resolved Ok({delivered}) instead of an error"),
2168 ));
2169 }
2170 Err(reason) => reason,
2171 };
2172 assert!(
2173 reason.starts_with("retryable:") && reason.contains("superseded by attempt 2"),
2174 "{reason}"
2175 );
2176
2177 pending.complete_activity(ActivityCompletion {
2180 workflow_id: workflow_id.clone(),
2181 activity_id: id.clone(),
2182 run_id: None,
2183 completion_token: fresh_token,
2184 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2185 ContentType::Json,
2186 br#""second attempt wins""#.to_vec(),
2187 )),
2188 })?;
2189 assert_eq!(
2190 fresh_rx.recv_timeout(Duration::from_millis(50)),
2191 Ok(Ok(r#""second attempt wins""#.to_owned()))
2192 );
2193 Ok(())
2194 }
2195
2196 #[test]
2200 fn a_lower_attempt_is_still_refused_by_name() -> Result<(), ServerError> {
2201 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2202 let workflow_id = WorkflowId::new_v4();
2203 let id = activity_id(4);
2204 let run_id = RunId::new_v4();
2205 let _held = pending.insert(workflow_id.clone(), &run_id, id.clone(), 3)?;
2206 let error = pending
2207 .insert(workflow_id, &run_id, id, 2)
2208 .err()
2209 .ok_or_else(|| {
2210 ServerError::worker_dispatch(
2211 "test",
2212 "stale-redispatch",
2213 "a lower attempt silently replaced a live responder",
2214 )
2215 })?;
2216 let text = error.to_string();
2217 assert!(
2218 text.contains("attempt 2 arrived while attempt 3 still holds"),
2219 "{text}"
2220 );
2221 Ok(())
2222 }
2223
2224 #[test]
2225 fn pending_complete_unknown_returns_false() {
2226 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2227 assert!(!pending.complete(
2228 &WorkflowId::new_v4(),
2229 &activity_id(99),
2230 None,
2231 Ok("orphan".to_owned())
2232 ));
2233 }
2234
2235 #[derive(Default)]
2236 struct RecordingOutboxCallback {
2237 completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
2238 failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
2239 live: bool,
2240 }
2241
2242 impl OutboxDeliveryCallback for RecordingOutboxCallback {
2243 fn deliver_completion(
2244 &self,
2245 workflow_id: &WorkflowId,
2246 activity_id: &ActivityId,
2247 run_id: Option<&RunId>,
2248 result: String,
2249 ) -> Result<bool, ServerError> {
2250 let _ = run_id;
2251 self.completions
2252 .lock()
2253 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2254 .push((workflow_id.clone(), activity_id.clone(), result));
2255 Ok(self.live)
2256 }
2257
2258 fn deliver_failure(
2259 &self,
2260 workflow_id: &WorkflowId,
2261 activity_id: &ActivityId,
2262 run_id: Option<&RunId>,
2263 reason: String,
2264 ) -> Result<bool, ServerError> {
2265 let _ = run_id;
2266 self.failures
2267 .lock()
2268 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2269 .push((workflow_id.clone(), activity_id.clone(), reason));
2270 Ok(self.live)
2271 }
2272 }
2273
2274 #[test]
2275 fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
2276 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2277 let callback = Arc::new(RecordingOutboxCallback {
2278 live: true,
2279 ..RecordingOutboxCallback::default()
2280 });
2281 pending.clone().set_outbox_delivery(callback.clone());
2283
2284 let workflow_id = WorkflowId::new_v4();
2285 let id = activity_id(7);
2286
2287 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
2290 let completions = callback
2291 .completions
2292 .lock()
2293 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
2294 assert_eq!(completions.len(), 1);
2295 assert_eq!(completions[0].0, workflow_id);
2296 assert_eq!(completions[0].1, id);
2297 assert_eq!(completions[0].2, "done");
2298 Ok(())
2299 }
2300
2301 #[test]
2302 fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
2303 -> Result<(), ServerError> {
2304 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2305 let callback = Arc::new(RecordingOutboxCallback::default());
2307 pending.set_outbox_delivery(callback.clone());
2308
2309 let workflow_id = WorkflowId::new_v4();
2310 let id = activity_id(8);
2311
2312 assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
2313 let failures = callback
2314 .failures
2315 .lock()
2316 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
2317 assert_eq!(failures.len(), 1);
2318 assert_eq!(failures[0].2, "retryable:boom");
2319 Ok(())
2320 }
2321
2322 #[test]
2323 fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
2324 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2327 assert!(!pending.complete(
2328 &WorkflowId::new_v4(),
2329 &activity_id(9),
2330 None,
2331 Ok("x".to_owned())
2332 ));
2333 }
2334
2335 #[test]
2336 fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
2337 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2338 let callback = Arc::new(RecordingOutboxCallback {
2339 live: true,
2340 ..RecordingOutboxCallback::default()
2341 });
2342 pending.set_outbox_delivery(callback.clone());
2343
2344 let workflow_id = WorkflowId::new_v4();
2345 let id = activity_id(10);
2346 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
2347
2348 assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
2349 assert_eq!(
2350 rx.recv_timeout(Duration::from_millis(50)),
2351 Ok(Ok("matched".to_owned()))
2352 );
2353 assert!(
2354 callback
2355 .completions
2356 .lock()
2357 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2358 .is_empty(),
2359 "a matched completion must deliver to its waiter, not the outbox callback"
2360 );
2361 Ok(())
2362 }
2363
2364 #[test]
2368 fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
2369 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2370 let workflow_id = WorkflowId::new_v4();
2371 let id = activity_id(11);
2372 let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
2373
2374 pending.park_activity(&workflow_id, &id)?;
2375 let result = rx
2376 .recv_timeout(Duration::from_millis(50))
2377 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2378 assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
2379 Ok(())
2380 }
2381
2382 #[test]
2386 fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
2387 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2388 let callback = Arc::new(RecordingOutboxCallback {
2389 live: true,
2390 ..RecordingOutboxCallback::default()
2391 });
2392 pending.set_outbox_delivery(callback.clone());
2393
2394 pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
2395
2396 assert!(
2397 callback
2398 .failures
2399 .lock()
2400 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2401 .is_empty(),
2402 "a park must never be delivered as an outbox failure"
2403 );
2404 assert!(
2405 callback
2406 .completions
2407 .lock()
2408 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2409 .is_empty(),
2410 "a park must never be delivered as an outbox completion"
2411 );
2412 Ok(())
2413 }
2414
2415 #[test]
2416 fn completion_sink_routes_success() -> Result<(), ServerError> {
2417 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2418 let workflow_id = WorkflowId::new_v4();
2419 let id = activity_id(2);
2420 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2421 let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
2422
2423 pending.complete_activity(ActivityCompletion {
2424 workflow_id,
2425 activity_id: id,
2426 run_id: None,
2427 completion_token,
2428 outcome: ActivityCompletionOutcome::Succeeded(payload),
2429 })?;
2430
2431 let result = rx
2432 .recv_timeout(Duration::from_millis(50))
2433 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2434 assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
2435 Ok(())
2436 }
2437
2438 #[test]
2439 fn accepted_hook_runs_before_the_result_is_observable() -> Result<(), ServerError> {
2440 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2441 let workflow_id = WorkflowId::new_v4();
2442 let id = activity_id(13);
2443 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2444 let accepted = AtomicBool::new(false);
2445
2446 pending.complete_activity_after_accept(
2447 ActivityCompletion {
2448 workflow_id,
2449 activity_id: id,
2450 run_id: None,
2451 completion_token,
2452 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2453 ContentType::Json,
2454 br#""ordered""#.to_vec(),
2455 )),
2456 },
2457 || {
2458 accepted.store(true, Ordering::SeqCst);
2459 Ok(())
2460 },
2461 )?;
2462
2463 let result = rx
2464 .recv_timeout(Duration::from_millis(50))
2465 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2466 assert_eq!(result, Ok(r#""ordered""#.to_owned()));
2467 assert!(
2468 accepted.load(Ordering::SeqCst),
2469 "the accepted-path cleanup hook must finish before the waiter can observe the result"
2470 );
2471 Ok(())
2472 }
2473
2474 #[test]
2475 fn a_non_publishing_settlement_restores_the_generation_for_the_true_resolver()
2476 -> Result<(), ServerError> {
2477 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2478 let workflow_id = WorkflowId::new_v4();
2479 let id = activity_id(16);
2480 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2481
2482 let settled = pending.complete_fenced_after_accept(
2487 &workflow_id,
2488 &id,
2489 None,
2490 &completion_token,
2491 Err("retryable:worker lost".to_owned()),
2492 || Ok(false),
2493 )?;
2494 assert!(!settled, "a non-publishing settlement must report false");
2495 assert!(
2496 rx.try_recv().is_err(),
2497 "a non-publishing settlement must not publish"
2498 );
2499
2500 let published = pending.complete_fenced_after_accept(
2503 &workflow_id,
2504 &id,
2505 None,
2506 &completion_token,
2507 Err("retryable:worker lost".to_owned()),
2508 || Ok(true),
2509 )?;
2510 assert!(
2511 published,
2512 "the true resolver must be accepted after a non-publishing settlement"
2513 );
2514 assert_eq!(
2515 rx.recv_timeout(Duration::from_millis(50))
2516 .map_err(|error| ServerError::worker_dispatch(
2517 "",
2518 "",
2519 format!("channel: {error}")
2520 ))?,
2521 Err("retryable:worker lost".to_owned())
2522 );
2523 Ok(())
2524 }
2525
2526 #[test]
2527 fn a_poisoned_tracker_fails_open_and_never_withholds_a_completion() {
2528 let tracker = HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW);
2529 tracker.poison_for_tests();
2530 let workflow_id = WorkflowId::new_v4();
2531 let id = activity_id(17);
2532 let was_tracked =
2533 clear_completed_task_tracking(&tracker, WorkerId::from_value(1), &workflow_id, &id);
2534 assert!(
2535 was_tracked,
2536 "a poisoned tracker must fail open (deliver), never withhold"
2537 );
2538 }
2539
2540 #[test]
2541 fn accepted_hook_failure_blocks_publication_and_restores_generation() -> Result<(), ServerError>
2542 {
2543 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2544 let workflow_id = WorkflowId::new_v4();
2545 let id = activity_id(15);
2546 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2547
2548 let rejected = pending.complete_activity_after_accept(
2549 ActivityCompletion {
2550 workflow_id: workflow_id.clone(),
2551 activity_id: id.clone(),
2552 run_id: None,
2553 completion_token: completion_token.clone(),
2554 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2555 ContentType::Json,
2556 br#""ordered""#.to_vec(),
2557 )),
2558 },
2559 || Err(ServerError::lock_poisoned("accepted-path settlement")),
2560 );
2561
2562 assert!(matches!(rejected, Err(ServerError::LockPoisoned { .. })));
2563 assert!(
2564 rx.try_recv().is_err(),
2565 "failed accepted-path settlement must not publish the completion"
2566 );
2567
2568 pending.complete_activity_after_accept(
2569 ActivityCompletion {
2570 workflow_id,
2571 activity_id: id,
2572 run_id: None,
2573 completion_token,
2574 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2575 ContentType::Json,
2576 br#""ordered""#.to_vec(),
2577 )),
2578 },
2579 || Ok(()),
2580 )?;
2581 assert_eq!(
2582 rx.recv_timeout(Duration::from_millis(50))
2583 .map_err(|error| ServerError::worker_dispatch(
2584 "",
2585 "",
2586 format!("channel: {error}")
2587 ))?,
2588 Ok(r#""ordered""#.to_owned())
2589 );
2590 Ok(())
2591 }
2592
2593 #[test]
2594 fn refused_generation_never_runs_the_accepted_hook() -> Result<(), ServerError> {
2595 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2596 let workflow_id = WorkflowId::new_v4();
2597 let id = activity_id(14);
2598 let (_completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2599 let accepted = AtomicBool::new(false);
2600
2601 let rejected = pending.complete_activity_after_accept(
2602 ActivityCompletion {
2603 workflow_id,
2604 activity_id: id,
2605 run_id: None,
2606 completion_token: CompletionToken::for_test(),
2607 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2608 ContentType::Json,
2609 br#""stale""#.to_vec(),
2610 )),
2611 },
2612 || {
2613 accepted.store(true, Ordering::SeqCst);
2614 Ok(())
2615 },
2616 );
2617
2618 assert!(matches!(
2619 rejected,
2620 Err(ServerError::ActivityCompletionRejected { .. })
2621 ));
2622 assert!(
2623 !accepted.load(Ordering::SeqCst),
2624 "a refused generation must preserve liveness by skipping accepted-path cleanup"
2625 );
2626 assert!(
2627 rx.try_recv().is_err(),
2628 "a refused generation must leave the current waiter unresolved"
2629 );
2630 Ok(())
2631 }
2632
2633 #[test]
2634 fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
2635 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2636 let workflow_id = WorkflowId::new_v4();
2637 let id = activity_id(12);
2638 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2639
2640 let malformed = pending.complete_activity(ActivityCompletion {
2641 workflow_id: workflow_id.clone(),
2642 activity_id: id.clone(),
2643 run_id: None,
2644 completion_token: completion_token.clone(),
2645 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2646 ContentType::Json,
2647 vec![0xff],
2648 )),
2649 });
2650 assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
2651 assert!(
2652 rx.try_recv().is_err(),
2653 "an invalid result must leave the waiter unresolved"
2654 );
2655
2656 pending.complete_activity(ActivityCompletion {
2657 workflow_id,
2658 activity_id: id,
2659 run_id: None,
2660 completion_token,
2661 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2662 ContentType::Json,
2663 br#""valid""#.to_vec(),
2664 )),
2665 })?;
2666 let result = rx
2667 .recv_timeout(Duration::from_millis(50))
2668 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2669 assert_eq!(result, Ok(r#""valid""#.to_owned()));
2670 Ok(())
2671 }
2672
2673 #[test]
2674 fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
2675 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2676 let workflow_id = WorkflowId::new_v4();
2677 let id = activity_id(3);
2678 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2679
2680 pending.complete_activity(ActivityCompletion {
2681 workflow_id,
2682 activity_id: id,
2683 run_id: None,
2684 completion_token,
2685 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2686 kind: ActivityErrorKind::Retryable,
2687 message: "temporary".to_owned(),
2688 details: None,
2689 }),
2690 })?;
2691
2692 let result = rx
2693 .recv_timeout(Duration::from_millis(50))
2694 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2695 assert_eq!(result, Err("retryable:temporary".to_owned()));
2696 Ok(())
2697 }
2698
2699 #[test]
2700 fn completion_sink_routes_policy_refusal_with_its_own_prefix() -> Result<(), ServerError> {
2701 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2702 let workflow_id = WorkflowId::new_v4();
2703 let id = activity_id(4);
2704 let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2705
2706 pending.complete_activity(ActivityCompletion {
2707 workflow_id,
2708 activity_id: id,
2709 run_id: None,
2710 completion_token,
2711 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2712 kind: ActivityErrorKind::PolicyRefused,
2713 message: "provider safety policy".to_owned(),
2714 details: None,
2715 }),
2716 })?;
2717
2718 let result = rx
2719 .recv_timeout(Duration::from_millis(50))
2720 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2721 assert_eq!(
2722 result,
2723 Err("policy_refused:provider safety policy".to_owned())
2724 );
2725 Ok(())
2726 }
2727
2728 #[test]
2737 fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
2738 -> Result<(), ServerError> {
2739 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2740 let post_restart_workflow = WorkflowId::new_v4();
2741 let pre_restart_workflow = WorkflowId::new_v4();
2742 let id = activity_id(1);
2744 let (completion_token, rx) = insert_attempt_one(&pending, &post_restart_workflow, &id)?;
2745
2746 let rejected = pending.complete_activity(ActivityCompletion {
2748 workflow_id: pre_restart_workflow,
2749 activity_id: id.clone(),
2750 run_id: None,
2751 completion_token: CompletionToken::for_test(),
2752 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2753 ContentType::Json,
2754 br#""stale""#.to_vec(),
2755 )),
2756 });
2757 assert!(matches!(
2758 rejected,
2759 Err(ServerError::ActivityCompletionRejected { .. })
2760 ));
2761 assert!(
2762 rx.try_recv().is_err(),
2763 "stale result for a different workflow must not complete this dispatch"
2764 );
2765
2766 pending.complete_activity(ActivityCompletion {
2768 workflow_id: post_restart_workflow,
2769 activity_id: id,
2770 run_id: None,
2771 completion_token,
2772 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2773 ContentType::Json,
2774 br#""fresh""#.to_vec(),
2775 )),
2776 })?;
2777 let result = rx
2778 .recv_timeout(Duration::from_millis(50))
2779 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2780 assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2781 Ok(())
2782 }
2783
2784 const TEST_HEARTBEAT_WINDOW: Duration = Duration::from_secs(5);
2791
2792 #[test]
2801 fn a_sink_derives_its_transport_budget_from_the_operator_window() {
2802 let window = Duration::from_secs(7);
2803 let sink = PendingActivities::new(window);
2804
2805 assert_eq!(
2806 sink.transport_losses().budget(),
2807 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2808 "the transport-loss budget must be the operator's heartbeat window times \
2809 TRANSPORT_LOSS_BUDGET_WINDOWS; a sink carrying any other budget is one no \
2810 operator declared"
2811 );
2812 }
2813
2814 #[test]
2822 fn a_dispatcher_carries_the_transport_budget_of_the_tracker_it_was_given() {
2823 let window = Duration::from_secs(11);
2824 let registry = ConnectedWorkerRegistry::default();
2825 let dispatcher =
2826 WorkerActivityDispatcher::new(registry, "default", HeartbeatTracker::new(window));
2827
2828 assert_eq!(
2829 dispatcher.pending().transport_losses().budget(),
2830 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2831 "the dispatcher's transport budget must be derived from the heartbeat window it \
2832 was constructed with; a zero or mismatched budget here means a worker loss is \
2833 declared transport-exhausted after ONE re-dispatch instead of after the \
2834 operator's window"
2835 );
2836 assert_ne!(
2837 dispatcher.pending().transport_losses().budget(),
2838 Duration::ZERO,
2839 "a ZERO budget is the specific regression this pin exists for: it grants exactly \
2840 one re-dispatchable loss and then declares the infrastructure flapping"
2841 );
2842 }
2843
2844 fn test_tracker() -> HeartbeatTracker {
2845 HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW)
2846 }
2847
2848 fn greet_request() -> ActivityDispatch {
2851 ActivityDispatch {
2852 namespace: "default".to_owned(),
2853 task_queue: "default".to_owned(),
2854 node: None,
2855 workflow_id: WorkflowId::new_v4(),
2856 run_id: RunId::new_v4(),
2857 activity_id: ActivityId::from_sequence_position(0),
2858 name: "greet".to_owned(),
2859 input: "{}".to_owned(),
2860 config: "{}".to_owned(),
2861 attempt: 1,
2862 labels: std::collections::BTreeMap::new(),
2863 advisory: false,
2864 }
2865 }
2866
2867 #[test]
2868 fn dispatcher_fails_immediately_when_draining_without_workers() {
2869 let registry = ConnectedWorkerRegistry::default();
2870 let drain = DrainState::default();
2871 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2872 .with_drain_state(drain.clone());
2873
2874 let _ = drain.begin();
2875
2876 let result = dispatcher.dispatch(greet_request());
2877
2878 assert!(result.is_err());
2879 let err = result.err().unwrap_or_default();
2880 assert!(
2881 err.contains("drain"),
2882 "expected drain rejection, got: {err}"
2883 );
2884 }
2885
2886 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2910 async fn a_retry_takes_over_a_timed_out_attempt_and_asks_its_worker_to_stop()
2911 -> Result<(), Box<dyn std::error::Error>> {
2912 let registry = ConnectedWorkerRegistry::default();
2913 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2914 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2915 let activity_types = [String::from("greet")];
2916 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2917 let worker_id = registration
2918 .worker_id()
2919 .ok_or("test worker registration missing id")?;
2920 let tracker = test_tracker();
2921 let dispatcher = Arc::new(
2922 WorkerActivityDispatcher::new(registry, "default", tracker.clone())
2923 .with_pending(pending.clone()),
2924 );
2925
2926 let request = greet_request();
2927 let workflow_id = request.workflow_id.clone();
2928 let activity_id = request.activity_id.clone();
2929
2930 let first_dispatcher = Arc::clone(&dispatcher);
2931 let first_request = request.clone();
2932 let first_attempt = tokio::spawn(futures::future::lazy(move |_| {
2933 first_dispatcher.dispatch(first_request)
2934 }));
2935 let Some(WorkerMessage::ActivityTask(first_task)) = worker_rx.recv().await else {
2936 return Err("expected attempt 1's task on the worker channel".into());
2937 };
2938 assert_eq!(first_task.attempt, 1);
2939
2940 let mut retry_request = request.clone();
2941 retry_request.attempt = 2;
2942 let retry_dispatcher = Arc::clone(&dispatcher);
2943 let retry_attempt = tokio::spawn(futures::future::lazy(move |_| {
2944 retry_dispatcher.dispatch(retry_request)
2945 }));
2946
2947 let Some(WorkerMessage::CancelActivity(cancel)) = worker_rx.recv().await else {
2950 return Err(
2951 "expected the superseded attempt's cancel frame before attempt 2's task".into(),
2952 );
2953 };
2954 assert_eq!(
2955 cancel
2956 .activity_id
2957 .map(|id| id.sequence_position)
2958 .unwrap_or_default(),
2959 activity_id.sequence_position()
2960 );
2961 let Some(WorkerMessage::ActivityTask(retry_task)) = worker_rx.recv().await else {
2962 return Err("expected attempt 2's task on the worker channel".into());
2963 };
2964 assert_eq!(retry_task.attempt, 2);
2965
2966 let in_flight = tracker.in_flight_for_workflow(&workflow_id)?;
2969 assert_eq!(
2970 in_flight
2971 .iter()
2972 .map(|task| task.attempt)
2973 .collect::<Vec<_>>(),
2974 vec![2],
2975 "exactly attempt 2 must remain tracked for worker {worker_id:?}"
2976 );
2977
2978 let retry_token = CompletionToken::from_wire(
2979 &workflow_id,
2980 &activity_id,
2981 retry_task.completion_token.clone(),
2982 )
2983 .map_err(|error| error.to_string())?;
2984 pending.complete_activity(ActivityCompletion {
2985 workflow_id: workflow_id.clone(),
2986 activity_id: activity_id.clone(),
2987 run_id: None,
2988 completion_token: retry_token,
2989 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2990 ContentType::Json,
2991 br#""attempt 2 delivered""#.to_vec(),
2992 )),
2993 })?;
2994
2995 let retry_result = retry_attempt.await.map_err(|error| error.to_string())?;
2996 assert_eq!(retry_result, Ok(r#""attempt 2 delivered""#.to_owned()));
2997 let first_result = first_attempt.await.map_err(|error| error.to_string())?;
2998 let first_reason = match first_result {
2999 Ok(delivered) => {
3000 return Err(
3001 format!("attempt 1 resolved Ok({delivered}) instead of superseded").into(),
3002 );
3003 }
3004 Err(reason) => reason,
3005 };
3006 assert!(
3007 first_reason.starts_with("retryable:")
3008 && first_reason.contains("superseded by attempt 2"),
3009 "{first_reason}"
3010 );
3011 registration.deregister()?;
3012 Ok(())
3013 }
3014
3015 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3016 async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
3017 -> Result<(), Box<dyn std::error::Error>> {
3018 let registry = ConnectedWorkerRegistry::default();
3019 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
3020 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
3021 let activity_types = [String::from("greet")];
3022 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
3023
3024 let sink = pending.clone();
3025 let echo_worker = tokio::spawn(async move {
3026 let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
3027 return Err("expected an activity task on the worker channel".to_owned());
3028 };
3029 let workflow_id = task
3030 .workflow_id
3031 .ok_or("task missing workflow id")
3032 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
3033 let activity_id = task
3034 .activity_id
3035 .map(ActivityId::from)
3036 .ok_or("task missing activity id")?;
3037 let completion_token =
3038 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
3039 .map_err(|error| error.to_string())?;
3040 sink.complete_activity(ActivityCompletion {
3041 workflow_id,
3042 activity_id,
3043 run_id: None,
3044 completion_token,
3045 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
3046 ContentType::Json,
3047 br#"{"greeting":"hello"}"#.to_vec(),
3048 )),
3049 })
3050 .map_err(|error| error.to_string())
3051 });
3052
3053 let dispatcher = Arc::new(
3054 WorkerActivityDispatcher::new(registry, "default", test_tracker())
3055 .with_pending(pending),
3056 );
3057 let started = Instant::now();
3058 let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
3061 dispatcher.dispatch(greet_request())
3062 }));
3063 let result = dispatch_task.await.map_err(|error| error.to_string())?;
3064 let elapsed = started.elapsed();
3065
3066 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
3067 assert!(
3068 elapsed < Duration::from_secs(5),
3069 "dispatch round trip took {elapsed:?}; task delivery must not \
3070 depend on the blocked dispatch thread"
3071 );
3072 echo_worker.await.map_err(|error| error.to_string())??;
3073 registration.deregister()?;
3074 Ok(())
3075 }
3076
3077 #[tokio::test]
3081 async fn dispatch_on_current_thread_runtime_fails_fast()
3082 -> Result<(), Box<dyn std::error::Error>> {
3083 let registry = ConnectedWorkerRegistry::default();
3084 let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
3085 let activity_types = [String::from("greet")];
3086 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
3087 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
3088
3089 let started = Instant::now();
3090 let result = dispatcher.dispatch(greet_request());
3091 let elapsed = started.elapsed();
3092
3093 let err = result.err().ok_or("expected dispatch to fail")?;
3094 assert!(
3095 err.contains("multi-thread tokio runtime"),
3096 "unexpected error: {err}"
3097 );
3098 assert!(
3099 elapsed < Duration::from_secs(5),
3100 "fail-fast path took {elapsed:?}"
3101 );
3102 registration.deregister()?;
3103 Ok(())
3104 }
3105
3106 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3112 async fn dispatch_pinned_to_node_reaches_only_that_node()
3113 -> Result<(), Box<dyn std::error::Error>> {
3114 let registry = ConnectedWorkerRegistry::default();
3115 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
3116 let activity_types = [String::from("greet")];
3117 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
3118 let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
3119 let on_n2 = registry.register_namespaces(
3128 [String::from("default")],
3129 "default",
3130 Some(String::from("n2")),
3131 activity_types.iter(),
3132 n2_tx,
3133 )?;
3134 let on_n1 = registry.register_namespaces(
3135 [String::from("default")],
3136 "default",
3137 Some(String::from("n1")),
3138 activity_types.iter(),
3139 n1_tx,
3140 )?;
3141
3142 let sink = pending.clone();
3146 let echo_n1 = tokio::spawn(async move {
3147 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
3148 return Err("expected an activity task on the n1 worker channel".to_owned());
3149 };
3150 let workflow_id = task
3151 .workflow_id
3152 .ok_or("task missing workflow id")
3153 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
3154 let activity_id = task
3155 .activity_id
3156 .map(ActivityId::from)
3157 .ok_or("task missing activity id")?;
3158 let completion_token =
3159 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
3160 .map_err(|error| error.to_string())?;
3161 sink.complete_activity(ActivityCompletion {
3162 workflow_id,
3163 activity_id,
3164 run_id: None,
3165 completion_token,
3166 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
3167 ContentType::Json,
3168 br#"{"greeting":"hello"}"#.to_vec(),
3169 )),
3170 })
3171 .map_err(|error| error.to_string())
3172 });
3173
3174 let dispatcher = Arc::new(
3175 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
3176 .with_pending(pending),
3177 );
3178
3179 let pinned = ActivityDispatch {
3180 node: Some(String::from("n1")),
3181 ..greet_request()
3182 };
3183 let started = Instant::now();
3184 let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
3185 .await
3186 .map_err(|error| error.to_string())?;
3187 let elapsed = started.elapsed();
3188
3189 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
3190 assert!(
3191 elapsed < Duration::from_secs(5),
3192 "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
3193 );
3194 echo_n1.await.map_err(|error| error.to_string())??;
3195
3196 assert!(
3198 n2_rx.try_recv().is_err(),
3199 "node=Some(\"n1\") dispatch must not reach the n2 worker"
3200 );
3201
3202 on_n1.deregister()?;
3203 on_n2.deregister()?;
3204 Ok(())
3205 }
3206
3207 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3210 async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
3211 -> Result<(), Box<dyn std::error::Error>> {
3212 let registry = ConnectedWorkerRegistry::default();
3213 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
3214 let activity_types = [String::from("greet")];
3215 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
3216 let on_n1 = registry.register_namespaces(
3217 [String::from("default")],
3218 "default",
3219 Some(String::from("n1")),
3220 activity_types.iter(),
3221 n1_tx,
3222 )?;
3223
3224 let sink = pending.clone();
3225 let echo = tokio::spawn(async move {
3226 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
3227 return Err("expected an activity task on the worker channel".to_owned());
3228 };
3229 let workflow_id = task
3230 .workflow_id
3231 .ok_or("task missing workflow id")
3232 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
3233 let activity_id = task
3234 .activity_id
3235 .map(ActivityId::from)
3236 .ok_or("task missing activity id")?;
3237 let completion_token =
3238 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
3239 .map_err(|error| error.to_string())?;
3240 sink.complete_activity(ActivityCompletion {
3241 workflow_id,
3242 activity_id,
3243 run_id: None,
3244 completion_token,
3245 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
3246 ContentType::Json,
3247 br#"{"greeting":"hello"}"#.to_vec(),
3248 )),
3249 })
3250 .map_err(|error| error.to_string())
3251 });
3252
3253 let dispatcher = Arc::new(
3254 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
3255 .with_pending(pending),
3256 );
3257
3258 let result = tokio::spawn(futures::future::lazy(move |_| {
3260 dispatcher.dispatch(greet_request())
3261 }))
3262 .await
3263 .map_err(|error| error.to_string())?;
3264
3265 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
3266 echo.await.map_err(|error| error.to_string())??;
3267 on_n1.deregister()?;
3268 Ok(())
3269 }
3270}