1use std::collections::BTreeMap;
4
5use aion_core::{ActivityError, ActivityId, Payload, RunId, WorkflowId};
6use aion_proto::{
7 ProtoActivityId, ProtoActivityResult, ProtoActivityTask, ProtoPayload, ProtoRunId,
8 ProtoWorkflowId, WireError, proto_activity_result,
9};
10
11use crate::error::ServerError;
12use crate::shutdown::DrainState;
13use crate::worker::envelope::{CompletionFences, CompletionToken, idempotency_key};
14use crate::worker::queue_service::declarations::QueueDeclarationSource;
15use crate::worker::queue_service::policy::QueueServiceConfig;
16use crate::worker::queue_service::state::QueueServiceState;
17use crate::worker::queue_service::taxonomy::{QueueServiceReason, ServiceAddress};
18use crate::worker::queue_service::wait::{
19 ServiceWait, clear_selection_miss, observe_selection_miss,
20};
21use crate::worker::registry::{ConnectedWorkerRegistry, WorkerMessage};
22use tracing::{Instrument, info_span};
23
24#[derive(Clone, Debug, Eq, PartialEq)]
26pub struct ScheduledActivity {
27 pub namespace: String,
30 pub task_queue: String,
34 pub activity_type: String,
37 pub node: Option<String>,
44 pub workflow_id: WorkflowId,
46 pub activity_id: ActivityId,
48 pub run_id: Option<RunId>,
50 pub input: Payload,
52 pub attempt: u32,
55 pub labels: BTreeMap<String, String>,
58}
59
60impl ScheduledActivity {
61 fn require_run_id(&self) -> Result<&RunId, ServerError> {
65 self.run_id.as_ref().ok_or_else(|| {
66 ServerError::worker_dispatch(
67 self.namespace.clone(),
68 self.activity_type.clone(),
69 "activity run id is missing; refusing unfenced external effect",
70 )
71 })
72 }
73
74 pub fn to_task(
81 &self,
82 completion_token: &CompletionToken,
83 ) -> Result<ProtoActivityTask, ServerError> {
84 let run_id = self.require_run_id()?;
85 Ok(ProtoActivityTask {
86 workflow_id: Some(ProtoWorkflowId::from(self.workflow_id.clone())),
87 activity_id: Some(ProtoActivityId::from(self.activity_id.clone())),
88 activity_type: self.activity_type.clone(),
89 input: Some(ProtoPayload::from(self.input.clone())),
90 attempt: self.attempt,
91 labels: self.labels.clone().into_iter().collect(),
92 run_id: Some(ProtoRunId::from(run_id.clone())),
93 completion_token: completion_token.as_str().to_owned(),
94 idempotency_key: idempotency_key(&self.workflow_id, run_id, &self.activity_id),
95 })
96 }
97}
98
99#[derive(Clone, Debug)]
101pub struct ActivityDispatcher {
102 registry: ConnectedWorkerRegistry,
103 drain_state: DrainState,
104 completion_fences: CompletionFences,
105 queue_declarations: QueueDeclarationSource,
114 queue_service_state: QueueServiceState,
115 queue_service_config: QueueServiceConfig,
116 cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
120}
121
122impl ActivityDispatcher {
123 #[must_use]
125 pub fn new(registry: ConnectedWorkerRegistry) -> Self {
126 Self {
127 registry,
128 drain_state: DrainState::default(),
129 completion_fences: CompletionFences::default(),
130 queue_declarations: QueueDeclarationSource::default(),
131 queue_service_state: QueueServiceState::default(),
132 queue_service_config: QueueServiceConfig::default(),
133 cluster_publisher: None,
134 }
135 }
136
137 #[must_use]
141 pub fn with_cluster_publisher(
142 mut self,
143 cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
144 ) -> Self {
145 self.cluster_publisher = Some(cluster_publisher);
146 self
147 }
148
149 #[must_use]
152 pub fn with_queue_service(
153 mut self,
154 declarations: QueueDeclarationSource,
155 state: QueueServiceState,
156 config: QueueServiceConfig,
157 ) -> Self {
158 self.queue_declarations = declarations;
159 self.queue_service_state = state;
160 self.queue_service_config = config;
161 self
162 }
163
164 #[must_use]
166 pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
167 self.drain_state = drain_state;
168 self
169 }
170
171 #[must_use]
173 pub fn with_completion_fences(mut self, completion_fences: CompletionFences) -> Self {
174 self.completion_fences = completion_fences;
175 self
176 }
177
178 pub async fn dispatch(&self, activity: &ScheduledActivity) -> Result<(), ServerError> {
185 let span = info_span!(
186 "activity_dispatch",
187 operation = "activity_dispatch",
188 namespace = %activity.namespace,
189 task_queue = %activity.task_queue,
190 node = activity.node.as_deref(),
191 workflow_id = %activity.workflow_id,
192 activity_id = %activity.activity_id,
193 activity_type = %activity.activity_type,
194 worker_id = tracing::field::Empty,
195 );
196 let span_fields = span.clone();
197
198 async {
199 self.dispatch_to_node(activity, activity.node.as_deref(), &span_fields)
200 .await
201 }
202 .instrument(span)
203 .await
204 .inspect_err(|error| {
205 log_dispatch_error("activity_dispatch", activity, error);
206 })
207 }
208
209 pub async fn dispatch_preferring(
237 &self,
238 activity: &ScheduledActivity,
239 preferred: &std::collections::BTreeSet<String>,
240 ) -> Result<(), ServerError> {
241 let tiers = crate::worker::preferred_node_order(&aion_store::NamespacePlacement::Prefer {
244 nodes: preferred.clone(),
245 });
246 self.dispatch_over_tiers(activity, &tiers).await
247 }
248
249 pub async fn dispatch_requiring(
278 &self,
279 activity: &ScheduledActivity,
280 required: &std::collections::BTreeSet<String>,
281 ) -> Result<(), ServerError> {
282 let span = info_span!(
283 "activity_dispatch",
284 operation = "activity_dispatch_requiring",
285 namespace = %activity.namespace,
286 task_queue = %activity.task_queue,
287 workflow_id = %activity.workflow_id,
288 activity_id = %activity.activity_id,
289 activity_type = %activity.activity_type,
290 worker_id = tracing::field::Empty,
291 );
292 let span_fields = span.clone();
293 async {
294 loop {
295 let arrival = self.registry.worker_arrival();
303 for label in required {
304 self.drain_state
305 .ensure_accepting(&activity.namespace, &activity.activity_type)?;
306 let candidates = self.registry.workers_for(
307 &activity.namespace,
308 &activity.task_queue,
309 &activity.activity_type,
310 Some(label.as_str()),
311 )?;
312 if let Some(()) = self
313 .send_to_candidates(activity, candidates, &span_fields)
314 .await?
315 {
316 return Ok(());
317 }
318 }
319 tracing::info!(
323 namespace = %activity.namespace,
324 task_queue = %activity.task_queue,
325 activity_type = %activity.activity_type,
326 workflow_id = %activity.workflow_id,
327 activity_id = %activity.activity_id,
328 "no worker on a required (Pinned) node; waiting — will NOT spill to any-node"
329 );
330 arrival.await;
331 }
332 }
333 .instrument(span)
334 .await
335 .inspect_err(|error| {
336 log_dispatch_error("activity_dispatch_requiring", activity, error);
337 })
338 }
339
340 async fn dispatch_over_tiers(
352 &self,
353 activity: &ScheduledActivity,
354 tiers: &[Option<String>],
355 ) -> Result<(), ServerError> {
356 let span = info_span!(
357 "activity_dispatch",
358 operation = "activity_dispatch_preferring",
359 namespace = %activity.namespace,
360 task_queue = %activity.task_queue,
361 workflow_id = %activity.workflow_id,
362 activity_id = %activity.activity_id,
363 activity_type = %activity.activity_type,
364 worker_id = tracing::field::Empty,
365 );
366 let span_fields = span.clone();
367 async {
368 for tier in tiers {
369 let Some(label) = tier else {
370 return self
373 .dispatch_to_node(activity, activity.node.as_deref(), &span_fields)
374 .await;
375 };
376 self.drain_state
377 .ensure_accepting(&activity.namespace, &activity.activity_type)?;
378 let candidates = self.registry.workers_for(
379 &activity.namespace,
380 &activity.task_queue,
381 &activity.activity_type,
382 Some(label.as_str()),
383 )?;
384 if let Some(()) = self
385 .send_to_candidates(activity, candidates, &span_fields)
386 .await?
387 {
388 return Ok(());
389 }
390 }
391 self.dispatch_to_node(activity, activity.node.as_deref(), &span_fields)
394 .await
395 }
396 .instrument(span)
397 .await
398 .inspect_err(|error| {
399 log_dispatch_error("activity_dispatch_preferring", activity, error);
400 })
401 }
402
403 async fn dispatch_to_node(
406 &self,
407 activity: &ScheduledActivity,
408 node: Option<&str>,
409 span_fields: &tracing::Span,
410 ) -> Result<(), ServerError> {
411 let address = ServiceAddress {
420 namespace: activity.namespace.clone(),
421 task_queue: activity.task_queue.clone(),
422 activity_type: activity.activity_type.clone(),
423 node: node.map(ToOwned::to_owned),
424 };
425 let wait = ServiceWait {
426 registry: &self.registry,
427 declarations: &self.queue_declarations,
428 config: &self.queue_service_config,
429 state: &self.queue_service_state,
430 address: &address,
431 workflow_id: &activity.workflow_id,
432 activity_id: &activity.activity_id,
433 publisher: self.cluster_publisher.as_ref(),
434 };
435 let policy = self
436 .queue_service_config
437 .policy_for(&activity.namespace, &activity.task_queue);
438 let started_at = std::time::Instant::now();
439 let mut reported: Option<QueueServiceReason> = None;
440 let workers = loop {
441 let arrival = self.registry.worker_arrival();
451 self.drain_state
452 .ensure_accepting(&activity.namespace, &activity.activity_type)
453 .inspect_err(|_| clear_selection_miss(&wait))?;
454 let candidates = self
455 .registry
456 .workers_for(
457 &activity.namespace,
458 &activity.task_queue,
459 &activity.activity_type,
460 node,
461 )
462 .inspect_err(|_| clear_selection_miss(&wait))?;
463 if !candidates.is_empty() {
464 if let Some(reason) = reported {
465 tracing::info!(
466 namespace = %activity.namespace,
467 task_queue = %activity.task_queue,
468 activity_type = %activity.activity_type,
469 workflow_id = %activity.workflow_id,
470 activity_id = %activity.activity_id,
471 queue_service_reason = reason.as_str(),
472 "queue service restored; the parked dispatch has a worker"
473 );
474 }
475 clear_selection_miss(&wait);
476 break candidates;
477 }
478 match observe_selection_miss(&wait, policy, None, started_at.elapsed(), reported) {
479 Ok(None) => {}
506 Ok(Some(observed)) => reported = Some(observed.reason),
507 Err(refusal) => {
508 clear_selection_miss(&wait);
509 return Err(ServerError::worker_dispatch(
510 activity.namespace.clone(),
511 activity.activity_type.clone(),
512 refusal.reason_string(),
513 ));
514 }
515 }
516 arrival.await;
517 };
518 match self
519 .send_to_candidates(activity, workers, span_fields)
520 .await?
521 {
522 Some(()) => Ok(()),
523 None => Err(ServerError::worker_dispatch(
524 activity.namespace.clone(),
525 activity.activity_type.clone(),
526 format!(
527 "all matching worker streams in task queue {} closed before task could be \
528 delivered",
529 activity.task_queue
530 ),
531 )),
532 }
533 }
534
535 async fn send_to_candidates(
540 &self,
541 activity: &ScheduledActivity,
542 candidates: Vec<crate::worker::registry::WorkerHandle>,
543 span_fields: &tracing::Span,
544 ) -> Result<Option<()>, ServerError> {
545 let run_id = activity.require_run_id()?;
546 let completion_token = self.completion_fences.issue(
551 &activity.workflow_id,
552 run_id,
553 &activity.activity_id,
554 activity.attempt,
555 )?;
556 let task = activity.to_task(&completion_token)?;
557 for worker in candidates {
558 if let Err(error) = self
559 .drain_state
560 .ensure_accepting(&activity.namespace, &activity.activity_type)
561 {
562 self.completion_fences.revoke(
566 &activity.workflow_id,
567 &activity.activity_id,
568 &completion_token,
569 )?;
570 return Err(error);
571 }
572 span_fields.record("worker_id", format!("{:?}", worker.id()));
573 if let Some(sender) = worker.sender()
578 && sender
579 .send(WorkerMessage::ActivityTask(Box::new(task.clone())))
580 .await
581 .is_ok()
582 {
583 return Ok(Some(()));
584 }
585 self.registry.deregister(worker.id())?;
586 }
587 self.completion_fences.revoke(
593 &activity.workflow_id,
594 &activity.activity_id,
595 &completion_token,
596 )?;
597 Ok(None)
598 }
599}
600
601fn log_dispatch_error(operation: &'static str, activity: &ScheduledActivity, error: &ServerError) {
602 let fields = error.trace_fields();
603 tracing::error!(
604 operation,
605 namespace = %activity.namespace,
606 task_queue = %activity.task_queue,
607 node = activity.node.as_deref(),
608 workflow_id = %activity.workflow_id,
609 activity_id = %activity.activity_id,
610 activity_type = %activity.activity_type,
611 error_type = %fields.error_type,
612 store_error_type = fields.store_error_type,
613 reason = %fields.reason,
614 "activity dispatch failed"
615 );
616}
617
618#[derive(Clone, Debug, Eq, PartialEq)]
620pub enum ActivityCompletionOutcome {
621 Succeeded(Payload),
623 Failed(ActivityError),
625 WorkerLost {
636 worker_id: crate::worker::registry::WorkerId,
638 },
639}
640
641#[derive(Clone, Debug, Eq, PartialEq)]
643pub struct ActivityCompletion {
644 pub workflow_id: WorkflowId,
646 pub activity_id: ActivityId,
648 pub run_id: Option<RunId>,
650 pub completion_token: CompletionToken,
652 pub outcome: ActivityCompletionOutcome,
654}
655
656impl TryFrom<ProtoActivityResult> for ActivityCompletion {
657 type Error = ServerError;
658
659 fn try_from(value: ProtoActivityResult) -> Result<Self, Self::Error> {
660 let workflow_id = value
661 .workflow_id
662 .ok_or_else(|| wire_error("activity result workflow id is missing"))
663 .and_then(|id| WorkflowId::try_from(id).map_err(ServerError::from))?;
664 let activity_id = value
665 .activity_id
666 .ok_or_else(|| wire_error("activity result activity id is missing"))
667 .map(ActivityId::from)?;
668 let run_id = value
669 .run_id
670 .ok_or_else(|| wire_error("activity result run id is missing"))
671 .and_then(|id| RunId::try_from(id).map_err(ServerError::from))?;
672 let completion_token =
673 CompletionToken::from_wire(&workflow_id, &activity_id, value.completion_token)?;
674 let outcome = match value.outcome {
675 Some(proto_activity_result::Outcome::Result(payload)) => {
676 ActivityCompletionOutcome::Succeeded(
677 Payload::try_from(payload).map_err(ServerError::from)?,
678 )
679 }
680 Some(proto_activity_result::Outcome::Error(error)) => {
681 ActivityCompletionOutcome::Failed(
682 ActivityError::try_from(error).map_err(ServerError::from)?,
683 )
684 }
685 None => return Err(wire_error("activity result outcome is missing")),
686 };
687
688 Ok(Self {
689 workflow_id,
690 activity_id,
691 run_id: Some(run_id),
692 completion_token,
693 outcome,
694 })
695 }
696}
697
698pub trait ActivityCompletionSink {
700 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError>;
706
707 fn park_activity(
724 &self,
725 workflow_id: &WorkflowId,
726 activity_id: &ActivityId,
727 ) -> Result<(), ServerError>;
728}
729
730pub fn handle_activity_result(
736 sink: &impl ActivityCompletionSink,
737 result: ProtoActivityResult,
738) -> Result<(), ServerError> {
739 sink.complete_activity(ActivityCompletion::try_from(result)?)
740}
741
742fn wire_error(message: &'static str) -> ServerError {
743 ServerError::Wire {
744 wire: WireError::backend(message),
745 }
746}
747
748#[cfg(test)]
749mod tests {
750 use std::sync::Mutex;
751
752 use aion_core::{ActivityErrorKind, ContentType};
753 use aion_proto::{ProtoActivityError, ProtoActivityErrorKind};
754 use serde_json::json;
755 use uuid::Uuid;
756
757 use crate::worker::queue_service::declarations::{QueueDeclaration, QueueDeclarations};
758 use crate::worker::registry::{ConnectedWorkerRegistry, WorkerRegistration};
759
760 use super::*;
761
762 fn workflow_id() -> WorkflowId {
763 WorkflowId::new(Uuid::nil())
764 }
765
766 fn activity_id() -> ActivityId {
767 ActivityId::from_sequence_position(42)
768 }
769
770 fn payload(value: &serde_json::Value) -> Result<Payload, Box<dyn std::error::Error>> {
771 Ok(Payload::from_json(value)?)
772 }
773
774 #[tokio::test]
775 async fn dispatch_pushes_activity_task_with_correlation()
776 -> Result<(), Box<dyn std::error::Error>> {
777 let registry = ConnectedWorkerRegistry::default();
778 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
779 let activity_types = [String::from("charge-card")];
780 let registration = registry.register("tenant-a", activity_types.iter(), tx)?;
781 let dispatcher = ActivityDispatcher::new(registry.clone());
782 let input = payload(&json!({"amount": 1200}))?;
783 let scheduled = ScheduledActivity {
784 namespace: String::from("tenant-a"),
785 task_queue: String::from("default"),
786 activity_type: String::from("charge-card"),
787 node: None,
788 workflow_id: workflow_id(),
789 activity_id: activity_id(),
790 run_id: Some(RunId::new_v4()),
791 input: input.clone(),
792 attempt: 1,
793 labels: std::collections::BTreeMap::new(),
794 };
795
796 dispatcher.dispatch(&scheduled).await?;
797 let message = rx.recv().await.ok_or("expected pushed activity task")?;
798 let WorkerMessage::ActivityTask(task) = message else {
799 return Err("expected activity task message".into());
800 };
801
802 assert_eq!(task.workflow_id, Some(ProtoWorkflowId::from(workflow_id())));
803 assert_eq!(task.activity_id, Some(ProtoActivityId::from(activity_id())));
804 assert_eq!(task.activity_type, "charge-card");
805 assert_eq!(task.input, Some(ProtoPayload::from(input)));
806 assert_eq!(task.attempt, 1, "wire task must carry the stamped attempt");
807
808 registration.deregister()?;
809 Ok(())
810 }
811
812 #[tokio::test]
813 async fn dispatch_waits_for_worker_then_delivers() -> Result<(), Box<dyn std::error::Error>> {
814 let registry = ConnectedWorkerRegistry::default();
815 let dispatcher = ActivityDispatcher::new(registry.clone());
816 let scheduled = ScheduledActivity {
817 namespace: String::from("tenant-a"),
818 task_queue: String::from("default"),
819 activity_type: String::from("charge-card"),
820 node: None,
821 workflow_id: workflow_id(),
822 activity_id: activity_id(),
823 run_id: Some(RunId::new_v4()),
824 input: Payload::new(ContentType::Json, b"{}".to_vec()),
825 attempt: 1,
826 labels: std::collections::BTreeMap::new(),
827 };
828
829 let dispatch_handle = tokio::spawn({
830 let dispatcher = dispatcher.clone();
831 let scheduled = scheduled.clone();
832 async move { dispatcher.dispatch(&scheduled).await }
833 });
834
835 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
836 assert!(!dispatch_handle.is_finished(), "dispatch should be waiting");
837
838 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
839 let activity_types = [String::from("charge-card")];
840 let _registration = registry.register("tenant-a", activity_types.iter(), tx)?;
841
842 dispatch_handle.await??;
843 assert!(rx.recv().await.is_some());
844 Ok(())
845 }
846
847 struct RegisterInsideTheSelectionWindow {
865 registry: ConnectedWorkerRegistry,
866 activity_types: Vec<String>,
867 delivery: tokio::sync::mpsc::Sender<WorkerMessage>,
868 placed: std::sync::OnceLock<Result<WorkerRegistration, ServerError>>,
873 }
874
875 impl QueueDeclarations for RegisterInsideTheSelectionWindow {
876 fn declaration_for(&self, _task_queue: &str) -> QueueDeclaration {
877 let placed = self.placed.get_or_init(|| {
882 self.registry.register(
883 "tenant-a",
884 self.activity_types.iter(),
885 self.delivery.clone(),
886 )
887 });
888 if let Err(error) = placed {
889 tracing::error!(%error, "T9 seam could not place its worker in the window");
890 }
891 QueueDeclaration::Declared
894 }
895 }
896
897 #[test]
931 fn a_registration_inside_the_selection_window_is_delivered_without_a_second_event()
932 -> Result<(), Box<dyn std::error::Error>> {
933 use std::future::Future;
934 use std::task::{Context, Poll, Waker};
935
936 let registry = ConnectedWorkerRegistry::default();
937 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
938 let seam = std::sync::Arc::new(RegisterInsideTheSelectionWindow {
939 registry: registry.clone(),
940 activity_types: vec![String::from("charge-card")],
941 delivery: tx,
942 placed: std::sync::OnceLock::new(),
943 });
944 let declarations = QueueDeclarationSource::default();
945 declarations.install(seam.clone());
946 let dispatcher = ActivityDispatcher::new(registry.clone()).with_queue_service(
947 declarations,
948 QueueServiceState::default(),
949 QueueServiceConfig::default(),
950 );
951 let scheduled = ScheduledActivity {
952 namespace: String::from("tenant-a"),
953 task_queue: String::from("default"),
954 activity_type: String::from("charge-card"),
955 node: None,
956 workflow_id: workflow_id(),
957 activity_id: activity_id(),
958 run_id: Some(RunId::new_v4()),
959 input: Payload::new(ContentType::Json, b"{}".to_vec()),
960 attempt: 1,
961 labels: std::collections::BTreeMap::new(),
962 };
963
964 assert!(
967 registry
968 .workers_for("tenant-a", "default", "charge-card", None)?
969 .is_empty(),
970 "the window is only a window if selection misses on the first pass"
971 );
972
973 let span = tracing::Span::none();
974 let mut dispatch = std::pin::pin!(dispatcher.dispatch_to_node(&scheduled, None, &span));
975 let mut context = Context::from_waker(Waker::noop());
976 let polled = dispatch.as_mut().poll(&mut context);
977
978 match seam.placed.get() {
981 Some(Ok(_)) => {}
982 Some(Err(error)) => {
983 return Err(format!("the seam's registration failed: {error}").into());
984 }
985 None => {
986 return Err(
987 "the seam was never consulted: the loop did not reach the census, \
988 so this test proved nothing about the park"
989 .into(),
990 );
991 }
992 }
993
994 assert!(
995 matches!(polled, Poll::Ready(Ok(()))),
996 "a registration that landed between the census and the park must be RETAINED: the \
997 loop holds a subscription taken before it looked, so it re-selects and delivers \
998 without any second event. Pending here is the finding — a dispatch parked past its \
999 own wake, with no probe, no verdict and no other registration left to free it."
1000 );
1001
1002 let message = rx.try_recv()?;
1003 let WorkerMessage::ActivityTask(task) = message else {
1004 return Err("expected the activity task to reach the window's worker".into());
1005 };
1006 assert_eq!(task.activity_type, "charge-card");
1007 Ok(())
1008 }
1009
1010 #[tokio::test]
1031 async fn a_dispatch_to_an_all_ineligible_pool_parks_until_eligibility_returns()
1032 -> Result<(), Box<dyn std::error::Error>> {
1033 let registry = ConnectedWorkerRegistry::default();
1034 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
1035 let activity_types = [String::from("charge-card")];
1036 let registration = registry.register("tenant-a", activity_types.iter(), tx)?;
1037 let worker_id = registration
1038 .worker_id()
1039 .ok_or("registration assigned no worker id")?;
1040 registry.set_dispatch_ineligible([worker_id].into_iter().collect())?;
1043
1044 let dispatcher = ActivityDispatcher::new(registry.clone());
1045 let scheduled = ScheduledActivity {
1046 namespace: String::from("tenant-a"),
1047 task_queue: String::from("default"),
1048 activity_type: String::from("charge-card"),
1049 node: None,
1050 workflow_id: workflow_id(),
1051 activity_id: activity_id(),
1052 run_id: Some(RunId::new_v4()),
1053 input: Payload::new(ContentType::Json, b"{}".to_vec()),
1054 attempt: 1,
1055 labels: std::collections::BTreeMap::new(),
1056 };
1057 let dispatch_handle = tokio::spawn({
1058 let dispatcher = dispatcher.clone();
1059 let scheduled = scheduled.clone();
1060 async move { dispatcher.dispatch(&scheduled).await }
1061 });
1062
1063 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1064 assert!(
1065 !dispatch_handle.is_finished(),
1066 "a worker the server cannot reach must not take the dispatch"
1067 );
1068
1069 registry.set_dispatch_ineligible(std::collections::BTreeSet::new())?;
1072
1073 dispatch_handle.await??;
1074 assert!(
1075 rx.recv().await.is_some(),
1076 "the parked dispatch delivers as soon as the pool has an eligible worker"
1077 );
1078
1079 registration.deregister()?;
1080 Ok(())
1081 }
1082
1083 #[tokio::test]
1095 async fn a_dispatch_with_no_worker_publishes_its_park_and_clears_on_arrival()
1096 -> Result<(), Box<dyn std::error::Error>> {
1097 let registry = ConnectedWorkerRegistry::default();
1098 let state = QueueServiceState::default();
1099 let dispatcher = ActivityDispatcher::new(registry.clone()).with_queue_service(
1100 QueueDeclarationSource::default(),
1101 state.clone(),
1102 QueueServiceConfig::default(),
1103 );
1104 let scheduled = ScheduledActivity {
1105 namespace: String::from("tenant-a"),
1106 task_queue: String::from("default"),
1107 activity_type: String::from("charge-card"),
1108 node: None,
1109 workflow_id: workflow_id(),
1110 activity_id: activity_id(),
1111 run_id: Some(RunId::new_v4()),
1112 input: Payload::new(ContentType::Json, b"{}".to_vec()),
1113 attempt: 1,
1114 labels: std::collections::BTreeMap::new(),
1115 };
1116
1117 assert!(
1120 state.unserved()?.is_empty(),
1121 "no dispatch has been made yet"
1122 );
1123
1124 let unwired = ActivityDispatcher::new(registry.clone());
1130 let unwired_handle = tokio::spawn({
1131 let scheduled = scheduled.clone();
1132 async move { unwired.dispatch(&scheduled).await }
1133 });
1134 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1135 assert!(
1136 state.unserved()?.is_empty(),
1137 "an unshared dispatcher must publish nothing HERE — that is the \
1138 defect this test exists to catch, reproduced on purpose"
1139 );
1140 unwired_handle.abort();
1141
1142 let dispatch_handle = tokio::spawn({
1143 let dispatcher = dispatcher.clone();
1144 let scheduled = scheduled.clone();
1145 async move { dispatcher.dispatch(&scheduled).await }
1146 });
1147 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1148 assert!(!dispatch_handle.is_finished(), "dispatch should be waiting");
1149
1150 let unserved = state.unserved()?;
1151 assert_eq!(
1152 unserved.len(),
1153 1,
1154 "the parked dispatch must be queryable, not just logged: {unserved:?}"
1155 );
1156 assert_eq!(unserved[0].key.task_queue, "default");
1157 assert_eq!(
1158 unserved[0].reason,
1159 QueueServiceReason::NoLivePollers,
1160 "an empty pool must be classified, not reported as a bare miss"
1161 );
1162 assert_eq!(
1163 state.parked_on_queue("default")?,
1164 1,
1165 "the run parked on the queue must be attributable to the queue"
1166 );
1167
1168 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
1171 let activity_types = [String::from("charge-card")];
1172 let _registration = registry.register("tenant-a", activity_types.iter(), tx)?;
1173
1174 dispatch_handle.await??;
1175 assert!(rx.recv().await.is_some(), "the task must be delivered");
1176 assert!(
1177 state.unserved()?.is_empty(),
1178 "a served dispatch must not be left published as unserved: {:?}",
1179 state.unserved()?
1180 );
1181 Ok(())
1182 }
1183
1184 #[tokio::test]
1185 async fn dispatch_skips_closed_worker_and_uses_next_match()
1186 -> Result<(), Box<dyn std::error::Error>> {
1187 let registry = ConnectedWorkerRegistry::default();
1188 let (closed_tx, closed_rx) = tokio::sync::mpsc::channel(1);
1189 let (live_tx, mut live_rx) = tokio::sync::mpsc::channel(1);
1190 let activity_types = [String::from("charge-card")];
1191 let closed_registration =
1192 registry.register("tenant-a", activity_types.iter(), closed_tx)?;
1193 let live_registration = registry.register("tenant-a", activity_types.iter(), live_tx)?;
1194 drop(closed_rx);
1195
1196 let dispatcher = ActivityDispatcher::new(registry.clone());
1197 let scheduled = ScheduledActivity {
1198 namespace: String::from("tenant-a"),
1199 task_queue: String::from("default"),
1200 activity_type: String::from("charge-card"),
1201 node: None,
1202 workflow_id: workflow_id(),
1203 activity_id: activity_id(),
1204 run_id: Some(RunId::new_v4()),
1205 input: Payload::new(ContentType::Json, b"{}".to_vec()),
1206 attempt: 1,
1207 labels: std::collections::BTreeMap::new(),
1208 };
1209
1210 dispatcher.dispatch(&scheduled).await?;
1211
1212 assert!(live_rx.recv().await.is_some());
1213 assert_eq!(
1214 registry
1215 .workers_for("tenant-a", "default", "charge-card", None)?
1216 .len(),
1217 1
1218 );
1219
1220 closed_registration.deregister()?;
1221 live_registration.deregister()?;
1222 Ok(())
1223 }
1224
1225 fn scheduled_unpinned() -> ScheduledActivity {
1226 ScheduledActivity {
1227 namespace: String::from("tenant-a"),
1228 task_queue: String::from("default"),
1229 activity_type: String::from("charge-card"),
1230 node: None,
1233 workflow_id: workflow_id(),
1234 activity_id: activity_id(),
1235 run_id: Some(RunId::new_v4()),
1236 input: Payload::new(ContentType::Json, b"{}".to_vec()),
1237 attempt: 1,
1238 labels: std::collections::BTreeMap::new(),
1239 }
1240 }
1241
1242 fn required(labels: &[&str]) -> std::collections::BTreeSet<String> {
1243 labels.iter().map(|l| (*l).to_owned()).collect()
1244 }
1245
1246 #[tokio::test]
1251 async fn dispatch_requiring_waits_and_never_spills_to_a_wrong_node_worker()
1252 -> Result<(), Box<dyn std::error::Error>> {
1253 let registry = ConnectedWorkerRegistry::default();
1254 let dispatcher = ActivityDispatcher::new(registry.clone());
1255 let scheduled = scheduled_unpinned();
1256 let types = [String::from("charge-card")];
1257
1258 let (wrong_tx, mut wrong_rx) = tokio::sync::mpsc::channel(1);
1261 let _wrong = registry.register_namespaces(
1262 [String::from("tenant-a")],
1263 "default",
1264 Some(String::from("n2")),
1265 types.iter(),
1266 wrong_tx,
1267 )?;
1268
1269 let handle = tokio::spawn({
1270 let dispatcher = dispatcher.clone();
1271 let scheduled = scheduled.clone();
1272 async move {
1273 dispatcher
1274 .dispatch_requiring(&scheduled, &required(&["n1"]))
1275 .await
1276 }
1277 });
1278
1279 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1281 assert!(
1282 !handle.is_finished(),
1283 "Pinned{{n1}} must WAIT rather than spill to the live n2 worker"
1284 );
1285 assert!(
1286 wrong_rx.try_recv().is_err(),
1287 "the wrong-node (n2) worker must never receive the task"
1288 );
1289
1290 let (right_tx, mut right_rx) = tokio::sync::mpsc::channel(1);
1292 let _right = registry.register_namespaces(
1293 [String::from("tenant-a")],
1294 "default",
1295 Some(String::from("n1")),
1296 types.iter(),
1297 right_tx,
1298 )?;
1299
1300 handle.await??;
1301 assert!(
1302 right_rx.recv().await.is_some(),
1303 "the required n1 worker receives the task once live"
1304 );
1305 assert!(
1306 wrong_rx.try_recv().is_err(),
1307 "the wrong-node worker still never received it"
1308 );
1309 Ok(())
1310 }
1311
1312 #[tokio::test]
1315 async fn dispatch_requiring_never_mutates_the_rows_node()
1316 -> Result<(), Box<dyn std::error::Error>> {
1317 let registry = ConnectedWorkerRegistry::default();
1318 let dispatcher = ActivityDispatcher::new(registry.clone());
1319 let scheduled = scheduled_unpinned();
1320 assert_eq!(scheduled.node, None, "precondition: the row is unpinned");
1321 let types = [String::from("charge-card")];
1322 let (tx, mut rx) = tokio::sync::mpsc::channel(1);
1323 let _right = registry.register_namespaces(
1324 [String::from("tenant-a")],
1325 "default",
1326 Some(String::from("n1")),
1327 types.iter(),
1328 tx,
1329 )?;
1330
1331 dispatcher
1332 .dispatch_requiring(&scheduled, &required(&["n1"]))
1333 .await?;
1334
1335 assert!(rx.recv().await.is_some(), "the n1 worker received the task");
1336 assert_eq!(
1337 scheduled.node, None,
1338 "the row's authored node MUST remain None through a Pinned dispatch \
1339 (the determinism invariant, CP-Phase-2 §2.4)"
1340 );
1341 Ok(())
1342 }
1343
1344 #[derive(Default)]
1345 struct RecordingSink {
1346 completions: Mutex<Vec<ActivityCompletion>>,
1347 }
1348
1349 impl ActivityCompletionSink for RecordingSink {
1350 fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
1351 self.completions
1352 .lock()
1353 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?
1354 .push(completion);
1355 Ok(())
1356 }
1357
1358 fn park_activity(
1359 &self,
1360 _workflow_id: &WorkflowId,
1361 _activity_id: &ActivityId,
1362 ) -> Result<(), ServerError> {
1363 Err(ServerError::worker_dispatch(
1364 "",
1365 "",
1366 "result-handoff tests never park a dispatch",
1367 ))
1368 }
1369 }
1370
1371 #[test]
1372 fn successful_activity_result_calls_completion_sink() -> Result<(), Box<dyn std::error::Error>>
1373 {
1374 let sink = RecordingSink::default();
1375 let output = payload(&json!({"ok": true}))?;
1376 let result = ProtoActivityResult {
1377 workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
1378 activity_id: Some(ProtoActivityId::from(activity_id())),
1379 run_id: Some(ProtoRunId::from(RunId::new_v4())),
1380 completion_token: String::from("generation-1"),
1381 outcome: Some(proto_activity_result::Outcome::Result(ProtoPayload::from(
1382 output.clone(),
1383 ))),
1384 };
1385
1386 handle_activity_result(&sink, result)?;
1387 let completions = sink
1388 .completions
1389 .lock()
1390 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1391
1392 assert_eq!(completions.len(), 1);
1393 assert_eq!(completions[0].workflow_id, workflow_id());
1394 assert_eq!(completions[0].activity_id, activity_id());
1395 assert_eq!(
1396 completions[0].outcome,
1397 ActivityCompletionOutcome::Succeeded(output)
1398 );
1399 Ok(())
1400 }
1401
1402 #[test]
1403 fn failed_activity_result_preserves_error_classification()
1404 -> Result<(), Box<dyn std::error::Error>> {
1405 let sink = RecordingSink::default();
1406 let error = ProtoActivityError {
1407 kind: ProtoActivityErrorKind::Retryable as i32,
1408 message: String::from("temporary outage"),
1409 details: Some(ProtoPayload::from(payload(
1410 &json!({"retry_after_ms": 500}),
1411 )?)),
1412 };
1413 let result = ProtoActivityResult {
1414 workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
1415 activity_id: Some(ProtoActivityId::from(activity_id())),
1416 run_id: Some(ProtoRunId::from(RunId::new_v4())),
1417 completion_token: String::from("generation-1"),
1418 outcome: Some(proto_activity_result::Outcome::Error(error)),
1419 };
1420
1421 handle_activity_result(&sink, result)?;
1422 let completions = sink
1423 .completions
1424 .lock()
1425 .map_err(|_| ServerError::lock_poisoned("recording completion sink"))?;
1426
1427 assert_eq!(completions.len(), 1);
1428 match &completions[0].outcome {
1429 ActivityCompletionOutcome::Failed(error) => {
1430 assert_eq!(error.kind, ActivityErrorKind::Retryable);
1431 assert!(error.is_retryable());
1432 }
1433 other => return Err(format!("expected failed outcome, got {other:?}").into()),
1434 }
1435 Ok(())
1436 }
1437}