1use std::collections::BTreeMap;
64use std::sync::{Arc, OnceLock};
65use std::time::{Duration, Instant};
66
67use aion::{ActivityDispatch, ActivityDispatcher};
68use aion_core::{ActivityId, ContentType, Payload, RunId, WorkflowId};
69use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
70use dashmap::DashMap;
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, Default)]
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]
348 pub fn with_heartbeat_window(mut self, heartbeat_window: std::time::Duration) -> Self {
349 self.transport_losses = super::transport_loss::TransportLossLedger::new(heartbeat_window);
350 self
351 }
352
353 #[must_use]
355 pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
356 &self.transport_losses
357 }
358
359 fn classify_worker_loss(
367 &self,
368 workflow_id: &WorkflowId,
369 activity_id: &ActivityId,
370 worker_id: crate::worker::registry::WorkerId,
371 ) -> String {
372 let detail = super::transport_loss::worker_lost_detail(worker_id);
373 match self
374 .transport_losses
375 .record_loss(workflow_id, activity_id, &detail)
376 {
377 Ok(verdict) => {
378 if verdict.exhausted {
379 tracing::error!(
380 operation = "activity_complete",
381 workflow_id = %workflow_id,
382 activity_id = %activity_id,
383 worker_id = ?worker_id,
384 error_type = "TransportExhausted",
385 losses = verdict.losses,
386 budget_ms = self.transport_losses.budget().as_millis(),
387 "activity abandoned: the transport kept losing its worker past the \
388 transport-loss budget"
389 );
390 } else {
391 tracing::warn!(
392 operation = "activity_complete",
393 workflow_id = %workflow_id,
394 activity_id = %activity_id,
395 worker_id = ?worker_id,
396 error_type = "WorkerLost",
397 losses = verdict.losses,
398 budget_ms = self.transport_losses.budget().as_millis(),
399 "worker lost before reporting an activity result; the activity never ran \
400 and will be re-dispatched attempt-neutrally"
401 );
402 }
403 verdict.reason
404 }
405 Err(error) => {
406 tracing::error!(
407 workflow_id = %workflow_id,
408 activity_id = %activity_id,
409 %error,
410 "transport-loss ledger is unreadable; abandoning the activity rather than \
411 re-dispatching it without a budget"
412 );
413 format!(
414 "{}{detail} (transport-loss budget unreadable: {error})",
415 super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
416 )
417 }
418 }
419 }
420}
421
422impl ActivityCompletionSink for PendingActivities {
423 fn park_activity(
434 &self,
435 workflow_id: &WorkflowId,
436 activity_id: &ActivityId,
437 ) -> Result<(), ServerError> {
438 self.completion_fences
439 .revoke_current(workflow_id, activity_id)?;
440 let matched = self
441 .pending
442 .remove(&(workflow_id.clone(), activity_id.clone()));
443 if let Some((_, sender)) = matched {
444 let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
448 }
449 Ok(())
450 }
451
452 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
453 let result = match completion.outcome {
454 ActivityCompletionOutcome::Succeeded(payload) => {
455 payload_to_string(&payload).map_err(|reason| {
456 tracing::error!(
457 operation = "activity_complete",
458 workflow_id = %completion.workflow_id,
459 activity_id = %completion.activity_id,
460 error_type = "ActivityResultDecode",
461 %reason,
462 "activity completion failed"
463 );
464 ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
465 })?
466 }
467 ActivityCompletionOutcome::Failed(error) => {
468 let prefix = if error.is_retryable() {
469 "retryable"
470 } else {
471 "terminal"
472 };
473 tracing::error!(
474 operation = "activity_complete",
475 workflow_id = %completion.workflow_id,
476 activity_id = %completion.activity_id,
477 error_type = "ActivityFailed",
478 error_kind = prefix,
479 reason = %error.message,
480 "activity completion failed"
481 );
482 Err(format!("{prefix}:{}", error.message))
483 }
484 ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
489 &completion.workflow_id,
490 &completion.activity_id,
491 worker_id,
492 )),
493 };
494 self.complete_fenced(
495 &completion.workflow_id,
496 &completion.activity_id,
497 completion.run_id.as_ref(),
498 &completion.completion_token,
499 result,
500 )?;
501 Ok(())
502 }
503}
504
505fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
506 match payload.content_type() {
507 ContentType::Json => String::from_utf8(payload.bytes().to_vec())
508 .map(Ok)
509 .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
510 }
511}
512
513pub struct WorkerActivityDispatcher {
521 registry: ConnectedWorkerRegistry,
522 namespace: String,
523 pending: PendingActivities,
524 heartbeat_tracker: HeartbeatTracker,
525 drain_state: DrainState,
526 tokio_handle: Option<tokio::runtime::Handle>,
527 attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
533 queue_service: QueueServiceConfig,
537 queue_declarations: QueueDeclarationSource,
540 queue_state: QueueServiceState,
542 cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
546}
547
548impl std::fmt::Debug for WorkerActivityDispatcher {
549 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
550 f.debug_struct("WorkerActivityDispatcher")
551 .field("namespace", &self.namespace)
552 .finish_non_exhaustive()
553 }
554}
555
556impl WorkerActivityDispatcher {
557 #[must_use]
565 pub fn new(
566 registry: ConnectedWorkerRegistry,
567 namespace: impl Into<String>,
568 heartbeat_tracker: HeartbeatTracker,
569 ) -> Self {
570 Self {
571 registry,
572 namespace: namespace.into(),
573 pending: PendingActivities::default(),
574 heartbeat_tracker,
575 drain_state: DrainState::default(),
576 tokio_handle: None,
577 attempt_owners: None,
578 queue_service: QueueServiceConfig::default(),
579 queue_declarations: QueueDeclarationSource::default(),
580 queue_state: QueueServiceState::default(),
581 cluster_publisher: None,
582 }
583 }
584
585 #[must_use]
589 pub fn with_cluster_publisher(
590 mut self,
591 cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
592 ) -> Self {
593 self.cluster_publisher = Some(cluster_publisher);
594 self
595 }
596
597 #[must_use]
600 pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
601 self.queue_service = queue_service;
602 self
603 }
604
605 #[must_use]
609 pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
610 self.queue_declarations = queue_declarations;
611 self
612 }
613
614 #[must_use]
617 pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
618 self.queue_state = queue_state;
619 self
620 }
621
622 #[must_use]
629 pub fn with_attempt_owners(
630 mut self,
631 attempt_owners: super::intervention::AttemptOwnerIndex,
632 ) -> Self {
633 self.attempt_owners = Some(attempt_owners);
634 self
635 }
636
637 #[must_use]
639 pub fn with_pending(mut self, pending: PendingActivities) -> Self {
640 self.pending = pending;
641 self
642 }
643
644 #[must_use]
646 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
647 self.drain_state = drain_state;
648 self
649 }
650
651 #[must_use]
653 pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
654 self.tokio_handle = Some(tokio_handle);
655 self
656 }
657}
658
659impl WorkerActivityDispatcher {
660 fn ensure_accepting(
684 &self,
685 namespace: &str,
686 activity_type: &str,
687 workflow_id: &WorkflowId,
688 activity_id: &ActivityId,
689 worker_id: Option<WorkerId>,
690 ) -> Result<(), String> {
691 self.drain_state
692 .ensure_accepting(namespace, activity_type)
693 .map_err(|error| {
694 log_worker_error(
695 "WorkerDispatch",
696 namespace,
697 activity_type,
698 workflow_id,
699 activity_id,
700 worker_id,
701 &error.to_string(),
702 );
703 aion::PARKED_ACTIVITY_REASON.to_owned()
704 })
705 }
706
707 fn select_worker_or_wait(
715 &self,
716 address: &ServiceAddress,
717 workflow_id: &WorkflowId,
718 activity_id: &ActivityId,
719 ) -> Result<WorkerHandle, String> {
720 let wait = ServiceWait {
721 registry: &self.registry,
722 declarations: &self.queue_declarations,
723 config: &self.queue_service,
724 state: &self.queue_state,
725 address,
726 workflow_id,
727 activity_id,
728 publisher: self.cluster_publisher.as_ref(),
729 };
730 let mut accepting = || {
731 self.ensure_accepting(
732 &address.namespace,
733 &address.activity_type,
734 workflow_id,
735 activity_id,
736 None,
737 )
738 };
739 let mut park = |budget: Option<Duration>| self.park_for_worker(budget);
740 select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
741 let reason = refusal.reason_string();
742 if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
745 let error_type = match refusal {
746 SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
747 _ => "WorkerRegistry",
748 };
749 log_worker_error(
750 error_type,
751 &address.namespace,
752 &address.activity_type,
753 workflow_id,
754 activity_id,
755 None,
756 &reason,
757 );
758 }
759 reason
760 })
761 }
762
763 fn park_for_worker(&self, budget: Option<Duration>) {
791 let handle = self
792 .tokio_handle
793 .clone()
794 .or_else(|| tokio::runtime::Handle::try_current().ok());
795 let Some(handle) = handle else {
796 std::thread::sleep(
801 budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
802 );
803 return;
804 };
805 handle.block_on(async {
806 let arrival = async {
807 tokio::select! {
808 () = self.registry.wait_for_worker() => {}
809 () = self.drain_state.wait_for_drain() => {}
810 }
811 };
812 match budget {
813 None => arrival.await,
814 Some(budget) => {
815 drop(tokio::time::timeout(budget, arrival).await);
818 }
819 }
820 });
821 }
822
823 fn track_worker_task(
824 &self,
825 worker_id: WorkerId,
826 activity_type: &str,
827 workflow_id: &WorkflowId,
828 activity_id: &ActivityId,
829 attempt: u32,
830 completion_token: CompletionToken,
831 ) -> Result<(), String> {
832 self.heartbeat_tracker
833 .track_task(
834 worker_id,
835 InFlightActivity {
836 workflow_id: workflow_id.clone(),
837 activity_id: activity_id.clone(),
838 attempt,
843 completion_token,
844 },
845 Instant::now(),
846 )
847 .map_err(|error| {
848 let reason = error.to_string();
849 log_worker_error(
850 "WorkerHeartbeatTracker",
851 &self.namespace,
852 activity_type,
853 workflow_id,
854 activity_id,
855 Some(worker_id),
856 &reason,
857 );
858 reason
859 })
860 }
861
862 fn cleanup_activity(
863 &self,
864 worker_id: WorkerId,
865 workflow_id: &WorkflowId,
866 activity_id: &ActivityId,
867 completion_token: &CompletionToken,
868 ) {
869 self.pending
870 .pending
871 .remove(&(workflow_id.clone(), activity_id.clone()));
872 if let Err(error) =
873 self.pending
874 .completion_fences
875 .revoke(workflow_id, activity_id, completion_token)
876 {
877 tracing::error!(
878 workflow_id = %workflow_id,
879 activity_id = %activity_id,
880 %error,
881 "failed to revoke undelivered activity generation"
882 );
883 }
884 let _ = self
885 .heartbeat_tracker
886 .complete_task(worker_id, workflow_id, activity_id);
887 self.drain_state.notify_activity_drained();
888 }
889
890 fn send_activity_task(
900 &self,
901 worker: &WorkerHandle,
902 task: ProtoActivityTask,
903 address: &ServiceAddress,
904 workflow_id: &WorkflowId,
905 activity_id: &ActivityId,
906 completion_token: &CompletionToken,
907 ) -> Result<(), String> {
908 match worker.delivery() {
909 WorkerDelivery::Grpc(sender) => {
910 let worker_id = worker.id();
911 let mut accepting = || {
912 self.ensure_accepting(
913 &address.namespace,
914 &address.activity_type,
915 workflow_id,
916 activity_id,
917 Some(worker_id),
918 )
919 };
920 let handed_over = deliver_within_schedule_to_start(
921 sender,
922 WorkerMessage::ActivityTask(Box::new(task)),
923 self.queue_service.schedule_to_start_timeout,
924 &mut accepting,
925 );
926 let Err(refusal) = handed_over else {
927 return Ok(());
928 };
929 self.cleanup_activity(worker_id, workflow_id, activity_id, completion_token);
930 let (error_type, reason) = self.hand_off_failure(&refusal, address);
931 if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
932 log_worker_error(
933 error_type,
934 &address.namespace,
935 &address.activity_type,
936 workflow_id,
937 activity_id,
938 Some(worker_id),
939 &reason,
940 );
941 }
942 Err(reason)
943 }
944 #[cfg(feature = "liminal-transport")]
951 WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
952 worker.id(),
953 delivery,
954 task,
955 &address.activity_type,
956 workflow_id,
957 activity_id,
958 ),
959 }
960 }
961
962 fn hand_off_failure(
969 &self,
970 refusal: &DeliveryRefusal,
971 address: &ServiceAddress,
972 ) -> (&'static str, String) {
973 match refusal {
974 DeliveryRefusal::Saturated { waited } => {
975 let census = self
978 .registry
979 .pool_census(
980 &address.namespace,
981 &address.task_queue,
982 &address.activity_type,
983 address.node.as_deref(),
984 )
985 .unwrap_or_else(|error| {
986 tracing::error!(
987 namespace = %address.namespace,
988 task_queue = %address.task_queue,
989 activity_type = %address.activity_type,
990 %error,
991 "poller census failed while reporting a saturated queue; \
992 the refusal carries an empty census"
993 );
994 PoolCensus::default()
995 });
996 let unavailable = WorkerUnavailable {
997 reason: QueueServiceReason::Saturated,
998 clock: Some(ExpiredClock::ScheduleToStart),
999 waited: *waited,
1000 address: address.clone(),
1001 census,
1002 };
1003 ("WorkerUnavailable", unavailable.reason_string())
1004 }
1005 DeliveryRefusal::Full => (
1006 "WorkerChannelClosed",
1007 "worker task channel full or closed: no available capacity".to_owned(),
1008 ),
1009 DeliveryRefusal::Closed => (
1010 "WorkerChannelClosed",
1011 "worker task channel full or closed: channel closed".to_owned(),
1012 ),
1013 DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
1014 }
1015 }
1016
1017 #[cfg(feature = "liminal-transport")]
1041 fn send_liminal_activity_task(
1042 &self,
1043 worker_id: WorkerId,
1044 delivery: &super::liminal_transport::LiminalWorkerDelivery,
1045 task: ProtoActivityTask,
1046 activity_type: &str,
1047 workflow_id: &WorkflowId,
1048 activity_id: &ActivityId,
1049 ) -> Result<(), String> {
1050 let completion_token =
1051 CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1052 .map_err(|error| error.to_string())?;
1053 let heartbeat_window_ms =
1054 u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1055 .unwrap_or(u64::MAX);
1056 let attempt = task.attempt;
1057 let run_id = task
1064 .run_id
1065 .map(RunId::try_from)
1066 .transpose()
1067 .map_err(|error| error.to_string())?
1068 .ok_or_else(|| {
1069 "activity task run id is missing; refusing to dispatch an unidentified run"
1070 .to_owned()
1071 })?;
1072 let request = super::liminal_transport::DispatchRequest {
1073 activity_type: activity_type.to_owned(),
1074 workflow_id: workflow_id.clone(),
1075 ordinal: activity_id.sequence_position(),
1076 run_id: Some(run_id.clone()),
1077 attempt,
1078 completion_token: task.completion_token,
1079 idempotency_key: task.idempotency_key,
1080 labels: task.labels.into_iter().collect(),
1081 heartbeat_window_ms,
1082 input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1083 };
1084 let awaiter = match delivery.push_dispatch(&request) {
1088 Ok(awaiter) => awaiter,
1089 Err(error) => {
1090 let reason = format!("worker liminal push failed: {error}");
1091 self.cleanup_activity(worker_id, workflow_id, activity_id, &completion_token);
1092 log_worker_error(
1093 "WorkerChannelClosed",
1094 &self.namespace,
1095 activity_type,
1096 workflow_id,
1097 activity_id,
1098 Some(worker_id),
1099 &reason,
1100 );
1101 return Err(reason);
1102 }
1103 };
1104 let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1108 super::liminal_transport::AttemptOwnerGuard::bind(
1109 owners.clone(),
1110 super::intervention::AttemptKey::new(
1111 workflow_id.clone(),
1112 run_id.clone(),
1113 activity_id.clone(),
1114 attempt,
1115 ),
1116 worker_id,
1117 )
1118 });
1119 self.spawn_liminal_reply_router(
1120 worker_id,
1121 awaiter,
1122 workflow_id,
1123 activity_id,
1124 &completion_token,
1125 owner_binding,
1126 );
1127 Ok(())
1128 }
1129
1130 #[cfg(feature = "liminal-transport")]
1170 fn spawn_liminal_reply_router(
1171 &self,
1172 worker_id: WorkerId,
1173 awaiter: liminal_server::server::connection::PushReplyAwaiter,
1174 workflow_id: &WorkflowId,
1175 activity_id: &ActivityId,
1176 completion_token: &CompletionToken,
1177 owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1178 ) {
1179 let pending = self.pending.clone();
1180 let heartbeat_tracker = self.heartbeat_tracker.clone();
1181 let drain_state = self.drain_state.clone();
1182 let workflow_id = workflow_id.clone();
1183 let activity_id = activity_id.clone();
1184 let completion_token = completion_token.clone();
1185 std::thread::spawn(move || {
1186 let _owner_binding = owner_binding;
1190 route_liminal_reply(
1191 &pending,
1192 &heartbeat_tracker,
1193 &drain_state,
1194 &awaiter,
1195 (worker_id, &workflow_id, &activity_id, &completion_token),
1196 );
1197 });
1198 }
1199
1200 fn await_activity_result(
1204 &self,
1205 context: &ActivityDispatchContext<'_>,
1206 rx: &SyncReceiver,
1207 ) -> Result<String, String> {
1208 match self.registry.is_registered(context.worker_id) {
1216 Ok(true) => {}
1217 Ok(false) => {
1218 if let Ok(result) = rx.try_recv() {
1222 return self.deliver_result(context, result);
1223 }
1224 self.cleanup_activity(
1225 context.worker_id,
1226 context.workflow_id,
1227 context.activity_id,
1228 &context.completion_token,
1229 );
1230 let reason = self.pending.classify_worker_loss(
1235 context.workflow_id,
1236 context.activity_id,
1237 context.worker_id,
1238 );
1239 log_worker_error(
1240 "WorkerLost",
1241 &self.namespace,
1242 context.activity_type,
1243 context.workflow_id,
1244 context.activity_id,
1245 Some(context.worker_id),
1246 &reason,
1247 );
1248 return Err(reason);
1249 }
1250 Err(error) => {
1251 self.cleanup_activity(
1252 context.worker_id,
1253 context.workflow_id,
1254 context.activity_id,
1255 &context.completion_token,
1256 );
1257 let reason = format!("worker registry inspection failed: {error}");
1258 log_worker_error(
1259 "WorkerRegistry",
1260 &self.namespace,
1261 context.activity_type,
1262 context.workflow_id,
1263 context.activity_id,
1264 Some(context.worker_id),
1265 &reason,
1266 );
1267 return Err(reason);
1268 }
1269 }
1270 if let Ok(result) = rx.recv() {
1271 return self.deliver_result(context, result);
1272 }
1273 self.cleanup_activity(
1276 context.worker_id,
1277 context.workflow_id,
1278 context.activity_id,
1279 &context.completion_token,
1280 );
1281 let reason = "activity response channel dropped".to_owned();
1282 log_worker_error(
1283 "WorkerChannelClosed",
1284 &self.namespace,
1285 context.activity_type,
1286 context.workflow_id,
1287 context.activity_id,
1288 Some(context.worker_id),
1289 &reason,
1290 );
1291 Err(reason)
1292 }
1293
1294 fn deliver_result(
1295 &self,
1296 context: &ActivityDispatchContext<'_>,
1297 result: Result<String, String>,
1298 ) -> Result<String, String> {
1299 self.pending
1300 .pending
1301 .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
1302 if let Err(reason) = &result
1307 && aion::is_parked_reason(reason)
1308 {
1309 tracing::info!(
1310 operation = "activity_dispatch",
1311 namespace = %self.namespace,
1312 workflow_id = %context.workflow_id,
1313 activity_id = %context.activity_id,
1314 activity_type = context.activity_type,
1315 worker_id = ?context.worker_id,
1316 "activity parked for restart recovery"
1317 );
1318 return result;
1319 }
1320 log_activity_completion(context, result.is_ok());
1321 result.inspect_err(|reason| {
1322 log_worker_error(
1323 "ActivityFailed",
1324 &self.namespace,
1325 context.activity_type,
1326 context.workflow_id,
1327 context.activity_id,
1328 Some(context.worker_id),
1329 reason,
1330 );
1331 })
1332 }
1333}
1334
1335impl ActivityDispatcher for WorkerActivityDispatcher {
1336 fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1337 match tokio::runtime::Handle::try_current() {
1338 Ok(handle) => match handle.runtime_flavor() {
1339 tokio::runtime::RuntimeFlavor::MultiThread => {
1340 tokio::task::block_in_place(|| self.dispatch_blocking(request))
1347 }
1348 flavor => Err(format!(
1349 "activity dispatch blocks the calling thread until the worker responds; \
1350 a {flavor:?} tokio runtime cannot host that wait because the worker \
1351 stream forwarder shares its only executor thread and the task could \
1352 never be delivered — run the engine on a multi-thread tokio runtime"
1353 )),
1354 },
1355 Err(_) => self.dispatch_blocking(request),
1359 }
1360 }
1361}
1362
1363impl WorkerActivityDispatcher {
1364 fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1381 let ActivityDispatch {
1382 namespace,
1383 task_queue,
1384 node,
1388 workflow_id,
1389 run_id,
1390 activity_id,
1391 name,
1392 input,
1393 config: _,
1394 attempt,
1395 labels,
1396 advisory: _,
1401 } = request;
1402 let started_at = Instant::now();
1403 self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1404 let address = ServiceAddress {
1405 namespace: namespace.clone(),
1406 task_queue: task_queue.clone(),
1407 activity_type: name.clone(),
1408 node: node.clone(),
1409 };
1410 let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1411 let worker_id = worker.id();
1412 let span = info_span!(
1413 "activity_dispatch",
1414 operation = "activity_dispatch",
1415 namespace = %namespace,
1416 task_queue = %task_queue,
1417 node = node.as_deref(),
1418 workflow_id = %workflow_id,
1419 activity_id = %activity_id,
1420 activity_type = %name,
1421 worker_id = ?worker_id,
1422 );
1423 let _span_guard = span.enter();
1424 self.ensure_accepting(
1425 &namespace,
1426 &name,
1427 &workflow_id,
1428 &activity_id,
1429 Some(worker_id),
1430 )?;
1431
1432 let (completion_token, rx) = self
1433 .pending
1434 .insert(workflow_id.clone(), activity_id.clone())
1435 .map_err(|error| error.to_string())?;
1436 let task = activity_task(
1437 &name,
1438 &input,
1439 (&workflow_id, &run_id, &activity_id),
1440 attempt,
1441 labels,
1442 &completion_token,
1443 );
1444 if let Err(error) = self.track_worker_task(
1445 worker_id,
1446 &name,
1447 &workflow_id,
1448 &activity_id,
1449 attempt,
1450 completion_token.clone(),
1451 ) {
1452 self.cleanup_activity(worker_id, &workflow_id, &activity_id, &completion_token);
1453 return Err(error);
1454 }
1455 self.send_activity_task(
1456 &worker,
1457 task,
1458 &address,
1459 &workflow_id,
1460 &activity_id,
1461 &completion_token,
1462 )?;
1463 let context = ActivityDispatchContext {
1464 namespace: &namespace,
1465 activity_type: &name,
1466 worker_id,
1467 workflow_id: &workflow_id,
1468 activity_id: &activity_id,
1469 completion_token,
1470 started_at,
1471 };
1472 self.await_activity_result(&context, &rx)
1473 }
1474}
1475
1476#[cfg(feature = "liminal-transport")]
1483fn route_liminal_reply(
1484 pending: &PendingActivities,
1485 heartbeat_tracker: &HeartbeatTracker,
1486 drain_state: &DrainState,
1487 awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1488 execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1489) {
1490 let (worker_id, workflow_id, activity_id, current_token) = execution;
1491 let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1496 heartbeat_tracker
1497 .is_tracked(worker_id, workflow_id, activity_id)
1498 .unwrap_or(false)
1499 });
1500 let (run_id, submitted_token, outcome, synthesized) = match waited {
1504 Ok(Some(response)) => {
1505 let submitted_token = match CompletionToken::from_wire(
1506 workflow_id,
1507 activity_id,
1508 response.completion_token,
1509 ) {
1510 Ok(token) => token,
1511 Err(error) => {
1512 tracing::warn!(
1513 worker_id = ?worker_id,
1514 workflow_id = %workflow_id,
1515 activity_id = %activity_id,
1516 %error,
1517 "liminal activity completion omitted its generation proof"
1518 );
1519 return;
1520 }
1521 };
1522 (response.run_id, submitted_token, response.outcome, false)
1523 }
1524 Ok(None) => {
1525 tracing::debug!(
1526 worker_id = ?worker_id,
1527 workflow_id = %workflow_id,
1528 activity_id = %activity_id,
1529 "liminal dispatch resolved by another path; abandoning reply wait"
1530 );
1531 return;
1532 }
1533 Err(error) if error.is_worker_connection_lost() => (
1534 None,
1535 current_token.clone(),
1536 Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1540 true,
1541 ),
1542 Err(error) => (
1543 None,
1544 current_token.clone(),
1545 Err(format!("retryable:worker liminal reply failed: {error}")),
1546 true,
1547 ),
1548 };
1549 if synthesized {
1554 let was_tracked =
1555 complete_liminal_tracking(heartbeat_tracker, worker_id, workflow_id, activity_id);
1556 if !was_tracked {
1557 tracing::debug!(
1560 worker_id = ?worker_id,
1561 workflow_id = %workflow_id,
1562 activity_id = %activity_id,
1563 "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1564 );
1565 return;
1566 }
1567 }
1568 if let Err(error) = pending.complete_fenced(
1569 workflow_id,
1570 activity_id,
1571 run_id.as_ref(),
1572 &submitted_token,
1573 outcome,
1574 ) {
1575 tracing::warn!(
1576 worker_id = ?worker_id,
1577 workflow_id = %workflow_id,
1578 activity_id = %activity_id,
1579 %error,
1580 "liminal activity completion handoff rejected"
1581 );
1582 return;
1583 }
1584 if !synthesized
1585 && let Err(error) = heartbeat_tracker.complete_task(worker_id, workflow_id, activity_id)
1586 {
1587 tracing::error!(
1588 worker_id = ?worker_id,
1589 workflow_id = %workflow_id,
1590 activity_id = %activity_id,
1591 %error,
1592 "failed to clear in-flight tracking for completed liminal activity"
1593 );
1594 }
1595 drain_state.notify_activity_drained();
1596}
1597
1598#[cfg(feature = "liminal-transport")]
1599fn complete_liminal_tracking(
1600 heartbeat_tracker: &HeartbeatTracker,
1601 worker_id: WorkerId,
1602 workflow_id: &WorkflowId,
1603 activity_id: &ActivityId,
1604) -> bool {
1605 heartbeat_tracker
1606 .complete_task(worker_id, workflow_id, activity_id)
1607 .unwrap_or_else(|error| {
1608 tracing::error!(
1609 worker_id = ?worker_id,
1610 workflow_id = %workflow_id,
1611 activity_id = %activity_id,
1612 %error,
1613 "failed to clear in-flight tracking for completed liminal activity"
1614 );
1615 true
1616 })
1617}
1618
1619struct ActivityDispatchContext<'a> {
1620 namespace: &'a str,
1621 activity_type: &'a str,
1622 worker_id: WorkerId,
1623 workflow_id: &'a WorkflowId,
1624 activity_id: &'a ActivityId,
1625 completion_token: CompletionToken,
1626 started_at: Instant,
1627}
1628
1629fn activity_task(
1630 activity_type: &str,
1631 input: &str,
1632 execution: (&WorkflowId, &RunId, &ActivityId),
1633 attempt: u32,
1634 labels: BTreeMap<String, String>,
1635 completion_token: &CompletionToken,
1636) -> ProtoActivityTask {
1637 let (workflow_id, run_id, activity_id) = execution;
1638 ProtoActivityTask {
1639 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1640 activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1641 activity_type: activity_type.to_owned(),
1642 input: Some(ProtoPayload {
1643 content_type: String::from("application/json"),
1644 bytes: input.as_bytes().to_vec(),
1645 }),
1646 attempt,
1647 labels: labels.into_iter().collect(),
1648 run_id: Some(run_id.clone().into()),
1649 completion_token: completion_token.as_str().to_owned(),
1650 idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
1651 }
1652}
1653
1654fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1655 let duration_ms = duration_ms(context.started_at.elapsed());
1656 tracing::info!(
1657 operation = "activity_complete",
1658 namespace = context.namespace,
1659 workflow_id = %context.workflow_id,
1660 activity_id = %context.activity_id,
1661 activity_type = context.activity_type,
1662 worker_id = ?context.worker_id,
1663 duration_ms,
1664 outcome = if succeeded { "succeeded" } else { "failed" },
1665 "activity completed"
1666 );
1667}
1668
1669fn duration_ms(duration: Duration) -> u64 {
1670 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1671}
1672
1673fn log_worker_error(
1674 error_type: &'static str,
1675 namespace: &str,
1676 activity_type: &str,
1677 workflow_id: &WorkflowId,
1678 activity_id: &ActivityId,
1679 worker_id: Option<super::registry::WorkerId>,
1680 reason: &str,
1681) {
1682 tracing::error!(
1683 operation = "activity_dispatch",
1684 namespace,
1685 workflow_id = %workflow_id,
1686 activity_id = %activity_id,
1687 activity_type,
1688 worker_id = ?worker_id,
1689 error_type,
1690 reason,
1691 "worker interaction failed"
1692 );
1693}
1694
1695#[cfg(test)]
1696mod tests {
1697 use std::sync::Mutex;
1698
1699 use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1700
1701 use super::*;
1702
1703 fn activity_id(pos: u64) -> ActivityId {
1704 ActivityId::from_sequence_position(pos)
1705 }
1706
1707 #[test]
1708 fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
1709 let pending = PendingActivities::default();
1710 let workflow_id = WorkflowId::new_v4();
1711 let id = activity_id(1);
1712 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1713
1714 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1715 assert_eq!(
1716 rx.recv_timeout(Duration::from_millis(50)),
1717 Ok(Ok("done".to_owned()))
1718 );
1719 Ok(())
1720 }
1721
1722 #[test]
1727 fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
1728 let pending = PendingActivities::default();
1729 let workflow_id = WorkflowId::new_v4();
1730 let id = activity_id(2);
1731 let _first = pending.insert(workflow_id.clone(), id.clone())?;
1732
1733 let error = pending.insert(workflow_id, id).err().ok_or_else(|| {
1734 ServerError::worker_dispatch(
1735 "test",
1736 "same-worker-collision",
1737 "a live responder collision was silently replaced",
1738 )
1739 })?;
1740
1741 assert!(
1742 error.to_string().contains("pending activity collision"),
1743 "{error}"
1744 );
1745 Ok(())
1746 }
1747
1748 #[test]
1749 fn pending_complete_unknown_returns_false() {
1750 let pending = PendingActivities::default();
1751 assert!(!pending.complete(
1752 &WorkflowId::new_v4(),
1753 &activity_id(99),
1754 None,
1755 Ok("orphan".to_owned())
1756 ));
1757 }
1758
1759 #[derive(Default)]
1760 struct RecordingOutboxCallback {
1761 completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1762 failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1763 live: bool,
1764 }
1765
1766 impl OutboxDeliveryCallback for RecordingOutboxCallback {
1767 fn deliver_completion(
1768 &self,
1769 workflow_id: &WorkflowId,
1770 activity_id: &ActivityId,
1771 run_id: Option<&RunId>,
1772 result: String,
1773 ) -> Result<bool, ServerError> {
1774 let _ = run_id;
1775 self.completions
1776 .lock()
1777 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1778 .push((workflow_id.clone(), activity_id.clone(), result));
1779 Ok(self.live)
1780 }
1781
1782 fn deliver_failure(
1783 &self,
1784 workflow_id: &WorkflowId,
1785 activity_id: &ActivityId,
1786 run_id: Option<&RunId>,
1787 reason: String,
1788 ) -> Result<bool, ServerError> {
1789 let _ = run_id;
1790 self.failures
1791 .lock()
1792 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1793 .push((workflow_id.clone(), activity_id.clone(), reason));
1794 Ok(self.live)
1795 }
1796 }
1797
1798 #[test]
1799 fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1800 let pending = PendingActivities::default();
1801 let callback = Arc::new(RecordingOutboxCallback {
1802 live: true,
1803 ..RecordingOutboxCallback::default()
1804 });
1805 pending.clone().set_outbox_delivery(callback.clone());
1807
1808 let workflow_id = WorkflowId::new_v4();
1809 let id = activity_id(7);
1810
1811 assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1814 let completions = callback
1815 .completions
1816 .lock()
1817 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1818 assert_eq!(completions.len(), 1);
1819 assert_eq!(completions[0].0, workflow_id);
1820 assert_eq!(completions[0].1, id);
1821 assert_eq!(completions[0].2, "done");
1822 Ok(())
1823 }
1824
1825 #[test]
1826 fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1827 -> Result<(), ServerError> {
1828 let pending = PendingActivities::default();
1829 let callback = Arc::new(RecordingOutboxCallback::default());
1831 pending.set_outbox_delivery(callback.clone());
1832
1833 let workflow_id = WorkflowId::new_v4();
1834 let id = activity_id(8);
1835
1836 assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1837 let failures = callback
1838 .failures
1839 .lock()
1840 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1841 assert_eq!(failures.len(), 1);
1842 assert_eq!(failures[0].2, "retryable:boom");
1843 Ok(())
1844 }
1845
1846 #[test]
1847 fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1848 let pending = PendingActivities::default();
1851 assert!(!pending.complete(
1852 &WorkflowId::new_v4(),
1853 &activity_id(9),
1854 None,
1855 Ok("x".to_owned())
1856 ));
1857 }
1858
1859 #[test]
1860 fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1861 let pending = PendingActivities::default();
1862 let callback = Arc::new(RecordingOutboxCallback {
1863 live: true,
1864 ..RecordingOutboxCallback::default()
1865 });
1866 pending.set_outbox_delivery(callback.clone());
1867
1868 let workflow_id = WorkflowId::new_v4();
1869 let id = activity_id(10);
1870 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1871
1872 assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1873 assert_eq!(
1874 rx.recv_timeout(Duration::from_millis(50)),
1875 Ok(Ok("matched".to_owned()))
1876 );
1877 assert!(
1878 callback
1879 .completions
1880 .lock()
1881 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1882 .is_empty(),
1883 "a matched completion must deliver to its waiter, not the outbox callback"
1884 );
1885 Ok(())
1886 }
1887
1888 #[test]
1892 fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
1893 let pending = PendingActivities::default();
1894 let workflow_id = WorkflowId::new_v4();
1895 let id = activity_id(11);
1896 let rx = pending.insert(workflow_id.clone(), id.clone())?.1;
1897
1898 pending.park_activity(&workflow_id, &id)?;
1899 let result = rx
1900 .recv_timeout(Duration::from_millis(50))
1901 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1902 assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
1903 Ok(())
1904 }
1905
1906 #[test]
1910 fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
1911 let pending = PendingActivities::default();
1912 let callback = Arc::new(RecordingOutboxCallback {
1913 live: true,
1914 ..RecordingOutboxCallback::default()
1915 });
1916 pending.set_outbox_delivery(callback.clone());
1917
1918 pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
1919
1920 assert!(
1921 callback
1922 .failures
1923 .lock()
1924 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1925 .is_empty(),
1926 "a park must never be delivered as an outbox failure"
1927 );
1928 assert!(
1929 callback
1930 .completions
1931 .lock()
1932 .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1933 .is_empty(),
1934 "a park must never be delivered as an outbox completion"
1935 );
1936 Ok(())
1937 }
1938
1939 #[test]
1940 fn completion_sink_routes_success() -> Result<(), ServerError> {
1941 let pending = PendingActivities::default();
1942 let workflow_id = WorkflowId::new_v4();
1943 let id = activity_id(2);
1944 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1945 let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
1946
1947 pending.complete_activity(ActivityCompletion {
1948 workflow_id,
1949 activity_id: id,
1950 run_id: None,
1951 completion_token,
1952 outcome: ActivityCompletionOutcome::Succeeded(payload),
1953 })?;
1954
1955 let result = rx
1956 .recv_timeout(Duration::from_millis(50))
1957 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
1958 assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
1959 Ok(())
1960 }
1961
1962 #[test]
1963 fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
1964 let pending = PendingActivities::default();
1965 let workflow_id = WorkflowId::new_v4();
1966 let id = activity_id(12);
1967 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
1968
1969 let malformed = pending.complete_activity(ActivityCompletion {
1970 workflow_id: workflow_id.clone(),
1971 activity_id: id.clone(),
1972 run_id: None,
1973 completion_token: completion_token.clone(),
1974 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1975 ContentType::Json,
1976 vec![0xff],
1977 )),
1978 });
1979 assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
1980 assert!(
1981 rx.try_recv().is_err(),
1982 "an invalid result must leave the waiter unresolved"
1983 );
1984
1985 pending.complete_activity(ActivityCompletion {
1986 workflow_id,
1987 activity_id: id,
1988 run_id: None,
1989 completion_token,
1990 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
1991 ContentType::Json,
1992 br#""valid""#.to_vec(),
1993 )),
1994 })?;
1995 let result = rx
1996 .recv_timeout(Duration::from_millis(50))
1997 .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
1998 assert_eq!(result, Ok(r#""valid""#.to_owned()));
1999 Ok(())
2000 }
2001
2002 #[test]
2003 fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
2004 let pending = PendingActivities::default();
2005 let workflow_id = WorkflowId::new_v4();
2006 let id = activity_id(3);
2007 let (completion_token, rx) = pending.insert(workflow_id.clone(), id.clone())?;
2008
2009 pending.complete_activity(ActivityCompletion {
2010 workflow_id,
2011 activity_id: id,
2012 run_id: None,
2013 completion_token,
2014 outcome: ActivityCompletionOutcome::Failed(ActivityError {
2015 kind: ActivityErrorKind::Retryable,
2016 message: "temporary".to_owned(),
2017 details: None,
2018 }),
2019 })?;
2020
2021 let result = rx
2022 .recv_timeout(Duration::from_millis(50))
2023 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2024 assert_eq!(result, Err("retryable:temporary".to_owned()));
2025 Ok(())
2026 }
2027
2028 #[test]
2037 fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
2038 -> Result<(), ServerError> {
2039 let pending = PendingActivities::default();
2040 let post_restart_workflow = WorkflowId::new_v4();
2041 let pre_restart_workflow = WorkflowId::new_v4();
2042 let id = activity_id(1);
2044 let (completion_token, rx) = pending.insert(post_restart_workflow.clone(), id.clone())?;
2045
2046 let rejected = pending.complete_activity(ActivityCompletion {
2048 workflow_id: pre_restart_workflow,
2049 activity_id: id.clone(),
2050 run_id: None,
2051 completion_token: CompletionToken::for_test(),
2052 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2053 ContentType::Json,
2054 br#""stale""#.to_vec(),
2055 )),
2056 });
2057 assert!(matches!(
2058 rejected,
2059 Err(ServerError::ActivityCompletionRejected { .. })
2060 ));
2061 assert!(
2062 rx.try_recv().is_err(),
2063 "stale result for a different workflow must not complete this dispatch"
2064 );
2065
2066 pending.complete_activity(ActivityCompletion {
2068 workflow_id: post_restart_workflow,
2069 activity_id: id,
2070 run_id: None,
2071 completion_token,
2072 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2073 ContentType::Json,
2074 br#""fresh""#.to_vec(),
2075 )),
2076 })?;
2077 let result = rx
2078 .recv_timeout(Duration::from_millis(50))
2079 .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2080 assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2081 Ok(())
2082 }
2083
2084 fn test_tracker() -> HeartbeatTracker {
2087 HeartbeatTracker::new(Duration::from_secs(5))
2088 }
2089
2090 fn greet_request() -> ActivityDispatch {
2093 ActivityDispatch {
2094 namespace: "default".to_owned(),
2095 task_queue: "default".to_owned(),
2096 node: None,
2097 workflow_id: WorkflowId::new_v4(),
2098 run_id: RunId::new_v4(),
2099 activity_id: ActivityId::from_sequence_position(0),
2100 name: "greet".to_owned(),
2101 input: "{}".to_owned(),
2102 config: "{}".to_owned(),
2103 attempt: 1,
2104 labels: std::collections::BTreeMap::new(),
2105 advisory: false,
2106 }
2107 }
2108
2109 #[test]
2110 fn dispatcher_fails_immediately_when_draining_without_workers() {
2111 let registry = ConnectedWorkerRegistry::default();
2112 let drain = DrainState::default();
2113 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2114 .with_drain_state(drain.clone());
2115
2116 let _ = drain.begin();
2117
2118 let result = dispatcher.dispatch(greet_request());
2119
2120 assert!(result.is_err());
2121 let err = result.err().unwrap_or_default();
2122 assert!(
2123 err.contains("drain"),
2124 "expected drain rejection, got: {err}"
2125 );
2126 }
2127
2128 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2144 async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
2145 -> Result<(), Box<dyn std::error::Error>> {
2146 let registry = ConnectedWorkerRegistry::default();
2147 let pending = PendingActivities::default();
2148 let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2149 let activity_types = [String::from("greet")];
2150 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2151
2152 let sink = pending.clone();
2153 let echo_worker = tokio::spawn(async move {
2154 let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
2155 return Err("expected an activity task on the worker channel".to_owned());
2156 };
2157 let workflow_id = task
2158 .workflow_id
2159 .ok_or("task missing workflow id")
2160 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2161 let activity_id = task
2162 .activity_id
2163 .map(ActivityId::from)
2164 .ok_or("task missing activity id")?;
2165 let completion_token =
2166 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2167 .map_err(|error| error.to_string())?;
2168 sink.complete_activity(ActivityCompletion {
2169 workflow_id,
2170 activity_id,
2171 run_id: None,
2172 completion_token,
2173 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2174 ContentType::Json,
2175 br#"{"greeting":"hello"}"#.to_vec(),
2176 )),
2177 })
2178 .map_err(|error| error.to_string())
2179 });
2180
2181 let dispatcher = Arc::new(
2182 WorkerActivityDispatcher::new(registry, "default", test_tracker())
2183 .with_pending(pending),
2184 );
2185 let started = Instant::now();
2186 let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
2189 dispatcher.dispatch(greet_request())
2190 }));
2191 let result = dispatch_task.await.map_err(|error| error.to_string())?;
2192 let elapsed = started.elapsed();
2193
2194 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2195 assert!(
2196 elapsed < Duration::from_secs(5),
2197 "dispatch round trip took {elapsed:?}; task delivery must not \
2198 depend on the blocked dispatch thread"
2199 );
2200 echo_worker.await.map_err(|error| error.to_string())??;
2201 registration.deregister()?;
2202 Ok(())
2203 }
2204
2205 #[tokio::test]
2209 async fn dispatch_on_current_thread_runtime_fails_fast()
2210 -> Result<(), Box<dyn std::error::Error>> {
2211 let registry = ConnectedWorkerRegistry::default();
2212 let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
2213 let activity_types = [String::from("greet")];
2214 let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2215 let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
2216
2217 let started = Instant::now();
2218 let result = dispatcher.dispatch(greet_request());
2219 let elapsed = started.elapsed();
2220
2221 let err = result.err().ok_or("expected dispatch to fail")?;
2222 assert!(
2223 err.contains("multi-thread tokio runtime"),
2224 "unexpected error: {err}"
2225 );
2226 assert!(
2227 elapsed < Duration::from_secs(5),
2228 "fail-fast path took {elapsed:?}"
2229 );
2230 registration.deregister()?;
2231 Ok(())
2232 }
2233
2234 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2240 async fn dispatch_pinned_to_node_reaches_only_that_node()
2241 -> Result<(), Box<dyn std::error::Error>> {
2242 let registry = ConnectedWorkerRegistry::default();
2243 let pending = PendingActivities::default();
2244 let activity_types = [String::from("greet")];
2245 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2246 let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
2247 let on_n2 = registry.register_namespaces(
2253 [String::from("default")],
2254 "default",
2255 Some(String::from("n2")),
2256 activity_types.iter(),
2257 n2_tx,
2258 )?;
2259 let on_n1 = registry.register_namespaces(
2260 [String::from("default")],
2261 "default",
2262 Some(String::from("n1")),
2263 activity_types.iter(),
2264 n1_tx,
2265 )?;
2266
2267 let sink = pending.clone();
2271 let echo_n1 = tokio::spawn(async move {
2272 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2273 return Err("expected an activity task on the n1 worker channel".to_owned());
2274 };
2275 let workflow_id = task
2276 .workflow_id
2277 .ok_or("task missing workflow id")
2278 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2279 let activity_id = task
2280 .activity_id
2281 .map(ActivityId::from)
2282 .ok_or("task missing activity id")?;
2283 let completion_token =
2284 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2285 .map_err(|error| error.to_string())?;
2286 sink.complete_activity(ActivityCompletion {
2287 workflow_id,
2288 activity_id,
2289 run_id: None,
2290 completion_token,
2291 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2292 ContentType::Json,
2293 br#"{"greeting":"hello"}"#.to_vec(),
2294 )),
2295 })
2296 .map_err(|error| error.to_string())
2297 });
2298
2299 let dispatcher = Arc::new(
2300 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2301 .with_pending(pending),
2302 );
2303
2304 let pinned = ActivityDispatch {
2305 node: Some(String::from("n1")),
2306 ..greet_request()
2307 };
2308 let started = Instant::now();
2309 let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
2310 .await
2311 .map_err(|error| error.to_string())?;
2312 let elapsed = started.elapsed();
2313
2314 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2315 assert!(
2316 elapsed < Duration::from_secs(5),
2317 "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
2318 );
2319 echo_n1.await.map_err(|error| error.to_string())??;
2320
2321 assert!(
2323 n2_rx.try_recv().is_err(),
2324 "node=Some(\"n1\") dispatch must not reach the n2 worker"
2325 );
2326
2327 on_n1.deregister()?;
2328 on_n2.deregister()?;
2329 Ok(())
2330 }
2331
2332 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2335 async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
2336 -> Result<(), Box<dyn std::error::Error>> {
2337 let registry = ConnectedWorkerRegistry::default();
2338 let pending = PendingActivities::default();
2339 let activity_types = [String::from("greet")];
2340 let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2341 let on_n1 = registry.register_namespaces(
2342 [String::from("default")],
2343 "default",
2344 Some(String::from("n1")),
2345 activity_types.iter(),
2346 n1_tx,
2347 )?;
2348
2349 let sink = pending.clone();
2350 let echo = tokio::spawn(async move {
2351 let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2352 return Err("expected an activity task on the worker channel".to_owned());
2353 };
2354 let workflow_id = task
2355 .workflow_id
2356 .ok_or("task missing workflow id")
2357 .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2358 let activity_id = task
2359 .activity_id
2360 .map(ActivityId::from)
2361 .ok_or("task missing activity id")?;
2362 let completion_token =
2363 CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2364 .map_err(|error| error.to_string())?;
2365 sink.complete_activity(ActivityCompletion {
2366 workflow_id,
2367 activity_id,
2368 run_id: None,
2369 completion_token,
2370 outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2371 ContentType::Json,
2372 br#"{"greeting":"hello"}"#.to_vec(),
2373 )),
2374 })
2375 .map_err(|error| error.to_string())
2376 });
2377
2378 let dispatcher = Arc::new(
2379 WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2380 .with_pending(pending),
2381 );
2382
2383 let result = tokio::spawn(futures::future::lazy(move |_| {
2385 dispatcher.dispatch(greet_request())
2386 }))
2387 .await
2388 .map_err(|error| error.to_string())?;
2389
2390 assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2391 echo.await.map_err(|error| error.to_string())??;
2392 on_n1.deregister()?;
2393 Ok(())
2394 }
2395}