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 read_history(
384 &self,
385 _: aion_proto::ProtoReadHistoryRequest,
386 ) -> Result<aion_proto::ProtoReadHistoryResponse, ClientError> {
387 Err(ClientError::unavailable("stub transport"))
388 }
389
390 async fn subscribe(
391 &self,
392 _: aion_proto::SubscriptionRequest,
393 resume_from_sequence: Option<u64>,
394 ) -> Result<SubscriptionAttempt, ClientError> {
395 self.resume_points.lock().await.push(resume_from_sequence);
396 if let Some(failure) = self.attach_failures.lock().await.pop_front() {
397 return Err(failure);
398 }
399 self.attempts
400 .lock()
401 .await
402 .pop_front()
403 .ok_or_else(|| ClientError::server("missing subscribe attempt"))
404 }
405 }
406
407 fn event(seq: u64, workflow_id: &WorkflowId) -> Event {
408 Event::WorkflowStarted {
409 envelope: EventEnvelope {
410 seq,
411 recorded_at: Utc::now(),
412 workflow_id: workflow_id.clone(),
413 },
414 workflow_type: String::from("checkout"),
415 input: Payload::new(ContentType::Json, Vec::new()),
416 run_id: aion_core::RunId::new(uuid::Uuid::from_u128(1)),
417 parent_run_id: None,
418 parent_workflow_id: None,
419 package_version: aion_core::PackageVersion::new("a".repeat(64)),
420 }
421 }
422
423 #[tokio::test]
424 async fn resumes_after_transient_disconnect_without_gaps_or_duplicates() {
425 let workflow_id = WorkflowId::new_v4();
426 let stub = Arc::new(SubscribeStub::default());
427 stub.attempts
428 .lock()
429 .await
430 .push_back(SubscriptionAttempt::new(
431 stream::iter(vec![
432 Ok(event(1, &workflow_id)),
433 Ok(event(2, &workflow_id)),
434 Err(ClientError::unavailable("transient disconnect")),
435 ])
436 .boxed(),
437 ));
438 stub.attempts
439 .lock()
440 .await
441 .push_back(SubscriptionAttempt::new(
442 stream::iter(vec![
443 Ok(event(2, &workflow_id)),
444 Ok(event(3, &workflow_id)),
445 Ok(event(4, &workflow_id)),
446 ])
447 .boxed(),
448 ));
449 let mut events = ResumingEventStream::new(
450 stub.clone(),
451 "tenant-a",
452 SubscribeTarget::Workflow {
453 workflow_id: workflow_id.clone(),
454 },
455 );
456
457 let mut seqs = Vec::new();
458 while let Some(item) = events.next().await {
459 let event = item
460 .map_err(|e| format!("unexpected stream error: {e}"))
461 .ok();
462 if let Some(event) = event {
463 seqs.push(event.seq());
464 }
465 }
466
467 assert_eq!(seqs, vec![1, 2, 3, 4]);
468 assert_eq!(*stub.resume_points.lock().await, vec![None, Some(3)]);
469 }
470
471 #[tokio::test]
472 async fn terminal_failure_is_yielded_before_end() {
473 let workflow_id = WorkflowId::new_v4();
474 let stub = Arc::new(SubscribeStub::default());
475 stub.attempts
476 .lock()
477 .await
478 .push_back(SubscriptionAttempt::new(
479 stream::iter(vec![Err(ClientError::unauthenticated("bad token"))]).boxed(),
480 ));
481 let mut events =
482 ResumingEventStream::new(stub, "tenant-a", SubscribeTarget::Workflow { workflow_id });
483
484 assert_eq!(
485 events.next().await,
486 Some(Err(ClientError::unauthenticated("bad token")))
487 );
488 assert_eq!(events.next().await, None);
489 }
490
491 #[tokio::test]
492 async fn namespace_denied_is_terminal_and_never_retried() {
493 let workflow_id = WorkflowId::new_v4();
494 let stub = Arc::new(SubscribeStub::default());
495 let denied =
496 ClientError::namespace_denied("namespace tenant-b is not granted to this caller");
497 stub.attempts
498 .lock()
499 .await
500 .push_back(SubscriptionAttempt::new(
501 stream::iter(vec![Err(denied.clone())]).boxed(),
502 ));
503 let mut events = ResumingEventStream::new(
504 stub.clone(),
505 "tenant-b",
506 SubscribeTarget::Workflow { workflow_id },
507 );
508
509 assert_eq!(events.next().await, Some(Err(denied)));
510 assert_eq!(events.next().await, None);
511 assert_eq!(stub.resume_points.lock().await.len(), 1);
512 }
513
514 #[tokio::test]
515 async fn from_sequence_passes_the_cursor_on_the_initial_attach() {
516 let workflow_id = WorkflowId::new_v4();
517 let stub = Arc::new(SubscribeStub::default());
518 stub.attempts
519 .lock()
520 .await
521 .push_back(SubscriptionAttempt::new(
522 stream::iter(vec![Ok(event(1, &workflow_id)), Ok(event(2, &workflow_id))]).boxed(),
523 ));
524 let Some(resume_from) = std::num::NonZeroU64::new(1) else {
525 unreachable!("1 is non-zero");
526 };
527 let mut events = super::ResumingEventStream::from_sequence(
528 stub.clone(),
529 "tenant-a",
530 workflow_id,
531 resume_from,
532 );
533
534 let mut seqs = Vec::new();
535 while let Some(item) = events.next().await {
536 if let Ok(event) = item {
537 seqs.push(event.seq());
538 }
539 }
540
541 assert_eq!(seqs, vec![1, 2]);
542 assert_eq!(
543 *stub.resume_points.lock().await,
544 vec![Some(1)],
545 "the initial attach must carry the explicit cursor"
546 );
547 }
548
549 #[tokio::test]
550 async fn live_only_streams_reconnect_only_before_any_delivery() {
551 let workflow_id = WorkflowId::new_v4();
554 let stub = Arc::new(SubscribeStub::default());
555 stub.attempts
556 .lock()
557 .await
558 .push_back(SubscriptionAttempt::new(
559 stream::iter(vec![Err(ClientError::unavailable("transient disconnect"))]).boxed(),
560 ));
561 stub.attempts
562 .lock()
563 .await
564 .push_back(SubscriptionAttempt::new(
565 stream::iter(vec![Ok(event(1, &workflow_id))]).boxed(),
566 ));
567 let mut events = ResumingEventStream::new(
568 stub.clone(),
569 "tenant-a",
570 SubscribeTarget::Filtered {
571 filter: aion_core::WorkflowFilter::default(),
572 },
573 );
574
575 let mut seqs = Vec::new();
576 while let Some(item) = events.next().await {
577 if let Ok(event) = item {
578 seqs.push(event.seq());
579 }
580 }
581
582 assert_eq!(seqs, vec![1]);
583 assert_eq!(
584 *stub.resume_points.lock().await,
585 vec![None, None],
586 "live-only streams never carry a resume cursor"
587 );
588 }
589
590 #[tokio::test]
591 async fn live_only_disconnect_after_delivery_is_honest_unavailable() {
592 for target in [
596 SubscribeTarget::Filtered {
597 filter: aion_core::WorkflowFilter::default(),
598 },
599 SubscribeTarget::Firehose,
600 ] {
601 let workflow_id = WorkflowId::new_v4();
602 let stub = Arc::new(SubscribeStub::default());
603 stub.attempts
604 .lock()
605 .await
606 .push_back(SubscriptionAttempt::new(
607 stream::iter(vec![
608 Ok(event(1, &workflow_id)),
609 Err(ClientError::unavailable("transient disconnect")),
610 ])
611 .boxed(),
612 ));
613 let mut events = ResumingEventStream::new(stub.clone(), "tenant-a", target);
614
615 let first = events.next().await;
616 assert!(matches!(first, Some(Ok(_))), "got {first:?}");
617 assert_eq!(
618 events.next().await,
619 Some(Err(ClientError::unavailable("transient disconnect")))
620 );
621 assert_eq!(events.next().await, None);
622 assert_eq!(
623 stub.resume_points.lock().await.len(),
624 1,
625 "no reattach may follow a post-delivery live-only disconnect"
626 );
627 }
628 }
629
630 #[tokio::test]
631 async fn live_only_streams_do_not_dedupe_sequence_numbers_across_workflows() {
632 let first_workflow = WorkflowId::new_v4();
635 let second_workflow = WorkflowId::new_v4();
636 let stub = Arc::new(SubscribeStub::default());
637 stub.attempts
638 .lock()
639 .await
640 .push_back(SubscriptionAttempt::new(
641 stream::iter(vec![
642 Ok(event(1, &first_workflow)),
643 Ok(event(1, &second_workflow)),
644 ])
645 .boxed(),
646 ));
647 let mut events = ResumingEventStream::new(stub, "tenant-a", SubscribeTarget::Firehose);
648
649 let mut delivered = Vec::new();
650 while let Some(item) = events.next().await {
651 if let Ok(event) = item {
652 delivered.push(event.envelope().workflow_id.clone());
653 }
654 }
655
656 assert_eq!(delivered, vec![first_workflow, second_workflow]);
657 }
658
659 #[tokio::test]
660 async fn not_found_is_terminal_and_never_retried() {
661 let workflow_id = WorkflowId::new_v4();
665 let stub = Arc::new(SubscribeStub::default());
666 stub.attempts
667 .lock()
668 .await
669 .push_back(SubscriptionAttempt::new(
670 stream::iter(vec![Err(ClientError::not_found("workflow was not found"))]).boxed(),
671 ));
672 let mut events = ResumingEventStream::new(
673 stub.clone(),
674 "tenant-a",
675 SubscribeTarget::Workflow { workflow_id },
676 );
677
678 assert_eq!(
679 events.next().await,
680 Some(Err(ClientError::not_found("workflow was not found")))
681 );
682 assert_eq!(events.next().await, None);
683 assert_eq!(stub.resume_points.lock().await.len(), 1);
684 }
685
686 #[tokio::test]
690 async fn unavailable_attach_failure_is_retried_until_attach_succeeds() -> Result<(), ClientError>
691 {
692 let workflow_id = WorkflowId::new_v4();
693 let stub = Arc::new(SubscribeStub::default());
694 stub.attach_failures
695 .lock()
696 .await
697 .push_back(ClientError::unavailable("connection refused"));
698 stub.attach_failures
699 .lock()
700 .await
701 .push_back(ClientError::unavailable("connection refused"));
702 stub.attempts
703 .lock()
704 .await
705 .push_back(SubscriptionAttempt::new(
706 stream::iter(vec![Ok(event(1, &workflow_id)), Ok(event(2, &workflow_id))]).boxed(),
707 ));
708 let mut events = ResumingEventStream::new(
709 stub.clone(),
710 "tenant-a",
711 SubscribeTarget::Workflow { workflow_id },
712 );
713
714 let mut seqs = Vec::new();
715 while let Some(item) = events.next().await {
716 seqs.push(item?.seq());
719 }
720
721 assert_eq!(seqs, vec![1, 2]);
722 assert_eq!(
723 *stub.resume_points.lock().await,
724 vec![None, None, None],
725 "every retried initial attach is still a live tail (no cursor)"
726 );
727 Ok(())
728 }
729
730 #[tokio::test]
733 async fn unavailable_reconnect_failure_retries_with_the_same_cursor() -> Result<(), ClientError>
734 {
735 let workflow_id = WorkflowId::new_v4();
736 let stub = Arc::new(SubscribeStub::default());
737 stub.attempts
738 .lock()
739 .await
740 .push_back(SubscriptionAttempt::new(
741 stream::iter(vec![
742 Ok(event(1, &workflow_id)),
743 Err(ClientError::unavailable("transient disconnect")),
744 ])
745 .boxed(),
746 ));
747 let mut events = ResumingEventStream::new(
748 stub.clone(),
749 "tenant-a",
750 SubscribeTarget::Workflow {
751 workflow_id: workflow_id.clone(),
752 },
753 );
754 let first = events.next().await;
755 assert!(matches!(first, Some(Ok(_))), "got {first:?}");
756 stub.attach_failures
758 .lock()
759 .await
760 .push_back(ClientError::unavailable("connection refused"));
761 stub.attempts
762 .lock()
763 .await
764 .push_back(SubscriptionAttempt::new(
765 stream::iter(vec![Ok(event(2, &workflow_id))]).boxed(),
766 ));
767
768 let mut seqs = vec![1];
769 while let Some(item) = events.next().await {
770 seqs.push(item?.seq());
773 }
774
775 assert_eq!(seqs, vec![1, 2]);
776 assert_eq!(
777 *stub.resume_points.lock().await,
778 vec![None, Some(2), Some(2)],
779 "the failed reconnect and the successful retry carry the same cursor"
780 );
781 Ok(())
782 }
783
784 #[tokio::test]
787 async fn non_retryable_attach_failure_is_terminal() {
788 let workflow_id = WorkflowId::new_v4();
789 let stub = Arc::new(SubscribeStub::default());
790 stub.attach_failures
791 .lock()
792 .await
793 .push_back(ClientError::unauthenticated("bad token"));
794 let mut events = ResumingEventStream::new(
795 stub.clone(),
796 "tenant-a",
797 SubscribeTarget::Workflow { workflow_id },
798 );
799
800 assert_eq!(
801 events.next().await,
802 Some(Err(ClientError::unauthenticated("bad token")))
803 );
804 assert_eq!(events.next().await, None);
805 assert_eq!(stub.resume_points.lock().await.len(), 1);
806 }
807
808 #[tokio::test]
811 async fn live_only_unavailable_attach_failure_is_retried_before_any_delivery() {
812 let workflow_id = WorkflowId::new_v4();
813 let stub = Arc::new(SubscribeStub::default());
814 stub.attach_failures
815 .lock()
816 .await
817 .push_back(ClientError::unavailable("connection refused"));
818 stub.attempts
819 .lock()
820 .await
821 .push_back(SubscriptionAttempt::new(
822 stream::iter(vec![Ok(event(1, &workflow_id))]).boxed(),
823 ));
824 let mut events =
825 ResumingEventStream::new(stub.clone(), "tenant-a", SubscribeTarget::Firehose);
826
827 let mut seqs = Vec::new();
828 while let Some(item) = events.next().await {
829 if let Ok(event) = item {
830 seqs.push(event.seq());
831 }
832 }
833
834 assert_eq!(seqs, vec![1]);
835 assert_eq!(*stub.resume_points.lock().await, vec![None, None]);
836 }
837}