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, 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 activity_id: ActivityId,
196 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
197 let completion_token = self.completion_fences.issue(&workflow_id, &activity_id)?;
198 let (tx, rx) = std::sync::mpsc::sync_channel(1);
199 let key = (workflow_id, activity_id);
200 match self.pending.entry(key) {
201 Entry::Vacant(entry) => {
202 entry.insert(tx);
203 Ok((completion_token, rx))
204 }
205 Entry::Occupied(entry) => {
206 let (workflow_id, activity_id) = entry.key();
207 Err(ServerError::PendingActivityCollision {
208 workflow_id: workflow_id.clone(),
209 activity_id: activity_id.clone(),
210 })
211 }
212 }
213 }
214
215 #[must_use]
217 pub fn completion_fences(&self) -> CompletionFences {
218 self.completion_fences.clone()
219 }
220
221 #[cfg(test)]
225 pub(crate) fn insert_for_test(
226 &self,
227 workflow_id: WorkflowId,
228 activity_id: ActivityId,
229 ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
230 self.insert(workflow_id, activity_id)
231 }
232
233 pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
239 if self.outbox_delivery.set(callback).is_err() {
240 tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
241 }
242 }
243
244 fn complete(
252 &self,
253 workflow_id: &WorkflowId,
254 activity_id: &ActivityId,
255 run_id: Option<&RunId>,
256 result: Result<String, String>,
257 ) -> bool {
258 let matched = self
261 .pending
262 .remove(&(workflow_id.clone(), activity_id.clone()));
263 if let Some((_, sender)) = matched {
264 return sender.send(result).is_ok();
265 }
266 let Some(callback) = self.outbox_delivery.get() else {
267 return false;
269 };
270 let outcome = match result {
271 Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
272 Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
273 };
274 match outcome {
275 Ok(true) => true,
276 Ok(false) => {
277 tracing::debug!(
279 workflow_id = %workflow_id,
280 activity_id = %activity_id,
281 "unmatched outbox completion for a workflow that is not currently live; \
282 recovery will re-arm it"
283 );
284 false
285 }
286 Err(error) => {
287 tracing::warn!(
288 workflow_id = %workflow_id,
289 activity_id = %activity_id,
290 %error,
291 "failed to deliver unmatched outbox completion to the live workflow"
292 );
293 false
294 }
295 }
296 }
297
298 fn complete_fenced(
301 &self,
302 workflow_id: &WorkflowId,
303 activity_id: &ActivityId,
304 run_id: Option<&RunId>,
305 completion_token: &CompletionToken,
306 result: Result<String, String>,
307 ) -> Result<bool, ServerError> {
308 self.completion_fences
309 .accept(workflow_id, activity_id, completion_token)
310 .inspect_err(|error| {
311 tracing::warn!(
312 workflow_id = %workflow_id,
313 activity_id = %activity_id,
314 %error,
315 "activity completion rejected by execution-generation fence"
316 );
317 })?;
318 let transport_domain = result
324 .as_ref()
325 .err()
326 .is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
327 if !transport_domain {
328 if let Err(error) = self.transport_losses.clear(workflow_id, activity_id) {
329 tracing::warn!(
330 workflow_id = %workflow_id,
331 activity_id = %activity_id,
332 %error,
333 "failed to retire the transport-loss budget for a resolved activity"
334 );
335 }
336 }
337 Ok(self.complete(workflow_id, activity_id, run_id, result))
338 }
339
340 #[must_use]
352 pub fn new(heartbeat_window: std::time::Duration) -> Self {
353 Self {
354 pending: Arc::default(),
355 completion_fences: CompletionFences::default(),
356 outbox_delivery: Arc::default(),
357 transport_losses: super::transport_loss::TransportLossLedger::new(heartbeat_window),
358 }
359 }
360
361 #[must_use]
363 pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
364 &self.transport_losses
365 }
366
367 fn classify_worker_loss(
375 &self,
376 workflow_id: &WorkflowId,
377 activity_id: &ActivityId,
378 worker_id: crate::worker::registry::WorkerId,
379 ) -> String {
380 let detail = super::transport_loss::worker_lost_detail(worker_id);
381 match self
382 .transport_losses
383 .record_loss(workflow_id, activity_id, &detail)
384 {
385 Ok(verdict) => {
386 if verdict.exhausted {
387 tracing::error!(
388 operation = "activity_complete",
389 workflow_id = %workflow_id,
390 activity_id = %activity_id,
391 worker_id = ?worker_id,
392 error_type = "TransportExhausted",
393 losses = verdict.losses,
394 budget_ms = self.transport_losses.budget().as_millis(),
395 "activity abandoned: the transport kept losing its worker past the \
396 transport-loss budget"
397 );
398 } else {
399 tracing::warn!(
400 operation = "activity_complete",
401 workflow_id = %workflow_id,
402 activity_id = %activity_id,
403 worker_id = ?worker_id,
404 error_type = "WorkerLost",
405 losses = verdict.losses,
406 budget_ms = self.transport_losses.budget().as_millis(),
407 "worker lost before reporting an activity result; the activity never ran \
408 and will be re-dispatched attempt-neutrally"
409 );
410 }
411 verdict.reason
412 }
413 Err(error) => {
414 tracing::error!(
415 workflow_id = %workflow_id,
416 activity_id = %activity_id,
417 %error,
418 "transport-loss ledger is unreadable; abandoning the activity rather than \
419 re-dispatching it without a budget"
420 );
421 format!(
422 "{}{detail} (transport-loss budget unreadable: {error})",
423 super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
424 )
425 }
426 }
427 }
428}
429
430impl ActivityCompletionSink for PendingActivities {
431 fn park_activity(
442 &self,
443 workflow_id: &WorkflowId,
444 activity_id: &ActivityId,
445 ) -> Result<(), ServerError> {
446 self.completion_fences
447 .revoke_current(workflow_id, activity_id)?;
448 let matched = self
449 .pending
450 .remove(&(workflow_id.clone(), activity_id.clone()));
451 if let Some((_, sender)) = matched {
452 let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
456 }
457 Ok(())
458 }
459
460 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
461 let result = match completion.outcome {
462 ActivityCompletionOutcome::Succeeded(payload) => {
463 payload_to_string(&payload).map_err(|reason| {
464 tracing::error!(
465 operation = "activity_complete",
466 workflow_id = %completion.workflow_id,
467 activity_id = %completion.activity_id,
468 error_type = "ActivityResultDecode",
469 %reason,
470 "activity completion failed"
471 );
472 ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
473 })?
474 }
475 ActivityCompletionOutcome::Failed(error) => {
476 let prefix = match error.kind {
477 ActivityErrorKind::Retryable => "retryable",
478 ActivityErrorKind::PolicyRefused => "policy_refused",
479 ActivityErrorKind::Terminal => "terminal",
480 };
481 tracing::error!(
482 operation = "activity_complete",
483 workflow_id = %completion.workflow_id,
484 activity_id = %completion.activity_id,
485 error_type = "ActivityFailed",
486 error_kind = prefix,
487 reason = %error.message,
488 "activity completion failed"
489 );
490 Err(format!("{prefix}:{}", error.message))
491 }
492 ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
497 &completion.workflow_id,
498 &completion.activity_id,
499 worker_id,
500 )),
501 };
502 self.complete_fenced(
503 &completion.workflow_id,
504 &completion.activity_id,
505 completion.run_id.as_ref(),
506 &completion.completion_token,
507 result,
508 )?;
509 Ok(())
510 }
511}
512
513fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
514 match payload.content_type() {
515 ContentType::Json => String::from_utf8(payload.bytes().to_vec())
516 .map(Ok)
517 .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
518 }
519}
520
521pub struct WorkerActivityDispatcher {
529 registry: ConnectedWorkerRegistry,
530 namespace: String,
531 pending: PendingActivities,
532 heartbeat_tracker: HeartbeatTracker,
533 drain_state: DrainState,
534 tokio_handle: Option<tokio::runtime::Handle>,
535 attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
541 queue_service: QueueServiceConfig,
545 queue_declarations: QueueDeclarationSource,
548 queue_state: QueueServiceState,
550 cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
554}
555
556impl std::fmt::Debug for WorkerActivityDispatcher {
557 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
558 f.debug_struct("WorkerActivityDispatcher")
559 .field("namespace", &self.namespace)
560 .finish_non_exhaustive()
561 }
562}
563
564impl WorkerActivityDispatcher {
565 #[must_use]
573 pub fn new(
574 registry: ConnectedWorkerRegistry,
575 namespace: impl Into<String>,
576 heartbeat_tracker: HeartbeatTracker,
577 ) -> Self {
578 Self {
579 registry,
580 namespace: namespace.into(),
581 pending: PendingActivities::new(heartbeat_tracker.heartbeat_window()),
585 heartbeat_tracker,
586 drain_state: DrainState::default(),
587 tokio_handle: None,
588 attempt_owners: None,
589 queue_service: QueueServiceConfig::default(),
590 queue_declarations: QueueDeclarationSource::default(),
591 queue_state: QueueServiceState::default(),
592 cluster_publisher: None,
593 }
594 }
595
596 #[must_use]
600 pub fn with_cluster_publisher(
601 mut self,
602 cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
603 ) -> Self {
604 self.cluster_publisher = Some(cluster_publisher);
605 self
606 }
607
608 #[must_use]
611 pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
612 self.queue_service = queue_service;
613 self
614 }
615
616 #[must_use]
620 pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
621 self.queue_declarations = queue_declarations;
622 self
623 }
624
625 #[must_use]
628 pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
629 self.queue_state = queue_state;
630 self
631 }
632
633 #[must_use]
640 pub fn with_attempt_owners(
641 mut self,
642 attempt_owners: super::intervention::AttemptOwnerIndex,
643 ) -> Self {
644 self.attempt_owners = Some(attempt_owners);
645 self
646 }
647
648 #[must_use]
653 pub const fn pending(&self) -> &PendingActivities {
654 &self.pending
655 }
656
657 #[must_use]
659 pub fn with_pending(mut self, pending: PendingActivities) -> Self {
660 self.pending = pending;
661 self
662 }
663
664 #[must_use]
666 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
667 self.drain_state = drain_state;
668 self
669 }
670
671 #[must_use]
673 pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
674 self.tokio_handle = Some(tokio_handle);
675 self
676 }
677}
678
679impl WorkerActivityDispatcher {
680 fn ensure_accepting(
704 &self,
705 namespace: &str,
706 activity_type: &str,
707 workflow_id: &WorkflowId,
708 activity_id: &ActivityId,
709 worker_id: Option<WorkerId>,
710 ) -> Result<(), String> {
711 self.drain_state
712 .ensure_accepting(namespace, activity_type)
713 .map_err(|error| {
714 log_worker_error(
715 "WorkerDispatch",
716 namespace,
717 activity_type,
718 workflow_id,
719 activity_id,
720 worker_id,
721 &error.to_string(),
722 );
723 aion::PARKED_ACTIVITY_REASON.to_owned()
724 })
725 }
726
727 fn select_worker_or_wait(
735 &self,
736 address: &ServiceAddress,
737 workflow_id: &WorkflowId,
738 activity_id: &ActivityId,
739 ) -> Result<WorkerHandle, String> {
740 let wait = ServiceWait {
741 registry: &self.registry,
742 declarations: &self.queue_declarations,
743 config: &self.queue_service,
744 state: &self.queue_state,
745 address,
746 workflow_id,
747 activity_id,
748 publisher: self.cluster_publisher.as_ref(),
749 };
750 let mut accepting = || {
751 self.ensure_accepting(
752 &address.namespace,
753 &address.activity_type,
754 workflow_id,
755 activity_id,
756 None,
757 )
758 };
759 let mut park = |budget: Option<Duration>| self.park_for_worker(budget);
760 select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
761 let reason = refusal.reason_string();
762 if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
765 let error_type = match refusal {
766 SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
767 _ => "WorkerRegistry",
768 };
769 log_worker_error(
770 error_type,
771 &address.namespace,
772 &address.activity_type,
773 workflow_id,
774 activity_id,
775 None,
776 &reason,
777 );
778 }
779 reason
780 })
781 }
782
783 fn park_for_worker(&self, budget: Option<Duration>) {
811 let handle = self
812 .tokio_handle
813 .clone()
814 .or_else(|| tokio::runtime::Handle::try_current().ok());
815 let Some(handle) = handle else {
816 std::thread::sleep(
821 budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
822 );
823 return;
824 };
825 handle.block_on(async {
826 let arrival = async {
827 tokio::select! {
828 () = self.registry.wait_for_worker() => {}
829 () = self.drain_state.wait_for_drain() => {}
830 }
831 };
832 match budget {
833 None => arrival.await,
834 Some(budget) => {
835 drop(tokio::time::timeout(budget, arrival).await);
838 }
839 }
840 });
841 }
842
843 fn track_worker_task(
844 &self,
845 worker_id: WorkerId,
846 activity_type: &str,
847 workflow_id: &WorkflowId,
848 activity_id: &ActivityId,
849 attempt: u32,
850 completion_token: CompletionToken,
851 ) -> Result<(), String> {
852 self.heartbeat_tracker
853 .track_task(
854 worker_id,
855 InFlightActivity {
856 workflow_id: workflow_id.clone(),
857 activity_id: activity_id.clone(),
858 attempt,
863 completion_token,
864 },
865 Instant::now(),
866 )
867 .map_err(|error| {
868 let reason = error.to_string();
869 log_worker_error(
870 "WorkerHeartbeatTracker",
871 &self.namespace,
872 activity_type,
873 workflow_id,
874 activity_id,
875 Some(worker_id),
876 &reason,
877 );
878 reason
879 })
880 }
881
882 fn cleanup_activity(
883 &self,
884 worker_id: WorkerId,
885 workflow_id: &WorkflowId,
886 activity_id: &ActivityId,
887 completion_token: &CompletionToken,
888 ) {
889 self.pending
890 .pending
891 .remove(&(workflow_id.clone(), activity_id.clone()));
892 if let Err(error) =
893 self.pending
894 .completion_fences
895 .revoke(workflow_id, activity_id, completion_token)
896 {
897 tracing::error!(
898 workflow_id = %workflow_id,
899 activity_id = %activity_id,
900 %error,
901 "failed to revoke undelivered activity generation"
902 );
903 }
904 let _ = self
905 .heartbeat_tracker
906 .complete_task(worker_id, workflow_id, activity_id);
907 self.drain_state.notify_activity_drained();
908 }
909
910 fn send_activity_task(
920 &self,
921 worker: &WorkerHandle,
922 task: ProtoActivityTask,
923 address: &ServiceAddress,
924 workflow_id: &WorkflowId,
925 activity_id: &ActivityId,
926 completion_token: &CompletionToken,
927 ) -> Result<(), String> {
928 match worker.delivery() {
929 WorkerDelivery::Grpc(sender) => {
930 let worker_id = worker.id();
931 let mut accepting = || {
932 self.ensure_accepting(
933 &address.namespace,
934 &address.activity_type,
935 workflow_id,
936 activity_id,
937 Some(worker_id),
938 )
939 };
940 let handed_over = deliver_within_schedule_to_start(
941 sender,
942 WorkerMessage::ActivityTask(Box::new(task)),
943 self.queue_service.schedule_to_start_timeout,
944 &mut accepting,
945 );
946 let Err(refusal) = handed_over else {
947 return Ok(());
948 };
949 self.cleanup_activity(worker_id, workflow_id, activity_id, completion_token);
950 let (error_type, reason) = self.hand_off_failure(&refusal, address);
951 if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
952 log_worker_error(
953 error_type,
954 &address.namespace,
955 &address.activity_type,
956 workflow_id,
957 activity_id,
958 Some(worker_id),
959 &reason,
960 );
961 }
962 Err(reason)
963 }
964 #[cfg(feature = "liminal-transport")]
971 WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
972 worker.id(),
973 delivery,
974 task,
975 &address.activity_type,
976 workflow_id,
977 activity_id,
978 ),
979 }
980 }
981
982 fn hand_off_failure(
989 &self,
990 refusal: &DeliveryRefusal,
991 address: &ServiceAddress,
992 ) -> (&'static str, String) {
993 match refusal {
994 DeliveryRefusal::Saturated { waited } => {
995 let census = self
998 .registry
999 .pool_census(
1000 &address.namespace,
1001 &address.task_queue,
1002 &address.activity_type,
1003 address.node.as_deref(),
1004 )
1005 .unwrap_or_else(|error| {
1006 tracing::error!(
1007 namespace = %address.namespace,
1008 task_queue = %address.task_queue,
1009 activity_type = %address.activity_type,
1010 %error,
1011 "poller census failed while reporting a saturated queue; \
1012 the refusal carries an empty census"
1013 );
1014 PoolCensus::default()
1015 });
1016 let unavailable = WorkerUnavailable {
1017 reason: QueueServiceReason::Saturated,
1018 clock: Some(ExpiredClock::ScheduleToStart),
1019 waited: *waited,
1020 address: address.clone(),
1021 census,
1022 };
1023 ("WorkerUnavailable", unavailable.reason_string())
1024 }
1025 DeliveryRefusal::Full => (
1026 "WorkerChannelClosed",
1027 "worker task channel full or closed: no available capacity".to_owned(),
1028 ),
1029 DeliveryRefusal::Closed => (
1030 "WorkerChannelClosed",
1031 "worker task channel full or closed: channel closed".to_owned(),
1032 ),
1033 DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
1034 }
1035 }
1036
1037 #[cfg(feature = "liminal-transport")]
1061 fn send_liminal_activity_task(
1062 &self,
1063 worker_id: WorkerId,
1064 delivery: &super::liminal_transport::LiminalWorkerDelivery,
1065 task: ProtoActivityTask,
1066 activity_type: &str,
1067 workflow_id: &WorkflowId,
1068 activity_id: &ActivityId,
1069 ) -> Result<(), String> {
1070 let completion_token =
1071 CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1072 .map_err(|error| error.to_string())?;
1073 let heartbeat_window_ms =
1074 u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1075 .unwrap_or(u64::MAX);
1076 let attempt = task.attempt;
1077 let run_id = task
1084 .run_id
1085 .map(RunId::try_from)
1086 .transpose()
1087 .map_err(|error| error.to_string())?
1088 .ok_or_else(|| {
1089 "activity task run id is missing; refusing to dispatch an unidentified run"
1090 .to_owned()
1091 })?;
1092 let request = super::liminal_transport::DispatchRequest {
1093 activity_type: activity_type.to_owned(),
1094 workflow_id: workflow_id.clone(),
1095 ordinal: activity_id.sequence_position(),
1096 run_id: Some(run_id.clone()),
1097 attempt,
1098 completion_token: task.completion_token,
1099 idempotency_key: task.idempotency_key,
1100 labels: task.labels.into_iter().collect(),
1101 heartbeat_window_ms,
1102 input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1103 };
1104 let awaiter = match delivery.push_dispatch(&request) {
1108 Ok(awaiter) => awaiter,
1109 Err(error) => {
1110 let reason = format!("worker liminal push failed: {error}");
1111 self.cleanup_activity(worker_id, workflow_id, activity_id, &completion_token);
1112 log_worker_error(
1113 "WorkerChannelClosed",
1114 &self.namespace,
1115 activity_type,
1116 workflow_id,
1117 activity_id,
1118 Some(worker_id),
1119 &reason,
1120 );
1121 return Err(reason);
1122 }
1123 };
1124 let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1128 super::liminal_transport::AttemptOwnerGuard::bind(
1129 owners.clone(),
1130 super::intervention::AttemptKey::new(
1131 workflow_id.clone(),
1132 run_id.clone(),
1133 activity_id.clone(),
1134 attempt,
1135 ),
1136 worker_id,
1137 )
1138 });
1139 self.spawn_liminal_reply_router(
1140 worker_id,
1141 awaiter,
1142 workflow_id,
1143 activity_id,
1144 &completion_token,
1145 owner_binding,
1146 );
1147 Ok(())
1148 }
1149
1150 #[cfg(feature = "liminal-transport")]
1190 fn spawn_liminal_reply_router(
1191 &self,
1192 worker_id: WorkerId,
1193 awaiter: liminal_server::server::connection::PushReplyAwaiter,
1194 workflow_id: &WorkflowId,
1195 activity_id: &ActivityId,
1196 completion_token: &CompletionToken,
1197 owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1198 ) {
1199 let pending = self.pending.clone();
1200 let heartbeat_tracker = self.heartbeat_tracker.clone();
1201 let drain_state = self.drain_state.clone();
1202 let workflow_id = workflow_id.clone();
1203 let activity_id = activity_id.clone();
1204 let completion_token = completion_token.clone();
1205 std::thread::spawn(move || {
1206 let _owner_binding = owner_binding;
1210 route_liminal_reply(
1211 &pending,
1212 &heartbeat_tracker,
1213 &drain_state,
1214 &awaiter,
1215 (worker_id, &workflow_id, &activity_id, &completion_token),
1216 );
1217 });
1218 }
1219
1220 fn await_activity_result(
1224 &self,
1225 context: &ActivityDispatchContext<'_>,
1226 rx: &SyncReceiver,
1227 ) -> Result<String, String> {
1228 match self.registry.is_registered(context.worker_id) {
1236 Ok(true) => {}
1237 Ok(false) => {
1238 if let Ok(result) = rx.try_recv() {
1242 return self.deliver_result(context, result);
1243 }
1244 self.cleanup_activity(
1245 context.worker_id,
1246 context.workflow_id,
1247 context.activity_id,
1248 &context.completion_token,
1249 );
1250 let reason = self.pending.classify_worker_loss(
1255 context.workflow_id,
1256 context.activity_id,
1257 context.worker_id,
1258 );
1259 log_worker_error(
1260 "WorkerLost",
1261 &self.namespace,
1262 context.activity_type,
1263 context.workflow_id,
1264 context.activity_id,
1265 Some(context.worker_id),
1266 &reason,
1267 );
1268 return Err(reason);
1269 }
1270 Err(error) => {
1271 self.cleanup_activity(
1272 context.worker_id,
1273 context.workflow_id,
1274 context.activity_id,
1275 &context.completion_token,
1276 );
1277 let reason = format!("worker registry inspection failed: {error}");
1278 log_worker_error(
1279 "WorkerRegistry",
1280 &self.namespace,
1281 context.activity_type,
1282 context.workflow_id,
1283 context.activity_id,
1284 Some(context.worker_id),
1285 &reason,
1286 );
1287 return Err(reason);
1288 }
1289 }
1290 if let Ok(result) = rx.recv() {
1291 return self.deliver_result(context, result);
1292 }
1293 self.cleanup_activity(
1296 context.worker_id,
1297 context.workflow_id,
1298 context.activity_id,
1299 &context.completion_token,
1300 );
1301 let reason = "activity response channel dropped".to_owned();
1302 log_worker_error(
1303 "WorkerChannelClosed",
1304 &self.namespace,
1305 context.activity_type,
1306 context.workflow_id,
1307 context.activity_id,
1308 Some(context.worker_id),
1309 &reason,
1310 );
1311 Err(reason)
1312 }
1313
1314 fn deliver_result(
1315 &self,
1316 context: &ActivityDispatchContext<'_>,
1317 result: Result<String, String>,
1318 ) -> Result<String, String> {
1319 self.pending
1320 .pending
1321 .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
1322 if let Err(reason) = &result
1327 && aion::is_parked_reason(reason)
1328 {
1329 tracing::info!(
1330 operation = "activity_dispatch",
1331 namespace = %self.namespace,
1332 workflow_id = %context.workflow_id,
1333 activity_id = %context.activity_id,
1334 activity_type = context.activity_type,
1335 worker_id = ?context.worker_id,
1336 "activity parked for restart recovery"
1337 );
1338 return result;
1339 }
1340 log_activity_completion(context, result.is_ok());
1341 result.inspect_err(|reason| {
1342 log_worker_error(
1343 "ActivityFailed",
1344 &self.namespace,
1345 context.activity_type,
1346 context.workflow_id,
1347 context.activity_id,
1348 Some(context.worker_id),
1349 reason,
1350 );
1351 })
1352 }
1353}
1354
1355impl ActivityDispatcher for WorkerActivityDispatcher {
1356 fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1357 match tokio::runtime::Handle::try_current() {
1358 Ok(handle) => match handle.runtime_flavor() {
1359 tokio::runtime::RuntimeFlavor::MultiThread => {
1360 tokio::task::block_in_place(|| self.dispatch_blocking(request))
1367 }
1368 flavor => Err(format!(
1369 "activity dispatch blocks the calling thread until the worker responds; \
1370 a {flavor:?} tokio runtime cannot host that wait because the worker \
1371 stream forwarder shares its only executor thread and the task could \
1372 never be delivered — run the engine on a multi-thread tokio runtime"
1373 )),
1374 },
1375 Err(_) => self.dispatch_blocking(request),
1379 }
1380 }
1381}
1382
1383impl WorkerActivityDispatcher {
1384 fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1401 let ActivityDispatch {
1402 namespace,
1403 task_queue,
1404 node,
1408 workflow_id,
1409 run_id,
1410 activity_id,
1411 name,
1412 input,
1413 config: _,
1414 attempt,
1415 labels,
1416 advisory: _,
1421 } = request;
1422 let started_at = Instant::now();
1423 self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1424 let address = ServiceAddress {
1425 namespace: namespace.clone(),
1426 task_queue: task_queue.clone(),
1427 activity_type: name.clone(),
1428 node: node.clone(),
1429 };
1430 let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1431 let worker_id = worker.id();
1432 let span = info_span!(
1433 "activity_dispatch",
1434 operation = "activity_dispatch",
1435 namespace = %namespace,
1436 task_queue = %task_queue,
1437 node = node.as_deref(),
1438 workflow_id = %workflow_id,
1439 activity_id = %activity_id,
1440 activity_type = %name,
1441 worker_id = ?worker_id,
1442 );
1443 let _span_guard = span.enter();
1444 self.ensure_accepting(
1445 &namespace,
1446 &name,
1447 &workflow_id,
1448 &activity_id,
1449 Some(worker_id),
1450 )?;
1451
1452 let (completion_token, rx) = self
1453 .pending
1454 .insert(workflow_id.clone(), activity_id.clone())
1455 .map_err(|error| error.to_string())?;
1456 let task = activity_task(
1457 &name,
1458 &input,
1459 (&workflow_id, &run_id, &activity_id),
1460 attempt,
1461 labels,
1462 &completion_token,
1463 );
1464 if let Err(error) = self.track_worker_task(
1465 worker_id,
1466 &name,
1467 &workflow_id,
1468 &activity_id,
1469 attempt,
1470 completion_token.clone(),
1471 ) {
1472 self.cleanup_activity(worker_id, &workflow_id, &activity_id, &completion_token);
1473 return Err(error);
1474 }
1475 self.send_activity_task(
1476 &worker,
1477 task,
1478 &address,
1479 &workflow_id,
1480 &activity_id,
1481 &completion_token,
1482 )?;
1483 let context = ActivityDispatchContext {
1484 namespace: &namespace,
1485 activity_type: &name,
1486 worker_id,
1487 workflow_id: &workflow_id,
1488 activity_id: &activity_id,
1489 completion_token,
1490 started_at,
1491 };
1492 self.await_activity_result(&context, &rx)
1493 }
1494}
1495
1496#[cfg(feature = "liminal-transport")]
1503fn route_liminal_reply(
1504 pending: &PendingActivities,
1505 heartbeat_tracker: &HeartbeatTracker,
1506 drain_state: &DrainState,
1507 awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1508 execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1509) {
1510 let (worker_id, workflow_id, activity_id, current_token) = execution;
1511 let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1516 heartbeat_tracker
1517 .is_tracked(worker_id, workflow_id, activity_id)
1518 .unwrap_or(false)
1519 });
1520 let (run_id, submitted_token, outcome, synthesized) = match waited {
1524 Ok(Some(response)) => {
1525 let submitted_token = match CompletionToken::from_wire(
1526 workflow_id,
1527 activity_id,
1528 response.completion_token,
1529 ) {
1530 Ok(token) => token,
1531 Err(error) => {
1532 tracing::warn!(
1533 worker_id = ?worker_id,
1534 workflow_id = %workflow_id,
1535 activity_id = %activity_id,
1536 %error,
1537 "liminal activity completion omitted its generation proof"
1538 );
1539 return;
1540 }
1541 };
1542 (response.run_id, submitted_token, response.outcome, false)
1543 }
1544 Ok(None) => {
1545 tracing::debug!(
1546 worker_id = ?worker_id,
1547 workflow_id = %workflow_id,
1548 activity_id = %activity_id,
1549 "liminal dispatch resolved by another path; abandoning reply wait"
1550 );
1551 return;
1552 }
1553 Err(error) if error.is_worker_connection_lost() => (
1554 None,
1555 current_token.clone(),
1556 Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1560 true,
1561 ),
1562 Err(error) => (
1563 None,
1564 current_token.clone(),
1565 Err(format!("retryable:worker liminal reply failed: {error}")),
1566 true,
1567 ),
1568 };
1569 if synthesized {
1574 let was_tracked =
1575 complete_liminal_tracking(heartbeat_tracker, worker_id, workflow_id, activity_id);
1576 if !was_tracked {
1577 tracing::debug!(
1580 worker_id = ?worker_id,
1581 workflow_id = %workflow_id,
1582 activity_id = %activity_id,
1583 "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1584 );
1585 return;
1586 }
1587 }
1588 if let Err(error) = pending.complete_fenced(
1589 workflow_id,
1590 activity_id,
1591 run_id.as_ref(),
1592 &submitted_token,
1593 outcome,
1594 ) {
1595 tracing::warn!(
1596 worker_id = ?worker_id,
1597 workflow_id = %workflow_id,
1598 activity_id = %activity_id,
1599 %error,
1600 "liminal activity completion handoff rejected"
1601 );
1602 return;
1603 }
1604 if !synthesized
1605 && let Err(error) = heartbeat_tracker.complete_task(worker_id, workflow_id, activity_id)
1606 {
1607 tracing::error!(
1608 worker_id = ?worker_id,
1609 workflow_id = %workflow_id,
1610 activity_id = %activity_id,
1611 %error,
1612 "failed to clear in-flight tracking for completed liminal activity"
1613 );
1614 }
1615 drain_state.notify_activity_drained();
1616}
1617
1618#[cfg(feature = "liminal-transport")]
1619fn complete_liminal_tracking(
1620 heartbeat_tracker: &HeartbeatTracker,
1621 worker_id: WorkerId,
1622 workflow_id: &WorkflowId,
1623 activity_id: &ActivityId,
1624) -> bool {
1625 heartbeat_tracker
1626 .complete_task(worker_id, workflow_id, activity_id)
1627 .unwrap_or_else(|error| {
1628 tracing::error!(
1629 worker_id = ?worker_id,
1630 workflow_id = %workflow_id,
1631 activity_id = %activity_id,
1632 %error,
1633 "failed to clear in-flight tracking for completed liminal activity"
1634 );
1635 true
1636 })
1637}
1638
1639struct ActivityDispatchContext<'a> {
1640 namespace: &'a str,
1641 activity_type: &'a str,
1642 worker_id: WorkerId,
1643 workflow_id: &'a WorkflowId,
1644 activity_id: &'a ActivityId,
1645 completion_token: CompletionToken,
1646 started_at: Instant,
1647}
1648
1649fn activity_task(
1650 activity_type: &str,
1651 input: &str,
1652 execution: (&WorkflowId, &RunId, &ActivityId),
1653 attempt: u32,
1654 labels: BTreeMap<String, String>,
1655 completion_token: &CompletionToken,
1656) -> ProtoActivityTask {
1657 let (workflow_id, run_id, activity_id) = execution;
1658 ProtoActivityTask {
1659 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1660 activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1661 activity_type: activity_type.to_owned(),
1662 input: Some(ProtoPayload {
1663 content_type: String::from("application/json"),
1664 bytes: input.as_bytes().to_vec(),
1665 }),
1666 attempt,
1667 labels: labels.into_iter().collect(),
1668 run_id: Some(run_id.clone().into()),
1669 completion_token: completion_token.as_str().to_owned(),
1670 idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
1671 }
1672}
1673
1674fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1675 let duration_ms = duration_ms(context.started_at.elapsed());
1676 tracing::info!(
1677 operation = "activity_complete",
1678 namespace = context.namespace,
1679 workflow_id = %context.workflow_id,
1680 activity_id = %context.activity_id,
1681 activity_type = context.activity_type,
1682 worker_id = ?context.worker_id,
1683 duration_ms,
1684 outcome = if succeeded { "succeeded" } else { "failed" },
1685 "activity completed"
1686 );
1687}
1688
1689fn duration_ms(duration: Duration) -> u64 {
1690 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1691}
1692
1693fn log_worker_error(
1694 error_type: &'static str,
1695 namespace: &str,
1696 activity_type: &str,
1697 workflow_id: &WorkflowId,
1698 activity_id: &ActivityId,
1699 worker_id: Option<super::registry::WorkerId>,
1700 reason: &str,
1701) {
1702 tracing::error!(
1703 operation = "activity_dispatch",
1704 namespace,
1705 workflow_id = %workflow_id,
1706 activity_id = %activity_id,
1707 activity_type,
1708 worker_id = ?worker_id,
1709 error_type,
1710 reason,
1711 "worker interaction failed"
1712 );
1713}
1714
1715#[cfg(test)]
1716mod tests {
1717 use std::sync::Mutex;
1718
1719 use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1720
1721 use super::*;
1722
1723 fn activity_id(pos: u64) -> ActivityId {
1724 ActivityId::from_sequence_position(pos)
1725 }
1726
1727 #[test]
1728 fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
1729 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1730 let workflow_id = WorkflowId::new_v4();
1731 let id = activity_id(1);
1732 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1733
1734 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1735 assert_eq!(
1736 rx.recv_timeout(Duration::from_millis(50)),
1737 Ok(Ok("done".to_owned()))
1738 );
1739 Ok(())
1740 }
1741
1742 #[test]
1747 fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
1748 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1749 let workflow_id = WorkflowId::new_v4();
1750 let id = activity_id(2);
1751 let _first = pending.insert(workflow_id.clone(), id.clone())?;
1752
1753 let error = pending.insert(workflow_id, id).err().ok_or_else(|| {
1754 ServerError::worker_dispatch(
1755 "test",
1756 "same-worker-collision",
1757 "a live responder collision was silently replaced",
1758 )
1759 })?;
1760
1761 assert!(
1762 error.to_string().contains("pending activity collision"),
1763 "{error}"
1764 );
1765 Ok(())
1766 }
1767
1768 #[test]
1769 fn pending_complete_unknown_returns_false() {
1770 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1771 assert!(!pending.complete(
1772 &WorkflowId::new_v4(),
1773 &activity_id(99),
1774 None,
1775 Ok("orphan".to_owned())
1776 ));
1777 }
1778
1779 #[derive(Default)]
1780 struct RecordingOutboxCallback {
1781 completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1782 failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1783 live: bool,
1784 }
1785
1786 impl OutboxDeliveryCallback for RecordingOutboxCallback {
1787 fn deliver_completion(
1788 &self,
1789 workflow_id: &WorkflowId,
1790 activity_id: &ActivityId,
1791 run_id: Option<&RunId>,
1792 result: String,
1793 ) -> Result<bool, ServerError> {
1794 let _ = run_id;
1795 self.completions
1796 .lock()
1797 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1798 .push((workflow_id.clone(), activity_id.clone(), result));
1799 Ok(self.live)
1800 }
1801
1802 fn deliver_failure(
1803 &self,
1804 workflow_id: &WorkflowId,
1805 activity_id: &ActivityId,
1806 run_id: Option<&RunId>,
1807 reason: String,
1808 ) -> Result<bool, ServerError> {
1809 let _ = run_id;
1810 self.failures
1811 .lock()
1812 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1813 .push((workflow_id.clone(), activity_id.clone(), reason));
1814 Ok(self.live)
1815 }
1816 }
1817
1818 #[test]
1819 fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1820 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1821 let callback = Arc::new(RecordingOutboxCallback {
1822 live: true,
1823 ..RecordingOutboxCallback::default()
1824 });
1825 pending.clone().set_outbox_delivery(callback.clone());
1827
1828 let workflow_id = WorkflowId::new_v4();
1829 let id = activity_id(7);
1830
1831 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1834 let completions = callback
1835 .completions
1836 .lock()
1837 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1838 assert_eq!(completions.len(), 1);
1839 assert_eq!(completions[0].0, workflow_id);
1840 assert_eq!(completions[0].1, id);
1841 assert_eq!(completions[0].2, "done");
1842 Ok(())
1843 }
1844
1845 #[test]
1846 fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1847 -> Result<(), ServerError> {
1848 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1849 let callback = Arc::new(RecordingOutboxCallback::default());
1851 pending.set_outbox_delivery(callback.clone());
1852
1853 let workflow_id = WorkflowId::new_v4();
1854 let id = activity_id(8);
1855
1856 assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1857 let failures = callback
1858 .failures
1859 .lock()
1860 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1861 assert_eq!(failures.len(), 1);
1862 assert_eq!(failures[0].2, "retryable:boom");
1863 Ok(())
1864 }
1865
1866 #[test]
1867 fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1868 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1871 assert!(!pending.complete(
1872 &WorkflowId::new_v4(),
1873 &activity_id(9),
1874 None,
1875 Ok("x".to_owned())
1876 ));
1877 }
1878
1879 #[test]
1880 fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1881 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1882 let callback = Arc::new(RecordingOutboxCallback {
1883 live: true,
1884 ..RecordingOutboxCallback::default()
1885 });
1886 pending.set_outbox_delivery(callback.clone());
1887
1888 let workflow_id = WorkflowId::new_v4();
1889 let id = activity_id(10);
1890 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1891
1892 assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1893 assert_eq!(
1894 rx.recv_timeout(Duration::from_millis(50)),
1895 Ok(Ok("matched".to_owned()))
1896 );
1897 assert!(
1898 callback
1899 .completions
1900 .lock()
1901 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1902 .is_empty(),
1903 "a matched completion must deliver to its waiter, not the outbox callback"
1904 );
1905 Ok(())
1906 }
1907
1908 #[test]
1912 fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
1913 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1914 let workflow_id = WorkflowId::new_v4();
1915 let id = activity_id(11);
1916 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1917
1918 pending.park_activity(&workflow_id, &id)?;
1919 let result = rx
1920 .recv_timeout(Duration::from_millis(50))
1921 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1922 assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
1923 Ok(())
1924 }
1925
1926 #[test]
1930 fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
1931 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1932 let callback = Arc::new(RecordingOutboxCallback {
1933 live: true,
1934 ..RecordingOutboxCallback::default()
1935 });
1936 pending.set_outbox_delivery(callback.clone());
1937
1938 pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
1939
1940 assert!(
1941 callback
1942 .failures
1943 .lock()
1944 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1945 .is_empty(),
1946 "a park must never be delivered as an outbox failure"
1947 );
1948 assert!(
1949 callback
1950 .completions
1951 .lock()
1952 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1953 .is_empty(),
1954 "a park must never be delivered as an outbox completion"
1955 );
1956 Ok(())
1957 }
1958
1959 #[test]
1960 fn completion_sink_routes_success() -> Result<(), ServerError> {
1961 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1962 let workflow_id = WorkflowId::new_v4();
1963 let id = activity_id(2);
1964 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1965 let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
1966
1967 pending.complete_activity(ActivityCompletion {
1968 workflow_id,
1969 activity_id: id,
1970 run_id: None,
1971 completion_token,
1972 outcome: ActivityCompletionOutcome::Succeeded(payload),
1973 })?;
1974
1975 let result = rx
1976 .recv_timeout(Duration::from_millis(50))
1977 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1978 assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
1979 Ok(())
1980 }
1981
1982 #[test]
1983 fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
1984 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1985 let workflow_id = WorkflowId::new_v4();
1986 let id = activity_id(12);
1987 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1988
1989 let malformed = pending.complete_activity(ActivityCompletion {
1990 workflow_id: workflow_id.clone(),
1991 activity_id: id.clone(),
1992 run_id: None,
1993 completion_token: completion_token.clone(),
1994 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1995 ContentType::Json,
1996 vec![0xff],
1997 )),
1998 });
1999 assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
2000 assert!(
2001 rx.try_recv().is_err(),
2002 "an invalid result must leave the waiter unresolved"
2003 );
2004
2005 pending.complete_activity(ActivityCompletion {
2006 workflow_id,
2007 activity_id: id,
2008 run_id: None,
2009 completion_token,
2010 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2011 ContentType::Json,
2012 br#""valid""#.to_vec(),
2013 )),
2014 })?;
2015 let result = rx
2016 .recv_timeout(Duration::from_millis(50))
2017 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2018 assert_eq!(result, Ok(r#""valid""#.to_owned()));
2019 Ok(())
2020 }
2021
2022 #[test]
2023 fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
2024 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2025 let workflow_id = WorkflowId::new_v4();
2026 let id = activity_id(3);
2027 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
2028
2029 pending.complete_activity(ActivityCompletion {
2030 workflow_id,
2031 activity_id: id,
2032 run_id: None,
2033 completion_token,
2034 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2035 kind: ActivityErrorKind::Retryable,
2036 message: "temporary".to_owned(),
2037 details: None,
2038 }),
2039 })?;
2040
2041 let result = rx
2042 .recv_timeout(Duration::from_millis(50))
2043 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2044 assert_eq!(result, Err("retryable:temporary".to_owned()));
2045 Ok(())
2046 }
2047
2048 #[test]
2049 fn completion_sink_routes_policy_refusal_with_its_own_prefix() -> Result<(), ServerError> {
2050 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2051 let workflow_id = WorkflowId::new_v4();
2052 let id = activity_id(4);
2053 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
2054
2055 pending.complete_activity(ActivityCompletion {
2056 workflow_id,
2057 activity_id: id,
2058 run_id: None,
2059 completion_token,
2060 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2061 kind: ActivityErrorKind::PolicyRefused,
2062 message: "provider safety policy".to_owned(),
2063 details: None,
2064 }),
2065 })?;
2066
2067 let result = rx
2068 .recv_timeout(Duration::from_millis(50))
2069 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2070 assert_eq!(
2071 result,
2072 Err("policy_refused:provider safety policy".to_owned())
2073 );
2074 Ok(())
2075 }
2076
2077 #[test]
2086 fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
2087 -> Result<(), ServerError> {
2088 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2089 let post_restart_workflow = WorkflowId::new_v4();
2090 let pre_restart_workflow = WorkflowId::new_v4();
2091 let id = activity_id(1);
2093 let (completion_token, rx) = pending.insert(post_restart_workflow.clone(), id.clone())?;
2094
2095 let rejected = pending.complete_activity(ActivityCompletion {
2097 workflow_id: pre_restart_workflow,
2098 activity_id: id.clone(),
2099 run_id: None,
2100 completion_token: CompletionToken::for_test(),
2101 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2102 ContentType::Json,
2103 br#""stale""#.to_vec(),
2104 )),
2105 });
2106 assert!(matches!(
2107 rejected,
2108 Err(ServerError::ActivityCompletionRejected { .. })
2109 ));
2110 assert!(
2111 rx.try_recv().is_err(),
2112 "stale result for a different workflow must not complete this dispatch"
2113 );
2114
2115 pending.complete_activity(ActivityCompletion {
2117 workflow_id: post_restart_workflow,
2118 activity_id: id,
2119 run_id: None,
2120 completion_token,
2121 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2122 ContentType::Json,
2123 br#""fresh""#.to_vec(),
2124 )),
2125 })?;
2126 let result = rx
2127 .recv_timeout(Duration::from_millis(50))
2128 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2129 assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2130 Ok(())
2131 }
2132
2133 const TEST_HEARTBEAT_WINDOW: Duration = Duration::from_secs(5);
2140
2141 #[test]
2150 fn a_sink_derives_its_transport_budget_from_the_operator_window() {
2151 let window = Duration::from_secs(7);
2152 let sink = PendingActivities::new(window);
2153
2154 assert_eq!(
2155 sink.transport_losses().budget(),
2156 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2157 "the transport-loss budget must be the operator's heartbeat window times \
2158 TRANSPORT_LOSS_BUDGET_WINDOWS; a sink carrying any other budget is one no \
2159 operator declared"
2160 );
2161 }
2162
2163 #[test]
2171 fn a_dispatcher_carries_the_transport_budget_of_the_tracker_it_was_given() {
2172 let window = Duration::from_secs(11);
2173 let registry = ConnectedWorkerRegistry::default();
2174 let dispatcher =
2175 WorkerActivityDispatcher::new(registry, "default", HeartbeatTracker::new(window));
2176
2177 assert_eq!(
2178 dispatcher.pending().transport_losses().budget(),
2179 window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2180 "the dispatcher's transport budget must be derived from the heartbeat window it \
2181 was constructed with; a zero or mismatched budget here means a worker loss is \
2182 declared transport-exhausted after ONE re-dispatch instead of after the \
2183 operator's window"
2184 );
2185 assert_ne!(
2186 dispatcher.pending().transport_losses().budget(),
2187 Duration::ZERO,
2188 "a ZERO budget is the specific regression this pin exists for: it grants exactly \
2189 one re-dispatchable loss and then declares the infrastructure flapping"
2190 );
2191 }
2192
2193 fn test_tracker() -> HeartbeatTracker {
2194 HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW)
2195 }
2196
2197 fn greet_request() -> ActivityDispatch {
2200 ActivityDispatch {
2201 namespace: "default".to_owned(),
2202 task_queue: "default".to_owned(),
2203 node: None,
2204 workflow_id: WorkflowId::new_v4(),
2205 run_id: RunId::new_v4(),
2206 activity_id: ActivityId::from_sequence_position(0),
2207 name: "greet".to_owned(),
2208 input: "{}".to_owned(),
2209 config: "{}".to_owned(),
2210 attempt: 1,
2211 labels: std::collections::BTreeMap::new(),
2212 advisory: false,
2213 }
2214 }
2215
2216 #[test]
2217 fn dispatcher_fails_immediately_when_draining_without_workers() {
2218 let registry = ConnectedWorkerRegistry::default();
2219 let drain = DrainState::default();
2220 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2221 .with_drain_state(drain.clone());
2222
2223 let _ = drain.begin();
2224
2225 let result = dispatcher.dispatch(greet_request());
2226
2227 assert!(result.is_err());
2228 let err = result.err().unwrap_or_default();
2229 assert!(
2230 err.contains("drain"),
2231 "expected drain rejection, got: {err}"
2232 );
2233 }
2234
2235 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2251 async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
2252 -> Result<(), Box<dyn std::error::Error>> {
2253 let registry = ConnectedWorkerRegistry::default();
2254 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2255 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2256 let activity_types = [String::from("greet")];
2257 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2258
2259 let sink = pending.clone();
2260 let echo_worker = tokio::spawn(async move {
2261 let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
2262 return Err("expected an activity task on the worker channel".to_owned());
2263 };
2264 let workflow_id = task
2265 .workflow_id
2266 .ok_or("task missing workflow id")
2267 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2268 let activity_id = task
2269 .activity_id
2270 .map(ActivityId::from)
2271 .ok_or("task missing activity id")?;
2272 let completion_token =
2273 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2274 .map_err(|error| error.to_string())?;
2275 sink.complete_activity(ActivityCompletion {
2276 workflow_id,
2277 activity_id,
2278 run_id: None,
2279 completion_token,
2280 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2281 ContentType::Json,
2282 br#"{"greeting":"hello"}"#.to_vec(),
2283 )),
2284 })
2285 .map_err(|error| error.to_string())
2286 });
2287
2288 let dispatcher = Arc::new(
2289 WorkerActivityDispatcher::new(registry, "default", test_tracker())
2290 .with_pending(pending),
2291 );
2292 let started = Instant::now();
2293 let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
2296 dispatcher.dispatch(greet_request())
2297 }));
2298 let result = dispatch_task.await.map_err(|error| error.to_string())?;
2299 let elapsed = started.elapsed();
2300
2301 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2302 assert!(
2303 elapsed < Duration::from_secs(5),
2304 "dispatch round trip took {elapsed:?}; task delivery must not \
2305 depend on the blocked dispatch thread"
2306 );
2307 echo_worker.await.map_err(|error| error.to_string())??;
2308 registration.deregister()?;
2309 Ok(())
2310 }
2311
2312 #[tokio::test]
2316 async fn dispatch_on_current_thread_runtime_fails_fast()
2317 -> Result<(), Box<dyn std::error::Error>> {
2318 let registry = ConnectedWorkerRegistry::default();
2319 let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
2320 let activity_types = [String::from("greet")];
2321 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2322 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
2323
2324 let started = Instant::now();
2325 let result = dispatcher.dispatch(greet_request());
2326 let elapsed = started.elapsed();
2327
2328 let err = result.err().ok_or("expected dispatch to fail")?;
2329 assert!(
2330 err.contains("multi-thread tokio runtime"),
2331 "unexpected error: {err}"
2332 );
2333 assert!(
2334 elapsed < Duration::from_secs(5),
2335 "fail-fast path took {elapsed:?}"
2336 );
2337 registration.deregister()?;
2338 Ok(())
2339 }
2340
2341 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2347 async fn dispatch_pinned_to_node_reaches_only_that_node()
2348 -> Result<(), Box<dyn std::error::Error>> {
2349 let registry = ConnectedWorkerRegistry::default();
2350 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2351 let activity_types = [String::from("greet")];
2352 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2353 let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
2354 let on_n2 = registry.register_namespaces(
2360 [String::from("default")],
2361 "default",
2362 Some(String::from("n2")),
2363 activity_types.iter(),
2364 n2_tx,
2365 )?;
2366 let on_n1 = registry.register_namespaces(
2367 [String::from("default")],
2368 "default",
2369 Some(String::from("n1")),
2370 activity_types.iter(),
2371 n1_tx,
2372 )?;
2373
2374 let sink = pending.clone();
2378 let echo_n1 = tokio::spawn(async move {
2379 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2380 return Err("expected an activity task on the n1 worker channel".to_owned());
2381 };
2382 let workflow_id = task
2383 .workflow_id
2384 .ok_or("task missing workflow id")
2385 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2386 let activity_id = task
2387 .activity_id
2388 .map(ActivityId::from)
2389 .ok_or("task missing activity id")?;
2390 let completion_token =
2391 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2392 .map_err(|error| error.to_string())?;
2393 sink.complete_activity(ActivityCompletion {
2394 workflow_id,
2395 activity_id,
2396 run_id: None,
2397 completion_token,
2398 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2399 ContentType::Json,
2400 br#"{"greeting":"hello"}"#.to_vec(),
2401 )),
2402 })
2403 .map_err(|error| error.to_string())
2404 });
2405
2406 let dispatcher = Arc::new(
2407 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2408 .with_pending(pending),
2409 );
2410
2411 let pinned = ActivityDispatch {
2412 node: Some(String::from("n1")),
2413 ..greet_request()
2414 };
2415 let started = Instant::now();
2416 let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
2417 .await
2418 .map_err(|error| error.to_string())?;
2419 let elapsed = started.elapsed();
2420
2421 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2422 assert!(
2423 elapsed < Duration::from_secs(5),
2424 "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
2425 );
2426 echo_n1.await.map_err(|error| error.to_string())??;
2427
2428 assert!(
2430 n2_rx.try_recv().is_err(),
2431 "node=Some(\"n1\") dispatch must not reach the n2 worker"
2432 );
2433
2434 on_n1.deregister()?;
2435 on_n2.deregister()?;
2436 Ok(())
2437 }
2438
2439 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2442 async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
2443 -> Result<(), Box<dyn std::error::Error>> {
2444 let registry = ConnectedWorkerRegistry::default();
2445 let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2446 let activity_types = [String::from("greet")];
2447 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2448 let on_n1 = registry.register_namespaces(
2449 [String::from("default")],
2450 "default",
2451 Some(String::from("n1")),
2452 activity_types.iter(),
2453 n1_tx,
2454 )?;
2455
2456 let sink = pending.clone();
2457 let echo = tokio::spawn(async move {
2458 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2459 return Err("expected an activity task on the worker channel".to_owned());
2460 };
2461 let workflow_id = task
2462 .workflow_id
2463 .ok_or("task missing workflow id")
2464 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2465 let activity_id = task
2466 .activity_id
2467 .map(ActivityId::from)
2468 .ok_or("task missing activity id")?;
2469 let completion_token =
2470 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2471 .map_err(|error| error.to_string())?;
2472 sink.complete_activity(ActivityCompletion {
2473 workflow_id,
2474 activity_id,
2475 run_id: None,
2476 completion_token,
2477 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2478 ContentType::Json,
2479 br#"{"greeting":"hello"}"#.to_vec(),
2480 )),
2481 })
2482 .map_err(|error| error.to_string())
2483 });
2484
2485 let dispatcher = Arc::new(
2486 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2487 .with_pending(pending),
2488 );
2489
2490 let result = tokio::spawn(futures::future::lazy(move |_| {
2492 dispatcher.dispatch(greet_request())
2493 }))
2494 .await
2495 .map_err(|error| error.to_string())?;
2496
2497 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2498 echo.await.map_err(|error| error.to_string())??;
2499 on_n1.deregister()?;
2500 Ok(())
2501 }
2502}