1use std::num::NonZeroU64;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::task::{Context, Poll};
7
8use aion_core::{Event, WorkflowFilter, WorkflowId};
9use aion_proto::{
10 FilteredSubscription, FirehoseSubscription, PerWorkflowSubscription, ProtoWorkflowId,
11 SubscriptionRequest, subscription_request,
12};
13use futures::Stream;
14use futures::future::BoxFuture;
15use futures::stream::BoxStream;
16
17use crate::error::ClientError;
18use crate::transport::{SubscriptionAttempt, WorkflowTransport};
19
20pub type EventStream = Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>;
22
23#[derive(Clone, Debug, PartialEq, Eq)]
25pub enum SubscribeTarget {
26 Workflow {
28 workflow_id: WorkflowId,
30 },
31 Filtered {
33 filter: WorkflowFilter,
35 },
36 Firehose,
38}
39
40impl SubscribeTarget {
41 pub(crate) fn request(&self, namespace: &str) -> SubscriptionRequest {
42 match self {
43 Self::Workflow { workflow_id } => SubscriptionRequest {
44 subscription: Some(subscription_request::Subscription::PerWorkflow(
45 PerWorkflowSubscription {
46 namespace: namespace.to_owned(),
47 workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
48 resume_from_seq: None,
49 },
50 )),
51 },
52 Self::Filtered { filter } => SubscriptionRequest {
53 subscription: Some(subscription_request::Subscription::Filtered(
54 FilteredSubscription {
55 namespace: namespace.to_owned(),
56 workflow_type: filter.workflow_type.clone(),
57 status: filter
58 .status
59 .map(|status| aion_proto::ProtoWorkflowStatus::from(status) as i32),
60 namespace_selector: None,
61 },
62 )),
63 },
64 Self::Firehose => SubscriptionRequest {
65 subscription: Some(subscription_request::Subscription::Firehose(
66 FirehoseSubscription {
67 namespace: namespace.to_owned(),
68 },
69 )),
70 },
71 }
72 }
73}
74
75pub struct ResumingEventStream {
94 transport: Arc<dyn WorkflowTransport>,
95 namespace: String,
96 target: SubscribeTarget,
97 last_seq: Option<u64>,
98 delivered_any: bool,
99 current: Option<BoxStream<'static, Result<Event, ClientError>>>,
100 pending_subscribe: Option<BoxFuture<'static, Result<SubscriptionAttempt, ClientError>>>,
101 terminal_error: Option<ClientError>,
102 finished: bool,
103}
104
105impl ResumingEventStream {
106 #[must_use]
108 pub fn new(
109 transport: Arc<dyn WorkflowTransport>,
110 namespace: impl Into<String>,
111 target: SubscribeTarget,
112 ) -> Self {
113 Self {
114 transport,
115 namespace: namespace.into(),
116 target,
117 last_seq: None,
118 delivered_any: false,
119 current: None,
120 pending_subscribe: None,
121 terminal_error: None,
122 finished: false,
123 }
124 }
125
126 #[must_use]
134 pub fn from_sequence(
135 transport: Arc<dyn WorkflowTransport>,
136 namespace: impl Into<String>,
137 workflow_id: WorkflowId,
138 resume_from: NonZeroU64,
139 ) -> Self {
140 let mut stream = Self::new(
141 transport,
142 namespace,
143 SubscribeTarget::Workflow { workflow_id },
144 );
145 stream.last_seq = Some(resume_from.get() - 1);
149 stream
150 }
151
152 fn is_per_workflow(&self) -> bool {
153 matches!(self.target, SubscribeTarget::Workflow { .. })
154 }
155
156 fn start_subscribe(&mut self) {
157 let transport = Arc::clone(&self.transport);
158 let request = self.target.request(&self.namespace);
159 let resume_from_sequence = if self.is_per_workflow() {
162 self.last_seq.map(|seq| seq.saturating_add(1))
163 } else {
164 None
165 };
166 self.pending_subscribe = Some(Box::pin(async move {
167 transport.subscribe(request, resume_from_sequence).await
168 }));
169 }
170}
171
172impl Stream for ResumingEventStream {
173 type Item = Result<Event, ClientError>;
174
175 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
176 let this = self.get_mut();
177 loop {
178 if this.finished {
179 return Poll::Ready(None);
180 }
181
182 if let Some(error) = this.terminal_error.take() {
183 this.finished = true;
184 return Poll::Ready(Some(Err(error)));
185 }
186
187 if this.current.is_none() && this.pending_subscribe.is_none() {
188 this.start_subscribe();
189 }
190
191 if let Some(pending) = this.pending_subscribe.as_mut() {
192 match pending.as_mut().poll(cx) {
193 Poll::Pending => return Poll::Pending,
194 Poll::Ready(Ok(attempt)) => {
195 this.pending_subscribe = None;
196 this.current = Some(attempt.events);
197 }
198 Poll::Ready(Err(error)) => {
199 this.pending_subscribe = None;
206 if is_retryable(&error) && (this.is_per_workflow() || !this.delivered_any) {
207 continue;
208 }
209 this.finished = true;
210 return Poll::Ready(Some(Err(error)));
211 }
212 }
213 }
214
215 let Some(current) = this.current.as_mut() else {
216 continue;
217 };
218 match current.as_mut().poll_next(cx) {
219 Poll::Pending => return Poll::Pending,
220 Poll::Ready(Some(Ok(event))) => {
221 if this.is_per_workflow() {
222 if this.last_seq.is_some_and(|seq| event.seq() <= seq) {
225 continue;
226 }
227 this.last_seq = Some(event.seq());
228 }
229 this.delivered_any = true;
230 return Poll::Ready(Some(Ok(event)));
231 }
232 Poll::Ready(Some(Err(error))) => {
233 this.current = None;
234 if is_retryable(&error) {
235 if this.is_per_workflow() {
236 continue;
237 }
238 if !this.delivered_any {
239 continue;
242 }
243 }
247 this.terminal_error = Some(error);
248 }
249 Poll::Ready(None) => {
250 this.current = None;
251 this.finished = true;
252 return Poll::Ready(None);
253 }
254 }
255 }
256 }
257}
258
259#[must_use]
261pub fn event_stream(
262 transport: Arc<dyn WorkflowTransport>,
263 namespace: impl Into<String>,
264 target: SubscribeTarget,
265) -> EventStream {
266 Box::pin(ResumingEventStream::new(transport, namespace, target))
267}
268
269#[must_use]
271pub fn event_stream_from(
272 transport: Arc<dyn WorkflowTransport>,
273 namespace: impl Into<String>,
274 workflow_id: WorkflowId,
275 resume_from: NonZeroU64,
276) -> EventStream {
277 Box::pin(ResumingEventStream::from_sequence(
278 transport,
279 namespace,
280 workflow_id,
281 resume_from,
282 ))
283}
284
285fn is_retryable(error: &ClientError) -> bool {
286 matches!(error, ClientError::Unavailable { .. })
287}
288
289#[cfg(test)]
290mod tests {
291 use std::collections::VecDeque;
292 use std::sync::Arc;
293
294 use aion_core::{ContentType, Event, EventEnvelope, Payload, WorkflowId};
295 use aion_proto::{
296 ProtoCancelResponse, ProtoDescribeWorkflowResponse, ProtoListWorkflowsResponse,
297 ProtoQueryResponse, ProtoSignalResponse, ProtoStartWorkflowResponse,
298 };
299 use async_trait::async_trait;
300 use chrono::Utc;
301 use futures::StreamExt;
302 use futures::stream;
303 use tokio::sync::Mutex;
304
305 use super::{ResumingEventStream, SubscribeTarget};
306 use crate::error::ClientError;
307 use crate::transport::{SubscriptionAttempt, WorkflowTransport};
308
309 #[derive(Default)]
310 struct SubscribeStub {
311 attach_failures: Mutex<VecDeque<ClientError>>,
314 attempts: Mutex<VecDeque<SubscriptionAttempt>>,
315 resume_points: Mutex<Vec<Option<u64>>>,
316 }
317
318 #[async_trait]
319 impl WorkflowTransport for SubscribeStub {
320 async fn start_workflow(
321 &self,
322 _: aion_proto::ProtoStartWorkflowRequest,
323 ) -> Result<ProtoStartWorkflowResponse, ClientError> {
324 Err(ClientError::unavailable("stub transport"))
325 }
326
327 async fn signal(
328 &self,
329 _: aion_proto::ProtoSignalRequest,
330 ) -> Result<ProtoSignalResponse, ClientError> {
331 Err(ClientError::unavailable("stub transport"))
332 }
333
334 async fn query(
335 &self,
336 _: aion_proto::ProtoQueryRequest,
337 ) -> Result<ProtoQueryResponse, ClientError> {
338 Err(ClientError::unavailable("stub transport"))
339 }
340
341 async fn cancel(
342 &self,
343 _: aion_proto::ProtoCancelRequest,
344 ) -> Result<ProtoCancelResponse, ClientError> {
345 Err(ClientError::unavailable("stub transport"))
346 }
347
348 async fn reopen(
349 &self,
350 _: aion_proto::ProtoReopenRequest,
351 ) -> Result<aion_proto::ProtoReopenResponse, ClientError> {
352 Err(ClientError::unavailable("stub transport"))
353 }
354
355 async fn pause(
356 &self,
357 _: aion_proto::ProtoPauseRequest,
358 ) -> Result<aion_proto::ProtoPauseResponse, ClientError> {
359 Err(ClientError::unavailable("stub transport"))
360 }
361
362 async fn resume(
363 &self,
364 _: aion_proto::ProtoResumeRequest,
365 ) -> Result<aion_proto::ProtoResumeResponse, ClientError> {
366 Err(ClientError::unavailable("stub transport"))
367 }
368
369 async fn list_workflows(
370 &self,
371 _: aion_proto::ProtoListWorkflowsRequest,
372 ) -> Result<ProtoListWorkflowsResponse, ClientError> {
373 Err(ClientError::unavailable("stub transport"))
374 }
375
376 async fn describe_workflow(
377 &self,
378 _: aion_proto::ProtoDescribeWorkflowRequest,
379 ) -> Result<ProtoDescribeWorkflowResponse, ClientError> {
380 Err(ClientError::unavailable("stub transport"))
381 }
382
383 async fn subscribe(
384 &self,
385 _: aion_proto::SubscriptionRequest,
386 resume_from_sequence: Option<u64>,
387 ) -> Result<SubscriptionAttempt, ClientError> {
388 self.resume_points.lock().await.push(resume_from_sequence);
389 if let Some(failure) = self.attach_failures.lock().await.pop_front() {
390 return Err(failure);
391 }
392 self.attempts
393 .lock()
394 .await
395 .pop_front()
396 .ok_or_else(|| ClientError::server("missing subscribe attempt"))
397 }
398 }
399
400 fn event(seq: u64, workflow_id: &WorkflowId) -> Event {
401 Event::WorkflowStarted {
402 envelope: EventEnvelope {
403 seq,
404 recorded_at: Utc::now(),
405 workflow_id: workflow_id.clone(),
406 },
407 workflow_type: String::from("checkout"),
408 input: Payload::new(ContentType::Json, Vec::new()),
409 run_id: aion_core::RunId::new(uuid::Uuid::from_u128(1)),
410 parent_run_id: None,
411 parent_workflow_id: None,
412 package_version: aion_core::PackageVersion::new("a".repeat(64)),
413 }
414 }
415
416 #[tokio::test]
417 async fn resumes_after_transient_disconnect_without_gaps_or_duplicates() {
418 let workflow_id = WorkflowId::new_v4();
419 let stub = Arc::new(SubscribeStub::default());
420 stub.attempts
421 .lock()
422 .await
423 .push_back(SubscriptionAttempt::new(
424 stream::iter(vec![
425 Ok(event(1, &workflow_id)),
426 Ok(event(2, &workflow_id)),
427 Err(ClientError::unavailable("transient disconnect")),
428 ])
429 .boxed(),
430 ));
431 stub.attempts
432 .lock()
433 .await
434 .push_back(SubscriptionAttempt::new(
435 stream::iter(vec![
436 Ok(event(2, &workflow_id)),
437 Ok(event(3, &workflow_id)),
438 Ok(event(4, &workflow_id)),
439 ])
440 .boxed(),
441 ));
442 let mut events = ResumingEventStream::new(
443 stub.clone(),
444 "tenant-a",
445 SubscribeTarget::Workflow {
446 workflow_id: workflow_id.clone(),
447 },
448 );
449
450 let mut seqs = Vec::new();
451 while let Some(item) = events.next().await {
452 let event = item
453 .map_err(|e| format!("unexpected stream error: {e}"))
454 .ok();
455 if let Some(event) = event {
456 seqs.push(event.seq());
457 }
458 }
459
460 assert_eq!(seqs, vec![1, 2, 3, 4]);
461 assert_eq!(*stub.resume_points.lock().await, vec![None, Some(3)]);
462 }
463
464 #[tokio::test]
465 async fn terminal_failure_is_yielded_before_end() {
466 let workflow_id = WorkflowId::new_v4();
467 let stub = Arc::new(SubscribeStub::default());
468 stub.attempts
469 .lock()
470 .await
471 .push_back(SubscriptionAttempt::new(
472 stream::iter(vec![Err(ClientError::unauthenticated("bad token"))]).boxed(),
473 ));
474 let mut events =
475 ResumingEventStream::new(stub, "tenant-a", SubscribeTarget::Workflow { workflow_id });
476
477 assert_eq!(
478 events.next().await,
479 Some(Err(ClientError::unauthenticated("bad token")))
480 );
481 assert_eq!(events.next().await, None);
482 }
483
484 #[tokio::test]
485 async fn namespace_denied_is_terminal_and_never_retried() {
486 let workflow_id = WorkflowId::new_v4();
487 let stub = Arc::new(SubscribeStub::default());
488 let denied =
489 ClientError::namespace_denied("namespace tenant-b is not granted to this caller");
490 stub.attempts
491 .lock()
492 .await
493 .push_back(SubscriptionAttempt::new(
494 stream::iter(vec![Err(denied.clone())]).boxed(),
495 ));
496 let mut events = ResumingEventStream::new(
497 stub.clone(),
498 "tenant-b",
499 SubscribeTarget::Workflow { workflow_id },
500 );
501
502 assert_eq!(events.next().await, Some(Err(denied)));
503 assert_eq!(events.next().await, None);
504 assert_eq!(stub.resume_points.lock().await.len(), 1);
505 }
506
507 #[tokio::test]
508 async fn from_sequence_passes_the_cursor_on_the_initial_attach() {
509 let workflow_id = WorkflowId::new_v4();
510 let stub = Arc::new(SubscribeStub::default());
511 stub.attempts
512 .lock()
513 .await
514 .push_back(SubscriptionAttempt::new(
515 stream::iter(vec![Ok(event(1, &workflow_id)), Ok(event(2, &workflow_id))]).boxed(),
516 ));
517 let Some(resume_from) = std::num::NonZeroU64::new(1) else {
518 unreachable!("1 is non-zero");
519 };
520 let mut events = super::ResumingEventStream::from_sequence(
521 stub.clone(),
522 "tenant-a",
523 workflow_id,
524 resume_from,
525 );
526
527 let mut seqs = Vec::new();
528 while let Some(item) = events.next().await {
529 if let Ok(event) = item {
530 seqs.push(event.seq());
531 }
532 }
533
534 assert_eq!(seqs, vec![1, 2]);
535 assert_eq!(
536 *stub.resume_points.lock().await,
537 vec![Some(1)],
538 "the initial attach must carry the explicit cursor"
539 );
540 }
541
542 #[tokio::test]
543 async fn live_only_streams_reconnect_only_before_any_delivery() {
544 let workflow_id = WorkflowId::new_v4();
547 let stub = Arc::new(SubscribeStub::default());
548 stub.attempts
549 .lock()
550 .await
551 .push_back(SubscriptionAttempt::new(
552 stream::iter(vec![Err(ClientError::unavailable("transient disconnect"))]).boxed(),
553 ));
554 stub.attempts
555 .lock()
556 .await
557 .push_back(SubscriptionAttempt::new(
558 stream::iter(vec![Ok(event(1, &workflow_id))]).boxed(),
559 ));
560 let mut events = ResumingEventStream::new(
561 stub.clone(),
562 "tenant-a",
563 SubscribeTarget::Filtered {
564 filter: aion_core::WorkflowFilter::default(),
565 },
566 );
567
568 let mut seqs = Vec::new();
569 while let Some(item) = events.next().await {
570 if let Ok(event) = item {
571 seqs.push(event.seq());
572 }
573 }
574
575 assert_eq!(seqs, vec![1]);
576 assert_eq!(
577 *stub.resume_points.lock().await,
578 vec![None, None],
579 "live-only streams never carry a resume cursor"
580 );
581 }
582
583 #[tokio::test]
584 async fn live_only_disconnect_after_delivery_is_honest_unavailable() {
585 for target in [
589 SubscribeTarget::Filtered {
590 filter: aion_core::WorkflowFilter::default(),
591 },
592 SubscribeTarget::Firehose,
593 ] {
594 let workflow_id = WorkflowId::new_v4();
595 let stub = Arc::new(SubscribeStub::default());
596 stub.attempts
597 .lock()
598 .await
599 .push_back(SubscriptionAttempt::new(
600 stream::iter(vec![
601 Ok(event(1, &workflow_id)),
602 Err(ClientError::unavailable("transient disconnect")),
603 ])
604 .boxed(),
605 ));
606 let mut events = ResumingEventStream::new(stub.clone(), "tenant-a", target);
607
608 let first = events.next().await;
609 assert!(matches!(first, Some(Ok(_))), "got {first:?}");
610 assert_eq!(
611 events.next().await,
612 Some(Err(ClientError::unavailable("transient disconnect")))
613 );
614 assert_eq!(events.next().await, None);
615 assert_eq!(
616 stub.resume_points.lock().await.len(),
617 1,
618 "no reattach may follow a post-delivery live-only disconnect"
619 );
620 }
621 }
622
623 #[tokio::test]
624 async fn live_only_streams_do_not_dedupe_sequence_numbers_across_workflows() {
625 let first_workflow = WorkflowId::new_v4();
628 let second_workflow = WorkflowId::new_v4();
629 let stub = Arc::new(SubscribeStub::default());
630 stub.attempts
631 .lock()
632 .await
633 .push_back(SubscriptionAttempt::new(
634 stream::iter(vec![
635 Ok(event(1, &first_workflow)),
636 Ok(event(1, &second_workflow)),
637 ])
638 .boxed(),
639 ));
640 let mut events = ResumingEventStream::new(stub, "tenant-a", SubscribeTarget::Firehose);
641
642 let mut delivered = Vec::new();
643 while let Some(item) = events.next().await {
644 if let Ok(event) = item {
645 delivered.push(event.envelope().workflow_id.clone());
646 }
647 }
648
649 assert_eq!(delivered, vec![first_workflow, second_workflow]);
650 }
651
652 #[tokio::test]
653 async fn not_found_is_terminal_and_never_retried() {
654 let workflow_id = WorkflowId::new_v4();
658 let stub = Arc::new(SubscribeStub::default());
659 stub.attempts
660 .lock()
661 .await
662 .push_back(SubscriptionAttempt::new(
663 stream::iter(vec![Err(ClientError::not_found("workflow was not found"))]).boxed(),
664 ));
665 let mut events = ResumingEventStream::new(
666 stub.clone(),
667 "tenant-a",
668 SubscribeTarget::Workflow { workflow_id },
669 );
670
671 assert_eq!(
672 events.next().await,
673 Some(Err(ClientError::not_found("workflow was not found")))
674 );
675 assert_eq!(events.next().await, None);
676 assert_eq!(stub.resume_points.lock().await.len(), 1);
677 }
678
679 #[tokio::test]
683 async fn unavailable_attach_failure_is_retried_until_attach_succeeds() -> Result<(), ClientError>
684 {
685 let workflow_id = WorkflowId::new_v4();
686 let stub = Arc::new(SubscribeStub::default());
687 stub.attach_failures
688 .lock()
689 .await
690 .push_back(ClientError::unavailable("connection refused"));
691 stub.attach_failures
692 .lock()
693 .await
694 .push_back(ClientError::unavailable("connection refused"));
695 stub.attempts
696 .lock()
697 .await
698 .push_back(SubscriptionAttempt::new(
699 stream::iter(vec![Ok(event(1, &workflow_id)), Ok(event(2, &workflow_id))]).boxed(),
700 ));
701 let mut events = ResumingEventStream::new(
702 stub.clone(),
703 "tenant-a",
704 SubscribeTarget::Workflow { workflow_id },
705 );
706
707 let mut seqs = Vec::new();
708 while let Some(item) = events.next().await {
709 seqs.push(item?.seq());
712 }
713
714 assert_eq!(seqs, vec![1, 2]);
715 assert_eq!(
716 *stub.resume_points.lock().await,
717 vec![None, None, None],
718 "every retried initial attach is still a live tail (no cursor)"
719 );
720 Ok(())
721 }
722
723 #[tokio::test]
726 async fn unavailable_reconnect_failure_retries_with_the_same_cursor() -> Result<(), ClientError>
727 {
728 let workflow_id = WorkflowId::new_v4();
729 let stub = Arc::new(SubscribeStub::default());
730 stub.attempts
731 .lock()
732 .await
733 .push_back(SubscriptionAttempt::new(
734 stream::iter(vec![
735 Ok(event(1, &workflow_id)),
736 Err(ClientError::unavailable("transient disconnect")),
737 ])
738 .boxed(),
739 ));
740 let mut events = ResumingEventStream::new(
741 stub.clone(),
742 "tenant-a",
743 SubscribeTarget::Workflow {
744 workflow_id: workflow_id.clone(),
745 },
746 );
747 let first = events.next().await;
748 assert!(matches!(first, Some(Ok(_))), "got {first:?}");
749 stub.attach_failures
751 .lock()
752 .await
753 .push_back(ClientError::unavailable("connection refused"));
754 stub.attempts
755 .lock()
756 .await
757 .push_back(SubscriptionAttempt::new(
758 stream::iter(vec![Ok(event(2, &workflow_id))]).boxed(),
759 ));
760
761 let mut seqs = vec![1];
762 while let Some(item) = events.next().await {
763 seqs.push(item?.seq());
766 }
767
768 assert_eq!(seqs, vec![1, 2]);
769 assert_eq!(
770 *stub.resume_points.lock().await,
771 vec![None, Some(2), Some(2)],
772 "the failed reconnect and the successful retry carry the same cursor"
773 );
774 Ok(())
775 }
776
777 #[tokio::test]
780 async fn non_retryable_attach_failure_is_terminal() {
781 let workflow_id = WorkflowId::new_v4();
782 let stub = Arc::new(SubscribeStub::default());
783 stub.attach_failures
784 .lock()
785 .await
786 .push_back(ClientError::unauthenticated("bad token"));
787 let mut events = ResumingEventStream::new(
788 stub.clone(),
789 "tenant-a",
790 SubscribeTarget::Workflow { workflow_id },
791 );
792
793 assert_eq!(
794 events.next().await,
795 Some(Err(ClientError::unauthenticated("bad token")))
796 );
797 assert_eq!(events.next().await, None);
798 assert_eq!(stub.resume_points.lock().await.len(), 1);
799 }
800
801 #[tokio::test]
804 async fn live_only_unavailable_attach_failure_is_retried_before_any_delivery() {
805 let workflow_id = WorkflowId::new_v4();
806 let stub = Arc::new(SubscribeStub::default());
807 stub.attach_failures
808 .lock()
809 .await
810 .push_back(ClientError::unavailable("connection refused"));
811 stub.attempts
812 .lock()
813 .await
814 .push_back(SubscriptionAttempt::new(
815 stream::iter(vec![Ok(event(1, &workflow_id))]).boxed(),
816 ));
817 let mut events =
818 ResumingEventStream::new(stub.clone(), "tenant-a", SubscribeTarget::Firehose);
819
820 let mut seqs = Vec::new();
821 while let Some(item) = events.next().await {
822 if let Ok(event) = item {
823 seqs.push(event.seq());
824 }
825 }
826
827 assert_eq!(seqs, vec![1]);
828 assert_eq!(*stub.resume_points.lock().await, vec![None, None]);
829 }
830}