1use std::{
11 any::{Any, TypeId},
12 collections::{BTreeMap, HashMap},
13 future::Future,
14 pin::Pin,
15 sync::{
16 atomic::{AtomicBool, Ordering},
17 Arc, Mutex, OnceLock,
18 },
19 task::{Context as TaskContext, Poll},
20 time::{Duration, Instant, SystemTime, UNIX_EPOCH},
21};
22
23use apache_avro::{from_avro_datum, from_value, to_avro_datum, to_value, Schema};
24use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
25use futures_util::{future::OptionFuture, task::noop_waker_ref};
26use serde::{de::DeserializeOwned, Deserialize, Serialize};
27pub use serde_json::{json, Value};
28use thiserror::Error;
29
30pub const WORKER_PROTOCOL_VERSION: &str = "1.2";
31pub const CONTROL_PLANE_VERSION: &str = "2";
32pub const DEFAULT_CODEC: &str = "avro";
33pub const JSON_CODEC: &str = "json";
34pub const SDK_VERSION: &str = concat!("durable-workflow-rust/", env!("CARGO_PKG_VERSION"));
35pub const QUERY_TASKS_CAPABILITY: &str = "query_tasks";
37pub const QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION: &str = "1.8";
39
40const MAX_LONG_POLL_TIMEOUT_SECONDS: u64 = 60;
41const WORKFLOW_TASK_WAITING_FOR_HISTORY_MESSAGE: &str =
42 "Workflow task waiting for scheduled history.";
43const WORKFLOW_TASK_WAITING_FOR_HISTORY_TYPE: &str = "WorkflowTaskWaitingForHistory";
44
45const QUERY_TASK_FINAL_REJECTION_REASONS: &[&str] = &[
46 "lease_expired",
47 "query_task_not_found",
48 "query_task_not_leased",
49 "query_task_timed_out",
50];
51
52const AVRO_PAYLOAD_SCHEMA_JSON: &str = r#"{"type":"record","name":"Payload","namespace":"durable_workflow","fields":[{"name":"json","type":"string"},{"name":"version","type":"int","default":1}]}"#;
53const AVRO_PAYLOAD_VERSION: i32 = 1;
54
55static AVRO_PAYLOAD_SCHEMA: OnceLock<std::result::Result<Schema, String>> = OnceLock::new();
56
57#[derive(Clone, Copy)]
58enum RequestProtocol {
59 ControlPlane,
60 Worker(&'static str),
61}
62
63impl RequestProtocol {
64 fn is_worker(self) -> bool {
65 matches!(self, Self::Worker(_))
66 }
67}
68
69pub type Result<T> = std::result::Result<T, Error>;
70
71#[derive(Debug, Error)]
72pub enum Error {
73 #[error("transport error: {0}")]
74 Transport(#[from] reqwest::Error),
75 #[error("json error: {0}")]
76 Json(#[from] serde_json::Error),
77 #[error("http {status}: {body}")]
78 Http {
79 status: reqwest::StatusCode,
80 body: String,
81 },
82 #[error("codec error: {0}")]
83 Codec(String),
84 #[error(transparent)]
85 QueryFailed(QueryFailure),
86 #[error(transparent)]
87 Protocol(ProtocolFailure),
88 #[error(transparent)]
89 NonDeterministicReplay(ReplayFailure),
90 #[error(transparent)]
91 ChildWorkflowFailed(ChildWorkflowFailure),
92 #[error("workflow handler {0:?} is not registered")]
93 WorkflowNotRegistered(String),
94 #[error("activity handler {0:?} is not registered")]
95 ActivityNotRegistered(String),
96 #[error("workflow future yielded without emitting a durable command")]
97 WorkflowYieldedWithoutCommand,
98 #[error("workflow state lock is poisoned")]
99 WorkflowStatePoisoned,
100 #[error("timer duration is too large for the worker protocol")]
101 TimerDurationOverflow,
102 #[error("operation timed out")]
103 Timeout,
104 #[error("worker loop error: {0}")]
105 WorkerLoop(String),
106 #[error("invalid child workflow options: {0}")]
107 InvalidChildWorkflowOptions(String),
108}
109
110#[derive(Clone, Copy, Debug, PartialEq, Eq)]
112pub enum ChildWorkflowFailureKind {
113 Failed,
114 Cancelled,
115 Terminated,
116}
117
118#[derive(Clone, Debug, Error)]
124#[error("child workflow failed ({reason}): {message}")]
125pub struct ChildWorkflowFailure {
126 pub kind: ChildWorkflowFailureKind,
127 pub reason: String,
128 pub message: String,
129 pub parent_workflow_id: Option<String>,
130 pub parent_workflow_run_id: Option<String>,
131 pub child_workflow_id: Option<String>,
132 pub child_workflow_run_id: Option<String>,
133 pub child_workflow_type: Option<String>,
134 pub failure_id: Option<String>,
135 pub failure_category: Option<String>,
136 pub exception_type: Option<String>,
137 pub exception_class: Option<String>,
138 pub non_retryable: bool,
139 pub code: Option<Value>,
140 pub exception: Option<Value>,
141}
142
143#[derive(Clone, Debug, PartialEq, Eq)]
145pub struct WorkflowIdentity {
146 pub workflow_id: Option<String>,
147 pub run_id: Option<String>,
148}
149
150#[derive(Clone, Debug, PartialEq)]
152pub struct ChildWorkflowResult {
153 pub parent: WorkflowIdentity,
154 pub child: WorkflowIdentity,
155 pub child_workflow_type: Option<String>,
156 pub result: Value,
157}
158
159#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
161pub enum ParentClosePolicy {
162 #[default]
163 Abandon,
164 RequestCancel,
165 Terminate,
166}
167
168impl ParentClosePolicy {
169 fn as_str(self) -> &'static str {
170 match self {
171 Self::Abandon => "abandon",
172 Self::RequestCancel => "request_cancel",
173 Self::Terminate => "terminate",
174 }
175 }
176}
177
178#[derive(Clone, Debug, Default, PartialEq, Eq)]
180pub struct ChildWorkflowRetryPolicy {
181 pub max_attempts: Option<u32>,
182 pub backoff_seconds: Vec<u64>,
183 pub non_retryable_error_types: Vec<String>,
184}
185
186#[derive(Clone, Debug, PartialEq, Eq)]
190pub struct ChildWorkflowOptions {
191 pub task_queue: String,
192 pub parent_close_policy: ParentClosePolicy,
193 pub retry_policy: Option<ChildWorkflowRetryPolicy>,
194 pub execution_timeout_seconds: Option<u64>,
195 pub run_timeout_seconds: Option<u64>,
196}
197
198impl ChildWorkflowOptions {
199 pub fn new(task_queue: impl Into<String>) -> Self {
200 Self {
201 task_queue: task_queue.into(),
202 parent_close_policy: ParentClosePolicy::Abandon,
203 retry_policy: None,
204 execution_timeout_seconds: None,
205 run_timeout_seconds: None,
206 }
207 }
208
209 pub fn parent_close_policy(mut self, policy: ParentClosePolicy) -> Self {
210 self.parent_close_policy = policy;
211 self
212 }
213
214 pub fn retry_policy(mut self, policy: ChildWorkflowRetryPolicy) -> Self {
215 self.retry_policy = Some(policy);
216 self
217 }
218
219 pub fn execution_timeout_seconds(mut self, seconds: u64) -> Self {
220 self.execution_timeout_seconds = Some(seconds);
221 self
222 }
223
224 pub fn run_timeout_seconds(mut self, seconds: u64) -> Self {
225 self.run_timeout_seconds = Some(seconds);
226 self
227 }
228}
229
230#[derive(Clone, Debug, Error)]
233#[error("non-deterministic workflow replay ({reason}) at sequence {sequence:?}: {message}")]
234pub struct ReplayFailure {
235 pub reason: String,
236 pub sequence: Option<u64>,
237 pub expected: Option<String>,
238 pub actual: Option<String>,
239 pub message: String,
240}
241
242impl ReplayFailure {
243 fn new(
244 reason: impl Into<String>,
245 sequence: Option<u64>,
246 expected: Option<String>,
247 actual: Option<String>,
248 message: impl Into<String>,
249 ) -> Self {
250 Self {
251 reason: reason.into(),
252 sequence,
253 expected,
254 actual,
255 message: message.into(),
256 }
257 }
258}
259
260#[derive(Clone, Debug, Error)]
262#[error("query failed ({reason}, HTTP {status}): {message}")]
263pub struct QueryFailure {
264 pub status: u16,
265 pub reason: String,
266 pub message: String,
267 pub body: Value,
268}
269
270#[derive(Clone, Debug, Error)]
272#[error("protocol rejected ({reason}, HTTP {status}): {message}")]
273pub struct ProtocolFailure {
274 pub status: u16,
275 pub reason: String,
276 pub message: String,
277 pub supported_version: Option<String>,
278 pub requested_version: Option<String>,
279 pub body: Value,
280}
281
282#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
283pub struct PayloadEnvelope {
284 pub codec: String,
285 pub blob: String,
286}
287
288impl PayloadEnvelope {
289 pub fn avro<T: Serialize>(value: &T) -> Result<Self> {
290 encode_payload(value, DEFAULT_CODEC)
291 }
292
293 pub fn json<T: Serialize>(value: &T) -> Result<Self> {
294 encode_payload(value, JSON_CODEC)
295 }
296}
297
298pub fn encode_payload<T: Serialize>(value: &T, codec: &str) -> Result<PayloadEnvelope> {
299 let value = serde_json::to_value(value)?;
300 let blob = encode_value_blob(&value, codec)?;
301
302 Ok(PayloadEnvelope {
303 codec: codec.to_string(),
304 blob,
305 })
306}
307
308pub fn decode_payload<T: DeserializeOwned>(envelope: &PayloadEnvelope) -> Result<T> {
309 let value = decode_blob(&envelope.blob, &envelope.codec)?;
310 Ok(serde_json::from_value(value)?)
311}
312
313fn encode_value_envelope(value: &Value, codec: &str) -> Result<Value> {
314 Ok(serde_json::to_value(encode_payload(value, codec)?)?)
315}
316
317fn encode_value_blob(value: &Value, codec: &str) -> Result<String> {
318 match codec {
319 JSON_CODEC => Ok(serde_json::to_string(value)?),
320 DEFAULT_CODEC => encode_avro_generic(value),
321 other => Err(Error::Codec(format!("unsupported payload codec {other:?}"))),
322 }
323}
324
325fn decode_wire_value(value: &Value, fallback_codec: &str) -> Result<Value> {
326 if value.is_null() {
327 return Ok(Value::Null);
328 }
329
330 if let Some(object) = value.as_object() {
331 if let (Some(codec), Some(blob)) = (
332 object.get("codec").and_then(Value::as_str),
333 object.get("blob").and_then(Value::as_str),
334 ) {
335 return decode_blob(blob, codec);
336 }
337 }
338
339 if let Some(blob) = value.as_str() {
340 return decode_blob(blob, fallback_codec);
341 }
342
343 Ok(value.clone())
344}
345
346fn decode_blob(blob: &str, codec: &str) -> Result<Value> {
347 match codec {
348 JSON_CODEC => Ok(serde_json::from_str(blob)?),
349 DEFAULT_CODEC => decode_avro_generic(blob),
350 other => Err(Error::Codec(format!("unsupported payload codec {other:?}"))),
351 }
352}
353
354fn encode_avro_generic(value: &Value) -> Result<String> {
355 let json = serde_json::to_string(value)?;
356 let datum = to_value(AvroPayload {
357 json,
358 version: AVRO_PAYLOAD_VERSION,
359 })
360 .map_err(|err| Error::Codec(format!("could not convert avro generic wrapper: {err}")))?;
361 let datum = to_avro_datum(avro_payload_schema()?, datum)
362 .map_err(|err| Error::Codec(format!("could not encode avro generic wrapper: {err}")))?;
363
364 let mut bytes = Vec::with_capacity(datum.len() + 1);
365 bytes.push(0x00);
366 bytes.extend_from_slice(&datum);
367 Ok(BASE64.encode(bytes))
368}
369
370fn decode_avro_generic(blob: &str) -> Result<Value> {
371 let bytes = BASE64
372 .decode(blob)
373 .map_err(|err| Error::Codec(format!("invalid avro base64 payload: {err}")))?;
374
375 if bytes.is_empty() {
376 return Err(Error::Codec("avro payload is empty".to_string()));
377 }
378
379 match bytes[0] {
380 0x00 => {}
381 0x01 => {
382 return Err(Error::Codec(
383 "typed avro payloads require a schema context; v1 supports the generic wrapper"
384 .to_string(),
385 ));
386 }
387 other => {
388 return Err(Error::Codec(format!(
389 "unknown avro payload prefix 0x{other:02x}"
390 )));
391 }
392 }
393
394 let mut datum = &bytes[1..];
395 let datum = from_avro_datum(avro_payload_schema()?, &mut datum, None)
396 .map_err(|err| Error::Codec(format!("could not decode avro generic wrapper: {err}")))?;
397 let payload: AvroPayload = from_value(&datum)
398 .map_err(|err| Error::Codec(format!("invalid avro generic wrapper record: {err}")))?;
399
400 if payload.version != AVRO_PAYLOAD_VERSION {
401 return Err(Error::Codec(format!(
402 "unsupported avro generic wrapper version {}",
403 payload.version
404 )));
405 }
406
407 Ok(serde_json::from_str(&payload.json)?)
408}
409
410#[derive(Debug, Serialize, Deserialize)]
411struct AvroPayload {
412 json: String,
413 version: i32,
414}
415
416fn avro_payload_schema() -> Result<&'static Schema> {
417 match AVRO_PAYLOAD_SCHEMA.get_or_init(|| {
418 Schema::parse_str(AVRO_PAYLOAD_SCHEMA_JSON)
419 .map_err(|err| format!("could not parse avro payload schema: {err}"))
420 }) {
421 Ok(schema) => Ok(schema),
422 Err(message) => Err(Error::Codec(message.clone())),
423 }
424}
425
426#[derive(Clone, Debug)]
427pub struct Client {
428 http: reqwest::Client,
429 base_url: String,
430 token: Option<String>,
431 control_token: Option<String>,
432 worker_token: Option<String>,
433 namespace: String,
434}
435
436impl Client {
437 pub fn new(base_url: impl Into<String>) -> Result<Self> {
438 Self::builder(base_url).build()
439 }
440
441 pub fn builder(base_url: impl Into<String>) -> ClientBuilder {
442 ClientBuilder {
443 base_url: base_url.into(),
444 token: None,
445 control_token: None,
446 worker_token: None,
447 namespace: "default".to_string(),
448 timeout: Duration::from_secs(60),
449 }
450 }
451
452 pub async fn health(&self) -> Result<Value> {
453 self.request_json(
454 reqwest::Method::GET,
455 "/health",
456 RequestProtocol::ControlPlane,
457 Option::<&Value>::None,
458 )
459 .await
460 }
461
462 pub async fn cluster_info(&self) -> Result<Value> {
463 self.request_json(
464 reqwest::Method::GET,
465 "/cluster/info",
466 RequestProtocol::ControlPlane,
467 Option::<&Value>::None,
468 )
469 .await
470 }
471
472 pub async fn start_workflow<T: Serialize>(
473 &self,
474 workflow_type: &str,
475 task_queue: &str,
476 workflow_id: &str,
477 input: T,
478 ) -> Result<WorkflowHandle> {
479 let input = serde_json::to_value(input)?;
480 let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
481 let body = json!({
482 "workflow_id": workflow_id,
483 "workflow_type": workflow_type,
484 "task_queue": task_queue,
485 "input": input_envelope,
486 "execution_timeout_seconds": 3600,
487 "run_timeout_seconds": 600
488 });
489
490 let data: Value = self
491 .request_json(
492 reqwest::Method::POST,
493 "/workflows",
494 RequestProtocol::ControlPlane,
495 Some(&body),
496 )
497 .await?;
498
499 Ok(WorkflowHandle {
500 client: self.clone(),
501 workflow_id: data
502 .get("workflow_id")
503 .and_then(Value::as_str)
504 .unwrap_or(workflow_id)
505 .to_string(),
506 run_id: data
507 .get("run_id")
508 .and_then(Value::as_str)
509 .map(str::to_string),
510 workflow_type: data
511 .get("workflow_type")
512 .and_then(Value::as_str)
513 .unwrap_or(workflow_type)
514 .to_string(),
515 })
516 }
517
518 pub async fn signal_workflow<T: Serialize>(
519 &self,
520 workflow_id: &str,
521 signal_name: &str,
522 input: T,
523 ) -> Result<Value> {
524 let input = serde_json::to_value(input)?;
525 let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
526 let body = json!({
527 "input": input_envelope
528 });
529 let path = format!("/workflows/{workflow_id}/signal/{signal_name}");
530 self.request_json(
531 reqwest::Method::POST,
532 &path,
533 RequestProtocol::ControlPlane,
534 Some(&body),
535 )
536 .await
537 }
538
539 pub async fn query_workflow<T: Serialize>(
545 &self,
546 workflow_id: &str,
547 query_name: &str,
548 input: T,
549 ) -> Result<Value> {
550 let input = serde_json::to_value(input)?;
551 let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
552 let body = json!({
553 "input": input_envelope
554 });
555 let path = format!("/workflows/{workflow_id}/query/{query_name}");
556 let response: Value = match self
557 .request_json(
558 reqwest::Method::POST,
559 &path,
560 RequestProtocol::ControlPlane,
561 Some(&body),
562 )
563 .await
564 {
565 Ok(response) => response,
566 Err(Error::Http { status, body }) => {
567 return Err(Error::QueryFailed(query_failure(status, body)));
568 }
569 Err(error) => return Err(error),
570 };
571
572 if let Some(envelope) = response
573 .get("result_envelope")
574 .filter(|envelope| !envelope.is_null())
575 {
576 return decode_wire_value(envelope, DEFAULT_CODEC);
577 }
578
579 Ok(response.get("result").cloned().unwrap_or(Value::Null))
580 }
581
582 pub async fn describe_workflow(&self, workflow_id: &str) -> Result<WorkflowDescription> {
583 let path = format!("/workflows/{workflow_id}");
584 let mut data: WorkflowDescription = self
585 .request_json(
586 reqwest::Method::GET,
587 &path,
588 RequestProtocol::ControlPlane,
589 Option::<&Value>::None,
590 )
591 .await?;
592 data.decode_payloads()?;
593 Ok(data)
594 }
595
596 pub async fn register_worker(
597 &self,
598 worker_id: &str,
599 task_queue: &str,
600 supported_workflow_types: Vec<String>,
601 supported_activity_types: Vec<String>,
602 max_concurrent_workflow_tasks: usize,
603 max_concurrent_activity_tasks: usize,
604 ) -> Result<RegisterWorkerResponse> {
605 self.register_worker_with_capabilities(
606 worker_id,
607 task_queue,
608 supported_workflow_types,
609 supported_activity_types,
610 max_concurrent_workflow_tasks,
611 max_concurrent_activity_tasks,
612 Vec::new(),
613 )
614 .await
615 }
616
617 pub async fn register_worker_with_capabilities(
619 &self,
620 worker_id: &str,
621 task_queue: &str,
622 supported_workflow_types: Vec<String>,
623 supported_activity_types: Vec<String>,
624 max_concurrent_workflow_tasks: usize,
625 max_concurrent_activity_tasks: usize,
626 capabilities: Vec<String>,
627 ) -> Result<RegisterWorkerResponse> {
628 let body = json!({
629 "worker_id": worker_id,
630 "task_queue": task_queue,
631 "runtime": "rust",
632 "sdk_version": SDK_VERSION,
633 "supported_workflow_types": supported_workflow_types,
634 "supported_activity_types": supported_activity_types,
635 "capabilities": capabilities,
636 "max_concurrent_workflow_tasks": max_concurrent_workflow_tasks,
637 "max_concurrent_activity_tasks": max_concurrent_activity_tasks
638 });
639
640 self.request_json(
641 reqwest::Method::POST,
642 "/worker/register",
643 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
644 Some(&body),
645 )
646 .await
647 }
648
649 pub async fn poll_query_task(
651 &self,
652 worker_id: &str,
653 task_queue: &str,
654 timeout: Duration,
655 ) -> Result<Option<QueryTask>> {
656 let timeout_seconds = long_poll_timeout_seconds(timeout);
657 let body = json!({
658 "worker_id": worker_id,
659 "task_queue": task_queue,
660 "poll_request_id": unique_request_id("rust-query-poll"),
661 "timeout_seconds": timeout_seconds,
662 });
663 let data: PollQueryTaskResponse = self
664 .request_json_with_timeout(
665 reqwest::Method::POST,
666 "/worker/query-tasks/poll",
667 RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
668 Some(&body),
669 timeout + Duration::from_secs(5),
670 )
671 .await?;
672 Ok(data.task)
673 }
674
675 pub async fn complete_query_task(
677 &self,
678 query_task_id: &str,
679 lease_owner: &str,
680 query_task_attempt: u64,
681 result: Value,
682 codec: &str,
683 ) -> Result<Value> {
684 let result_envelope = encode_value_envelope(&result, codec)?;
685 self.complete_query_task_with_envelope(
686 query_task_id,
687 lease_owner,
688 query_task_attempt,
689 result,
690 result_envelope,
691 )
692 .await
693 }
694
695 async fn complete_query_task_with_envelope(
696 &self,
697 query_task_id: &str,
698 lease_owner: &str,
699 query_task_attempt: u64,
700 result: Value,
701 result_envelope: Value,
702 ) -> Result<Value> {
703 let body = json!({
704 "lease_owner": lease_owner,
705 "query_task_attempt": query_task_attempt,
706 "result": result,
707 "result_envelope": result_envelope,
708 });
709 let path = format!("/worker/query-tasks/{query_task_id}/complete");
710 let response = self
711 .request_json(
712 reqwest::Method::POST,
713 &path,
714 RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
715 Some(&body),
716 )
717 .await;
718 query_task_response(response)
719 }
720
721 pub async fn fail_query_task(
723 &self,
724 query_task_id: &str,
725 lease_owner: &str,
726 query_task_attempt: u64,
727 message: impl Into<String>,
728 reason: impl Into<String>,
729 failure_type: impl Into<String>,
730 ) -> Result<Value> {
731 let body = json!({
732 "lease_owner": lease_owner,
733 "query_task_attempt": query_task_attempt,
734 "failure": {
735 "message": message.into(),
736 "reason": reason.into(),
737 "type": failure_type.into(),
738 }
739 });
740 let path = format!("/worker/query-tasks/{query_task_id}/fail");
741 let response = self
742 .request_json(
743 reqwest::Method::POST,
744 &path,
745 RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
746 Some(&body),
747 )
748 .await;
749 query_task_response(response)
750 }
751
752 pub async fn heartbeat_worker(
753 &self,
754 worker_id: &str,
755 workflow_available: usize,
756 activity_available: usize,
757 ) -> Result<Value> {
758 let body = json!({
759 "worker_id": worker_id,
760 "task_slots": {
761 "workflow_available": workflow_available,
762 "activity_available": activity_available
763 },
764 "process_metrics": {
765 "process_id": std::process::id(),
766 "process_uptime_seconds": 0
767 }
768 });
769
770 self.request_json(
771 reqwest::Method::POST,
772 "/worker/heartbeat",
773 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
774 Some(&body),
775 )
776 .await
777 }
778
779 pub async fn poll_workflow_task(
780 &self,
781 worker_id: &str,
782 task_queue: &str,
783 timeout: Duration,
784 ) -> Result<Option<WorkflowTask>> {
785 Ok(self
786 .poll_workflow_task_response(worker_id, task_queue, timeout)
787 .await?
788 .task)
789 }
790
791 pub async fn poll_workflow_task_response(
792 &self,
793 worker_id: &str,
794 task_queue: &str,
795 timeout: Duration,
796 ) -> Result<PollWorkflowTaskResponse> {
797 let body = json!({
798 "worker_id": worker_id,
799 "task_queue": task_queue,
800 "timeout_seconds": long_poll_timeout_seconds(timeout),
801 });
802 let mut data: PollWorkflowTaskResponse = self
803 .request_json_with_timeout(
804 reqwest::Method::POST,
805 "/worker/workflow-tasks/poll",
806 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
807 Some(&body),
808 timeout + Duration::from_secs(5),
809 )
810 .await?;
811
812 if let Some(task) = data.task.as_mut() {
813 self.fetch_remaining_workflow_history(worker_id, task)
814 .await?;
815 }
816
817 Ok(data)
818 }
819
820 async fn fetch_remaining_workflow_history(
821 &self,
822 worker_id: &str,
823 task: &mut WorkflowTask,
824 ) -> Result<()> {
825 let mut next_token = task.next_history_page_token.clone();
826
827 while let Some(token) = next_token.take().filter(|token| !token.is_empty()) {
828 let lease_owner = task
829 .lease_owner
830 .clone()
831 .unwrap_or_else(|| worker_id.to_string());
832 let page = self
833 .workflow_task_history_page(
834 &task.task_id,
835 &lease_owner,
836 task.workflow_task_attempt,
837 &token,
838 )
839 .await?;
840
841 task.append_history_page(page);
842
843 if task.next_history_page_token.as_deref() == Some(token.as_str()) {
844 return Err(Error::Codec(
845 "workflow history pagination returned the same page token".to_string(),
846 ));
847 }
848
849 next_token = task.next_history_page_token.clone();
850 }
851
852 Ok(())
853 }
854
855 async fn workflow_task_history_page(
856 &self,
857 task_id: &str,
858 lease_owner: &str,
859 workflow_task_attempt: u64,
860 next_history_page_token: &str,
861 ) -> Result<WorkflowTaskHistoryPage> {
862 let body = json!({
863 "lease_owner": lease_owner,
864 "workflow_task_attempt": workflow_task_attempt,
865 "next_history_page_token": next_history_page_token
866 });
867 let path = format!("/worker/workflow-tasks/{task_id}/history");
868
869 self.request_json(
870 reqwest::Method::POST,
871 &path,
872 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
873 Some(&body),
874 )
875 .await
876 }
877
878 pub async fn complete_workflow_task(
879 &self,
880 task_id: &str,
881 lease_owner: &str,
882 workflow_task_attempt: u64,
883 commands: Vec<Value>,
884 ) -> Result<Value> {
885 let body = json!({
886 "lease_owner": lease_owner,
887 "workflow_task_attempt": workflow_task_attempt,
888 "commands": commands
889 });
890 let path = format!("/worker/workflow-tasks/{task_id}/complete");
891 self.request_json(
892 reqwest::Method::POST,
893 &path,
894 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
895 Some(&body),
896 )
897 .await
898 }
899
900 pub async fn fail_workflow_task(
901 &self,
902 task_id: &str,
903 lease_owner: &str,
904 workflow_task_attempt: u64,
905 message: impl Into<String>,
906 ) -> Result<Value> {
907 self.fail_workflow_task_with_type(
908 task_id,
909 lease_owner,
910 workflow_task_attempt,
911 message,
912 "RustWorkflowTaskFailure",
913 )
914 .await
915 }
916
917 async fn fail_workflow_task_with_type(
918 &self,
919 task_id: &str,
920 lease_owner: &str,
921 workflow_task_attempt: u64,
922 message: impl Into<String>,
923 failure_type: &str,
924 ) -> Result<Value> {
925 let body = json!({
926 "lease_owner": lease_owner,
927 "workflow_task_attempt": workflow_task_attempt,
928 "failure": {
929 "message": message.into(),
930 "type": failure_type
931 }
932 });
933 let path = format!("/worker/workflow-tasks/{task_id}/fail");
934 self.request_json(
935 reqwest::Method::POST,
936 &path,
937 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
938 Some(&body),
939 )
940 .await
941 }
942
943 pub async fn poll_activity_task(
944 &self,
945 worker_id: &str,
946 task_queue: &str,
947 timeout: Duration,
948 ) -> Result<Option<ActivityTask>> {
949 let body = json!({
950 "worker_id": worker_id,
951 "task_queue": task_queue,
952 "timeout_seconds": long_poll_timeout_seconds(timeout),
953 });
954 let data: PollActivityTaskResponse = self
955 .request_json_with_timeout(
956 reqwest::Method::POST,
957 "/worker/activity-tasks/poll",
958 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
959 Some(&body),
960 timeout + Duration::from_secs(5),
961 )
962 .await?;
963 Ok(data.task)
964 }
965
966 pub async fn complete_activity_task(
967 &self,
968 task_id: &str,
969 activity_attempt_id: &str,
970 lease_owner: &str,
971 result: Value,
972 codec: &str,
973 ) -> Result<Value> {
974 let result = encode_value_envelope(&result, codec)?;
975 let body = json!({
976 "activity_attempt_id": activity_attempt_id,
977 "lease_owner": lease_owner,
978 "result": result
979 });
980 let path = format!("/worker/activity-tasks/{task_id}/complete");
981 self.request_json(
982 reqwest::Method::POST,
983 &path,
984 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
985 Some(&body),
986 )
987 .await
988 }
989
990 pub async fn fail_activity_task(
991 &self,
992 task_id: &str,
993 activity_attempt_id: &str,
994 lease_owner: &str,
995 message: impl Into<String>,
996 non_retryable: bool,
997 ) -> Result<Value> {
998 let body = json!({
999 "activity_attempt_id": activity_attempt_id,
1000 "lease_owner": lease_owner,
1001 "failure": {
1002 "message": message.into(),
1003 "type": "RustActivityFailure",
1004 "non_retryable": non_retryable
1005 }
1006 });
1007 let path = format!("/worker/activity-tasks/{task_id}/fail");
1008 self.request_json(
1009 reqwest::Method::POST,
1010 &path,
1011 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
1012 Some(&body),
1013 )
1014 .await
1015 }
1016
1017 pub async fn heartbeat_activity_task(
1018 &self,
1019 task_id: &str,
1020 activity_attempt_id: &str,
1021 lease_owner: &str,
1022 details: Value,
1023 ) -> Result<ActivityHeartbeatResponse> {
1024 let body = json!({
1025 "activity_attempt_id": activity_attempt_id,
1026 "lease_owner": lease_owner,
1027 "details": details
1028 });
1029 let path = format!("/worker/activity-tasks/{task_id}/heartbeat");
1030 self.request_json(
1031 reqwest::Method::POST,
1032 &path,
1033 RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
1034 Some(&body),
1035 )
1036 .await
1037 }
1038
1039 async fn request_json<T: DeserializeOwned, B: Serialize + ?Sized>(
1040 &self,
1041 method: reqwest::Method,
1042 path: &str,
1043 protocol: RequestProtocol,
1044 body: Option<&B>,
1045 ) -> Result<T> {
1046 self.request_json_with_timeout(method, path, protocol, body, Duration::from_secs(60))
1047 .await
1048 }
1049
1050 async fn request_json_with_timeout<T: DeserializeOwned, B: Serialize + ?Sized>(
1051 &self,
1052 method: reqwest::Method,
1053 path: &str,
1054 protocol: RequestProtocol,
1055 body: Option<&B>,
1056 timeout: Duration,
1057 ) -> Result<T> {
1058 let mut request = self
1059 .http
1060 .request(method, format!("{}/api{}", self.base_url, path))
1061 .timeout(timeout)
1062 .header(reqwest::header::ACCEPT, "application/json")
1063 .header(reqwest::header::CONTENT_TYPE, "application/json")
1064 .header("X-Namespace", &self.namespace);
1065
1066 match protocol {
1067 RequestProtocol::Worker(version) => {
1068 request = request.header("X-Durable-Workflow-Protocol-Version", version);
1069 }
1070 RequestProtocol::ControlPlane => {
1071 request = request.header(
1072 "X-Durable-Workflow-Control-Plane-Version",
1073 CONTROL_PLANE_VERSION,
1074 );
1075 }
1076 }
1077
1078 if let Some(token) = self.auth_token(protocol.is_worker()) {
1079 request = request.bearer_auth(token);
1080 }
1081
1082 if let Some(body) = body {
1083 request = request.json(body);
1084 }
1085
1086 let response = request.send().await?;
1087 let status = response.status();
1088 let bytes = response.bytes().await?;
1089
1090 if !status.is_success() {
1091 let body = String::from_utf8_lossy(&bytes).to_string();
1092 if let Some(protocol) = protocol_failure(status, &body) {
1093 return Err(Error::Protocol(protocol));
1094 }
1095 return Err(Error::Http { status, body });
1096 }
1097
1098 if bytes.is_empty() {
1099 return Ok(serde_json::from_value(Value::Null)?);
1100 }
1101
1102 Ok(serde_json::from_slice(&bytes)?)
1103 }
1104
1105 fn auth_token(&self, worker: bool) -> Option<&str> {
1106 if worker {
1107 self.worker_token
1108 .as_deref()
1109 .or(self.token.as_deref())
1110 .or(self.control_token.as_deref())
1111 } else {
1112 self.control_token
1113 .as_deref()
1114 .or(self.token.as_deref())
1115 .or(self.worker_token.as_deref())
1116 }
1117 }
1118}
1119
1120fn query_failure(status: reqwest::StatusCode, raw_body: String) -> QueryFailure {
1121 let body = serde_json::from_str(&raw_body).unwrap_or_else(|_| json!({"message": raw_body}));
1122 let reason = body
1123 .get("reason")
1124 .and_then(Value::as_str)
1125 .unwrap_or("query_rejected")
1126 .to_string();
1127 let message = body
1128 .get("message")
1129 .or_else(|| body.get("error"))
1130 .and_then(Value::as_str)
1131 .unwrap_or("workflow query was rejected")
1132 .to_string();
1133
1134 QueryFailure {
1135 status: status.as_u16(),
1136 reason,
1137 message,
1138 body,
1139 }
1140}
1141
1142fn query_task_response(response: Result<Value>) -> Result<Value> {
1143 match response {
1144 Err(Error::Http { status, body }) => Err(Error::QueryFailed(query_failure(status, body))),
1145 response => response,
1146 }
1147}
1148
1149fn query_task_rejection_is_final(error: &Error) -> bool {
1150 matches!(
1151 error,
1152 Error::QueryFailed(failure)
1153 if QUERY_TASK_FINAL_REJECTION_REASONS.contains(&failure.reason.as_str())
1154 )
1155}
1156
1157fn protocol_failure(status: reqwest::StatusCode, raw_body: &str) -> Option<ProtocolFailure> {
1158 let body: Value = serde_json::from_str(raw_body).ok()?;
1159 let reason = body.get("reason")?.as_str()?;
1160 if !matches!(
1161 reason,
1162 "missing_protocol_version"
1163 | "unsupported_protocol_version"
1164 | "missing_control_plane_version"
1165 | "unsupported_control_plane_version"
1166 ) {
1167 return None;
1168 }
1169
1170 Some(ProtocolFailure {
1171 status: status.as_u16(),
1172 reason: reason.to_string(),
1173 message: body
1174 .get("message")
1175 .or_else(|| body.get("error"))
1176 .and_then(Value::as_str)
1177 .unwrap_or("protocol version rejected")
1178 .to_string(),
1179 supported_version: body
1180 .get("supported_version")
1181 .and_then(Value::as_str)
1182 .map(str::to_string),
1183 requested_version: body
1184 .get("requested_version")
1185 .and_then(Value::as_str)
1186 .map(str::to_string),
1187 body,
1188 })
1189}
1190
1191fn long_poll_timeout_seconds(timeout: Duration) -> u64 {
1192 timeout
1193 .as_secs()
1194 .saturating_add(u64::from(timeout.subsec_nanos() > 0))
1195 .min(MAX_LONG_POLL_TIMEOUT_SECONDS)
1196}
1197
1198fn worker_operation_is_retryable(error: &Error) -> bool {
1199 match error {
1200 Error::Transport(error) => {
1201 error.is_timeout() || error.is_connect() || error.is_request() || error.is_body()
1202 }
1203 Error::Http { status, .. } => {
1204 matches!(
1205 *status,
1206 reqwest::StatusCode::REQUEST_TIMEOUT | reqwest::StatusCode::TOO_MANY_REQUESTS
1207 ) || status.is_server_error()
1208 }
1209 _ => false,
1210 }
1211}
1212
1213fn worker_retry_delay(policy: WorkerRetryPolicy, retry: usize) -> Duration {
1214 let exponent = retry.saturating_sub(1).min(31) as u32;
1215 policy
1216 .initial_backoff
1217 .saturating_mul(1_u32 << exponent)
1218 .min(policy.max_backoff)
1219}
1220
1221#[derive(Debug)]
1222pub struct ClientBuilder {
1223 base_url: String,
1224 token: Option<String>,
1225 control_token: Option<String>,
1226 worker_token: Option<String>,
1227 namespace: String,
1228 timeout: Duration,
1229}
1230
1231impl ClientBuilder {
1232 pub fn token(mut self, token: Option<String>) -> Self {
1233 self.token = token;
1234 self
1235 }
1236
1237 pub fn control_token(mut self, token: Option<String>) -> Self {
1238 self.control_token = token;
1239 self
1240 }
1241
1242 pub fn worker_token(mut self, token: Option<String>) -> Self {
1243 self.worker_token = token;
1244 self
1245 }
1246
1247 pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
1248 self.namespace = namespace.into();
1249 self
1250 }
1251
1252 pub fn timeout(mut self, timeout: Duration) -> Self {
1253 self.timeout = timeout;
1254 self
1255 }
1256
1257 pub fn build(self) -> Result<Client> {
1258 Ok(Client {
1259 http: reqwest::Client::builder().timeout(self.timeout).build()?,
1260 base_url: self.base_url.trim_end_matches('/').to_string(),
1261 token: self.token,
1262 control_token: self.control_token,
1263 worker_token: self.worker_token,
1264 namespace: self.namespace,
1265 })
1266 }
1267}
1268
1269#[derive(Clone, Debug)]
1270pub struct WorkflowHandle {
1271 client: Client,
1272 pub workflow_id: String,
1273 pub run_id: Option<String>,
1274 pub workflow_type: String,
1275}
1276
1277impl WorkflowHandle {
1278 pub async fn describe(&self) -> Result<WorkflowDescription> {
1279 self.client.describe_workflow(&self.workflow_id).await
1280 }
1281
1282 pub async fn signal<T: Serialize>(&self, signal_name: &str, input: T) -> Result<Value> {
1283 self.client
1284 .signal_workflow(&self.workflow_id, signal_name, input)
1285 .await
1286 }
1287
1288 pub async fn query<T: Serialize>(&self, query_name: &str, input: T) -> Result<Value> {
1290 self.client
1291 .query_workflow(&self.workflow_id, query_name, input)
1292 .await
1293 }
1294
1295 pub async fn result(&self, options: WorkflowResultOptions) -> Result<Value> {
1296 let started = Instant::now();
1297
1298 loop {
1299 let description = self.describe().await?;
1300 if description.is_completed() {
1301 return Ok(description.output.unwrap_or(Value::Null));
1302 }
1303
1304 if description.is_terminal() {
1305 return Err(Error::Codec(format!(
1306 "workflow {} closed with status {:?}",
1307 self.workflow_id, description.status
1308 )));
1309 }
1310
1311 if started.elapsed() >= options.timeout {
1312 return Err(Error::Timeout);
1313 }
1314
1315 tokio::time::sleep(options.poll_interval).await;
1316 }
1317 }
1318}
1319
1320#[derive(Clone, Copy, Debug)]
1321pub struct WorkflowResultOptions {
1322 pub poll_interval: Duration,
1323 pub timeout: Duration,
1324}
1325
1326impl Default for WorkflowResultOptions {
1327 fn default() -> Self {
1328 Self {
1329 poll_interval: Duration::from_millis(500),
1330 timeout: Duration::from_secs(30),
1331 }
1332 }
1333}
1334
1335#[derive(Clone, Debug, Deserialize)]
1336pub struct WorkflowDescription {
1337 pub workflow_id: Option<String>,
1338 pub run_id: Option<String>,
1339 pub workflow_type: Option<String>,
1340 pub status: Option<String>,
1341 #[serde(default)]
1342 pub output: Option<Value>,
1343 #[serde(default)]
1344 pub output_envelope: Option<Value>,
1345 #[serde(flatten)]
1346 pub raw: HashMap<String, Value>,
1347}
1348
1349impl WorkflowDescription {
1350 pub fn is_completed(&self) -> bool {
1351 matches!(self.status.as_deref(), Some("completed" | "Completed"))
1352 }
1353
1354 pub fn is_terminal(&self) -> bool {
1355 matches!(
1356 self.status.as_deref(),
1357 Some(
1358 "completed"
1359 | "Completed"
1360 | "failed"
1361 | "Failed"
1362 | "cancelled"
1363 | "Cancelled"
1364 | "terminated"
1365 | "Terminated"
1366 | "timed_out"
1367 | "TimedOut",
1368 )
1369 )
1370 }
1371
1372 fn decode_payloads(&mut self) -> Result<()> {
1373 if let Some(envelope) = &self.output_envelope {
1374 self.output = Some(decode_wire_value(envelope, DEFAULT_CODEC)?);
1375 }
1376
1377 Ok(())
1378 }
1379}
1380
1381#[derive(Clone, Debug, Deserialize)]
1382pub struct RegisterWorkerResponse {
1383 pub worker_id: String,
1384 pub registered: bool,
1385 #[serde(default)]
1386 pub heartbeat_interval_seconds: Option<u64>,
1387 #[serde(default)]
1388 pub protocol_version: Option<String>,
1389 #[serde(default)]
1390 pub server_capabilities: Option<Value>,
1391}
1392
1393#[derive(Clone, Debug, Deserialize)]
1394pub struct PollWorkflowTaskResponse {
1395 #[serde(default)]
1396 pub task: Option<WorkflowTask>,
1397 #[serde(default)]
1398 pub poll_status: Option<String>,
1399 #[serde(default)]
1400 pub reason: Option<String>,
1401 #[serde(default)]
1402 pub protocol_version: Option<String>,
1403 #[serde(default)]
1404 pub server_capabilities: Option<Value>,
1405}
1406
1407#[derive(Clone, Debug, Deserialize)]
1408struct PollActivityTaskResponse {
1409 #[serde(default)]
1410 task: Option<ActivityTask>,
1411}
1412
1413#[derive(Clone, Debug, Deserialize)]
1414struct PollQueryTaskResponse {
1415 #[serde(default)]
1416 task: Option<QueryTask>,
1417}
1418
1419#[derive(Clone, Debug, Deserialize)]
1421pub struct QueryTask {
1422 pub query_task_id: String,
1423 #[serde(default = "default_workflow_task_attempt")]
1424 pub query_task_attempt: u64,
1425 #[serde(default)]
1426 pub lease_owner: Option<String>,
1427 #[serde(default)]
1428 pub workflow_id: Option<String>,
1429 #[serde(default)]
1430 pub run_id: Option<String>,
1431 pub workflow_type: String,
1432 pub query_name: String,
1433 #[serde(default = "default_payload_codec")]
1434 pub payload_codec: String,
1435 #[serde(default)]
1436 pub workflow_arguments: Option<Value>,
1437 #[serde(default)]
1438 pub query_arguments: Option<Value>,
1439 #[serde(default)]
1440 pub history_events: Vec<HistoryEvent>,
1441 #[serde(default)]
1442 pub history_export: Option<Value>,
1443 #[serde(default)]
1444 pub run_status: Option<String>,
1445}
1446
1447#[derive(Clone, Debug, Deserialize)]
1448pub struct WorkflowTask {
1449 pub task_id: String,
1450 #[serde(default)]
1451 pub workflow_id: Option<String>,
1452 #[serde(default)]
1453 pub run_id: Option<String>,
1454 pub workflow_type: String,
1455 #[serde(default = "default_payload_codec")]
1456 pub payload_codec: String,
1457 #[serde(default)]
1458 pub arguments: Option<Value>,
1459 #[serde(default)]
1460 pub history_events: Vec<HistoryEvent>,
1461 #[serde(default)]
1462 pub total_history_events: Option<u64>,
1463 #[serde(default)]
1464 pub next_history_page_token: Option<String>,
1465 #[serde(default = "default_workflow_task_attempt")]
1466 pub workflow_task_attempt: u64,
1467 #[serde(default)]
1468 pub workflow_signal_id: Option<String>,
1469 #[serde(default)]
1470 pub signal_name: Option<String>,
1471 #[serde(default)]
1472 pub signal_arguments: Option<Value>,
1473 #[serde(default)]
1474 pub lease_owner: Option<String>,
1475}
1476
1477impl WorkflowTask {
1478 fn append_history_page(&mut self, page: WorkflowTaskHistoryPage) {
1479 self.history_events.extend(page.history_events);
1480
1481 if page.total_history_events.is_some() {
1482 self.total_history_events = page.total_history_events;
1483 }
1484
1485 self.next_history_page_token = page
1486 .next_history_page_token
1487 .filter(|token| !token.is_empty());
1488 }
1489}
1490
1491#[derive(Clone, Debug, Deserialize)]
1492struct WorkflowTaskHistoryPage {
1493 #[serde(default)]
1494 history_events: Vec<HistoryEvent>,
1495 #[serde(default)]
1496 total_history_events: Option<u64>,
1497 #[serde(default)]
1498 next_history_page_token: Option<String>,
1499}
1500
1501#[derive(Clone, Debug, Deserialize)]
1502pub struct ActivityTask {
1503 pub task_id: String,
1504 #[serde(default)]
1505 pub activity_attempt_id: Option<String>,
1506 #[serde(default)]
1507 pub attempt_id: Option<String>,
1508 pub activity_type: String,
1509 #[serde(default = "default_payload_codec")]
1510 pub payload_codec: String,
1511 #[serde(default)]
1512 pub arguments: Option<Value>,
1513 #[serde(default = "default_attempt_number")]
1514 pub attempt_number: u64,
1515 #[serde(default)]
1516 pub lease_owner: Option<String>,
1517}
1518
1519#[derive(Clone, Debug, Deserialize)]
1520pub struct HistoryEvent {
1521 #[serde(alias = "type")]
1522 pub event_type: String,
1523 #[serde(default)]
1524 pub payload: Value,
1525 #[serde(flatten)]
1526 pub raw: HashMap<String, Value>,
1527}
1528
1529#[derive(Clone, Debug, PartialEq)]
1531pub struct QuerySignal {
1532 pub id: Option<String>,
1533 pub name: String,
1534 pub arguments: Vec<Value>,
1535 pub workflow_sequence: Option<u64>,
1536}
1537
1538#[derive(Clone, Debug)]
1544pub struct QueryContext {
1545 pub workflow_id: Option<String>,
1546 pub run_id: Option<String>,
1547 pub workflow_type: String,
1548 pub run_status: Option<String>,
1549 workflow_input: Value,
1550 history_events: Arc<Vec<HistoryEvent>>,
1551 signal_events: Arc<Vec<QuerySignal>>,
1552}
1553
1554impl QueryContext {
1555 pub fn workflow_input(&self) -> &Value {
1557 &self.workflow_input
1558 }
1559
1560 pub fn history_events(&self) -> &[HistoryEvent] {
1562 self.history_events.as_slice()
1563 }
1564
1565 pub fn signal_events(&self) -> &[QuerySignal] {
1567 self.signal_events.as_slice()
1568 }
1569
1570 pub fn signals(&self, signal_name: &str) -> Vec<Vec<Value>> {
1572 self.signal_events
1573 .iter()
1574 .filter(|signal| signal.name == signal_name)
1575 .map(|signal| signal.arguments.clone())
1576 .collect()
1577 }
1578}
1579
1580#[derive(Clone, Debug, Deserialize)]
1581pub struct ActivityHeartbeatResponse {
1582 #[serde(default)]
1583 pub cancel_requested: bool,
1584 #[serde(default)]
1585 pub heartbeat_recorded: bool,
1586}
1587
1588fn default_payload_codec() -> String {
1589 DEFAULT_CODEC.to_string()
1590}
1591
1592fn default_workflow_task_attempt() -> u64 {
1593 1
1594}
1595
1596fn default_attempt_number() -> u64 {
1597 1
1598}
1599
1600type WorkflowFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1601type WorkflowHandler = Arc<dyn Fn(WorkflowContext, Value) -> WorkflowFuture + Send + Sync>;
1602type ErasedWorkflowState = Arc<dyn Any + Send + Sync>;
1603type WorkflowStateSnapshot = Arc<dyn Fn() -> Result<ErasedWorkflowState> + Send + Sync>;
1604type ReplayedWorkflowHandler =
1605 Arc<dyn Fn(WorkflowContext, Value) -> ReplayedWorkflowInvocation + Send + Sync>;
1606type ActivityFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1607type ActivityHandler = Arc<dyn Fn(ActivityContext, Value) -> ActivityFuture + Send + Sync>;
1608type QueryFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1609type QueryHandler = Arc<dyn Fn(QueryContext, Value) -> QueryFuture + Send + Sync>;
1610type ReplayedQueryHandler = Arc<
1611 dyn Fn(QueryContext, ErasedWorkflowState, Value) -> std::result::Result<QueryFuture, String>
1612 + Send
1613 + Sync,
1614>;
1615type WorkerHeartbeatObserver = Arc<dyn Fn(&WorkerHeartbeatObservation) + Send + Sync>;
1616
1617struct ReplayedWorkflowInvocation {
1618 future: WorkflowFuture,
1619 snapshot: WorkflowStateSnapshot,
1620}
1621
1622#[derive(Clone)]
1623struct RegisteredWorkflow {
1624 execute: WorkflowHandler,
1625 replay: Option<ReplayedWorkflowHandler>,
1626 state_type: Option<TypeId>,
1627}
1628
1629#[derive(Clone)]
1630enum RegisteredQuery {
1631 Snapshot(QueryHandler),
1632 Replayed {
1633 state_type: TypeId,
1634 handler: ReplayedQueryHandler,
1635 },
1636}
1637
1638#[derive(Clone, Debug)]
1639pub struct WorkerHeartbeatObservation {
1640 pub worker_id: String,
1641 pub task_queue: String,
1642 pub acknowledged_at_unix_millis: u64,
1643 pub acknowledgement: Value,
1644}
1645
1646#[derive(Clone, Copy, Debug)]
1653pub struct WorkerRetryPolicy {
1654 pub max_retries: usize,
1656 pub initial_backoff: Duration,
1658 pub max_backoff: Duration,
1660}
1661
1662impl Default for WorkerRetryPolicy {
1663 fn default() -> Self {
1664 Self {
1665 max_retries: 5,
1666 initial_backoff: Duration::from_millis(100),
1667 max_backoff: Duration::from_secs(5),
1668 }
1669 }
1670}
1671
1672#[derive(Clone)]
1673pub struct Worker {
1674 client: Client,
1675 worker_id: String,
1676 task_queue: String,
1677 workflows: HashMap<String, RegisteredWorkflow>,
1678 activities: HashMap<String, ActivityHandler>,
1679 queries: HashMap<String, HashMap<String, RegisteredQuery>>,
1680 max_concurrent_workflow_tasks: usize,
1681 max_concurrent_activity_tasks: usize,
1682 poll_timeout: Duration,
1683 heartbeat_interval: Duration,
1684 retry_policy: WorkerRetryPolicy,
1685 heartbeat_observer: Option<WorkerHeartbeatObserver>,
1686}
1687
1688impl Worker {
1689 pub fn new(client: Client, task_queue: impl Into<String>) -> Self {
1690 Self {
1691 client,
1692 worker_id: default_worker_id(),
1693 task_queue: task_queue.into(),
1694 workflows: HashMap::new(),
1695 activities: HashMap::new(),
1696 queries: HashMap::new(),
1697 max_concurrent_workflow_tasks: 10,
1698 max_concurrent_activity_tasks: 10,
1699 poll_timeout: Duration::from_secs(30),
1700 heartbeat_interval: Duration::from_secs(60),
1701 retry_policy: WorkerRetryPolicy::default(),
1702 heartbeat_observer: None,
1703 }
1704 }
1705
1706 pub fn worker_id(mut self, worker_id: impl Into<String>) -> Self {
1707 self.worker_id = worker_id.into();
1708 self
1709 }
1710
1711 pub fn poll_timeout(mut self, timeout: Duration) -> Self {
1712 self.poll_timeout = timeout;
1713 self
1714 }
1715
1716 pub fn heartbeat_interval(mut self, interval: Duration) -> Self {
1717 self.heartbeat_interval = interval;
1718 self
1719 }
1720
1721 pub fn retry_policy(mut self, policy: WorkerRetryPolicy) -> Self {
1723 self.retry_policy = policy;
1724 self
1725 }
1726
1727 pub fn on_worker_heartbeat<F>(mut self, observer: F) -> Self
1728 where
1729 F: Fn(&WorkerHeartbeatObservation) + Send + Sync + 'static,
1730 {
1731 self.heartbeat_observer = Some(Arc::new(observer));
1732 self
1733 }
1734
1735 pub fn max_concurrent_workflow_tasks(mut self, count: usize) -> Self {
1736 self.max_concurrent_workflow_tasks = count.max(1);
1737 self
1738 }
1739
1740 pub fn max_concurrent_activity_tasks(mut self, count: usize) -> Self {
1741 self.max_concurrent_activity_tasks = count.max(1);
1742 self
1743 }
1744
1745 pub fn register_workflow<F, Fut>(&mut self, workflow_type: impl Into<String>, handler: F)
1746 where
1747 F: Fn(WorkflowContext, Value) -> Fut + Send + Sync + 'static,
1748 Fut: Future<Output = Result<Value>> + Send + 'static,
1749 {
1750 self.workflows.insert(
1751 workflow_type.into(),
1752 RegisteredWorkflow {
1753 execute: Arc::new(move |ctx, input| Box::pin(handler(ctx, input))),
1754 replay: None,
1755 state_type: None,
1756 },
1757 );
1758 }
1759
1760 pub fn register_replayed_workflow<S, Factory, F, Fut>(
1768 &mut self,
1769 workflow_type: impl Into<String>,
1770 state_factory: Factory,
1771 handler: F,
1772 ) where
1773 S: Clone + Send + Sync + 'static,
1774 Factory: Fn() -> S + Send + Sync + 'static,
1775 F: Fn(WorkflowContext, Value, WorkflowInstance<S>) -> Fut + Send + Sync + 'static,
1776 Fut: Future<Output = Result<Value>> + Send + 'static,
1777 {
1778 let state_factory = Arc::new(state_factory);
1779 let handler = Arc::new(handler);
1780
1781 let execute_factory = Arc::clone(&state_factory);
1782 let execute_handler = Arc::clone(&handler);
1783 let execute = Arc::new(move |ctx: WorkflowContext, input: Value| {
1784 let state = WorkflowInstance::new(execute_factory());
1785 let future = execute_handler(ctx, input, state);
1786 Box::pin(future) as WorkflowFuture
1787 });
1788
1789 let replay = Arc::new(move |ctx: WorkflowContext, input: Value| {
1790 let state = WorkflowInstance::new(state_factory());
1791 let snapshot_state = state.clone();
1792 let snapshot: WorkflowStateSnapshot =
1793 Arc::new(move || Ok(Arc::new(snapshot_state.snapshot()?) as ErasedWorkflowState));
1794 let future = handler(ctx, input, state);
1795 ReplayedWorkflowInvocation {
1796 future: Box::pin(future),
1797 snapshot,
1798 }
1799 });
1800
1801 self.workflows.insert(
1802 workflow_type.into(),
1803 RegisteredWorkflow {
1804 execute,
1805 replay: Some(replay),
1806 state_type: Some(TypeId::of::<S>()),
1807 },
1808 );
1809 }
1810
1811 pub fn register_activity<F, Fut>(&mut self, activity_type: impl Into<String>, handler: F)
1812 where
1813 F: Fn(ActivityContext, Value) -> Fut + Send + Sync + 'static,
1814 Fut: Future<Output = Result<Value>> + Send + 'static,
1815 {
1816 self.activities.insert(
1817 activity_type.into(),
1818 Arc::new(move |ctx, args| Box::pin(handler(ctx, args))),
1819 );
1820 }
1821
1822 pub fn register_query<F, Fut>(
1828 &mut self,
1829 workflow_type: impl Into<String>,
1830 query_name: impl Into<String>,
1831 handler: F,
1832 ) where
1833 F: Fn(QueryContext, Value) -> Fut + Send + Sync + 'static,
1834 Fut: Future<Output = Result<Value>> + Send + 'static,
1835 {
1836 self.queries
1837 .entry(workflow_type.into())
1838 .or_default()
1839 .insert(
1840 query_name.into(),
1841 RegisteredQuery::Snapshot(Arc::new(move |ctx, args| Box::pin(handler(ctx, args)))),
1842 );
1843 }
1844
1845 pub fn register_replayed_query<S, F, Fut>(
1852 &mut self,
1853 workflow_type: impl Into<String>,
1854 query_name: impl Into<String>,
1855 handler: F,
1856 ) where
1857 S: Clone + Send + Sync + 'static,
1858 F: Fn(QueryContext, Arc<S>, Value) -> Fut + Send + Sync + 'static,
1859 Fut: Future<Output = Result<Value>> + Send + 'static,
1860 {
1861 let handler = Arc::new(handler);
1862 let erased_handler: ReplayedQueryHandler = Arc::new(move |ctx, state, args| {
1863 let state = state.downcast::<S>().map_err(|_| {
1864 "registered query state type does not match the replayed workflow state".to_string()
1865 })?;
1866 Ok(Box::pin(handler(ctx, state, args)))
1867 });
1868
1869 self.queries
1870 .entry(workflow_type.into())
1871 .or_default()
1872 .insert(
1873 query_name.into(),
1874 RegisteredQuery::Replayed {
1875 state_type: TypeId::of::<S>(),
1876 handler: erased_handler,
1877 },
1878 );
1879 }
1880
1881 pub async fn register(&self) -> Result<RegisterWorkerResponse> {
1882 self.client
1883 .register_worker_with_capabilities(
1884 &self.worker_id,
1885 &self.task_queue,
1886 self.workflows.keys().cloned().collect(),
1887 self.activities.keys().cloned().collect(),
1888 self.max_concurrent_workflow_tasks,
1889 self.max_concurrent_activity_tasks,
1890 (!self.queries.is_empty())
1891 .then(|| QUERY_TASKS_CAPABILITY.to_string())
1892 .into_iter()
1893 .collect(),
1894 )
1895 .await
1896 }
1897
1898 pub async fn run(&self) -> Result<()> {
1904 self.run_until(std::future::pending::<()>()).await
1905 }
1906
1907 pub async fn run_until<F>(&self, shutdown: F) -> Result<()>
1911 where
1912 F: Future<Output = ()>,
1913 {
1914 let registration = self.register().await?;
1915 let mut heartbeat = tokio::time::interval(Duration::from_secs(
1916 registration
1917 .heartbeat_interval_seconds
1918 .unwrap_or(self.heartbeat_interval.as_secs().max(1)),
1919 ));
1920 tokio::pin!(shutdown);
1921 let stop = Arc::new(AtomicBool::new(false));
1922 let mut workflow_poller = (!self.workflows.is_empty()).then(|| {
1926 let worker = self.clone();
1927 let stop = Arc::clone(&stop);
1928 tokio::spawn(async move { worker.poll_workflows_until_stopped(stop).await })
1929 });
1930 let mut activity_poller = (!self.activities.is_empty()).then(|| {
1931 let worker = self.clone();
1932 let stop = Arc::clone(&stop);
1933 tokio::spawn(async move { worker.poll_activities_until_stopped(stop).await })
1934 });
1935 let mut query_poller = (!self.queries.is_empty()).then(|| {
1936 let worker = self.clone();
1937 let stop = Arc::clone(&stop);
1938 tokio::spawn(async move { worker.poll_queries_until_stopped(stop).await })
1939 });
1940
1941 loop {
1942 tokio::select! {
1943 _ = &mut shutdown => {
1944 stop.store(true, Ordering::SeqCst);
1945 break;
1946 }
1947 _ = heartbeat.tick() => {
1948 match self.retry_worker_operation(|| {
1949 self.client.heartbeat_worker(
1950 &self.worker_id,
1951 self.max_concurrent_workflow_tasks,
1952 self.max_concurrent_activity_tasks,
1953 )
1954 }).await
1955 {
1956 Ok(acknowledgement) => {
1957 if let Some(observer) = &self.heartbeat_observer {
1958 observer(&WorkerHeartbeatObservation {
1959 worker_id: self.worker_id.clone(),
1960 task_queue: self.task_queue.clone(),
1961 acknowledged_at_unix_millis: SystemTime::now()
1962 .duration_since(UNIX_EPOCH)
1963 .unwrap_or_default()
1964 .as_millis()
1965 .min(u64::MAX as u128)
1966 as u64,
1967 acknowledgement,
1968 });
1969 }
1970 }
1971 Err(error) => {
1972 stop.store(true, Ordering::SeqCst);
1973 join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await?;
1974 return Err(error);
1975 }
1976 }
1977 }
1978 result = OptionFuture::from(workflow_poller.as_mut()), if workflow_poller.is_some() => {
1979 workflow_poller = None;
1980 stop.store(true, Ordering::SeqCst);
1981 let poller_result = optional_poller_result("workflow", result);
1982 let join_result =
1983 join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
1984 poller_result?;
1985 join_result?;
1986 return Err(Error::WorkerLoop(
1987 "workflow poller stopped unexpectedly".to_string(),
1988 ));
1989 }
1990 result = OptionFuture::from(activity_poller.as_mut()), if activity_poller.is_some() => {
1991 activity_poller = None;
1992 stop.store(true, Ordering::SeqCst);
1993 let poller_result = optional_poller_result("activity", result);
1994 let join_result =
1995 join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
1996 poller_result?;
1997 join_result?;
1998 return Err(Error::WorkerLoop(
1999 "activity poller stopped unexpectedly".to_string(),
2000 ));
2001 }
2002 result = OptionFuture::from(query_poller.as_mut()), if query_poller.is_some() => {
2003 query_poller = None;
2004 stop.store(true, Ordering::SeqCst);
2005 let poller_result = optional_poller_result("query", result);
2006 let join_result =
2007 join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
2008 poller_result?;
2009 join_result?;
2010 return Err(Error::WorkerLoop(
2011 "query poller stopped unexpectedly".to_string(),
2012 ));
2013 }
2014 }
2015 }
2016
2017 join_pollers(
2018 workflow_poller.take(),
2019 activity_poller.take(),
2020 query_poller.take(),
2021 )
2022 .await
2023 }
2024
2025 pub async fn run_once(&self) -> Result<usize> {
2026 let mut handled = 0;
2027 if self.poll_workflow_once().await? {
2028 handled += 1;
2029 }
2030 if self.poll_activity_once().await? {
2031 handled += 1;
2032 }
2033 if !self.queries.is_empty() && self.poll_query_once().await? {
2034 handled += 1;
2035 }
2036 Ok(handled)
2037 }
2038
2039 async fn poll_workflow_once(&self) -> Result<bool> {
2040 let Some(task) = self
2041 .retry_worker_operation(|| {
2042 self.client
2043 .poll_workflow_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2044 })
2045 .await?
2046 else {
2047 return Ok(false);
2048 };
2049
2050 let task_id = task.task_id.clone();
2051 let attempt = task.workflow_task_attempt;
2052 let lease_owner = task
2053 .lease_owner
2054 .clone()
2055 .unwrap_or_else(|| self.worker_id.clone());
2056
2057 match self.execute_workflow_task(task) {
2058 Ok(commands) if commands.is_empty() => {
2059 self.client
2065 .fail_workflow_task_with_type(
2066 &task_id,
2067 &lease_owner,
2068 attempt,
2069 WORKFLOW_TASK_WAITING_FOR_HISTORY_MESSAGE,
2070 WORKFLOW_TASK_WAITING_FOR_HISTORY_TYPE,
2071 )
2072 .await?;
2073 }
2074 Ok(commands) => {
2075 self.client
2076 .complete_workflow_task(&task_id, &lease_owner, attempt, commands)
2077 .await?;
2078 }
2079 Err(error) => {
2080 self.client
2081 .fail_workflow_task(&task_id, &lease_owner, attempt, error.to_string())
2082 .await?;
2083 }
2084 }
2085
2086 Ok(true)
2087 }
2088
2089 async fn poll_workflows_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2090 while !stop.load(Ordering::SeqCst) {
2091 self.poll_workflow_once().await?;
2092 }
2093
2094 Ok(())
2095 }
2096
2097 async fn poll_activity_once(&self) -> Result<bool> {
2098 let Some(task) = self
2099 .retry_worker_operation(|| {
2100 self.client
2101 .poll_activity_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2102 })
2103 .await?
2104 else {
2105 return Ok(false);
2106 };
2107
2108 let task_id = task.task_id.clone();
2109 let attempt_id = task
2110 .activity_attempt_id
2111 .clone()
2112 .or(task.attempt_id.clone())
2113 .unwrap_or_default();
2114 let lease_owner = task
2115 .lease_owner
2116 .clone()
2117 .unwrap_or_else(|| self.worker_id.clone());
2118 let codec = task.payload_codec.clone();
2119 let result = self.execute_activity_task(task).await;
2120 match result {
2121 Ok(value) => {
2122 self.client
2123 .complete_activity_task(&task_id, &attempt_id, &lease_owner, value, &codec)
2124 .await?;
2125 }
2126 Err(error) => {
2127 self.client
2128 .fail_activity_task(
2129 &task_id,
2130 &attempt_id,
2131 &lease_owner,
2132 error.to_string(),
2133 false,
2134 )
2135 .await?;
2136 }
2137 }
2138
2139 Ok(true)
2140 }
2141
2142 async fn poll_activities_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2143 while !stop.load(Ordering::SeqCst) {
2144 self.poll_activity_once().await?;
2145 }
2146
2147 Ok(())
2148 }
2149
2150 async fn poll_query_once(&self) -> Result<bool> {
2151 let Some(task) = self
2152 .retry_worker_operation(|| {
2153 self.client
2154 .poll_query_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2155 })
2156 .await?
2157 else {
2158 return Ok(false);
2159 };
2160
2161 let query_task_id = task.query_task_id.clone();
2162 let attempt = task.query_task_attempt;
2163 let lease_owner = task
2164 .lease_owner
2165 .clone()
2166 .unwrap_or_else(|| self.worker_id.clone());
2167 let codec = task.payload_codec.clone();
2168
2169 match self.execute_query_task(task).await {
2170 Ok(value) => {
2171 let result_envelope = match encode_value_envelope(&value, &codec) {
2172 Ok(result_envelope) => result_envelope,
2173 Err(error) => {
2174 let failure = self
2175 .client
2176 .fail_query_task(
2177 &query_task_id,
2178 &lease_owner,
2179 attempt,
2180 error.to_string(),
2181 "query_result_encode_failed",
2182 "QueryResultEncodeFailed",
2183 )
2184 .await;
2185 if let Err(error) = failure {
2186 if !query_task_rejection_is_final(&error) {
2187 return Err(error);
2188 }
2189 }
2190 return Ok(true);
2191 }
2192 };
2193
2194 if let Err(error) = self
2195 .client
2196 .complete_query_task_with_envelope(
2197 &query_task_id,
2198 &lease_owner,
2199 attempt,
2200 value,
2201 result_envelope,
2202 )
2203 .await
2204 {
2205 if !query_task_rejection_is_final(&error) {
2206 return Err(error);
2207 }
2208 }
2209 }
2210 Err(failure) => {
2211 let result = self
2212 .client
2213 .fail_query_task(
2214 &query_task_id,
2215 &lease_owner,
2216 attempt,
2217 failure.message,
2218 failure.reason,
2219 failure.failure_type,
2220 )
2221 .await;
2222 if let Err(error) = result {
2223 if !query_task_rejection_is_final(&error) {
2224 return Err(error);
2225 }
2226 }
2227 }
2228 }
2229
2230 Ok(true)
2231 }
2232
2233 async fn poll_queries_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2234 while !stop.load(Ordering::SeqCst) {
2235 self.poll_query_once().await?;
2236 }
2237
2238 Ok(())
2239 }
2240
2241 async fn retry_worker_operation<T, F, Fut>(&self, mut operation: F) -> Result<T>
2242 where
2243 F: FnMut() -> Fut,
2244 Fut: Future<Output = Result<T>>,
2245 {
2246 let mut retries = 0;
2247
2248 loop {
2249 match operation().await {
2250 Err(error)
2251 if worker_operation_is_retryable(&error)
2252 && retries < self.retry_policy.max_retries =>
2253 {
2254 retries += 1;
2255 tokio::time::sleep(worker_retry_delay(self.retry_policy, retries)).await;
2256 }
2257 result => return result,
2258 }
2259 }
2260 }
2261
2262 async fn execute_query_task(
2263 &self,
2264 mut task: QueryTask,
2265 ) -> std::result::Result<Value, QueryTaskExecutionFailure> {
2266 if !matches!(task.payload_codec.as_str(), DEFAULT_CODEC | JSON_CODEC) {
2267 return Err(QueryTaskExecutionFailure::new(
2268 "query_payload_decode_failed",
2269 format!(
2270 "cannot decode query payload with unsupported codec {:?}",
2271 task.payload_codec
2272 ),
2273 "QueryPayloadDecodeFailed",
2274 ));
2275 }
2276
2277 if !self.workflows.contains_key(&task.workflow_type) {
2278 return Err(QueryTaskExecutionFailure::new(
2279 "query_workflow_type_not_registered",
2280 format!("no workflow registered for type {:?}", task.workflow_type),
2281 "WorkflowTypeNotRegistered",
2282 ));
2283 }
2284
2285 let Some(handlers) = self.queries.get(&task.workflow_type) else {
2286 return Err(QueryTaskExecutionFailure::new(
2287 "query_handler_unavailable",
2288 format!(
2289 "query handlers are unavailable for workflow type {:?}",
2290 task.workflow_type
2291 ),
2292 "QueryHandlerUnavailable",
2293 ));
2294 };
2295 let Some(query) = handlers.get(&task.query_name) else {
2296 return Err(QueryTaskExecutionFailure::new(
2297 "rejected_unknown_query",
2298 format!("unknown query {:?}", task.query_name),
2299 "QueryFailed",
2300 ));
2301 };
2302
2303 let args = decode_task_arguments(task.query_arguments.as_ref(), &task.payload_codec)
2304 .map_err(|error| {
2305 QueryTaskExecutionFailure::new(
2306 "query_payload_decode_failed",
2307 format!("cannot decode query arguments: {error}"),
2308 "QueryPayloadDecodeFailed",
2309 )
2310 })?;
2311 let workflow_input =
2312 decode_task_arguments(task.workflow_arguments.as_ref(), &task.payload_codec).map_err(
2313 |error| {
2314 QueryTaskExecutionFailure::new(
2315 "query_workflow_state_unavailable",
2316 format!("cannot decode workflow start input: {error}"),
2317 "QueryWorkflowStateUnavailable",
2318 )
2319 },
2320 )?;
2321 hydrate_query_history_from_export(&mut task).map_err(|error| {
2322 QueryTaskExecutionFailure::new(
2323 "query_workflow_state_unavailable",
2324 format!("cannot restore query history snapshot: {error}"),
2325 "QueryWorkflowStateUnavailable",
2326 )
2327 })?;
2328 enrich_query_history_from_export(&mut task).map_err(|error| {
2329 QueryTaskExecutionFailure::new(
2330 "query_workflow_state_unavailable",
2331 format!("cannot restore compact query history payloads: {error}"),
2332 "QueryWorkflowStateUnavailable",
2333 )
2334 })?;
2335 let signal_events = query_signal_events(&task).map_err(|error| {
2336 QueryTaskExecutionFailure::new(
2337 "query_workflow_state_unavailable",
2338 format!("cannot decode committed workflow signals: {error}"),
2339 "QueryWorkflowStateUnavailable",
2340 )
2341 })?;
2342 let history_events = Arc::new(std::mem::take(&mut task.history_events));
2343 let context = QueryContext {
2344 workflow_id: task.workflow_id,
2345 run_id: task.run_id,
2346 workflow_type: task.workflow_type.clone(),
2347 run_status: task.run_status,
2348 workflow_input,
2349 history_events: Arc::clone(&history_events),
2350 signal_events: Arc::new(signal_events),
2351 };
2352
2353 let future = match query {
2354 RegisteredQuery::Snapshot(handler) => handler(context, args),
2355 RegisteredQuery::Replayed {
2356 state_type,
2357 handler,
2358 } => {
2359 let workflow = self
2360 .workflows
2361 .get(&task.workflow_type)
2362 .expect("workflow registration was checked above");
2363 if workflow.state_type != Some(*state_type) {
2364 return Err(QueryTaskExecutionFailure::new(
2365 "query_workflow_state_unavailable",
2366 "replayed query state type does not match its workflow registration",
2367 "QueryWorkflowStateUnavailable",
2368 ));
2369 }
2370 let replay = workflow.replay.as_ref().ok_or_else(|| {
2371 QueryTaskExecutionFailure::new(
2372 "query_workflow_state_unavailable",
2373 format!(
2374 "workflow type {:?} is not registered for instance-state replay",
2375 task.workflow_type
2376 ),
2377 "QueryWorkflowStateUnavailable",
2378 )
2379 })?;
2380 let workflow_state = Arc::new(Mutex::new(
2381 WorkflowState::new_with_identity(
2382 history_events.as_ref().clone(),
2383 context.workflow_id.clone(),
2384 context.run_id.clone(),
2385 self.task_queue.clone(),
2386 task.payload_codec,
2387 None,
2388 )
2389 .map_err(|error| {
2390 QueryTaskExecutionFailure::new(
2391 "query_workflow_state_unavailable",
2392 format!("workflow replay failed before query: {error}"),
2393 "QueryWorkflowStateUnavailable",
2394 )
2395 })?,
2396 ));
2397 let workflow_context = WorkflowContext {
2398 state: workflow_state,
2399 };
2400 let mut invocation =
2401 replay(workflow_context.clone(), context.workflow_input.clone());
2402 let mut cx = TaskContext::from_waker(noop_waker_ref());
2403 match invocation.future.as_mut().poll(&mut cx) {
2404 Poll::Ready(Ok(_)) => {
2405 workflow_context
2406 .ensure_history_consumed()
2407 .map_err(|error| {
2408 QueryTaskExecutionFailure::new(
2409 "query_workflow_state_unavailable",
2410 format!("workflow replay failed before query: {error}"),
2411 "QueryWorkflowStateUnavailable",
2412 )
2413 })?;
2414 }
2415 Poll::Ready(Err(error)) => {
2416 return Err(QueryTaskExecutionFailure::new(
2417 "query_workflow_state_unavailable",
2418 format!("workflow replay failed before query: {error}"),
2419 "QueryWorkflowStateUnavailable",
2420 ));
2421 }
2422 Poll::Pending => {
2423 let commands = workflow_context.take_commands().map_err(|error| {
2424 QueryTaskExecutionFailure::new(
2425 "query_workflow_state_unavailable",
2426 format!("workflow replay failed before query: {error}"),
2427 "QueryWorkflowStateUnavailable",
2428 )
2429 })?;
2430 if commands.is_empty()
2431 && !workflow_context
2432 .matched_recorded_pending()
2433 .map_err(|error| {
2434 QueryTaskExecutionFailure::new(
2435 "query_workflow_state_unavailable",
2436 format!("workflow replay failed before query: {error}"),
2437 "QueryWorkflowStateUnavailable",
2438 )
2439 })?
2440 {
2441 return Err(QueryTaskExecutionFailure::new(
2442 "query_workflow_state_unavailable",
2443 "workflow replay yielded without a durable command",
2444 "QueryWorkflowStateUnavailable",
2445 ));
2446 }
2447 }
2448 }
2449 let state = (invocation.snapshot)().map_err(|error| {
2450 QueryTaskExecutionFailure::new(
2451 "query_workflow_state_unavailable",
2452 format!("cannot snapshot replayed workflow state: {error}"),
2453 "QueryWorkflowStateUnavailable",
2454 )
2455 })?;
2456 handler(context, state, args).map_err(|message| {
2457 QueryTaskExecutionFailure::new(
2458 "query_workflow_state_unavailable",
2459 message,
2460 "QueryWorkflowStateUnavailable",
2461 )
2462 })?
2463 }
2464 };
2465
2466 future.await.map_err(|error| {
2467 QueryTaskExecutionFailure::new("query_rejected", error.to_string(), "QueryFailed")
2468 })
2469 }
2470
2471 fn execute_workflow_task(&self, task: WorkflowTask) -> Result<Vec<Value>> {
2472 let workflow = self
2473 .workflows
2474 .get(&task.workflow_type)
2475 .ok_or_else(|| Error::WorkflowNotRegistered(task.workflow_type.clone()))?;
2476 let input = decode_task_arguments(task.arguments.as_ref(), &task.payload_codec)?;
2477 let resume_signal = decode_resume_signal(&task)?;
2478 let state = Arc::new(Mutex::new(WorkflowState::new_with_identity(
2479 task.history_events,
2480 task.workflow_id,
2481 task.run_id,
2482 self.task_queue.clone(),
2483 task.payload_codec.clone(),
2484 resume_signal,
2485 )?));
2486 let ctx = WorkflowContext { state };
2487 let mut future = (workflow.execute)(ctx.clone(), input);
2488 let mut cx = TaskContext::from_waker(noop_waker_ref());
2489
2490 match future.as_mut().poll(&mut cx) {
2491 Poll::Ready(Ok(result)) => {
2492 ctx.ensure_history_consumed()?;
2493 let result = encode_value_envelope(&result, &task.payload_codec)?;
2494 Ok(vec![json!({
2495 "type": "complete_workflow",
2496 "result": result
2497 })])
2498 }
2499 Poll::Ready(Err(error @ Error::ChildWorkflowFailed(_))) => {
2500 Ok(vec![workflow_failure_command(&error)])
2501 }
2502 Poll::Ready(Err(error)) => Err(error),
2503 Poll::Pending => {
2504 let commands = ctx.take_commands()?;
2505 if commands.is_empty() && !ctx.matched_recorded_pending()? {
2506 Err(Error::WorkflowYieldedWithoutCommand)
2507 } else {
2508 Ok(commands)
2509 }
2510 }
2511 }
2512 }
2513
2514 async fn execute_activity_task(&self, task: ActivityTask) -> Result<Value> {
2515 let handler = self
2516 .activities
2517 .get(&task.activity_type)
2518 .ok_or_else(|| Error::ActivityNotRegistered(task.activity_type.clone()))?;
2519 let args = decode_task_arguments(task.arguments.as_ref(), &task.payload_codec)?;
2520 let attempt_id = task
2521 .activity_attempt_id
2522 .clone()
2523 .or(task.attempt_id.clone())
2524 .unwrap_or_default();
2525 let lease_owner = task
2526 .lease_owner
2527 .clone()
2528 .unwrap_or_else(|| self.worker_id.clone());
2529 let ctx = ActivityContext {
2530 client: self.client.clone(),
2531 task_id: task.task_id,
2532 activity_attempt_id: attempt_id,
2533 lease_owner,
2534 activity_type: task.activity_type,
2535 attempt_number: task.attempt_number,
2536 task_queue: self.task_queue.clone(),
2537 worker_id: self.worker_id.clone(),
2538 };
2539
2540 handler(ctx, args).await
2541 }
2542}
2543
2544fn poller_result(
2545 kind: &str,
2546 result: std::result::Result<Result<()>, tokio::task::JoinError>,
2547) -> Result<()> {
2548 match result {
2549 Ok(result) => result,
2550 Err(error) => Err(Error::WorkerLoop(format!(
2551 "{kind} poller join error: {error}"
2552 ))),
2553 }
2554}
2555
2556fn optional_poller_result(
2557 kind: &str,
2558 result: Option<std::result::Result<Result<()>, tokio::task::JoinError>>,
2559) -> Result<()> {
2560 match result {
2561 Some(result) => poller_result(kind, result),
2562 None => Ok(()),
2563 }
2564}
2565
2566async fn join_pollers(
2567 workflow_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2568 activity_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2569 query_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2570) -> Result<()> {
2571 let mut first_error = None;
2572
2573 if let Some(handle) = workflow_poller {
2574 if let Err(error) = poller_result("workflow", handle.await) {
2575 first_error.get_or_insert(error);
2576 }
2577 }
2578
2579 if let Some(handle) = activity_poller {
2580 if let Err(error) = poller_result("activity", handle.await) {
2581 first_error.get_or_insert(error);
2582 }
2583 }
2584
2585 if let Some(handle) = query_poller {
2586 if let Err(error) = poller_result("query", handle.await) {
2587 first_error.get_or_insert(error);
2588 }
2589 }
2590
2591 if let Some(error) = first_error {
2592 Err(error)
2593 } else {
2594 Ok(())
2595 }
2596}
2597
2598fn default_worker_id() -> String {
2599 let millis = SystemTime::now()
2600 .duration_since(UNIX_EPOCH)
2601 .unwrap_or_default()
2602 .as_millis();
2603 format!("rust-worker-{}-{millis}", std::process::id())
2604}
2605
2606fn unique_request_id(prefix: &str) -> String {
2607 let nanos = SystemTime::now()
2608 .duration_since(UNIX_EPOCH)
2609 .unwrap_or_default()
2610 .as_nanos();
2611 format!("{prefix}-{}-{nanos}", std::process::id())
2612}
2613
2614#[derive(Debug)]
2615struct QueryTaskExecutionFailure {
2616 reason: String,
2617 message: String,
2618 failure_type: String,
2619}
2620
2621impl QueryTaskExecutionFailure {
2622 fn new(
2623 reason: impl Into<String>,
2624 message: impl Into<String>,
2625 failure_type: impl Into<String>,
2626 ) -> Self {
2627 Self {
2628 reason: reason.into(),
2629 message: message.into(),
2630 failure_type: failure_type.into(),
2631 }
2632 }
2633}
2634
2635#[derive(Clone, Debug)]
2641pub struct WorkflowInstance<S> {
2642 state: Arc<Mutex<S>>,
2643}
2644
2645impl<S> WorkflowInstance<S> {
2646 fn new(state: S) -> Self {
2647 Self {
2648 state: Arc::new(Mutex::new(state)),
2649 }
2650 }
2651
2652 pub fn read<R>(&self, reader: impl FnOnce(&S) -> R) -> Result<R> {
2654 let state = self
2655 .state
2656 .lock()
2657 .map_err(|_| Error::WorkflowStatePoisoned)?;
2658 Ok(reader(&state))
2659 }
2660
2661 pub fn update<R>(&self, transition: impl FnOnce(&mut S) -> R) -> Result<R> {
2663 let mut state = self
2664 .state
2665 .lock()
2666 .map_err(|_| Error::WorkflowStatePoisoned)?;
2667 Ok(transition(&mut state))
2668 }
2669}
2670
2671impl<S: Clone> WorkflowInstance<S> {
2672 fn snapshot(&self) -> Result<S> {
2673 self.read(Clone::clone)
2674 }
2675}
2676
2677#[derive(Clone, Debug)]
2678pub struct WorkflowContext {
2679 state: Arc<Mutex<WorkflowState>>,
2680}
2681
2682impl WorkflowContext {
2683 pub fn workflow_identity(&self) -> Result<WorkflowIdentity> {
2685 let state = self
2686 .state
2687 .lock()
2688 .map_err(|_| Error::WorkflowStatePoisoned)?;
2689 Ok(WorkflowIdentity {
2690 workflow_id: state.workflow_id.clone(),
2691 run_id: state.run_id.clone(),
2692 })
2693 }
2694
2695 pub fn activity<T: Serialize>(
2696 &self,
2697 activity_type: impl Into<String>,
2698 args: T,
2699 ) -> ActivityCall {
2700 self.activity_on_queue(activity_type, None::<String>, args)
2701 }
2702
2703 pub fn activity_on_queue<T, Q>(
2704 &self,
2705 activity_type: impl Into<String>,
2706 task_queue: Option<Q>,
2707 args: T,
2708 ) -> ActivityCall
2709 where
2710 T: Serialize,
2711 Q: Into<String>,
2712 {
2713 ActivityCall {
2714 ctx: self.clone(),
2715 activity_type: activity_type.into(),
2716 task_queue: task_queue.map(Into::into),
2717 args: Some(serde_json::to_value(args).map_err(Error::from)),
2718 scheduled: false,
2719 }
2720 }
2721
2722 pub fn wait_signal(&self, signal_name: impl Into<String>) -> SignalCall {
2723 SignalCall {
2724 ctx: self.clone(),
2725 signal_name: signal_name.into(),
2726 opened_wait: false,
2727 matched_pending: false,
2728 }
2729 }
2730
2731 pub fn sleep(&self, duration: Duration) -> TimerCall {
2752 let delay_seconds = duration
2753 .as_secs()
2754 .checked_add(u64::from(duration.subsec_nanos() > 0));
2755 TimerCall {
2756 ctx: self.clone(),
2757 delay_seconds,
2758 scheduled: false,
2759 matched_pending: false,
2760 }
2761 }
2762
2763 pub fn start_timer(&self, duration: Duration) -> TimerCall {
2765 self.sleep(duration)
2766 }
2767
2768 pub fn start_child_workflow<T: Serialize>(
2795 &self,
2796 workflow_type: impl Into<String>,
2797 options: ChildWorkflowOptions,
2798 args: T,
2799 ) -> ChildWorkflowCall {
2800 ChildWorkflowCall {
2801 ctx: self.clone(),
2802 workflow_type: workflow_type.into(),
2803 options,
2804 args: Some(serde_json::to_value(args).map_err(Error::from)),
2805 scheduled: false,
2806 matched_pending: false,
2807 }
2808 }
2809
2810 fn take_commands(&self) -> Result<Vec<Value>> {
2811 let mut state = self
2812 .state
2813 .lock()
2814 .map_err(|_| Error::WorkflowStatePoisoned)?;
2815 Ok(std::mem::take(&mut state.commands))
2816 }
2817
2818 fn matched_recorded_pending(&self) -> Result<bool> {
2819 let state = self
2820 .state
2821 .lock()
2822 .map_err(|_| Error::WorkflowStatePoisoned)?;
2823 Ok(state.matched_recorded_pending)
2824 }
2825
2826 fn ensure_history_consumed(&self) -> Result<()> {
2827 let state = self
2828 .state
2829 .lock()
2830 .map_err(|_| Error::WorkflowStatePoisoned)?;
2831 if let Some(command) = state.recorded_commands.get(state.command_cursor) {
2832 return Err(Error::NonDeterministicReplay(ReplayFailure::new(
2833 "recorded_commands_unconsumed",
2834 Some(command.sequence()),
2835 Some(command.shape().to_string()),
2836 Some("workflow completion".to_string()),
2837 "workflow completed before consuming all recorded durable commands",
2838 )));
2839 }
2840 Ok(())
2841 }
2842}
2843
2844#[derive(Debug)]
2845struct WorkflowState {
2846 history: Vec<HistoryEvent>,
2847 workflow_id: Option<String>,
2848 run_id: Option<String>,
2849 task_queue: String,
2850 payload_codec: String,
2851 resume_signal: Option<ResumeSignal>,
2852 recorded_commands: Vec<RecordedCommand>,
2853 command_cursor: usize,
2854 matched_recorded_pending: bool,
2855 signal_cursors: HashMap<String, usize>,
2856 commands: Vec<Value>,
2857}
2858
2859impl WorkflowState {
2860 #[cfg(test)]
2861 fn new(
2862 history: Vec<HistoryEvent>,
2863 task_queue: String,
2864 payload_codec: String,
2865 resume_signal: Option<ResumeSignal>,
2866 ) -> Result<Self> {
2867 Self::new_with_identity(
2868 history,
2869 None,
2870 None,
2871 task_queue,
2872 payload_codec,
2873 resume_signal,
2874 )
2875 }
2876
2877 fn new_with_identity(
2878 history: Vec<HistoryEvent>,
2879 workflow_id: Option<String>,
2880 run_id: Option<String>,
2881 task_queue: String,
2882 payload_codec: String,
2883 resume_signal: Option<ResumeSignal>,
2884 ) -> Result<Self> {
2885 let recorded_commands = recorded_commands(
2886 &history,
2887 &payload_codec,
2888 WorkflowIdentity {
2889 workflow_id: workflow_id.clone(),
2890 run_id: run_id.clone(),
2891 },
2892 )?;
2893 Ok(Self {
2894 history,
2895 workflow_id,
2896 run_id,
2897 task_queue,
2898 payload_codec,
2899 resume_signal,
2900 recorded_commands,
2901 command_cursor: 0,
2902 matched_recorded_pending: false,
2903 signal_cursors: HashMap::new(),
2904 commands: Vec::new(),
2905 })
2906 }
2907}
2908
2909#[derive(Clone, Debug)]
2910enum RecordedCommand {
2911 Activity {
2912 sequence: u64,
2913 activity_type: Option<String>,
2914 result: Option<Value>,
2915 },
2916 Timer {
2917 sequence: u64,
2918 delay_seconds: u64,
2919 fired: bool,
2920 },
2921 ChildWorkflow {
2922 sequence: u64,
2923 workflow_type: Option<String>,
2924 outcome: Option<ChildWorkflowOutcome>,
2925 },
2926 SignalWait {
2927 sequence: u64,
2928 signal_name: Option<String>,
2929 },
2930}
2931
2932impl RecordedCommand {
2933 fn sequence(&self) -> u64 {
2934 match self {
2935 Self::Activity { sequence, .. }
2936 | Self::Timer { sequence, .. }
2937 | Self::ChildWorkflow { sequence, .. }
2938 | Self::SignalWait { sequence, .. } => *sequence,
2939 }
2940 }
2941
2942 fn shape(&self) -> &'static str {
2943 match self {
2944 Self::Activity { .. } => "activity",
2945 Self::Timer { .. } => "timer",
2946 Self::ChildWorkflow { .. } => "child workflow",
2947 Self::SignalWait { .. } => "signal wait",
2948 }
2949 }
2950}
2951
2952#[derive(Clone, Debug)]
2953struct ResumeSignal {
2954 signal_id: Option<String>,
2955 signal_name: String,
2956 arguments: Vec<Value>,
2957}
2958
2959pub struct ActivityCall {
2960 ctx: WorkflowContext,
2961 activity_type: String,
2962 task_queue: Option<String>,
2963 args: Option<Result<Value>>,
2964 scheduled: bool,
2965}
2966
2967impl Future for ActivityCall {
2968 type Output = Result<Value>;
2969
2970 fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
2971 let ctx = self.ctx.clone();
2972 let mut state = match ctx.state.lock() {
2973 Ok(state) => state,
2974 Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
2975 };
2976
2977 if self.scheduled {
2978 return Poll::Pending;
2979 }
2980
2981 if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
2982 let sequence = recorded.sequence();
2983 match recorded {
2984 RecordedCommand::Activity {
2985 activity_type,
2986 result,
2987 ..
2988 } => {
2989 if let Some(recorded_type) = activity_type {
2990 if recorded_type != self.activity_type {
2991 return Poll::Ready(Err(Error::NonDeterministicReplay(
2992 ReplayFailure::new(
2993 "recorded_command_detail_mismatch",
2994 Some(sequence),
2995 Some(format!("activity:{recorded_type}")),
2996 Some(format!("activity:{}", self.activity_type)),
2997 "recorded activity type differs from the current workflow command",
2998 ),
2999 )));
3000 }
3001 }
3002 state.command_cursor += 1;
3003 if let Some(value) = result {
3004 return Poll::Ready(Ok(value));
3005 }
3006 state.matched_recorded_pending = true;
3007 self.scheduled = true;
3008 return Poll::Pending;
3009 }
3010 other => {
3011 return Poll::Ready(Err(command_mismatch(
3012 &other,
3013 format!("activity:{}", self.activity_type),
3014 )));
3015 }
3016 }
3017 }
3018
3019 if !self.scheduled {
3020 let task_queue = self
3021 .task_queue
3022 .clone()
3023 .unwrap_or_else(|| state.task_queue.clone());
3024 let args = match self.args.take().unwrap_or(Ok(Value::Null)) {
3025 Ok(args) => args,
3026 Err(error) => return Poll::Ready(Err(error)),
3027 };
3028 let arguments = normalize_arguments(args);
3029 let envelope = match encode_value_envelope(&arguments, &state.payload_codec) {
3030 Ok(envelope) => envelope,
3031 Err(error) => return Poll::Ready(Err(error)),
3032 };
3033
3034 state.commands.push(json!({
3035 "type": "schedule_activity",
3036 "activity_type": self.activity_type.clone(),
3037 "queue": task_queue,
3038 "arguments": envelope
3039 }));
3040 self.scheduled = true;
3041 }
3042
3043 Poll::Pending
3044 }
3045}
3046
3047pub struct TimerCall {
3049 ctx: WorkflowContext,
3050 delay_seconds: Option<u64>,
3051 scheduled: bool,
3052 matched_pending: bool,
3053}
3054
3055impl Future for TimerCall {
3056 type Output = Result<()>;
3057
3058 fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3059 if self.matched_pending {
3060 return Poll::Pending;
3061 }
3062
3063 let ctx = self.ctx.clone();
3064 let Some(requested_delay) = self.delay_seconds else {
3065 return Poll::Ready(Err(Error::TimerDurationOverflow));
3066 };
3067 let mut state = match ctx.state.lock() {
3068 Ok(state) => state,
3069 Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3070 };
3071
3072 if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3073 match recorded {
3074 RecordedCommand::Timer {
3075 sequence,
3076 delay_seconds,
3077 fired,
3078 ..
3079 } => {
3080 if delay_seconds != requested_delay {
3081 return Poll::Ready(Err(Error::NonDeterministicReplay(
3082 ReplayFailure::new(
3083 "timer_delay_mismatch",
3084 Some(sequence),
3085 Some(format!("timer:{delay_seconds}s")),
3086 Some(format!("timer:{requested_delay}s")),
3087 "recorded timer delay differs from the current workflow command",
3088 ),
3089 )));
3090 }
3091 state.command_cursor += 1;
3092 if fired {
3093 return Poll::Ready(Ok(()));
3094 }
3095 state.matched_recorded_pending = true;
3096 self.scheduled = true;
3097 self.matched_pending = true;
3098 return Poll::Pending;
3099 }
3100 other => return Poll::Ready(Err(command_mismatch(&other, "timer"))),
3101 }
3102 }
3103
3104 if !self.scheduled {
3105 state.commands.push(json!({
3106 "type": "start_timer",
3107 "delay_seconds": requested_delay,
3108 }));
3109 self.scheduled = true;
3110 }
3111
3112 Poll::Pending
3113 }
3114}
3115
3116pub struct ChildWorkflowCall {
3118 ctx: WorkflowContext,
3119 workflow_type: String,
3120 options: ChildWorkflowOptions,
3121 args: Option<Result<Value>>,
3122 scheduled: bool,
3123 matched_pending: bool,
3124}
3125
3126impl Future for ChildWorkflowCall {
3127 type Output = Result<ChildWorkflowResult>;
3128
3129 fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3130 if self.matched_pending {
3131 return Poll::Pending;
3132 }
3133
3134 let ctx = self.ctx.clone();
3135 let mut state = match ctx.state.lock() {
3136 Ok(state) => state,
3137 Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3138 };
3139
3140 if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3141 let sequence = recorded.sequence();
3142 match recorded {
3143 RecordedCommand::ChildWorkflow {
3144 workflow_type,
3145 outcome,
3146 ..
3147 } => {
3148 if let Some(recorded_type) = workflow_type {
3149 if recorded_type != self.workflow_type {
3150 return Poll::Ready(Err(Error::NonDeterministicReplay(
3151 ReplayFailure::new(
3152 "recorded_command_detail_mismatch",
3153 Some(sequence),
3154 Some(format!("child workflow:{recorded_type}")),
3155 Some(format!("child workflow:{}", self.workflow_type)),
3156 "recorded child workflow type differs from the current workflow command",
3157 ),
3158 )));
3159 }
3160 }
3161 state.command_cursor += 1;
3162 if let Some(outcome) = outcome {
3163 return Poll::Ready(outcome.map_err(Error::ChildWorkflowFailed));
3164 }
3165 state.matched_recorded_pending = true;
3166 self.scheduled = true;
3167 self.matched_pending = true;
3168 return Poll::Pending;
3169 }
3170 other => {
3171 return Poll::Ready(Err(command_mismatch(
3172 &other,
3173 format!("child workflow:{}", self.workflow_type),
3174 )));
3175 }
3176 }
3177 }
3178
3179 if !self.scheduled {
3180 if self.options.task_queue.trim().is_empty() {
3181 return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3182 "task_queue must not be empty".to_string(),
3183 )));
3184 }
3185 for (name, value) in [
3186 (
3187 "execution_timeout_seconds",
3188 self.options.execution_timeout_seconds,
3189 ),
3190 ("run_timeout_seconds", self.options.run_timeout_seconds),
3191 ] {
3192 if value == Some(0) {
3193 return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(format!(
3194 "{name} must be at least 1"
3195 ))));
3196 }
3197 }
3198
3199 let args = match self.args.take().unwrap_or(Ok(Value::Null)) {
3200 Ok(args) => args,
3201 Err(error) => return Poll::Ready(Err(error)),
3202 };
3203 let arguments =
3204 match encode_value_envelope(&normalize_arguments(args), &state.payload_codec) {
3205 Ok(arguments) => arguments,
3206 Err(error) => return Poll::Ready(Err(error)),
3207 };
3208 let mut command = json!({
3209 "type": "start_child_workflow",
3210 "workflow_type": self.workflow_type,
3211 "queue": self.options.task_queue,
3212 "parent_close_policy": self.options.parent_close_policy.as_str(),
3213 "arguments": arguments,
3214 });
3215 let object = command
3216 .as_object_mut()
3217 .expect("child workflow command is always an object");
3218 if let Some(policy) = &self.options.retry_policy {
3219 let mut retry_policy = serde_json::Map::new();
3220 if let Some(max_attempts) = policy.max_attempts {
3221 if max_attempts == 0 {
3222 return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3223 "retry_policy.max_attempts must be at least 1".to_string(),
3224 )));
3225 }
3226 retry_policy.insert("max_attempts".to_string(), json!(max_attempts));
3227 }
3228 if !policy.backoff_seconds.is_empty() {
3229 retry_policy
3230 .insert("backoff_seconds".to_string(), json!(policy.backoff_seconds));
3231 }
3232 if !policy.non_retryable_error_types.is_empty() {
3233 retry_policy.insert(
3234 "non_retryable_error_types".to_string(),
3235 json!(policy.non_retryable_error_types),
3236 );
3237 }
3238 if retry_policy.is_empty() {
3239 return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3240 "retry_policy must configure at least one field".to_string(),
3241 )));
3242 }
3243 object.insert("retry_policy".to_string(), Value::Object(retry_policy));
3244 }
3245 if let Some(seconds) = self.options.execution_timeout_seconds {
3246 object.insert("execution_timeout_seconds".to_string(), json!(seconds));
3247 }
3248 if let Some(seconds) = self.options.run_timeout_seconds {
3249 object.insert("run_timeout_seconds".to_string(), json!(seconds));
3250 }
3251 state.commands.push(command);
3252 self.scheduled = true;
3253 }
3254
3255 Poll::Pending
3256 }
3257}
3258
3259fn command_mismatch(recorded: &RecordedCommand, actual: impl Into<String>) -> Error {
3260 Error::NonDeterministicReplay(ReplayFailure::new(
3261 "recorded_command_mismatch",
3262 Some(recorded.sequence()),
3263 Some(recorded.shape().to_string()),
3264 Some(actual.into()),
3265 "current workflow command does not match the recorded durable command sequence",
3266 ))
3267}
3268
3269pub struct SignalCall {
3270 ctx: WorkflowContext,
3271 signal_name: String,
3272 opened_wait: bool,
3273 matched_pending: bool,
3274}
3275
3276impl Future for SignalCall {
3277 type Output = Result<Vec<Value>>;
3278
3279 fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3280 if self.matched_pending {
3281 return Poll::Pending;
3282 }
3283
3284 let ctx = self.ctx.clone();
3285 let mut state = match ctx.state.lock() {
3286 Ok(state) => state,
3287 Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3288 };
3289
3290 let signals = match signal_values(
3291 &state.history,
3292 &self.signal_name,
3293 &state.payload_codec,
3294 state.resume_signal.as_ref(),
3295 ) {
3296 Ok(signals) => signals,
3297 Err(error) => return Poll::Ready(Err(error)),
3298 };
3299 let cursor = *state.signal_cursors.get(&self.signal_name).unwrap_or(&0);
3300
3301 if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3302 match recorded {
3303 RecordedCommand::SignalWait {
3304 sequence,
3305 signal_name,
3306 } => {
3307 if let Some(recorded_name) = signal_name {
3308 if recorded_name != self.signal_name {
3309 return Poll::Ready(Err(Error::NonDeterministicReplay(
3310 ReplayFailure::new(
3311 "recorded_command_detail_mismatch",
3312 Some(sequence),
3313 Some(format!("signal wait:{recorded_name}")),
3314 Some(format!("signal wait:{}", self.signal_name)),
3315 "recorded signal name differs from the current workflow command",
3316 ),
3317 )));
3318 }
3319 }
3320
3321 state.command_cursor += 1;
3322 if cursor < signals.len() {
3323 state
3324 .signal_cursors
3325 .insert(self.signal_name.clone(), cursor + 1);
3326 return Poll::Ready(Ok(signals[cursor].clone()));
3327 }
3328
3329 state.matched_recorded_pending = true;
3330 self.opened_wait = true;
3331 self.matched_pending = true;
3332 return Poll::Pending;
3333 }
3334 other => {
3335 return Poll::Ready(Err(command_mismatch(
3336 &other,
3337 format!("signal wait:{}", self.signal_name),
3338 )));
3339 }
3340 }
3341 }
3342
3343 if cursor < signals.len() {
3344 state
3345 .signal_cursors
3346 .insert(self.signal_name.clone(), cursor + 1);
3347 return Poll::Ready(Ok(signals[cursor].clone()));
3348 }
3349
3350 if !self.opened_wait {
3351 state.commands.push(json!({
3352 "type": "open_condition_wait",
3353 "condition_key": format!("signal:{}", self.signal_name)
3354 }));
3355 self.opened_wait = true;
3356 }
3357
3358 Poll::Pending
3359 }
3360}
3361
3362#[derive(Clone, Debug)]
3363pub struct ActivityContext {
3364 client: Client,
3365 pub task_id: String,
3366 pub activity_attempt_id: String,
3367 pub lease_owner: String,
3368 pub activity_type: String,
3369 pub attempt_number: u64,
3370 pub task_queue: String,
3371 pub worker_id: String,
3372}
3373
3374impl ActivityContext {
3375 pub async fn heartbeat<T: Serialize>(&self, details: T) -> Result<ActivityHeartbeatResponse> {
3376 self.client
3377 .heartbeat_activity_task(
3378 &self.task_id,
3379 &self.activity_attempt_id,
3380 &self.lease_owner,
3381 serde_json::to_value(details)?,
3382 )
3383 .await
3384 }
3385}
3386
3387fn decode_task_arguments(value: Option<&Value>, codec: &str) -> Result<Value> {
3388 match value {
3389 Some(value) => Ok(normalize_arguments(decode_wire_value(value, codec)?)),
3390 None => Ok(Value::Array(Vec::new())),
3391 }
3392}
3393
3394fn decode_resume_signal(task: &WorkflowTask) -> Result<Option<ResumeSignal>> {
3395 let Some(signal_name) = task
3396 .signal_name
3397 .as_deref()
3398 .filter(|value| !value.is_empty())
3399 else {
3400 return Ok(None);
3401 };
3402 let Some(arguments) = task.signal_arguments.as_ref() else {
3403 return Ok(None);
3404 };
3405
3406 let decoded = normalize_arguments(decode_wire_value(arguments, &task.payload_codec)?);
3407 let Value::Array(arguments) = decoded else {
3408 unreachable!("normalize_arguments always returns an array");
3409 };
3410
3411 Ok(Some(ResumeSignal {
3412 signal_id: task.workflow_signal_id.clone(),
3413 signal_name: signal_name.to_string(),
3414 arguments,
3415 }))
3416}
3417
3418fn normalize_arguments(value: Value) -> Value {
3419 match value {
3420 Value::Null => Value::Array(Vec::new()),
3421 Value::Array(_) => value,
3422 other => Value::Array(vec![other]),
3423 }
3424}
3425
3426fn recorded_commands(
3427 events: &[HistoryEvent],
3428 fallback_codec: &str,
3429 parent: WorkflowIdentity,
3430) -> Result<Vec<RecordedCommand>> {
3431 let mut events_by_sequence: BTreeMap<u64, Vec<&HistoryEvent>> = BTreeMap::new();
3432
3433 for event in events {
3434 let is_activity = matches!(
3435 event.event_type.as_str(),
3436 "ActivityScheduled"
3437 | "ActivityStarted"
3438 | "ActivityHeartbeatRecorded"
3439 | "ActivityRetryScheduled"
3440 | "ActivityCompleted"
3441 | "ActivityFailed"
3442 | "ActivityCancelled"
3443 | "ActivityTimedOut"
3444 );
3445 let is_workflow_timer = matches!(
3446 event.event_type.as_str(),
3447 "TimerScheduled" | "TimerCancelled" | "TimerFired"
3448 ) && !is_internal_timer_event(event);
3449 let is_child_workflow = matches!(
3450 event.event_type.as_str(),
3451 "ChildWorkflowScheduled"
3452 | "ChildRunCompleted"
3453 | "ChildRunFailed"
3454 | "ChildRunCancelled"
3455 | "ChildRunTerminated"
3456 );
3457 let is_signal_wait = is_recorded_signal_wait_event(event);
3458 if !is_activity && !is_workflow_timer && !is_child_workflow && !is_signal_wait {
3459 continue;
3460 }
3461
3462 let sequence = durable_event_sequence(event).ok_or_else(|| {
3463 Error::NonDeterministicReplay(ReplayFailure::new(
3464 "durable_command_sequence_missing",
3465 None,
3466 Some("positive workflow sequence".to_string()),
3467 Some(event.event_type.clone()),
3468 "durable command history event has no workflow sequence",
3469 ))
3470 })?;
3471 if sequence == 0 {
3472 return Err(Error::NonDeterministicReplay(ReplayFailure::new(
3473 "durable_command_sequence_invalid",
3474 Some(sequence),
3475 Some("positive workflow sequence".to_string()),
3476 Some(sequence.to_string()),
3477 "durable command history uses an invalid workflow sequence",
3478 )));
3479 }
3480 events_by_sequence.entry(sequence).or_default().push(event);
3481 }
3482
3483 events_by_sequence
3484 .into_iter()
3485 .map(|(sequence, sequence_events)| {
3486 let activity_events: Vec<_> = sequence_events
3487 .iter()
3488 .copied()
3489 .filter(|event| event.event_type.starts_with("Activity"))
3490 .collect();
3491 let timer_events: Vec<_> = sequence_events
3492 .iter()
3493 .copied()
3494 .filter(|event| event.event_type.starts_with("Timer"))
3495 .collect();
3496 let child_events: Vec<_> = sequence_events
3497 .iter()
3498 .copied()
3499 .filter(|event| {
3500 event.event_type == "ChildWorkflowScheduled"
3501 || event.event_type.starts_with("ChildRun")
3502 })
3503 .collect();
3504 let signal_wait_events: Vec<_> = sequence_events
3505 .iter()
3506 .copied()
3507 .filter(|event| is_recorded_signal_wait_event(event))
3508 .collect();
3509
3510 let command_kind_count = usize::from(!activity_events.is_empty())
3511 + usize::from(!timer_events.is_empty())
3512 + usize::from(!child_events.is_empty())
3513 + usize::from(!signal_wait_events.is_empty());
3514 if command_kind_count > 1 {
3515 let actual = [
3516 (!activity_events.is_empty()).then_some("activity"),
3517 (!timer_events.is_empty()).then_some("timer"),
3518 (!child_events.is_empty()).then_some("child workflow"),
3519 (!signal_wait_events.is_empty()).then_some("signal wait"),
3520 ]
3521 .into_iter()
3522 .flatten()
3523 .collect::<Vec<_>>()
3524 .join(" and ");
3525 return Err(invalid_recorded_history(
3526 "durable_command_sequence_collision",
3527 sequence,
3528 "one durable command kind",
3529 &actual,
3530 "one workflow sequence records more than one durable command kind",
3531 ));
3532 }
3533
3534 if !activity_events.is_empty() {
3535 let scheduled_count = activity_events
3536 .iter()
3537 .filter(|event| event.event_type == "ActivityScheduled")
3538 .count();
3539 if scheduled_count > 1 {
3540 return Err(invalid_recorded_history(
3541 "duplicate_activity_schedule",
3542 sequence,
3543 "at most one ActivityScheduled event",
3544 "multiple ActivityScheduled events",
3545 "activity history schedules more than one command at one workflow sequence",
3546 ));
3547 }
3548 let activity_type = activity_events.iter().find_map(|event| {
3549 event
3550 .payload
3551 .get("activity_type")
3552 .or_else(|| event.payload.get("activity_name"))
3553 .and_then(Value::as_str)
3554 .map(str::to_string)
3555 });
3556 if activity_events.iter().filter_map(|event| {
3557 event
3558 .payload
3559 .get("activity_type")
3560 .or_else(|| event.payload.get("activity_name"))
3561 .and_then(Value::as_str)
3562 }).any(|candidate| Some(candidate) != activity_type.as_deref()) {
3563 return Err(invalid_recorded_history(
3564 "activity_identity_mismatch",
3565 sequence,
3566 activity_type.as_deref().unwrap_or("one activity identity"),
3567 "conflicting activity identities",
3568 "activity lifecycle events at one workflow sequence disagree on identity",
3569 ));
3570 }
3571 let completed: Vec<_> = activity_events
3572 .iter()
3573 .copied()
3574 .filter(|event| event.event_type == "ActivityCompleted")
3575 .collect();
3576 if completed.len() > 1 {
3577 return Err(invalid_recorded_history(
3578 "duplicate_activity_completion",
3579 sequence,
3580 "one ActivityCompleted event",
3581 "multiple ActivityCompleted events",
3582 "activity history contains more than one completion for a workflow sequence",
3583 ));
3584 }
3585 let result = completed
3586 .first()
3587 .map(|event| {
3588 let codec = event
3589 .payload
3590 .get("payload_codec")
3591 .and_then(Value::as_str)
3592 .unwrap_or(fallback_codec);
3593 decode_wire_value(
3594 event.payload.get("result").unwrap_or(&Value::Null),
3595 codec,
3596 )
3597 })
3598 .transpose()?;
3599 return Ok(RecordedCommand::Activity {
3600 sequence,
3601 activity_type,
3602 result,
3603 });
3604 }
3605
3606 if !child_events.is_empty() {
3607 let scheduled: Vec<_> = child_events
3608 .iter()
3609 .copied()
3610 .filter(|event| event.event_type == "ChildWorkflowScheduled")
3611 .collect();
3612 if scheduled.len() != 1 {
3613 return Err(invalid_recorded_history(
3614 "child_workflow_schedule_missing_or_duplicate",
3615 sequence,
3616 "one ChildWorkflowScheduled event",
3617 &format!("{} ChildWorkflowScheduled events", scheduled.len()),
3618 "child workflow replay requires exactly one recorded schedule event",
3619 ));
3620 }
3621 let workflow_type = child_events.iter().find_map(|event| {
3622 event
3623 .payload
3624 .get("child_workflow_type")
3625 .or_else(|| event.payload.get("workflow_type"))
3626 .and_then(Value::as_str)
3627 .filter(|value| !value.is_empty())
3628 .map(str::to_string)
3629 });
3630 if child_events
3631 .iter()
3632 .filter_map(|event| {
3633 event
3634 .payload
3635 .get("child_workflow_type")
3636 .or_else(|| event.payload.get("workflow_type"))
3637 .and_then(Value::as_str)
3638 })
3639 .any(|candidate| Some(candidate) != workflow_type.as_deref())
3640 {
3641 return Err(invalid_recorded_history(
3642 "child_workflow_identity_mismatch",
3643 sequence,
3644 workflow_type
3645 .as_deref()
3646 .unwrap_or("one child workflow type"),
3647 "conflicting child workflow types",
3648 "child workflow lifecycle events at one sequence disagree on type",
3649 ));
3650 }
3651 let mut outcomes = child_workflow_outcomes(
3652 &child_events.iter().map(|event| (*event).clone()).collect::<Vec<_>>(),
3653 fallback_codec,
3654 parent.clone(),
3655 )?;
3656 if outcomes.len() > 1 {
3657 return Err(invalid_recorded_history(
3658 "duplicate_child_workflow_terminal_event",
3659 sequence,
3660 "at most one terminal child event",
3661 "multiple terminal child events",
3662 "child workflow history settles one command more than once",
3663 ));
3664 }
3665 return Ok(RecordedCommand::ChildWorkflow {
3666 sequence,
3667 workflow_type,
3668 outcome: outcomes.pop(),
3669 });
3670 }
3671
3672 if !signal_wait_events.is_empty() {
3673 let opened_count = signal_wait_events
3674 .iter()
3675 .filter(|event| {
3676 matches!(
3677 event.event_type.as_str(),
3678 "SignalWaitOpened" | "ConditionWaitOpened"
3679 )
3680 })
3681 .count();
3682 if opened_count > 1 {
3683 return Err(invalid_recorded_history(
3684 "duplicate_signal_wait_open",
3685 sequence,
3686 "at most one signal wait open event",
3687 "multiple signal wait open events",
3688 "signal history opens more than one durable wait at one workflow sequence",
3689 ));
3690 }
3691
3692 let signal_name = signal_wait_events
3693 .iter()
3694 .find_map(|event| recorded_signal_wait_name(event));
3695 if signal_wait_events
3696 .iter()
3697 .filter_map(|event| recorded_signal_wait_name(event))
3698 .any(|candidate| Some(candidate.as_str()) != signal_name.as_deref())
3699 {
3700 return Err(invalid_recorded_history(
3701 "signal_wait_identity_mismatch",
3702 sequence,
3703 signal_name.as_deref().unwrap_or("one signal name"),
3704 "conflicting signal names",
3705 "signal wait lifecycle events at one workflow sequence disagree on identity",
3706 ));
3707 }
3708 return Ok(RecordedCommand::SignalWait {
3709 sequence,
3710 signal_name,
3711 });
3712 }
3713
3714 let scheduled: Vec<_> = timer_events
3715 .iter()
3716 .copied()
3717 .filter(|event| event.event_type == "TimerScheduled")
3718 .collect();
3719 let fired: Vec<_> = timer_events
3720 .iter()
3721 .copied()
3722 .filter(|event| event.event_type == "TimerFired")
3723 .collect();
3724 if scheduled.len() != 1 {
3725 return Err(invalid_recorded_history(
3726 "timer_schedule_missing_or_duplicate",
3727 sequence,
3728 "one TimerScheduled event",
3729 &format!("{} TimerScheduled events", scheduled.len()),
3730 "timer replay requires exactly one recorded schedule event",
3731 ));
3732 }
3733 if fired.len() > 1 {
3734 return Err(invalid_recorded_history(
3735 "duplicate_timer_fire",
3736 sequence,
3737 "at most one TimerFired event",
3738 "multiple TimerFired events",
3739 "timer history contains more than one fire event for a workflow sequence",
3740 ));
3741 }
3742
3743 let scheduled = scheduled[0];
3744 let timer_id = required_history_string(scheduled, "timer_id", sequence)?;
3745 let delay_seconds = required_history_u64(scheduled, "delay_seconds", sequence)?;
3746 if let Some(fired) = fired.first() {
3747 let fired_timer_id = required_history_string(fired, "timer_id", sequence)?;
3748 if fired_timer_id != timer_id {
3749 return Err(invalid_recorded_history(
3750 "timer_identity_mismatch",
3751 sequence,
3752 &timer_id,
3753 &fired_timer_id,
3754 "TimerFired does not correspond to the recorded TimerScheduled event",
3755 ));
3756 }
3757 let fired_delay = required_history_u64(fired, "delay_seconds", sequence)?;
3758 if fired_delay != delay_seconds {
3759 return Err(invalid_recorded_history(
3760 "timer_history_delay_mismatch",
3761 sequence,
3762 &delay_seconds.to_string(),
3763 &fired_delay.to_string(),
3764 "TimerScheduled and TimerFired record different delays",
3765 ));
3766 }
3767 }
3768
3769 Ok(RecordedCommand::Timer {
3770 sequence,
3771 delay_seconds,
3772 fired: !fired.is_empty(),
3773 })
3774 })
3775 .collect()
3776}
3777
3778fn durable_event_sequence(event: &HistoryEvent) -> Option<u64> {
3779 event
3780 .payload
3781 .get("sequence")
3782 .or_else(|| event.payload.get("workflow_sequence"))
3783 .or_else(|| event.raw.get("sequence"))
3784 .or_else(|| event.raw.get("workflow_sequence"))
3785 .and_then(value_as_u64)
3786}
3787
3788fn is_internal_timer_event(event: &HistoryEvent) -> bool {
3789 matches!(
3790 event
3791 .payload
3792 .get("timer_kind")
3793 .or_else(|| event.raw.get("timer_kind"))
3794 .and_then(Value::as_str),
3795 Some("condition_timeout" | "signal_timeout")
3796 )
3797}
3798
3799fn recorded_signal_wait_name(event: &HistoryEvent) -> Option<String> {
3800 match event.event_type.as_str() {
3801 "SignalWaitOpened" | "SignalApplied" => event
3802 .payload
3803 .get("signal_name")
3804 .or_else(|| event.raw.get("signal_name"))
3805 .and_then(Value::as_str)
3806 .filter(|value| !value.is_empty())
3807 .map(str::to_string),
3808 "ConditionWaitOpened" | "ConditionWaitSatisfied" | "ConditionWaitTimedOut" => event
3809 .payload
3810 .get("condition_key")
3811 .or_else(|| event.raw.get("condition_key"))
3812 .and_then(Value::as_str)
3813 .and_then(|key| key.strip_prefix("signal:"))
3814 .filter(|value| !value.is_empty())
3815 .map(str::to_string),
3816 _ => None,
3817 }
3818}
3819
3820fn is_recorded_signal_wait_event(event: &HistoryEvent) -> bool {
3821 match event.event_type.as_str() {
3822 "SignalWaitOpened" | "SignalApplied" => true,
3823 "ConditionWaitOpened" | "ConditionWaitSatisfied" | "ConditionWaitTimedOut" => event
3824 .payload
3825 .get("condition_key")
3826 .or_else(|| event.raw.get("condition_key"))
3827 .and_then(Value::as_str)
3828 .is_some_and(|key| key.starts_with("signal:")),
3829 _ => false,
3830 }
3831}
3832
3833fn required_history_string(event: &HistoryEvent, field: &str, sequence: u64) -> Result<String> {
3834 event
3835 .payload
3836 .get(field)
3837 .and_then(Value::as_str)
3838 .filter(|value| !value.is_empty())
3839 .map(str::to_string)
3840 .ok_or_else(|| {
3841 invalid_recorded_history(
3842 "timer_history_field_missing",
3843 sequence,
3844 field,
3845 &event.event_type,
3846 "timer history is missing a required identity field",
3847 )
3848 })
3849}
3850
3851fn required_history_u64(event: &HistoryEvent, field: &str, sequence: u64) -> Result<u64> {
3852 event
3853 .payload
3854 .get(field)
3855 .and_then(value_as_u64)
3856 .ok_or_else(|| {
3857 invalid_recorded_history(
3858 "timer_history_field_missing",
3859 sequence,
3860 field,
3861 &event.event_type,
3862 "timer history is missing a required numeric field",
3863 )
3864 })
3865}
3866
3867fn invalid_recorded_history(
3868 reason: &str,
3869 sequence: u64,
3870 expected: &str,
3871 actual: &str,
3872 message: &str,
3873) -> Error {
3874 Error::NonDeterministicReplay(ReplayFailure::new(
3875 reason,
3876 Some(sequence),
3877 Some(expected.to_string()),
3878 Some(actual.to_string()),
3879 message,
3880 ))
3881}
3882
3883type ChildWorkflowOutcome = std::result::Result<ChildWorkflowResult, ChildWorkflowFailure>;
3884
3885fn child_workflow_outcomes(
3886 events: &[HistoryEvent],
3887 fallback_codec: &str,
3888 parent: WorkflowIdentity,
3889) -> Result<Vec<ChildWorkflowOutcome>> {
3890 let mut outcomes = Vec::new();
3891
3892 for event in events {
3893 let kind = match event.event_type.as_str() {
3894 "ChildRunCompleted" => None,
3895 "ChildRunFailed" => Some((
3896 ChildWorkflowFailureKind::Failed,
3897 "child_workflow",
3898 "child workflow failed",
3899 )),
3900 "ChildRunCancelled" => Some((
3901 ChildWorkflowFailureKind::Cancelled,
3902 "cancelled",
3903 "child workflow was cancelled",
3904 )),
3905 "ChildRunTerminated" => Some((
3906 ChildWorkflowFailureKind::Terminated,
3907 "terminated",
3908 "child workflow was terminated",
3909 )),
3910 _ => continue,
3911 };
3912 let payload = &event.payload;
3913 let child_workflow_id = payload_string(payload, "child_workflow_instance_id");
3914 let child_workflow_run_id = payload_string(payload, "child_workflow_run_id");
3915 let child_workflow_type = payload_string(payload, "child_workflow_type");
3916
3917 if let Some((kind, reason, fallback_message)) = kind {
3918 let exception = payload
3919 .get("exception")
3920 .filter(|value| !value.is_null())
3921 .cloned();
3922 let message = payload_string(payload, "message")
3923 .or_else(|| {
3924 exception
3925 .as_ref()
3926 .and_then(|value| payload_string(value, "message"))
3927 })
3928 .unwrap_or_else(|| fallback_message.to_string());
3929 let exception_type = payload_string(payload, "exception_type").or_else(|| {
3930 exception
3931 .as_ref()
3932 .and_then(|value| payload_string(value, "type"))
3933 });
3934 let exception_class = payload_string(payload, "exception_class").or_else(|| {
3935 exception
3936 .as_ref()
3937 .and_then(|value| payload_string(value, "class"))
3938 });
3939 outcomes.push(Err(ChildWorkflowFailure {
3940 kind,
3941 reason: reason.to_string(),
3942 message,
3943 parent_workflow_id: parent.workflow_id.clone(),
3944 parent_workflow_run_id: parent.run_id.clone(),
3945 child_workflow_id,
3946 child_workflow_run_id,
3947 child_workflow_type,
3948 failure_id: payload_string(payload, "failure_id"),
3949 failure_category: payload_string(payload, "failure_category"),
3950 exception_type,
3951 exception_class,
3952 non_retryable: payload
3953 .get("non_retryable")
3954 .and_then(Value::as_bool)
3955 .unwrap_or(false),
3956 code: payload
3957 .get("code")
3958 .filter(|value| !value.is_null())
3959 .cloned(),
3960 exception,
3961 }));
3962 continue;
3963 }
3964
3965 let codec = payload
3966 .get("payload_codec")
3967 .and_then(Value::as_str)
3968 .unwrap_or(fallback_codec);
3969 let result = payload
3970 .get("result")
3971 .or_else(|| payload.get("output"))
3972 .unwrap_or(&Value::Null);
3973 outcomes.push(Ok(ChildWorkflowResult {
3974 parent: parent.clone(),
3975 child: WorkflowIdentity {
3976 workflow_id: child_workflow_id,
3977 run_id: child_workflow_run_id,
3978 },
3979 child_workflow_type,
3980 result: decode_wire_value(result, codec)?,
3981 }));
3982 }
3983
3984 Ok(outcomes)
3985}
3986
3987fn payload_string(payload: &Value, key: &str) -> Option<String> {
3988 payload
3989 .get(key)
3990 .and_then(Value::as_str)
3991 .filter(|value| !value.is_empty())
3992 .map(str::to_string)
3993}
3994
3995fn workflow_failure_command(error: &Error) -> Value {
3996 let (exception_type, exception_class, properties) = match error {
3997 Error::ChildWorkflowFailed(failure) => (
3998 match failure.kind {
3999 ChildWorkflowFailureKind::Failed => "ChildWorkflowFailed",
4000 ChildWorkflowFailureKind::Cancelled => "ChildWorkflowCancelled",
4001 ChildWorkflowFailureKind::Terminated => "ChildWorkflowTerminated",
4002 },
4003 "durable_workflow::ChildWorkflowFailure",
4004 json!({
4005 "reason": failure.reason,
4006 "parent_workflow_id": failure.parent_workflow_id,
4007 "parent_workflow_run_id": failure.parent_workflow_run_id,
4008 "child_workflow_id": failure.child_workflow_id,
4009 "child_workflow_run_id": failure.child_workflow_run_id,
4010 "child_workflow_type": failure.child_workflow_type,
4011 "failure_id": failure.failure_id,
4012 "failure_category": failure.failure_category,
4013 "child_exception_type": failure.exception_type,
4014 "child_exception_class": failure.exception_class,
4015 "child_non_retryable": failure.non_retryable,
4016 "child_code": failure.code,
4017 "child_exception": failure.exception,
4018 }),
4019 ),
4020 Error::NonDeterministicReplay(_) => (
4021 "NonDeterministicReplay",
4022 "durable_workflow::Error",
4023 Value::Null,
4024 ),
4025 _ => ("RustWorkflowError", "durable_workflow::Error", Value::Null),
4026 };
4027 let non_retryable = match error {
4028 Error::ChildWorkflowFailed(failure) => failure.non_retryable,
4029 Error::NonDeterministicReplay(_) => true,
4030 _ => false,
4031 };
4032
4033 json!({
4034 "type": "fail_workflow",
4035 "message": error.to_string(),
4036 "exception_type": exception_type,
4037 "exception_class": exception_class,
4038 "non_retryable": non_retryable,
4039 "exception": {
4040 "type": exception_type,
4041 "class": exception_class,
4042 "message": error.to_string(),
4043 "properties": properties,
4044 }
4045 })
4046}
4047
4048fn signal_values(
4049 events: &[HistoryEvent],
4050 signal_name: &str,
4051 fallback_codec: &str,
4052 resume_signal: Option<&ResumeSignal>,
4053) -> Result<Vec<Vec<Value>>> {
4054 let mut signals = Vec::new();
4055
4056 for event in events {
4057 if event.event_type != "SignalApplied" && event.event_type != "SignalReceived" {
4058 continue;
4059 }
4060
4061 if event.payload.get("signal_name").and_then(Value::as_str) != Some(signal_name) {
4062 continue;
4063 }
4064
4065 let codec = event
4066 .payload
4067 .get("payload_codec")
4068 .and_then(Value::as_str)
4069 .unwrap_or(fallback_codec);
4070 let raw = event
4071 .payload
4072 .get("value")
4073 .or_else(|| event.payload.get("input"))
4074 .or_else(|| event.payload.get("arguments"));
4075 let decoded = match raw.filter(|value| !value.is_null()) {
4076 Some(value) => decode_wire_value(value, codec)?,
4077 None => resume_signal
4078 .filter(|signal| resume_signal_matches_event(signal, event, signal_name))
4079 .map(|signal| Value::Array(signal.arguments.clone()))
4080 .unwrap_or_else(|| Value::Array(Vec::new())),
4081 };
4082 let args = match normalize_arguments(decoded) {
4083 Value::Array(values) => values,
4084 _ => unreachable!("normalize_arguments always returns an array"),
4085 };
4086 signals.push(args);
4087 }
4088
4089 Ok(signals)
4090}
4091
4092fn hydrate_query_history_from_export(task: &mut QueryTask) -> Result<()> {
4093 let Some(export_events) = task
4094 .history_export
4095 .as_ref()
4096 .and_then(|export| export.get("history_events"))
4097 .and_then(Value::as_array)
4098 else {
4099 return Ok(());
4100 };
4101
4102 if export_events.len() > task.history_events.len() {
4103 task.history_events = serde_json::from_value(Value::Array(export_events.clone()))?;
4104 }
4105
4106 Ok(())
4107}
4108
4109fn enrich_query_history_from_export(task: &mut QueryTask) -> Result<()> {
4110 let Some(export) = task.history_export.as_ref() else {
4111 return Ok(());
4112 };
4113 let signals = export
4114 .get("signals")
4115 .and_then(Value::as_array)
4116 .cloned()
4117 .unwrap_or_default();
4118 let activities = export
4119 .get("activities")
4120 .and_then(Value::as_array)
4121 .cloned()
4122 .unwrap_or_default();
4123 let export_codec = export
4124 .get("payloads")
4125 .and_then(|payloads| payloads.get("codec"))
4126 .and_then(Value::as_str)
4127 .unwrap_or(&task.payload_codec)
4128 .to_string();
4129 let mut signal_name_offsets: HashMap<String, usize> = HashMap::new();
4130
4131 for event in &mut task.history_events {
4132 if event.event_type == "ActivityCompleted" {
4133 let sequence = event
4134 .payload
4135 .get("sequence")
4136 .or_else(|| event.payload.get("workflow_sequence"))
4137 .and_then(value_as_u64);
4138 let Some(activity) = sequence.and_then(|sequence| {
4139 activities.iter().find(|activity| {
4140 activity.get("sequence").and_then(value_as_u64) == Some(sequence)
4141 })
4142 }) else {
4143 continue;
4144 };
4145 let Some(payload) = event.payload.as_object_mut() else {
4146 continue;
4147 };
4148 if missing_payload(payload.get("result")) {
4149 if let Some(result) = activity
4150 .get("result")
4151 .filter(|value| !missing_payload(Some(value)))
4152 {
4153 payload.insert("result".to_string(), result.clone());
4154 }
4155 }
4156 for field in ["payload_codec", "activity_type"] {
4157 if payload
4158 .get(field)
4159 .and_then(Value::as_str)
4160 .unwrap_or_default()
4161 .is_empty()
4162 {
4163 if let Some(value) = activity.get(field) {
4164 payload.insert(field.to_string(), value.clone());
4165 }
4166 }
4167 }
4168 continue;
4169 }
4170
4171 if event.event_type != "SignalReceived" && event.event_type != "SignalApplied" {
4172 continue;
4173 }
4174 let signal_id = event.payload.get("signal_id").and_then(Value::as_str);
4175 let command_id = event
4176 .payload
4177 .get("workflow_command_id")
4178 .or_else(|| event.raw.get("workflow_command_id"))
4179 .and_then(Value::as_str);
4180 let signal_name = event
4181 .payload
4182 .get("signal_name")
4183 .and_then(Value::as_str)
4184 .unwrap_or_default()
4185 .to_string();
4186 let matched = signals
4187 .iter()
4188 .find(|signal| {
4189 signal_id.is_some() && signal.get("id").and_then(Value::as_str) == signal_id
4190 })
4191 .or_else(|| {
4192 signals.iter().find(|signal| {
4193 command_id.is_some()
4194 && signal.get("command_id").and_then(Value::as_str) == command_id
4195 })
4196 })
4197 .or_else(|| {
4198 let offset = signal_name_offsets.entry(signal_name.clone()).or_default();
4199 let signal = signals
4200 .iter()
4201 .filter(|signal| {
4202 signal.get("name").and_then(Value::as_str) == Some(signal_name.as_str())
4203 })
4204 .nth(*offset);
4205 if signal.is_some() {
4206 *offset += 1;
4207 }
4208 signal
4209 });
4210 let Some(signal) = matched else {
4211 continue;
4212 };
4213 let signal_codec = signal
4214 .get("payload_codec")
4215 .and_then(Value::as_str)
4216 .unwrap_or(&export_codec);
4217 let Some(payload) = event.payload.as_object_mut() else {
4218 continue;
4219 };
4220 if missing_payload(payload.get("arguments")) {
4221 if let Some(arguments) = signal
4222 .get("arguments")
4223 .filter(|value| !missing_payload(Some(value)))
4224 {
4225 let envelope = match arguments {
4226 Value::String(blob) => json!({"codec": signal_codec, "blob": blob}),
4227 other => other.clone(),
4228 };
4229 payload.insert("arguments".to_string(), envelope);
4230 }
4231 }
4232 if payload
4233 .get("payload_codec")
4234 .and_then(Value::as_str)
4235 .unwrap_or_default()
4236 .is_empty()
4237 {
4238 payload.insert("payload_codec".to_string(), json!(signal_codec));
4239 }
4240 }
4241
4242 Ok(())
4243}
4244
4245fn missing_payload(value: Option<&Value>) -> bool {
4246 match value {
4247 None | Some(Value::Null) => true,
4248 Some(Value::String(value)) => value.is_empty(),
4249 Some(_) => false,
4250 }
4251}
4252
4253fn query_signal_events(task: &QueryTask) -> Result<Vec<QuerySignal>> {
4254 let export_signals = task
4255 .history_export
4256 .as_ref()
4257 .and_then(|export| export.get("signals"))
4258 .and_then(Value::as_array)
4259 .cloned()
4260 .unwrap_or_default();
4261 let export_codec = task
4262 .history_export
4263 .as_ref()
4264 .and_then(|export| export.get("payloads"))
4265 .and_then(|payloads| payloads.get("codec"))
4266 .and_then(Value::as_str)
4267 .unwrap_or(&task.payload_codec);
4268 let mut name_offsets: HashMap<String, usize> = HashMap::new();
4269 let mut signals = Vec::new();
4270
4271 for event in &task.history_events {
4272 if event.event_type != "SignalApplied" && event.event_type != "SignalReceived" {
4273 continue;
4274 }
4275
4276 let name = event
4277 .payload
4278 .get("signal_name")
4279 .and_then(Value::as_str)
4280 .unwrap_or_default();
4281 if name.is_empty() {
4282 continue;
4283 }
4284 let signal_id = event.payload.get("signal_id").and_then(Value::as_str);
4285 let command_id = event
4286 .payload
4287 .get("workflow_command_id")
4288 .or_else(|| event.raw.get("workflow_command_id"))
4289 .and_then(Value::as_str);
4290 let matched_export = export_signals
4291 .iter()
4292 .find(|candidate| {
4293 signal_id.is_some() && candidate.get("id").and_then(Value::as_str) == signal_id
4294 })
4295 .or_else(|| {
4296 export_signals.iter().find(|candidate| {
4297 command_id.is_some()
4298 && candidate.get("command_id").and_then(Value::as_str) == command_id
4299 })
4300 })
4301 .or_else(|| {
4302 let offset = name_offsets.entry(name.to_string()).or_default();
4303 let candidate = export_signals
4304 .iter()
4305 .filter(|candidate| candidate.get("name").and_then(Value::as_str) == Some(name))
4306 .nth(*offset);
4307 if candidate.is_some() {
4308 *offset += 1;
4309 }
4310 candidate
4311 });
4312 let codec = event
4313 .payload
4314 .get("payload_codec")
4315 .and_then(Value::as_str)
4316 .or_else(|| {
4317 matched_export
4318 .and_then(|signal| signal.get("payload_codec"))
4319 .and_then(Value::as_str)
4320 })
4321 .unwrap_or(export_codec);
4322 let raw_arguments = event
4323 .payload
4324 .get("value")
4325 .or_else(|| event.payload.get("input"))
4326 .or_else(|| event.payload.get("arguments"))
4327 .filter(|value| !value.is_null())
4328 .or_else(|| matched_export.and_then(|signal| signal.get("arguments")));
4329 let arguments = decode_query_signal_arguments(raw_arguments, codec)?;
4330 let workflow_sequence = event
4331 .payload
4332 .get("workflow_sequence")
4333 .and_then(value_as_u64)
4334 .or_else(|| {
4335 matched_export
4336 .and_then(|signal| signal.get("workflow_sequence"))
4337 .and_then(value_as_u64)
4338 });
4339
4340 signals.push(QuerySignal {
4341 id: signal_id.map(str::to_string).or_else(|| {
4342 matched_export
4343 .and_then(|signal| signal.get("id"))
4344 .and_then(Value::as_str)
4345 .map(str::to_string)
4346 }),
4347 name: name.to_string(),
4348 arguments,
4349 workflow_sequence,
4350 });
4351 }
4352
4353 if signals.is_empty() {
4354 for signal in export_signals {
4355 if signal.get("status").and_then(Value::as_str) == Some("rejected") {
4356 continue;
4357 }
4358 let Some(name) = signal.get("name").and_then(Value::as_str) else {
4359 continue;
4360 };
4361 let codec = signal
4362 .get("payload_codec")
4363 .and_then(Value::as_str)
4364 .unwrap_or(export_codec);
4365 let arguments = decode_query_signal_arguments(signal.get("arguments"), codec)?;
4366 signals.push(QuerySignal {
4367 id: signal.get("id").and_then(Value::as_str).map(str::to_string),
4368 name: name.to_string(),
4369 arguments,
4370 workflow_sequence: signal.get("workflow_sequence").and_then(value_as_u64),
4371 });
4372 }
4373 signals.sort_by_key(|signal| signal.workflow_sequence.unwrap_or(u64::MAX));
4374 }
4375
4376 Ok(signals)
4377}
4378
4379fn decode_query_signal_arguments(raw: Option<&Value>, codec: &str) -> Result<Vec<Value>> {
4380 let decoded = match raw.filter(|value| !value.is_null()) {
4381 Some(value) => decode_wire_value(value, codec)?,
4382 None => Value::Array(Vec::new()),
4383 };
4384 let Value::Array(arguments) = normalize_arguments(decoded) else {
4385 unreachable!("normalize_arguments always returns an array");
4386 };
4387 Ok(arguments)
4388}
4389
4390fn value_as_u64(value: &Value) -> Option<u64> {
4391 value
4392 .as_u64()
4393 .or_else(|| value.as_str().and_then(|value| value.parse().ok()))
4394}
4395
4396fn resume_signal_matches_event(
4397 resume_signal: &ResumeSignal,
4398 event: &HistoryEvent,
4399 signal_name: &str,
4400) -> bool {
4401 if resume_signal.signal_name != signal_name {
4402 return false;
4403 }
4404
4405 match (
4406 resume_signal.signal_id.as_deref(),
4407 event.payload.get("signal_id").and_then(Value::as_str),
4408 ) {
4409 (Some(resume_id), Some(event_id)) => resume_id == event_id,
4410 _ => true,
4411 }
4412}
4413
4414#[cfg(test)]
4415mod tests {
4416 use super::*;
4417 use std::{
4418 io::{Read, Write},
4419 net::{SocketAddr, TcpListener, TcpStream},
4420 thread,
4421 };
4422
4423 #[derive(Clone, Debug, Default, PartialEq)]
4424 struct ReplayCounterState {
4425 loaded: Option<String>,
4426 count: i64,
4427 finished: bool,
4428 }
4429
4430 fn replay_counter_worker() -> Worker {
4431 let client = Client::new("http://127.0.0.1:8080").expect("client");
4432 let mut worker = Worker::new(client, "rust-workers");
4433 worker.register_replayed_workflow(
4434 "replay-counter",
4435 ReplayCounterState::default,
4436 |ctx, _input, state| async move {
4437 let loaded = ctx.activity("load-counter", json!([])).await?;
4438 state.update(|current| {
4439 current.loaded = loaded.as_str().map(str::to_string);
4440 })?;
4441 for _ in 0..2 {
4442 let signal = ctx.wait_signal("increment").await?;
4443 let amount = signal.first().and_then(Value::as_i64).unwrap_or_default();
4444 state.update(|current| current.count += amount)?;
4445 }
4446 state.update(|current| current.finished = true)?;
4447 state.read(|current| Ok(json!(current.count)))?
4448 },
4449 );
4450 worker.register_replayed_query::<ReplayCounterState, _, _>(
4451 "replay-counter",
4452 "current",
4453 |_ctx, state, _args| async move {
4454 Ok(json!({
4455 "loaded": state.loaded,
4456 "count": state.count,
4457 "finished": state.finished,
4458 }))
4459 },
4460 );
4461 worker.register_replayed_query::<ReplayCounterState, _, _>(
4462 "replay-counter",
4463 "detached-mutation",
4464 |_ctx, state, _args| async move {
4465 let mut detached = (*state).clone();
4466 detached.count = 999;
4467 Ok(json!(detached.count))
4468 },
4469 );
4470 worker.register_replayed_query::<ReplayCounterState, _, _>(
4471 "replay-counter",
4472 "failed-mutation",
4473 |_ctx, state, _args| async move {
4474 let mut detached = (*state).clone();
4475 detached.count = 999;
4476 Err(Error::WorkerLoop("query refused".to_string()))
4477 },
4478 );
4479 worker
4480 }
4481
4482 fn replay_counter_query(
4483 query_name: &str,
4484 history_events: Value,
4485 run_status: &str,
4486 ) -> QueryTask {
4487 serde_json::from_value(json!({
4488 "query_task_id": format!("query-{query_name}"),
4489 "workflow_type": "replay-counter",
4490 "query_name": query_name,
4491 "payload_codec": "json",
4492 "workflow_arguments": {"codec": "json", "blob": "[]"},
4493 "query_arguments": {"codec": "json", "blob": "[]"},
4494 "history_events": history_events,
4495 "run_status": run_status,
4496 }))
4497 .expect("query task")
4498 }
4499
4500 fn workflow_context(history: Vec<HistoryEvent>) -> WorkflowContext {
4501 WorkflowContext {
4502 state: Arc::new(Mutex::new(
4503 WorkflowState::new_with_identity(
4504 history,
4505 None,
4506 None,
4507 "rust-workers".to_string(),
4508 JSON_CODEC.to_string(),
4509 None,
4510 )
4511 .expect("valid workflow history"),
4512 )),
4513 }
4514 }
4515
4516 fn history_event(event_type: &str, payload: Value) -> HistoryEvent {
4517 HistoryEvent {
4518 event_type: event_type.to_string(),
4519 payload,
4520 raw: HashMap::new(),
4521 }
4522 }
4523
4524 #[test]
4525 fn avro_generic_wrapper_round_trips_json_values() {
4526 let value = json!({"greeting": "hello", "count": 3, "ok": true});
4527 let envelope = PayloadEnvelope::avro(&value).expect("encode");
4528 assert_eq!(envelope.codec, DEFAULT_CODEC);
4529 assert_eq!(decode_payload::<Value>(&envelope).expect("decode"), value);
4530 }
4531
4532 #[test]
4533 fn json_codec_remains_plain_json() {
4534 let value = json!({"greeting": "hello", "count": 3, "ok": true});
4535 let envelope = PayloadEnvelope::json(&value).expect("encode");
4536
4537 assert_eq!(envelope.codec, JSON_CODEC);
4538 assert_eq!(envelope.blob, serde_json::to_string(&value).expect("json"));
4539 assert_eq!(decode_payload::<Value>(&envelope).expect("decode"), value);
4540 }
4541
4542 #[test]
4543 fn typed_avro_payload_without_schema_context_keeps_diagnostic() {
4544 let envelope = PayloadEnvelope {
4545 codec: DEFAULT_CODEC.to_string(),
4546 blob: BASE64.encode([0x01]),
4547 };
4548
4549 let error = decode_payload::<Value>(&envelope).expect_err("typed payload must fail");
4550 assert_eq!(
4551 error.to_string(),
4552 "codec error: typed avro payloads require a schema context; v1 supports the generic wrapper"
4553 );
4554 }
4555
4556 #[test]
4557 fn workflow_context_schedules_activity_until_completion_is_in_history() {
4558 let ctx = WorkflowContext {
4559 state: Arc::new(Mutex::new(
4560 WorkflowState::new_with_identity(
4561 Vec::new(),
4562 Some("wf-parent".to_string()),
4563 Some("run-parent".to_string()),
4564 "rust-workers".to_string(),
4565 DEFAULT_CODEC.to_string(),
4566 None,
4567 )
4568 .expect("workflow state"),
4569 )),
4570 };
4571
4572 let mut call = Box::pin(ctx.activity("hello.activity", json!(["Ada"])));
4573 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4574 assert!(matches!(
4575 call.as_mut().poll(&mut task_context),
4576 Poll::Pending
4577 ));
4578
4579 let commands = ctx.take_commands().expect("commands");
4580 assert_eq!(commands[0]["type"], "schedule_activity");
4581 assert_eq!(commands[0]["activity_type"], "hello.activity");
4582 }
4583
4584 #[test]
4585 fn workflow_sleep_emits_one_durable_timer_and_rounds_up() {
4586 let ctx = workflow_context(Vec::new());
4587 let mut sleep = Box::pin(ctx.sleep(Duration::from_millis(1_001)));
4588 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4589
4590 assert!(matches!(
4591 sleep.as_mut().poll(&mut task_context),
4592 Poll::Pending
4593 ));
4594 assert!(matches!(
4595 sleep.as_mut().poll(&mut task_context),
4596 Poll::Pending
4597 ));
4598
4599 let commands = ctx.take_commands().expect("timer command");
4600 assert_eq!(
4601 commands,
4602 vec![json!({
4603 "type": "start_timer",
4604 "delay_seconds": 2,
4605 })]
4606 );
4607 }
4608
4609 #[test]
4610 fn workflow_sleep_replays_matching_schedule_and_fire_without_a_command() {
4611 let history = vec![
4612 history_event(
4613 "TimerScheduled",
4614 json!({
4615 "sequence": 1,
4616 "timer_id": "timer-1",
4617 "delay_seconds": 5,
4618 "fire_at": "2026-07-11T12:00:05Z",
4619 }),
4620 ),
4621 history_event(
4622 "TimerFired",
4623 json!({
4624 "sequence": 1,
4625 "timer_id": "timer-1",
4626 "delay_seconds": 5,
4627 "fire_at": "2026-07-11T12:00:05Z",
4628 "fired_at": "2026-07-11T12:00:05Z",
4629 }),
4630 ),
4631 ];
4632
4633 for _restart in 0..2 {
4634 let ctx = workflow_context(history.clone());
4635 let mut sleep = Box::pin(ctx.sleep(Duration::from_secs(5)));
4636 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4637 assert!(matches!(
4638 sleep.as_mut().poll(&mut task_context),
4639 Poll::Ready(Ok(()))
4640 ));
4641 assert!(ctx.take_commands().expect("commands").is_empty());
4642 ctx.ensure_history_consumed().expect("history consumed");
4643 }
4644 }
4645
4646 #[test]
4647 fn workflow_sleep_rejects_changed_delay_during_replay() {
4648 let ctx = workflow_context(vec![
4649 history_event(
4650 "TimerScheduled",
4651 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4652 ),
4653 history_event(
4654 "TimerFired",
4655 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4656 ),
4657 ]);
4658 let mut sleep = Box::pin(ctx.sleep(Duration::from_secs(500)));
4659 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4660
4661 let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4662 sleep.as_mut().poll(&mut task_context)
4663 else {
4664 panic!("changed timer delay must be rejected");
4665 };
4666 assert_eq!(failure.reason, "timer_delay_mismatch");
4667 assert_eq!(failure.sequence, Some(1));
4668 }
4669
4670 #[test]
4671 fn workflow_history_rejects_unpaired_or_mismatched_timer_events() {
4672 let lone_fire = WorkflowState::new(
4673 vec![history_event(
4674 "TimerFired",
4675 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4676 )],
4677 "rust-workers".to_string(),
4678 JSON_CODEC.to_string(),
4679 None,
4680 )
4681 .expect_err("TimerFired requires TimerScheduled");
4682 assert!(matches!(
4683 lone_fire,
4684 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4685 if reason == "timer_schedule_missing_or_duplicate"
4686 ));
4687
4688 let wrong_identity = WorkflowState::new(
4689 vec![
4690 history_event(
4691 "TimerScheduled",
4692 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4693 ),
4694 history_event(
4695 "TimerFired",
4696 json!({"sequence": 1, "timer_id": "timer-2", "delay_seconds": 5}),
4697 ),
4698 ],
4699 "rust-workers".to_string(),
4700 JSON_CODEC.to_string(),
4701 None,
4702 )
4703 .expect_err("fire must match scheduled timer identity");
4704 assert!(matches!(
4705 wrong_identity,
4706 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4707 if reason == "timer_identity_mismatch"
4708 ));
4709
4710 let duplicate_fire = WorkflowState::new(
4711 vec![
4712 history_event(
4713 "TimerScheduled",
4714 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4715 ),
4716 history_event(
4717 "TimerFired",
4718 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4719 ),
4720 history_event(
4721 "TimerFired",
4722 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4723 ),
4724 ],
4725 "rust-workers".to_string(),
4726 JSON_CODEC.to_string(),
4727 None,
4728 )
4729 .expect_err("a durable timer cannot fire twice");
4730 assert!(matches!(
4731 duplicate_fire,
4732 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4733 if reason == "duplicate_timer_fire"
4734 ));
4735
4736 let wrong_fired_delay = WorkflowState::new(
4737 vec![
4738 history_event(
4739 "TimerScheduled",
4740 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4741 ),
4742 history_event(
4743 "TimerFired",
4744 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 6}),
4745 ),
4746 ],
4747 "rust-workers".to_string(),
4748 JSON_CODEC.to_string(),
4749 None,
4750 )
4751 .expect_err("timer schedule and fire delays must agree");
4752 assert!(matches!(
4753 wrong_fired_delay,
4754 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4755 if reason == "timer_history_delay_mismatch"
4756 ));
4757 }
4758
4759 #[test]
4760 fn replay_rejects_activity_moved_before_recorded_timer() {
4761 let ctx = workflow_context(vec![
4762 history_event(
4763 "TimerScheduled",
4764 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4765 ),
4766 history_event(
4767 "TimerFired",
4768 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4769 ),
4770 history_event(
4771 "ActivityCompleted",
4772 json!({
4773 "sequence": 2,
4774 "activity_type": "after-timer",
4775 "payload_codec": "json",
4776 "result": {"codec": "json", "blob": "\"done\""},
4777 }),
4778 ),
4779 ]);
4780 let mut activity = Box::pin(ctx.activity("after-timer", json!([])));
4781 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4782
4783 let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4784 activity.as_mut().poll(&mut task_context)
4785 else {
4786 panic!("reordered durable command must be rejected");
4787 };
4788 assert_eq!(failure.reason, "recorded_command_mismatch");
4789 assert_eq!(failure.sequence, Some(1));
4790 assert_eq!(failure.expected.as_deref(), Some("timer"));
4791 assert_eq!(failure.actual.as_deref(), Some("activity:after-timer"));
4792 }
4793
4794 #[test]
4795 fn replay_orders_signal_waits_and_timers_in_one_command_stream() {
4796 let signal_then_timer = vec![
4797 history_event(
4798 "ConditionWaitOpened",
4799 json!({"sequence": 1, "condition_key": "signal:go"}),
4800 ),
4801 history_event(
4802 "SignalReceived",
4803 json!({
4804 "signal_name": "go",
4805 "arguments": ["now"],
4806 }),
4807 ),
4808 history_event(
4809 "TimerScheduled",
4810 json!({"sequence": 2, "timer_id": "timer-2", "delay_seconds": 5}),
4811 ),
4812 history_event(
4813 "TimerFired",
4814 json!({"sequence": 2, "timer_id": "timer-2", "delay_seconds": 5}),
4815 ),
4816 ];
4817
4818 let ctx = workflow_context(signal_then_timer.clone());
4819 let mut signal = Box::pin(ctx.wait_signal("go"));
4820 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4821 assert!(matches!(
4822 signal.as_mut().poll(&mut task_context),
4823 Poll::Ready(Ok(arguments)) if arguments == vec![json!("now")]
4824 ));
4825 let mut timer = Box::pin(ctx.sleep(Duration::from_secs(5)));
4826 assert!(matches!(
4827 timer.as_mut().poll(&mut task_context),
4828 Poll::Ready(Ok(()))
4829 ));
4830 ctx.ensure_history_consumed()
4831 .expect("signal and timer history consumed in order");
4832
4833 let reordered = workflow_context(signal_then_timer);
4834 let mut timer_first = Box::pin(reordered.sleep(Duration::from_secs(5)));
4835 let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4836 timer_first.as_mut().poll(&mut task_context)
4837 else {
4838 panic!("timer cannot consume signal-wait-first history");
4839 };
4840 assert_eq!(failure.reason, "recorded_command_mismatch");
4841 assert_eq!(failure.sequence, Some(1));
4842 assert_eq!(failure.expected.as_deref(), Some("signal wait"));
4843
4844 let timer_then_signal = vec![
4845 history_event(
4846 "TimerScheduled",
4847 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4848 ),
4849 history_event(
4850 "TimerFired",
4851 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4852 ),
4853 history_event(
4854 "ConditionWaitOpened",
4855 json!({"sequence": 2, "condition_key": "signal:go"}),
4856 ),
4857 history_event(
4858 "SignalReceived",
4859 json!({"signal_name": "go", "arguments": []}),
4860 ),
4861 ];
4862 let reordered = workflow_context(timer_then_signal);
4863 let mut signal_first = Box::pin(reordered.wait_signal("go"));
4864 let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4865 signal_first.as_mut().poll(&mut task_context)
4866 else {
4867 panic!("signal wait cannot consume timer-first history");
4868 };
4869 assert_eq!(failure.reason, "recorded_command_mismatch");
4870 assert_eq!(failure.sequence, Some(1));
4871 assert_eq!(failure.expected.as_deref(), Some("timer"));
4872 }
4873
4874 #[test]
4875 fn workflow_history_rejects_duplicate_or_colliding_command_sequences() {
4876 let duplicate_timer = WorkflowState::new(
4877 vec![
4878 history_event(
4879 "TimerScheduled",
4880 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4881 ),
4882 history_event(
4883 "TimerScheduled",
4884 json!({"sequence": 1, "timer_id": "timer-2", "delay_seconds": 5}),
4885 ),
4886 ],
4887 "rust-workers".to_string(),
4888 JSON_CODEC.to_string(),
4889 None,
4890 )
4891 .expect_err("one workflow sequence cannot schedule two timers");
4892 assert!(matches!(
4893 duplicate_timer,
4894 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4895 if reason == "timer_schedule_missing_or_duplicate"
4896 ));
4897
4898 let colliding_kinds = WorkflowState::new(
4899 vec![
4900 history_event(
4901 "TimerScheduled",
4902 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4903 ),
4904 history_event(
4905 "ActivityCompleted",
4906 json!({"sequence": 1, "activity_type": "same-sequence"}),
4907 ),
4908 ],
4909 "rust-workers".to_string(),
4910 JSON_CODEC.to_string(),
4911 None,
4912 )
4913 .expect_err("one workflow sequence cannot identify two command kinds");
4914 assert!(matches!(
4915 colliding_kinds,
4916 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4917 if reason == "durable_command_sequence_collision"
4918 ));
4919
4920 let duplicate_signal_wait = WorkflowState::new(
4921 vec![
4922 history_event(
4923 "SignalWaitOpened",
4924 json!({"sequence": 1, "signal_name": "go"}),
4925 ),
4926 history_event(
4927 "SignalWaitOpened",
4928 json!({"sequence": 1, "signal_name": "go"}),
4929 ),
4930 ],
4931 "rust-workers".to_string(),
4932 JSON_CODEC.to_string(),
4933 None,
4934 )
4935 .expect_err("one workflow sequence cannot open two signal waits");
4936 assert!(matches!(
4937 duplicate_signal_wait,
4938 Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4939 if reason == "duplicate_signal_wait_open"
4940 ));
4941 }
4942
4943 #[test]
4944 fn workflow_sleep_rejects_unrepresentable_rounded_duration() {
4945 let ctx = workflow_context(Vec::new());
4946 let mut sleep = Box::pin(ctx.start_timer(Duration::new(u64::MAX, 1)));
4947 let mut task_context = TaskContext::from_waker(noop_waker_ref());
4948 assert!(matches!(
4949 sleep.as_mut().poll(&mut task_context),
4950 Poll::Ready(Err(Error::TimerDurationOverflow))
4951 ));
4952 assert!(ctx.take_commands().expect("commands").is_empty());
4953 }
4954
4955 #[test]
4956 fn workflow_task_replay_completes_without_rescheduling_recorded_commands() {
4957 let client = Client::new("http://127.0.0.1:8080").expect("client");
4958 let mut worker = Worker::new(client, "rust-workers");
4959 worker.register_workflow("rust.timer", |ctx, _input| async move {
4960 ctx.sleep(Duration::from_secs(5)).await?;
4961 ctx.activity("after-timer", json!([])).await
4962 });
4963
4964 let task = |history_events| WorkflowTask {
4965 task_id: "wft-rust-timer-1".to_string(),
4966 workflow_id: Some("wf-rust-timer".to_string()),
4967 run_id: Some("run-rust-timer".to_string()),
4968 workflow_type: "rust.timer".to_string(),
4969 payload_codec: JSON_CODEC.to_string(),
4970 arguments: Some(json!({"codec": "json", "blob": "[]"})),
4971 history_events,
4972 total_history_events: None,
4973 next_history_page_token: None,
4974 workflow_task_attempt: 1,
4975 workflow_signal_id: None,
4976 signal_name: None,
4977 signal_arguments: None,
4978 lease_owner: Some("rust-worker".to_string()),
4979 };
4980
4981 let initial = worker
4982 .execute_workflow_task(task(Vec::new()))
4983 .expect("initial timer task");
4984 assert_eq!(
4985 initial,
4986 vec![json!({"type": "start_timer", "delay_seconds": 5})]
4987 );
4988
4989 let replayed = worker
4990 .execute_workflow_task(task(vec![
4991 history_event(
4992 "TimerScheduled",
4993 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4994 ),
4995 history_event(
4996 "TimerFired",
4997 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4998 ),
4999 history_event(
5000 "ActivityCompleted",
5001 json!({
5002 "sequence": 2,
5003 "activity_type": "after-timer",
5004 "payload_codec": "json",
5005 "result": {"codec": "json", "blob": "\"done\""},
5006 }),
5007 ),
5008 ]))
5009 .expect("replayed workflow task");
5010 assert_eq!(replayed.len(), 1);
5011 assert_eq!(replayed[0]["type"], "complete_workflow");
5012 assert_eq!(
5013 decode_wire_value(&replayed[0]["result"], JSON_CODEC).expect("result"),
5014 json!("done")
5015 );
5016 }
5017
5018 #[test]
5019 fn workflow_task_replay_keeps_recorded_unfired_timer_pending_without_rescheduling() {
5020 let client = Client::new("http://127.0.0.1:8080").expect("client");
5021 let mut worker = Worker::new(client, "rust-workers");
5022 worker.register_workflow("rust.timer.pending", |ctx, _input| async move {
5023 ctx.sleep(Duration::from_secs(5)).await?;
5024 Ok(json!({"status": "timer fired"}))
5025 });
5026
5027 let task = WorkflowTask {
5028 task_id: "wft-rust-timer-pending".to_string(),
5029 workflow_id: Some("wf-rust-timer".to_string()),
5030 run_id: Some("run-rust-timer".to_string()),
5031 workflow_type: "rust.timer.pending".to_string(),
5032 payload_codec: JSON_CODEC.to_string(),
5033 arguments: Some(json!({"codec": "json", "blob": "[]"})),
5034 history_events: vec![history_event(
5035 "TimerScheduled",
5036 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5037 )],
5038 total_history_events: Some(1),
5039 next_history_page_token: None,
5040 workflow_task_attempt: 1,
5041 workflow_signal_id: None,
5042 signal_name: None,
5043 signal_arguments: None,
5044 lease_owner: Some("rust-worker".to_string()),
5045 };
5046
5047 for _redelivery_or_restart in 0..2 {
5048 let commands = worker
5049 .execute_workflow_task(task.clone())
5050 .expect("recorded timer remains pending");
5051 assert!(
5052 commands.is_empty(),
5053 "recorded timer must not be rescheduled"
5054 );
5055 }
5056 }
5057
5058 #[test]
5059 fn workflow_task_rejects_recorded_command_removed_from_workflow_code() {
5060 let client = Client::new("http://127.0.0.1:8080").expect("client");
5061 let mut worker = Worker::new(client, "rust-workers");
5062 worker.register_workflow("rust.timer.removed", |_ctx, _input| async move {
5063 Ok(json!({"status": "completed"}))
5064 });
5065 let task = WorkflowTask {
5066 task_id: "wft-rust-timer-removed".to_string(),
5067 workflow_id: Some("wf-rust-timer".to_string()),
5068 run_id: Some("run-rust-timer".to_string()),
5069 workflow_type: "rust.timer.removed".to_string(),
5070 payload_codec: JSON_CODEC.to_string(),
5071 arguments: Some(json!({"codec": "json", "blob": "[]"})),
5072 history_events: vec![
5073 history_event(
5074 "TimerScheduled",
5075 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5076 ),
5077 history_event(
5078 "TimerFired",
5079 json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5080 ),
5081 ],
5082 total_history_events: Some(2),
5083 next_history_page_token: None,
5084 workflow_task_attempt: 1,
5085 workflow_signal_id: None,
5086 signal_name: None,
5087 signal_arguments: None,
5088 lease_owner: Some("rust-worker".to_string()),
5089 };
5090
5091 let Error::NonDeterministicReplay(failure) = worker
5092 .execute_workflow_task(task)
5093 .expect_err("removed timer must fail replay")
5094 else {
5095 panic!("expected typed replay failure");
5096 };
5097 assert_eq!(failure.reason, "recorded_commands_unconsumed");
5098 assert_eq!(failure.sequence, Some(1));
5099 }
5100
5101 #[test]
5102 fn workflow_context_emits_explicit_child_workflow_contract() {
5103 let ctx = WorkflowContext {
5104 state: Arc::new(Mutex::new(
5105 WorkflowState::new_with_identity(
5106 Vec::new(),
5107 Some("wf-parent".to_string()),
5108 Some("run-parent".to_string()),
5109 "parent-workers".to_string(),
5110 JSON_CODEC.to_string(),
5111 None,
5112 )
5113 .expect("workflow state"),
5114 )),
5115 };
5116 let options = ChildWorkflowOptions::new("python-workers")
5117 .parent_close_policy(ParentClosePolicy::RequestCancel)
5118 .retry_policy(ChildWorkflowRetryPolicy {
5119 max_attempts: Some(3),
5120 backoff_seconds: vec![1, 5],
5121 non_retryable_error_types: vec!["ValidationError".to_string()],
5122 })
5123 .execution_timeout_seconds(600)
5124 .run_timeout_seconds(120);
5125 let mut call = Box::pin(ctx.start_child_workflow(
5126 "python.fulfil-order",
5127 options,
5128 json!([{"order_id": "order-42"}]),
5129 ));
5130 let mut task_context = TaskContext::from_waker(noop_waker_ref());
5131
5132 assert!(matches!(
5133 call.as_mut().poll(&mut task_context),
5134 Poll::Pending
5135 ));
5136 let commands = ctx.take_commands().expect("commands");
5137 assert_eq!(commands.len(), 1);
5138 let command = &commands[0];
5139 assert_eq!(command["type"], "start_child_workflow");
5140 assert_eq!(command["workflow_type"], "python.fulfil-order");
5141 assert_eq!(command["queue"], "python-workers");
5142 assert_eq!(command["parent_close_policy"], "request_cancel");
5143 assert_eq!(command["retry_policy"]["max_attempts"], 3);
5144 assert_eq!(command["execution_timeout_seconds"], 600);
5145 assert_eq!(command["run_timeout_seconds"], 120);
5146 assert_eq!(
5147 decode_wire_value(&command["arguments"], JSON_CODEC).expect("child args"),
5148 json!([{"order_id": "order-42"}])
5149 );
5150 }
5151
5152 fn child_parent_worker() -> Worker {
5153 let client = Client::new("http://127.0.0.1:8080").expect("client");
5154 let mut worker = Worker::new(client, "rust-parent-workers");
5155 worker.register_workflow("rust.parent", |ctx, _input| async move {
5156 let child = ctx
5157 .start_child_workflow(
5158 "python.child",
5159 ChildWorkflowOptions::new("python-child-workers")
5160 .parent_close_policy(ParentClosePolicy::Terminate),
5161 json!([{"codec_probe": [1, true, "rust"]}]),
5162 )
5163 .await?;
5164 Ok(json!({
5165 "parent_workflow_id": child.parent.workflow_id,
5166 "parent_run_id": child.parent.run_id,
5167 "child_workflow_id": child.child.workflow_id,
5168 "child_run_id": child.child.run_id,
5169 "child_workflow_type": child.child_workflow_type,
5170 "result": child.result,
5171 }))
5172 });
5173 worker
5174 }
5175
5176 fn child_parent_task(event_type: &str, payload: Value) -> WorkflowTask {
5177 WorkflowTask {
5178 task_id: "wft-child-parent".to_string(),
5179 workflow_id: Some("wf-parent".to_string()),
5180 run_id: Some("run-parent".to_string()),
5181 workflow_type: "rust.parent".to_string(),
5182 payload_codec: JSON_CODEC.to_string(),
5183 arguments: Some(encode_value_envelope(&json!([]), JSON_CODEC).expect("input")),
5184 history_events: vec![
5185 HistoryEvent {
5186 event_type: "ChildWorkflowScheduled".to_string(),
5187 payload: json!({
5188 "sequence": 1,
5189 "child_call_id": "call-child",
5190 "child_workflow_instance_id": "wf-child",
5191 "child_workflow_run_id": "run-child",
5192 "child_workflow_type": "python.child",
5193 }),
5194 raw: HashMap::new(),
5195 },
5196 HistoryEvent {
5197 event_type: event_type.to_string(),
5198 payload,
5199 raw: HashMap::new(),
5200 },
5201 ],
5202 total_history_events: Some(2),
5203 next_history_page_token: None,
5204 workflow_task_attempt: 1,
5205 workflow_signal_id: None,
5206 signal_name: None,
5207 signal_arguments: None,
5208 lease_owner: Some("rust-worker".to_string()),
5209 }
5210 }
5211
5212 #[test]
5213 fn committed_child_result_replays_without_starting_a_duplicate() {
5214 let worker = child_parent_worker();
5215 let task = child_parent_task(
5216 "ChildRunCompleted",
5217 json!({
5218 "sequence": 1,
5219 "child_call_id": "call-child",
5220 "child_workflow_instance_id": "wf-child",
5221 "child_workflow_run_id": "run-child",
5222 "child_workflow_type": "python.child",
5223 "payload_codec": "json",
5224 "result": {"codec": "json", "blob": "{\"from\":\"python\",\"ok\":true}"},
5225 }),
5226 );
5227
5228 for _restart in 0..2 {
5229 let commands = worker
5230 .execute_workflow_task(task.clone())
5231 .expect("replayed parent task");
5232 assert_eq!(commands.len(), 1);
5233 assert_eq!(commands[0]["type"], "complete_workflow");
5234 assert!(!commands
5235 .iter()
5236 .any(|command| command["type"] == "start_child_workflow"));
5237 let output =
5238 decode_wire_value(&commands[0]["result"], JSON_CODEC).expect("parent output");
5239 assert_eq!(output["parent_workflow_id"], "wf-parent");
5240 assert_eq!(output["parent_run_id"], "run-parent");
5241 assert_eq!(output["child_workflow_id"], "wf-child");
5242 assert_eq!(output["child_run_id"], "run-child");
5243 assert_eq!(output["result"], json!({"from": "python", "ok": true}));
5244 }
5245 }
5246
5247 #[test]
5248 fn pending_child_replays_after_restart_without_starting_a_duplicate() {
5249 let worker = child_parent_worker();
5250 let mut task = child_parent_task("unused", Value::Null);
5251 task.history_events.truncate(1);
5252 task.total_history_events = Some(1);
5253
5254 for _redelivery_or_restart in 0..2 {
5255 let commands = worker
5256 .execute_workflow_task(task.clone())
5257 .expect("recorded child remains pending");
5258 assert!(
5259 commands.is_empty(),
5260 "recorded pending child must not be started again"
5261 );
5262 }
5263 }
5264
5265 #[test]
5266 fn child_cancellation_becomes_stable_parent_failure_command() {
5267 let worker = child_parent_worker();
5268 let task = child_parent_task(
5269 "ChildRunCancelled",
5270 json!({
5271 "sequence": 1,
5272 "child_workflow_instance_id": "wf-child",
5273 "child_workflow_run_id": "run-child",
5274 "child_workflow_type": "python.child",
5275 "failure_id": "failure-child",
5276 "failure_category": "cancelled",
5277 "message": "cancelled by parent-close policy",
5278 }),
5279 );
5280
5281 let commands = worker
5282 .execute_workflow_task(task)
5283 .expect("parent settlement");
5284 assert_eq!(commands.len(), 1);
5285 assert_eq!(commands[0]["type"], "fail_workflow");
5286 assert_eq!(commands[0]["exception_type"], "ChildWorkflowCancelled");
5287 assert_eq!(
5288 commands[0]["exception"]["properties"]["reason"],
5289 "cancelled"
5290 );
5291 assert_eq!(
5292 commands[0]["exception"]["properties"]["child_workflow_run_id"],
5293 "run-child"
5294 );
5295 }
5296
5297 #[test]
5298 fn workflow_can_handle_typed_child_failure() {
5299 let client = Client::new("http://127.0.0.1:8080").expect("client");
5300 let mut worker = Worker::new(client, "rust-parent-workers");
5301 worker.register_workflow("rust.handled-parent", |ctx, _input| async move {
5302 match ctx
5303 .start_child_workflow(
5304 "python.child",
5305 ChildWorkflowOptions::new("python-child-workers"),
5306 json!([]),
5307 )
5308 .await
5309 {
5310 Err(Error::ChildWorkflowFailed(failure)) => Ok(json!({
5311 "reason": failure.reason,
5312 "failure_id": failure.failure_id,
5313 "exception_class": failure.exception_class,
5314 "child_run_id": failure.child_workflow_run_id,
5315 })),
5316 Err(error) => Err(error),
5317 Ok(_) => Err(Error::WorkerLoop(
5318 "child unexpectedly succeeded".to_string(),
5319 )),
5320 }
5321 });
5322 let mut task = child_parent_task(
5323 "ChildRunFailed",
5324 json!({
5325 "sequence": 1,
5326 "child_workflow_instance_id": "wf-child",
5327 "child_workflow_run_id": "run-child",
5328 "child_workflow_type": "python.child",
5329 "failure_id": "failure-child",
5330 "failure_category": "child_workflow",
5331 "message": "payment rejected",
5332 "exception": {
5333 "type": "PaymentRejected",
5334 "class": "payments.PaymentRejected",
5335 "message": "payment rejected"
5336 }
5337 }),
5338 );
5339 task.workflow_type = "rust.handled-parent".to_string();
5340
5341 let commands = worker.execute_workflow_task(task).expect("handled failure");
5342 assert_eq!(commands[0]["type"], "complete_workflow");
5343 let output = decode_wire_value(&commands[0]["result"], JSON_CODEC).expect("parent output");
5344 assert_eq!(output["reason"], "child_workflow");
5345 assert_eq!(output["failure_id"], "failure-child");
5346 assert_eq!(output["exception_class"], "payments.PaymentRejected");
5347 assert_eq!(output["child_run_id"], "run-child");
5348 }
5349
5350 #[test]
5351 fn rust_hello_world_uses_signal_arguments_from_resume_payload() {
5352 let client = Client::new("http://127.0.0.1:8080").expect("client");
5353 let mut worker = Worker::new(client, "rust-workers");
5354
5355 worker.register_workflow("rust.hello_workflow", |ctx, _input| async move {
5356 let signal = ctx.wait_signal("start").await?;
5357 let name = signal
5358 .first()
5359 .and_then(|value| value.as_str())
5360 .unwrap_or("world");
5361 let greeting = ctx.activity("rust.hello_activity", json!([name])).await?;
5362 Ok(json!({
5363 "greeting": greeting,
5364 "language": "rust"
5365 }))
5366 });
5367
5368 let signal_arguments =
5369 encode_value_envelope(&json!(["Rust"]), DEFAULT_CODEC).expect("signal arguments");
5370 let task = WorkflowTask {
5371 task_id: "wft-rust-signal-1".to_string(),
5372 workflow_id: Some("wf-rust-hello".to_string()),
5373 run_id: Some("run-rust-hello".to_string()),
5374 workflow_type: "rust.hello_workflow".to_string(),
5375 payload_codec: DEFAULT_CODEC.to_string(),
5376 arguments: Some(encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("input")),
5377 history_events: vec![HistoryEvent {
5378 event_type: "SignalReceived".to_string(),
5379 payload: json!({
5380 "signal_id": "sig-rust-1",
5381 "signal_name": "start"
5382 }),
5383 raw: HashMap::new(),
5384 }],
5385 total_history_events: Some(1),
5386 next_history_page_token: None,
5387 workflow_task_attempt: 1,
5388 workflow_signal_id: Some("sig-rust-1".to_string()),
5389 signal_name: Some("start".to_string()),
5390 signal_arguments: Some(signal_arguments),
5391 lease_owner: Some("rust-worker".to_string()),
5392 };
5393
5394 let commands = worker.execute_workflow_task(task).expect("workflow task");
5395
5396 assert_eq!(commands.len(), 1);
5397 assert_eq!(commands[0]["type"], "schedule_activity");
5398 assert_eq!(commands[0]["activity_type"], "rust.hello_activity");
5399 assert_eq!(
5400 decode_wire_value(&commands[0]["arguments"], DEFAULT_CODEC).expect("activity args"),
5401 json!(["Rust"])
5402 );
5403 }
5404
5405 #[test]
5406 fn workflow_task_appends_paginated_history_events() {
5407 let mut task = WorkflowTask {
5408 task_id: "wft-rust-pages-1".to_string(),
5409 workflow_id: Some("wf-rust-pages".to_string()),
5410 run_id: Some("run-rust-pages".to_string()),
5411 workflow_type: "rust.hello_workflow".to_string(),
5412 payload_codec: DEFAULT_CODEC.to_string(),
5413 arguments: Some(encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("input")),
5414 history_events: vec![HistoryEvent {
5415 event_type: "WorkflowStarted".to_string(),
5416 payload: json!({}),
5417 raw: HashMap::new(),
5418 }],
5419 total_history_events: Some(3),
5420 next_history_page_token: Some("MQ==".to_string()),
5421 workflow_task_attempt: 1,
5422 workflow_signal_id: None,
5423 signal_name: None,
5424 signal_arguments: None,
5425 lease_owner: Some("rust-worker".to_string()),
5426 };
5427
5428 task.append_history_page(WorkflowTaskHistoryPage {
5429 history_events: vec![
5430 HistoryEvent {
5431 event_type: "SignalReceived".to_string(),
5432 payload: json!({
5433 "signal_id": "sig-rust-1",
5434 "signal_name": "start",
5435 "arguments": encode_value_envelope(&json!(["Rust"]), DEFAULT_CODEC)
5436 .expect("signal arguments")
5437 }),
5438 raw: HashMap::new(),
5439 },
5440 HistoryEvent {
5441 event_type: "MarkerRecorded".to_string(),
5442 payload: json!({"sequence": 3}),
5443 raw: HashMap::new(),
5444 },
5445 ],
5446 total_history_events: Some(3),
5447 next_history_page_token: None,
5448 });
5449
5450 assert_eq!(task.history_events.len(), 3);
5451 assert_eq!(task.total_history_events, Some(3));
5452 assert_eq!(task.next_history_page_token, None);
5453
5454 let signals =
5455 signal_values(&task.history_events, "start", DEFAULT_CODEC, None).expect("signals");
5456 assert_eq!(signals, vec![vec![json!("Rust")]]);
5457 }
5458
5459 #[tokio::test]
5460 async fn query_handler_reads_ordered_cross_codec_signals_without_commands() {
5461 let client = Client::new("http://127.0.0.1:8080").expect("client");
5462 let mut worker = Worker::new(client, "rust-workers");
5463 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5464 worker.register_query("counter", "current", |ctx, _args| async move {
5465 let mut count = 0_i64;
5466 for signal in ctx.signal_events() {
5467 let value = signal
5468 .arguments
5469 .first()
5470 .and_then(Value::as_i64)
5471 .unwrap_or_default();
5472 match signal.name.as_str() {
5473 "increment" => count += value,
5474 "set" => count = value,
5475 _ => {}
5476 }
5477 }
5478 Ok(json!(count))
5479 });
5480
5481 let task = QueryTask {
5482 query_task_id: "query-rust-counter".to_string(),
5483 query_task_attempt: 1,
5484 lease_owner: Some("rust-worker".to_string()),
5485 workflow_id: Some("counter-1".to_string()),
5486 run_id: Some("run-counter-1".to_string()),
5487 workflow_type: "counter".to_string(),
5488 query_name: "current".to_string(),
5489 payload_codec: DEFAULT_CODEC.to_string(),
5490 workflow_arguments: Some(
5491 encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("workflow input"),
5492 ),
5493 query_arguments: Some(
5494 encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("query arguments"),
5495 ),
5496 history_events: vec![
5497 HistoryEvent {
5498 event_type: "SignalReceived".to_string(),
5499 payload: json!({
5500 "signal_id": "php-signal-1",
5501 "signal_name": "increment",
5502 "workflow_sequence": 1,
5503 "payload_codec": DEFAULT_CODEC,
5504 "arguments": encode_value_envelope(&json!([3]), DEFAULT_CODEC).expect("php avro signal")
5505 }),
5506 raw: HashMap::new(),
5507 },
5508 HistoryEvent {
5509 event_type: "SignalReceived".to_string(),
5510 payload: json!({
5511 "signal_id": "python-signal-2",
5512 "signal_name": "increment",
5513 "workflow_sequence": 2,
5514 "payload_codec": JSON_CODEC,
5515 "arguments": encode_value_envelope(&json!([5]), JSON_CODEC).expect("python json signal")
5516 }),
5517 raw: HashMap::new(),
5518 },
5519 HistoryEvent {
5520 event_type: "SignalReceived".to_string(),
5521 payload: json!({
5522 "signal_id": "rust-signal-3",
5523 "signal_name": "set",
5524 "workflow_sequence": 3,
5525 "payload_codec": DEFAULT_CODEC,
5526 "arguments": encode_value_envelope(&json!([0]), DEFAULT_CODEC).expect("rust avro signal")
5527 }),
5528 raw: HashMap::new(),
5529 },
5530 ],
5531 history_export: None,
5532 run_status: Some("completed".to_string()),
5533 };
5534
5535 let result = worker.execute_query_task(task).await.expect("query result");
5536 assert_eq!(result, json!(0));
5537 }
5538
5539 #[tokio::test]
5540 async fn replayed_queries_read_running_completed_and_cold_restarted_instance_state() {
5541 let worker = replay_counter_worker();
5542 let running_history = json!([
5543 {
5544 "type": "ActivityCompleted",
5545 "payload": {
5546 "sequence": 1,
5547 "activity_type": "load-counter",
5548 "payload_codec": "json",
5549 "result": {"codec": "json", "blob": "\"loaded\""}
5550 }
5551 },
5552 {
5553 "type": "SignalReceived",
5554 "payload": {
5555 "signal_id": "signal-3",
5556 "signal_name": "increment",
5557 "payload_codec": "json",
5558 "arguments": {"codec": "json", "blob": "[3]"}
5559 }
5560 }
5561 ]);
5562
5563 let running = worker
5564 .execute_query_task(replay_counter_query(
5565 "current",
5566 running_history.clone(),
5567 "running",
5568 ))
5569 .await
5570 .expect("running replay query");
5571 assert_eq!(
5572 running,
5573 json!({"loaded": "loaded", "count": 3, "finished": false})
5574 );
5575
5576 let detached = worker
5577 .execute_query_task(replay_counter_query(
5578 "detached-mutation",
5579 running_history.clone(),
5580 "running",
5581 ))
5582 .await
5583 .expect("query mutates only its detached state clone");
5584 assert_eq!(detached, json!(999));
5585 let failed = worker
5586 .execute_query_task(replay_counter_query(
5587 "failed-mutation",
5588 running_history.clone(),
5589 "running",
5590 ))
5591 .await
5592 .expect_err("failed query");
5593 assert_eq!(failed.reason, "query_rejected");
5594 let unchanged = worker
5595 .execute_query_task(replay_counter_query("current", running_history, "running"))
5596 .await
5597 .expect("later query reconstructs unchanged state");
5598 assert_eq!(unchanged, running);
5599
5600 let restarted_worker = replay_counter_worker();
5601 let restarted_task: QueryTask = serde_json::from_value(json!({
5602 "query_task_id": "query-after-restart",
5603 "workflow_id": "counter-1",
5604 "run_id": "run-counter-1",
5605 "workflow_type": "replay-counter",
5606 "query_name": "current",
5607 "payload_codec": "json",
5608 "workflow_arguments": {"codec": "json", "blob": "[]"},
5609 "query_arguments": {"codec": "json", "blob": "[]"},
5610 "history_events": [],
5611 "history_export": {
5612 "payloads": {"codec": "json"},
5613 "history_events": [
5614 {
5615 "type": "ActivityCompleted",
5616 "payload": {
5617 "sequence": 1,
5618 "activity_type": "load-counter",
5619 "payload_codec": "json",
5620 "result": null
5621 }
5622 },
5623 {
5624 "type": "SignalReceived",
5625 "payload": {"signal_id": "signal-3", "signal_name": "increment"}
5626 },
5627 {
5628 "type": "SignalReceived",
5629 "payload": {"signal_id": "signal-5", "signal_name": "increment"}
5630 }
5631 ],
5632 "activities": [{
5633 "sequence": 1,
5634 "activity_type": "load-counter",
5635 "payload_codec": "json",
5636 "result": {"codec": "json", "blob": "\"loaded\""}
5637 }],
5638 "signals": [
5639 {
5640 "id": "signal-3",
5641 "name": "increment",
5642 "payload_codec": "json",
5643 "arguments": "[3]"
5644 },
5645 {
5646 "id": "signal-5",
5647 "name": "increment",
5648 "payload_codec": "json",
5649 "arguments": "[5]"
5650 }
5651 ]
5652 },
5653 "run_status": "completed"
5654 }))
5655 .expect("cold replay query task");
5656 let completed = restarted_worker
5657 .execute_query_task(restarted_task)
5658 .await
5659 .expect("completed cold replay query");
5660 assert_eq!(
5661 completed,
5662 json!({"loaded": "loaded", "count": 8, "finished": true})
5663 );
5664 }
5665
5666 #[tokio::test]
5667 async fn replayed_query_replay_failures_are_machine_readable() {
5668 let worker = replay_counter_worker();
5669 let task = replay_counter_query(
5670 "current",
5671 json!([{
5672 "type": "ActivityCompleted",
5673 "payload": {
5674 "sequence": 1,
5675 "payload_codec": "json",
5676 "result": {"codec": "json", "blob": "{"}
5677 }
5678 }]),
5679 "running",
5680 );
5681 let failure = worker
5682 .execute_query_task(task)
5683 .await
5684 .expect_err("invalid replay history payload");
5685 assert_eq!(failure.reason, "query_workflow_state_unavailable");
5686 assert_eq!(failure.failure_type, "QueryWorkflowStateUnavailable");
5687 }
5688
5689 #[tokio::test]
5690 async fn query_task_restores_compact_history_from_export() {
5691 let client = Client::new("http://127.0.0.1:8080").expect("client");
5692 let mut worker = Worker::new(client, "rust-workers");
5693 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5694 worker.register_query("counter", "current", |ctx, _args| async move {
5695 Ok(json!(ctx.signals("increment")[0][0]))
5696 });
5697 let task: QueryTask = serde_json::from_value(json!({
5698 "query_task_id": "query-export",
5699 "workflow_type": "counter",
5700 "query_name": "current",
5701 "payload_codec": "json",
5702 "workflow_arguments": {"codec": "json", "blob": "[]"},
5703 "query_arguments": {"codec": "json", "blob": "[]"},
5704 "history_events": [],
5705 "history_export": {
5706 "payloads": {"codec": "json"},
5707 "history_events": [{
5708 "type": "SignalReceived",
5709 "payload": {"signal_id": "signal-export", "signal_name": "increment"}
5710 }],
5711 "signals": [{
5712 "id": "signal-export",
5713 "name": "increment",
5714 "status": "applied",
5715 "workflow_sequence": 1,
5716 "payload_codec": "json",
5717 "arguments": "[9]"
5718 }]
5719 }
5720 }))
5721 .expect("query task");
5722
5723 let result = worker.execute_query_task(task).await.expect("query result");
5724 assert_eq!(result, json!(9));
5725 }
5726
5727 #[tokio::test]
5728 async fn query_task_failures_have_stable_reasons() {
5729 let client = Client::new("http://127.0.0.1:8080").expect("client");
5730 let mut worker = Worker::new(client, "rust-workers");
5731 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5732 worker.register_query(
5733 "counter",
5734 "current",
5735 |_ctx, _args| async move { Ok(json!(0)) },
5736 );
5737
5738 let base_task = QueryTask {
5739 query_task_id: "query-errors".to_string(),
5740 query_task_attempt: 1,
5741 lease_owner: None,
5742 workflow_id: Some("counter-errors".to_string()),
5743 run_id: Some("run-errors".to_string()),
5744 workflow_type: "counter".to_string(),
5745 query_name: "missing".to_string(),
5746 payload_codec: JSON_CODEC.to_string(),
5747 workflow_arguments: Some(json!({"codec": "json", "blob": "[]"})),
5748 query_arguments: Some(json!({"codec": "json", "blob": "[]"})),
5749 history_events: Vec::new(),
5750 history_export: None,
5751 run_status: Some("running".to_string()),
5752 };
5753
5754 let unknown = worker
5755 .execute_query_task(base_task.clone())
5756 .await
5757 .expect_err("unknown query");
5758 assert_eq!(unknown.reason, "rejected_unknown_query");
5759
5760 let mut malformed = base_task;
5761 malformed.query_name = "current".to_string();
5762 malformed.query_arguments = Some(json!({"codec": "json", "blob": "{"}));
5763 let malformed = worker
5764 .execute_query_task(malformed)
5765 .await
5766 .expect_err("malformed payload");
5767 assert_eq!(malformed.reason, "query_payload_decode_failed");
5768
5769 let client = Client::new("http://127.0.0.1:8080").expect("client");
5770 let mut unavailable_worker = Worker::new(client, "rust-workers");
5771 unavailable_worker
5772 .register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5773 let unavailable_task: QueryTask = serde_json::from_value(json!({
5774 "query_task_id": "query-unavailable",
5775 "workflow_type": "counter",
5776 "query_name": "current",
5777 "payload_codec": "json",
5778 "workflow_arguments": {"codec": "json", "blob": "[]"},
5779 "query_arguments": {"codec": "json", "blob": "[]"}
5780 }))
5781 .expect("query task");
5782 let unavailable = unavailable_worker
5783 .execute_query_task(unavailable_task)
5784 .await
5785 .expect_err("query handler unavailable");
5786 assert_eq!(unavailable.reason, "query_handler_unavailable");
5787 }
5788
5789 #[tokio::test]
5790 async fn client_query_decodes_result_and_typed_failure() {
5791 let server = MockWorkerServer::start();
5792 let client = Client::builder(server.base_url())
5793 .timeout(Duration::from_secs(2))
5794 .build()
5795 .expect("client");
5796
5797 let result = client
5798 .query_workflow("counter-1", "current", json!([]))
5799 .await
5800 .expect("query result");
5801 assert_eq!(result, json!({"count": 8}));
5802
5803 let error = client
5804 .query_workflow("counter-1", "missing", json!([]))
5805 .await
5806 .expect_err("unknown query");
5807 let Error::QueryFailed(failure) = error else {
5808 panic!("expected typed query failure");
5809 };
5810 assert_eq!(failure.status, 404);
5811 assert_eq!(failure.reason, "rejected_unknown_query");
5812 }
5813
5814 #[tokio::test]
5815 async fn baseline_worker_endpoints_send_the_baseline_protocol() {
5816 let server = MockWorkerServer::start();
5817 let client = Client::builder(server.base_url())
5818 .timeout(Duration::from_secs(2))
5819 .build()
5820 .expect("client");
5821
5822 client
5823 .register_worker("capture-worker", "capture", vec![], vec![], 1, 1)
5824 .await
5825 .expect("register");
5826 client
5827 .heartbeat_worker("capture-worker", 1, 1)
5828 .await
5829 .expect("heartbeat");
5830 client
5831 .poll_workflow_task("capture-worker", "capture", Duration::from_millis(10))
5832 .await
5833 .expect("workflow poll");
5834 client
5835 .poll_activity_task("capture-worker", "capture", Duration::from_millis(10))
5836 .await
5837 .expect("activity poll");
5838
5839 for path in [
5840 "/api/worker/register",
5841 "/api/worker/heartbeat",
5842 "/api/worker/workflow-tasks/poll",
5843 "/api/worker/activity-tasks/poll",
5844 ] {
5845 assert_eq!(
5846 server.worker_protocol_for(path).as_deref(),
5847 Some(WORKER_PROTOCOL_VERSION),
5848 "unexpected protocol for {path}"
5849 );
5850 }
5851
5852 assert_eq!(
5853 server.request_body("/api/worker/workflow-tasks/poll")["timeout_seconds"],
5854 1
5855 );
5856 assert_eq!(
5857 server.request_body("/api/worker/activity-tasks/poll")["timeout_seconds"],
5858 1
5859 );
5860 }
5861
5862 #[tokio::test]
5863 async fn query_task_endpoints_send_the_query_feature_protocol() {
5864 let server = MockWorkerServer::start();
5865 let client = Client::builder(server.base_url())
5866 .timeout(Duration::from_secs(2))
5867 .build()
5868 .expect("client");
5869
5870 client
5871 .poll_query_task("capture-worker", "capture", Duration::from_millis(10))
5872 .await
5873 .expect("query poll");
5874 client
5875 .complete_query_task("query-capture", "capture-worker", 1, json!(8), JSON_CODEC)
5876 .await
5877 .expect("query complete");
5878 client
5879 .fail_query_task(
5880 "query-capture",
5881 "capture-worker",
5882 1,
5883 "failed",
5884 "query_rejected",
5885 "QueryFailed",
5886 )
5887 .await
5888 .expect("query fail");
5889
5890 for path in [
5891 "/api/worker/query-tasks/poll",
5892 "/api/worker/query-tasks/query-capture/complete",
5893 "/api/worker/query-tasks/query-capture/fail",
5894 ] {
5895 assert_eq!(
5896 server.worker_protocol_for(path).as_deref(),
5897 Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
5898 "unexpected protocol for {path}"
5899 );
5900 }
5901
5902 assert_eq!(
5903 server.request_body("/api/worker/query-tasks/poll")["timeout_seconds"],
5904 1
5905 );
5906 }
5907
5908 #[tokio::test]
5909 async fn query_protocol_rejection_from_older_server_is_typed() {
5910 let server = MockWorkerServer::reject_query_protocol();
5911 let client = Client::builder(server.base_url())
5912 .timeout(Duration::from_secs(2))
5913 .build()
5914 .expect("client");
5915
5916 let error = client
5917 .poll_query_task("capture-worker", "capture", Duration::from_millis(10))
5918 .await
5919 .expect_err("server below query protocol floor must reject");
5920 let Error::Protocol(failure) = error else {
5921 panic!("expected typed protocol failure");
5922 };
5923
5924 assert_eq!(failure.status, 400);
5925 assert_eq!(failure.reason, "unsupported_protocol_version");
5926 assert_eq!(failure.supported_version.as_deref(), Some("1.7"));
5927 assert_eq!(
5928 failure.requested_version.as_deref(),
5929 Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION)
5930 );
5931 assert_eq!(
5932 server
5933 .worker_protocol_for("/api/worker/query-tasks/poll")
5934 .as_deref(),
5935 Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION)
5936 );
5937 }
5938
5939 #[tokio::test]
5940 async fn run_once_without_query_handlers_keeps_pre_query_server_compatibility() {
5941 let server = MockWorkerServer::reject_query_protocol();
5942 let client = Client::builder(server.base_url())
5943 .timeout(Duration::from_secs(2))
5944 .build()
5945 .expect("client");
5946 let mut worker = Worker::new(client, "rust-workers")
5947 .worker_id("baseline-worker")
5948 .poll_timeout(Duration::from_millis(10));
5949
5950 worker.register_workflow("baseline.workflow", |_ctx, _input| async move {
5951 Ok(Value::Null)
5952 });
5953
5954 assert_eq!(worker.run_once().await.expect("baseline run once"), 0);
5955 assert_eq!(
5956 server
5957 .worker_protocol_for("/api/worker/workflow-tasks/poll")
5958 .as_deref(),
5959 Some(WORKER_PROTOCOL_VERSION)
5960 );
5961 assert_eq!(
5962 server.worker_protocol_for("/api/worker/query-tasks/poll"),
5963 None,
5964 "a worker without query handlers must not use the query-task endpoint"
5965 );
5966 }
5967
5968 #[tokio::test]
5969 async fn completion_time_query_rejection_is_typed_without_stopping_worker() {
5970 let server = MockWorkerServer::reject_query_completion();
5971 let client = Client::builder(server.base_url())
5972 .timeout(Duration::from_secs(2))
5973 .build()
5974 .expect("client");
5975
5976 let error = client
5977 .complete_query_task("query-late", "late-worker", 1, json!(8), JSON_CODEC)
5978 .await
5979 .expect_err("expired completion must be rejected");
5980 let Error::QueryFailed(failure) = error else {
5981 panic!("expected typed query failure");
5982 };
5983 assert_eq!(failure.status, 409);
5984 assert_eq!(failure.reason, "query_task_timed_out");
5985
5986 let mut worker = Worker::new(client, "rust-workers")
5987 .worker_id("late-worker")
5988 .poll_timeout(Duration::from_millis(10));
5989 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5990 worker.register_query(
5991 "counter",
5992 "current",
5993 |_ctx, _args| async move { Ok(json!(8)) },
5994 );
5995
5996 assert_eq!(worker.run_once().await.expect("late task is handled"), 1);
5997 assert_eq!(
5998 worker
5999 .run_once()
6000 .await
6001 .expect("worker continues after late completion"),
6002 0
6003 );
6004 assert_eq!(
6005 server.request_count("/api/worker/query-tasks/query-late/complete"),
6006 2
6007 );
6008 assert_eq!(
6009 server.request_count("/api/worker/query-tasks/query-late/fail"),
6010 0,
6011 "a server completion rejection must not be reported as an encoding failure"
6012 );
6013 }
6014
6015 #[tokio::test]
6016 async fn activity_only_worker_can_shutdown_without_workflow_poller() {
6017 let server = MockWorkerServer::start();
6018 let client = Client::builder(server.base_url())
6019 .timeout(Duration::from_secs(2))
6020 .build()
6021 .expect("client");
6022 let mut worker = Worker::new(client, "rust-workers")
6023 .worker_id("activity-only-worker")
6024 .poll_timeout(Duration::from_millis(10));
6025
6026 worker.register_activity(
6027 "activity.only",
6028 |_ctx, _args| async move { Ok(Value::Null) },
6029 );
6030
6031 worker.run_until(async {}).await.expect("run worker");
6032 }
6033
6034 #[tokio::test]
6035 async fn workflow_only_worker_can_shutdown_without_activity_poller() {
6036 let server = MockWorkerServer::start();
6037 let client = Client::builder(server.base_url())
6038 .timeout(Duration::from_secs(2))
6039 .build()
6040 .expect("client");
6041 let mut worker = Worker::new(client, "rust-workers")
6042 .worker_id("workflow-only-worker")
6043 .poll_timeout(Duration::from_millis(10));
6044
6045 worker.register_workflow(
6046 "workflow.only",
6047 |_ctx, _input| async move { Ok(Value::Null) },
6048 );
6049
6050 worker.run_until(async {}).await.expect("run worker");
6051 }
6052
6053 #[tokio::test]
6054 async fn worker_heartbeat_observer_receives_server_acknowledgements() {
6055 let server = MockWorkerServer::start();
6056 let client = Client::builder(server.base_url())
6057 .timeout(Duration::from_secs(2))
6058 .build()
6059 .expect("client");
6060 let observations = Arc::new(Mutex::new(Vec::new()));
6061 let observed = Arc::clone(&observations);
6062 let mut worker = Worker::new(client, "rust-workers")
6063 .worker_id("observed-heartbeat-worker")
6064 .poll_timeout(Duration::from_millis(10))
6065 .on_worker_heartbeat(move |observation| {
6066 observed
6067 .lock()
6068 .expect("heartbeat observations")
6069 .push(observation.clone());
6070 });
6071
6072 worker.register_workflow("workflow.observed", |_ctx, _input| async move {
6073 Ok(Value::Null)
6074 });
6075 worker
6076 .run_until(tokio::time::sleep(Duration::from_millis(20)))
6077 .await
6078 .expect("run worker");
6079
6080 let observations = observations.lock().expect("heartbeat observations");
6081 let first = observations.first().expect("heartbeat acknowledgement");
6082 assert_eq!(first.worker_id, "observed-heartbeat-worker");
6083 assert_eq!(first.task_queue, "rust-workers");
6084 assert!(first.acknowledged_at_unix_millis > 0);
6085 assert_eq!(first.acknowledgement, json!({}));
6086 }
6087
6088 #[tokio::test]
6089 async fn query_enabled_worker_stays_live_when_signal_replay_emits_no_commands() {
6090 let server = MockWorkerServer::waiting_query_worker();
6091 let client = Client::builder(server.base_url())
6092 .timeout(Duration::from_secs(2))
6093 .build()
6094 .expect("client");
6095 let observations = Arc::new(Mutex::new(Vec::new()));
6096 let observed = Arc::clone(&observations);
6097 let mut worker = Worker::new(client, "rust-snapshot-workers")
6098 .worker_id("rust-snapshot-worker")
6099 .poll_timeout(Duration::from_millis(10))
6100 .on_worker_heartbeat(move |observation| {
6101 observed
6102 .lock()
6103 .expect("heartbeat observations")
6104 .push(observation.clone());
6105 });
6106
6107 worker.register_workflow("snapshot", |ctx, _input| async move {
6108 ctx.wait_signal("finish").await?;
6109 Ok(json!({"status": "finished"}))
6110 });
6111 worker.register_query("snapshot", "current", |ctx, _args| async move {
6112 let current = ctx
6113 .signals("increment")
6114 .iter()
6115 .filter_map(|arguments| arguments.first().and_then(Value::as_i64))
6116 .sum::<i64>();
6117 Ok(json!(current))
6118 });
6119
6120 worker
6121 .run_until(tokio::time::sleep(Duration::from_millis(3_200)))
6122 .await
6123 .expect("pending workflow and query poller must remain live until shutdown");
6124
6125 assert!(
6126 observations.lock().expect("heartbeat observations").len() >= 4,
6127 "the immediate heartbeat and at least three advertised one-second intervals must be acknowledged"
6128 );
6129 assert!(
6130 server.request_count("/api/worker/workflow-tasks/poll") >= 3,
6131 "workflow polling must continue after empty replay acknowledgements"
6132 );
6133 assert!(
6134 server.request_count("/api/worker/query-tasks/poll") >= 2,
6135 "query polling must continue after serving the current query"
6136 );
6137 assert_eq!(
6138 server.request_body("/api/worker/register")["capabilities"],
6139 json!([QUERY_TASKS_CAPABILITY])
6140 );
6141
6142 for task_id in ["snapshot-wait-3", "snapshot-wait-5"] {
6143 let fail_path = format!("/api/worker/workflow-tasks/{task_id}/fail");
6144 let completion_path = format!("/api/worker/workflow-tasks/{task_id}/complete");
6145 let failure = server.request_body(&fail_path);
6146 assert_eq!(
6147 failure["failure"]["type"],
6148 WORKFLOW_TASK_WAITING_FOR_HISTORY_TYPE
6149 );
6150 assert_eq!(server.request_count(&completion_path), 0);
6151 }
6152
6153 let query_completion =
6154 server.request_body("/api/worker/query-tasks/snapshot-current/complete");
6155 assert_eq!(query_completion["result"], json!(8));
6156 }
6157
6158 #[tokio::test]
6159 async fn worker_retries_poll_and_heartbeat_transport_failures_independently() {
6160 let server = MockWorkerServer::transient_worker_failures();
6161 let client = Client::builder(server.base_url())
6162 .timeout(Duration::from_secs(2))
6163 .build()
6164 .expect("client");
6165 let mut worker = Worker::new(client, "rust-workers")
6166 .worker_id("retry-worker")
6167 .poll_timeout(Duration::from_millis(10))
6168 .retry_policy(WorkerRetryPolicy {
6169 max_retries: 2,
6170 initial_backoff: Duration::from_millis(1),
6171 max_backoff: Duration::from_millis(1),
6172 });
6173 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6174 worker.register_activity(
6175 "counter.activity",
6176 |_ctx, _input| async move { Ok(Value::Null) },
6177 );
6178 worker.register_query(
6179 "counter",
6180 "current",
6181 |_ctx, _args| async move { Ok(json!(8)) },
6182 );
6183
6184 worker
6185 .run_until(tokio::time::sleep(Duration::from_millis(75)))
6186 .await
6187 .expect("transient failures must not stop the worker");
6188
6189 for path in [
6190 "/api/worker/heartbeat",
6191 "/api/worker/workflow-tasks/poll",
6192 "/api/worker/activity-tasks/poll",
6193 "/api/worker/query-tasks/poll",
6194 ] {
6195 assert!(
6196 server.request_count(path) >= 2,
6197 "{path} must continue after its transient failure"
6198 );
6199 }
6200 }
6201
6202 #[tokio::test]
6203 async fn worker_bounds_transport_retries() {
6204 let server = MockWorkerServer::unavailable_polls();
6205 let client = Client::builder(server.base_url())
6206 .timeout(Duration::from_secs(2))
6207 .build()
6208 .expect("client");
6209 let mut worker = Worker::new(client, "rust-workers")
6210 .worker_id("bounded-retry-worker")
6211 .poll_timeout(Duration::from_millis(10))
6212 .retry_policy(WorkerRetryPolicy {
6213 max_retries: 2,
6214 initial_backoff: Duration::from_millis(1),
6215 max_backoff: Duration::from_millis(1),
6216 });
6217 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6218
6219 let error = worker.run().await.expect_err("retry bound must terminate");
6220 assert!(matches!(error, Error::Transport(_)));
6221 assert_eq!(
6222 server.request_count("/api/worker/workflow-tasks/poll"),
6223 3,
6224 "one initial request plus two retries"
6225 );
6226 }
6227
6228 #[tokio::test]
6229 async fn worker_does_not_retry_authentication_failures() {
6230 let server = MockWorkerServer::unauthorized_polls();
6231 let client = Client::builder(server.base_url())
6232 .timeout(Duration::from_secs(2))
6233 .build()
6234 .expect("client");
6235 let mut worker = Worker::new(client, "rust-workers")
6236 .worker_id("unauthorized-worker")
6237 .poll_timeout(Duration::from_millis(10));
6238 worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6239
6240 let error = worker
6241 .run()
6242 .await
6243 .expect_err("authentication must terminate");
6244 let Error::Http { status, body } = error else {
6245 panic!("expected stable HTTP authentication error");
6246 };
6247 assert_eq!(status, reqwest::StatusCode::UNAUTHORIZED);
6248 assert!(body.contains("authentication_failed"));
6249 assert_eq!(
6250 server.request_count("/api/worker/workflow-tasks/poll"),
6251 1,
6252 "authentication failures must not be retried"
6253 );
6254 }
6255
6256 #[derive(Clone, Debug)]
6257 struct CapturedRequest {
6258 path: String,
6259 worker_protocol: Option<String>,
6260 body: String,
6261 }
6262
6263 struct MockWorkerServer {
6264 addr: SocketAddr,
6265 stop: Arc<AtomicBool>,
6266 requests: Arc<Mutex<Vec<CapturedRequest>>>,
6267 thread: Option<thread::JoinHandle<()>>,
6268 }
6269
6270 #[derive(Clone, Copy, Default)]
6271 struct MockWorkerBehavior {
6272 reject_query_protocol: bool,
6273 reject_query_completion: bool,
6274 waiting_query_worker: bool,
6275 poll_failures_per_path: usize,
6276 heartbeat_failures: usize,
6277 unauthorized_polls: bool,
6278 }
6279
6280 impl MockWorkerServer {
6281 fn start() -> Self {
6282 Self::start_with_behavior(MockWorkerBehavior::default())
6283 }
6284
6285 fn reject_query_protocol() -> Self {
6286 Self::start_with_behavior(MockWorkerBehavior {
6287 reject_query_protocol: true,
6288 ..MockWorkerBehavior::default()
6289 })
6290 }
6291
6292 fn reject_query_completion() -> Self {
6293 Self::start_with_behavior(MockWorkerBehavior {
6294 reject_query_completion: true,
6295 ..MockWorkerBehavior::default()
6296 })
6297 }
6298
6299 fn waiting_query_worker() -> Self {
6300 Self::start_with_behavior(MockWorkerBehavior {
6301 waiting_query_worker: true,
6302 ..MockWorkerBehavior::default()
6303 })
6304 }
6305
6306 fn transient_worker_failures() -> Self {
6307 Self::start_with_behavior(MockWorkerBehavior {
6308 poll_failures_per_path: 1,
6309 heartbeat_failures: 1,
6310 ..MockWorkerBehavior::default()
6311 })
6312 }
6313
6314 fn unavailable_polls() -> Self {
6315 Self::start_with_behavior(MockWorkerBehavior {
6316 poll_failures_per_path: usize::MAX,
6317 ..MockWorkerBehavior::default()
6318 })
6319 }
6320
6321 fn unauthorized_polls() -> Self {
6322 Self::start_with_behavior(MockWorkerBehavior {
6323 unauthorized_polls: true,
6324 ..MockWorkerBehavior::default()
6325 })
6326 }
6327
6328 fn start_with_behavior(behavior: MockWorkerBehavior) -> Self {
6329 let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock server");
6330 listener
6331 .set_nonblocking(true)
6332 .expect("configure mock listener");
6333 let addr = listener.local_addr().expect("mock server address");
6334 let stop = Arc::new(AtomicBool::new(false));
6335 let server_stop = Arc::clone(&stop);
6336 let requests = Arc::new(Mutex::new(Vec::new()));
6337 let server_requests = Arc::clone(&requests);
6338 let thread = thread::spawn(move || {
6339 while !server_stop.load(Ordering::SeqCst) {
6340 match listener.accept() {
6341 Ok((mut stream, _)) => {
6342 handle_mock_worker_request(&mut stream, &server_requests, behavior)
6343 }
6344 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
6345 thread::sleep(Duration::from_millis(5));
6346 }
6347 Err(_) => break,
6348 }
6349 }
6350 });
6351
6352 Self {
6353 addr,
6354 stop,
6355 requests,
6356 thread: Some(thread),
6357 }
6358 }
6359
6360 fn base_url(&self) -> String {
6361 format!("http://{}", self.addr)
6362 }
6363
6364 fn worker_protocol_for(&self, path: &str) -> Option<String> {
6365 self.requests
6366 .lock()
6367 .expect("captured requests")
6368 .iter()
6369 .find(|request| request.path == path)
6370 .and_then(|request| request.worker_protocol.clone())
6371 }
6372
6373 fn request_count(&self, path: &str) -> usize {
6374 self.requests
6375 .lock()
6376 .expect("captured requests")
6377 .iter()
6378 .filter(|request| request.path == path)
6379 .count()
6380 }
6381
6382 fn request_body(&self, path: &str) -> Value {
6383 let requests = self.requests.lock().expect("captured requests");
6384 let body = &requests
6385 .iter()
6386 .find(|request| request.path == path)
6387 .unwrap_or_else(|| panic!("missing request for {path}"))
6388 .body;
6389 serde_json::from_str(body).unwrap_or_else(|error| {
6390 panic!("invalid JSON request body for {path}: {error}: {body:?}")
6391 })
6392 }
6393 }
6394
6395 impl Drop for MockWorkerServer {
6396 fn drop(&mut self) {
6397 self.stop.store(true, Ordering::SeqCst);
6398 let _ = TcpStream::connect(self.addr);
6399
6400 if let Some(thread) = self.thread.take() {
6401 thread.join().expect("join mock server");
6402 }
6403 }
6404 }
6405
6406 fn handle_mock_worker_request(
6407 stream: &mut TcpStream,
6408 requests: &Arc<Mutex<Vec<CapturedRequest>>>,
6409 behavior: MockWorkerBehavior,
6410 ) {
6411 let _ = stream.set_read_timeout(Some(Duration::from_millis(200)));
6412 let mut buffer = [0_u8; 8192];
6413 let mut request = Vec::new();
6414
6415 loop {
6416 match stream.read(&mut buffer) {
6417 Ok(0) => break,
6418 Ok(read) => {
6419 request.extend_from_slice(&buffer[..read]);
6420 if mock_request_is_complete(&request) {
6421 break;
6422 }
6423 }
6424 Err(error)
6425 if matches!(
6426 error.kind(),
6427 std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
6428 ) =>
6429 {
6430 break;
6431 }
6432 Err(_) => return,
6433 }
6434 }
6435
6436 let request = String::from_utf8_lossy(&request);
6437 let body = request
6438 .split_once("\r\n\r\n")
6439 .map(|(_, body)| body)
6440 .unwrap_or_default();
6441 let path = request
6442 .lines()
6443 .next()
6444 .and_then(|line| line.split_whitespace().nth(1))
6445 .unwrap_or_default();
6446 let worker_protocol = request.lines().find_map(|line| {
6447 let (name, value) = line.split_once(':')?;
6448 name.eq_ignore_ascii_case("X-Durable-Workflow-Protocol-Version")
6449 .then(|| value.trim().to_string())
6450 });
6451 let request_number = {
6452 let mut requests = requests.lock().expect("captured requests");
6453 requests.push(CapturedRequest {
6454 path: path.to_string(),
6455 worker_protocol: worker_protocol.clone(),
6456 body: body.to_string(),
6457 });
6458 requests
6459 .iter()
6460 .filter(|request| request.path == path)
6461 .count()
6462 };
6463
6464 let is_poll = matches!(
6465 path,
6466 "/api/worker/workflow-tasks/poll"
6467 | "/api/worker/activity-tasks/poll"
6468 | "/api/worker/query-tasks/poll"
6469 );
6470 if is_poll && request_number <= behavior.poll_failures_per_path {
6471 return;
6472 }
6473 if path == "/api/worker/heartbeat" && request_number <= behavior.heartbeat_failures {
6474 return;
6475 }
6476 if behavior.unauthorized_polls && is_poll {
6477 write_mock_response(
6478 stream,
6479 "401 Unauthorized",
6480 r#"{"reason":"authentication_failed","message":"invalid worker token"}"#,
6481 );
6482 return;
6483 }
6484
6485 if behavior.reject_query_protocol && path.starts_with("/api/worker/query-tasks/") {
6486 let requested_version = worker_protocol.as_deref().unwrap_or("missing");
6487 let body = format!(
6488 r#"{{"reason":"unsupported_protocol_version","message":"Query tasks require worker protocol 1.8 or newer.","supported_version":"1.7","requested_version":"{requested_version}"}}"#
6489 );
6490 write_mock_response(stream, "400 Bad Request", &body);
6491 return;
6492 }
6493
6494 if behavior.reject_query_completion && path == "/api/worker/query-tasks/query-late/complete"
6495 {
6496 write_mock_response(
6497 stream,
6498 "409 Conflict",
6499 r#"{"reason":"query_task_timed_out","message":"query task timed out before completion"}"#,
6500 );
6501 return;
6502 }
6503
6504 if behavior.waiting_query_worker {
6505 if path == "/api/worker/workflow-tasks/poll" && request_number <= 2 {
6506 let amounts = if request_number == 1 {
6507 vec![3]
6508 } else {
6509 vec![3, 5]
6510 };
6511 let task_id = if request_number == 1 {
6512 "snapshot-wait-3"
6513 } else {
6514 "snapshot-wait-5"
6515 };
6516 let history_events = std::iter::once(json!({
6517 "event_type": "SignalWaitOpened",
6518 "payload": {"sequence": 1, "signal_name": "finish"}
6519 }))
6520 .chain(amounts.iter().enumerate().map(|(index, amount)| {
6521 json!({
6522 "event_type": "SignalReceived",
6523 "payload": {
6524 "signal_id": format!("increment-{amount}"),
6525 "signal_name": "increment",
6526 "workflow_sequence": index + 2,
6527 "payload_codec": DEFAULT_CODEC,
6528 "arguments": encode_value_envelope(&json!([amount]), DEFAULT_CODEC)
6529 .expect("Avro signal envelope")
6530 }
6531 })
6532 }))
6533 .collect::<Vec<_>>();
6534 let body = json!({
6535 "task": {
6536 "task_id": task_id,
6537 "workflow_id": "snapshot-1",
6538 "run_id": "snapshot-run-1",
6539 "workflow_type": "snapshot",
6540 "payload_codec": DEFAULT_CODEC,
6541 "arguments": encode_value_envelope(&json!([]), DEFAULT_CODEC)
6542 .expect("Avro workflow arguments"),
6543 "history_events": history_events,
6544 "workflow_task_attempt": 1,
6545 "workflow_signal_id": format!("increment-{}", amounts.last().expect("amount")),
6546 "signal_name": "increment",
6547 "signal_arguments": encode_value_envelope(
6548 &json!([amounts.last().expect("amount")]),
6549 DEFAULT_CODEC,
6550 )
6551 .expect("Avro resume signal"),
6552 "lease_owner": "rust-snapshot-worker"
6553 }
6554 })
6555 .to_string();
6556 write_mock_response(stream, "200 OK", &body);
6557 return;
6558 }
6559
6560 if path == "/api/worker/query-tasks/poll" && request_number == 1 {
6561 let history_events = [3, 5]
6562 .into_iter()
6563 .enumerate()
6564 .map(|(index, amount)| {
6565 json!({
6566 "event_type": "SignalReceived",
6567 "payload": {
6568 "signal_id": format!("increment-{amount}"),
6569 "signal_name": "increment",
6570 "workflow_sequence": index + 2,
6571 "payload_codec": DEFAULT_CODEC,
6572 "arguments": encode_value_envelope(&json!([amount]), DEFAULT_CODEC)
6573 .expect("Avro query signal envelope")
6574 }
6575 })
6576 })
6577 .collect::<Vec<_>>();
6578 let body = json!({
6579 "task": {
6580 "query_task_id": "snapshot-current",
6581 "query_task_attempt": 1,
6582 "lease_owner": "rust-snapshot-worker",
6583 "workflow_id": "snapshot-1",
6584 "run_id": "snapshot-run-1",
6585 "workflow_type": "snapshot",
6586 "query_name": "current",
6587 "payload_codec": DEFAULT_CODEC,
6588 "workflow_arguments": encode_value_envelope(&json!([]), DEFAULT_CODEC)
6589 .expect("Avro workflow arguments"),
6590 "query_arguments": encode_value_envelope(&json!([]), DEFAULT_CODEC)
6591 .expect("Avro query arguments"),
6592 "history_events": history_events,
6593 "run_status": "waiting"
6594 }
6595 })
6596 .to_string();
6597 write_mock_response(stream, "200 OK", &body);
6598 return;
6599 }
6600
6601 if path == "/api/worker/workflow-tasks/snapshot-wait-3/fail"
6602 || path == "/api/worker/workflow-tasks/snapshot-wait-5/fail"
6603 {
6604 write_mock_response(
6605 stream,
6606 "200 OK",
6607 r#"{"outcome":"waiting_for_history","recorded":true}"#,
6608 );
6609 return;
6610 }
6611
6612 if path == "/api/worker/query-tasks/snapshot-current/complete" {
6613 write_mock_response(stream, "200 OK", r#"{"outcome":"completed"}"#);
6614 return;
6615 }
6616 }
6617
6618 let (status, body) = match path {
6619 "/api/worker/register" if behavior.waiting_query_worker => (
6620 "200 OK",
6621 r#"{"worker_id":"rust-snapshot-worker","registered":true,"heartbeat_interval_seconds":1}"#,
6622 ),
6623 "/api/worker/register" => (
6624 "200 OK",
6625 r#"{"worker_id":"mock-worker","registered":true,"heartbeat_interval_seconds":3600}"#,
6626 ),
6627 "/api/worker/heartbeat" => ("200 OK", "{}"),
6628 "/api/worker/activity-tasks/poll" | "/api/worker/workflow-tasks/poll" => {
6629 ("200 OK", r#"{"task":null}"#)
6630 }
6631 "/api/worker/query-tasks/poll"
6632 if behavior.reject_query_completion && request_number == 1 =>
6633 {
6634 (
6635 "200 OK",
6636 r#"{"task":{"query_task_id":"query-late","query_task_attempt":1,"lease_owner":"late-worker","workflow_id":"counter-late","run_id":"run-late","workflow_type":"counter","query_name":"current","payload_codec":"json","workflow_arguments":{"codec":"json","blob":"[]"},"query_arguments":{"codec":"json","blob":"[]"},"history_events":[],"run_status":"running"}}"#,
6637 )
6638 }
6639 "/api/worker/query-tasks/poll" => ("200 OK", r#"{"task":null}"#),
6640 "/api/worker/query-tasks/query-capture/complete"
6641 | "/api/worker/query-tasks/query-capture/fail" => ("200 OK", "{}"),
6642 "/api/workflows/counter-1/query/current" => (
6643 "200 OK",
6644 r#"{"workflow_id":"counter-1","query_name":"current","result":{"count":8},"result_envelope":{"codec":"json","blob":"{\"count\":8}"}}"#,
6645 ),
6646 "/api/workflows/counter-1/query/missing" => (
6647 "404 Not Found",
6648 r#"{"workflow_id":"counter-1","query_name":"missing","reason":"rejected_unknown_query","message":"unknown query"}"#,
6649 ),
6650 _ => ("404 Not Found", r#"{"message":"not found"}"#),
6651 };
6652 write_mock_response(stream, status, body);
6653 }
6654
6655 fn mock_request_is_complete(request: &[u8]) -> bool {
6656 let Some(header_end) = request
6657 .windows(4)
6658 .position(|window| window == b"\r\n\r\n")
6659 .map(|position| position + 4)
6660 else {
6661 return false;
6662 };
6663 let headers = String::from_utf8_lossy(&request[..header_end]);
6664 let content_length = headers.lines().find_map(|line| {
6665 let (name, value) = line.split_once(':')?;
6666 name.eq_ignore_ascii_case("content-length")
6667 .then(|| value.trim().parse::<usize>().ok())
6668 .flatten()
6669 });
6670
6671 request.len() >= header_end + content_length.unwrap_or(0)
6672 }
6673
6674 fn write_mock_response(stream: &mut TcpStream, status: &str, body: &str) {
6675 let response = format!(
6676 "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
6677 body.len()
6678 );
6679
6680 let _ = stream.write_all(response.as_bytes());
6681 let _ = stream.flush();
6682 }
6683}