1use std::{
2 collections::VecDeque,
3 fmt,
4 sync::{
5 Arc, Condvar, Mutex, MutexGuard, Weak,
6 atomic::{AtomicBool, Ordering},
7 },
8 time::{Duration, Instant},
9};
10
11use futures::channel::oneshot;
12
13use crate::stream::{
14 BoxStream, Cancellable, NotUsed, Sink, Source, StreamCancellation, StreamCompletion,
15 StreamError, StreamResult,
16};
17
18use super::{Actor, ActorProcessingErr, ActorRef, ActorResult, Message, block_on_ractor_runtime};
19
20const DEFAULT_STREAM_REF_BUFFER_CAPACITY: usize = 32;
21const DEFAULT_STREAM_REF_SUBSCRIPTION_TIMEOUT: Duration = Duration::from_secs(30);
22const DEFAULT_STREAM_REF_DEMAND_REDELIVERY: Duration = Duration::from_secs(1);
23const STREAM_REF_WAIT_POLL: Duration = Duration::from_millis(1);
24
25mod stream_ref_actor_bound {
26 pub trait Bound {}
27
28 impl<T> Bound for T {}
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct StreamRefSettings {
38 buffer_capacity: usize,
39 subscription_timeout: Duration,
40 demand_redelivery_interval: Duration,
41}
42
43impl Default for StreamRefSettings {
44 fn default() -> Self {
45 Self {
46 buffer_capacity: DEFAULT_STREAM_REF_BUFFER_CAPACITY,
47 subscription_timeout: DEFAULT_STREAM_REF_SUBSCRIPTION_TIMEOUT,
48 demand_redelivery_interval: DEFAULT_STREAM_REF_DEMAND_REDELIVERY,
49 }
50 }
51}
52
53impl StreamRefSettings {
54 #[must_use]
55 pub fn buffer_capacity(&self) -> usize {
56 self.buffer_capacity
57 }
58
59 #[must_use]
60 pub fn subscription_timeout(&self) -> Duration {
61 self.subscription_timeout
62 }
63
64 #[must_use]
65 pub fn demand_redelivery_interval(&self) -> Duration {
66 self.demand_redelivery_interval
67 }
68
69 #[must_use]
70 pub fn with_buffer_capacity(mut self, capacity: usize) -> Self {
71 assert!(
72 capacity > 0,
73 "StreamRef buffer capacity must be greater than zero"
74 );
75 self.buffer_capacity = capacity;
76 self
77 }
78
79 #[must_use]
80 pub fn with_subscription_timeout(mut self, timeout: Duration) -> Self {
81 self.subscription_timeout = timeout;
82 self
83 }
84
85 #[must_use]
86 pub fn with_demand_redelivery_interval(mut self, interval: Duration) -> Self {
87 self.demand_redelivery_interval = interval;
88 self
89 }
90}
91
92pub struct StreamRefs;
99
100impl StreamRefs {
101 #[must_use]
102 pub fn source_ref<T>() -> Sink<T, SourceRef<T>>
103 where
104 T: Send + 'static,
105 {
106 Self::source_ref_with_settings(StreamRefSettings::default())
107 }
108
109 #[must_use]
110 pub fn source_ref_with_settings<T>(settings: StreamRefSettings) -> Sink<T, SourceRef<T>>
111 where
112 T: Send + 'static,
113 {
114 stream_ref_source_sink(settings)
115 }
116
117 #[must_use]
118 pub fn sink_ref<T>() -> Source<T, SinkRef<T>>
119 where
120 T: Send + 'static,
121 {
122 Self::sink_ref_with_settings(StreamRefSettings::default())
123 }
124
125 #[must_use]
126 pub fn sink_ref_with_settings<T>(settings: StreamRefSettings) -> Source<T, SinkRef<T>>
127 where
128 T: Send + 'static,
129 {
130 stream_ref_sink_source(settings)
131 }
132}
133
134pub struct SourceRef<T> {
139 inner: Arc<SourceRefInner<T>>,
140}
141
142struct SourceRefInner<T> {
143 #[allow(dead_code)]
144 producer: LazyProducerStatus<T>,
145 direct: Arc<SourceRefDirectState<T>>,
146 settings: StreamRefSettings,
147 subscribed: AtomicBool,
148 timeout: Option<Cancellable>,
149}
150
151impl<T> Clone for SourceRef<T> {
152 fn clone(&self) -> Self {
153 Self {
154 inner: Arc::clone(&self.inner),
155 }
156 }
157}
158
159impl<T> fmt::Debug for SourceRef<T> {
160 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161 f.debug_struct("SourceRef").finish_non_exhaustive()
162 }
163}
164
165impl<T> SourceRef<T>
166where
167 T: stream_ref_actor_bound::Bound + Send + 'static,
168{
169 #[must_use]
170 pub fn source(&self) -> Source<T, NotUsed> {
171 let inner = Arc::clone(&self.inner);
172 Source::from_materialized_factory(move |_materializer| {
173 if inner.subscribed.swap(true, Ordering::SeqCst) {
174 return Ok((
175 failed_once("source ref has already been materialized"),
176 NotUsed,
177 ));
178 }
179
180 if let Some(timeout) = &inner.timeout {
181 timeout.cancel();
182 }
183 match inner.direct.claim_input() {
184 Ok(input) => Ok((
185 Box::new(SourceRefDirectStream::new(
186 input,
187 inner.settings,
188 Arc::clone(&inner),
189 )) as BoxStream<T>,
190 NotUsed,
191 )),
192 Err(error) => Ok((failed_stream(error), NotUsed)),
193 }
194 })
195 }
196}
197
198pub(super) fn proto_source<T>(source_ref: &SourceRef<T>) -> Source<T, NotUsed>
199where
200 T: Send + 'static,
201{
202 let inner = Arc::clone(&source_ref.inner);
203 Source::from_materialized_factory(move |_materializer| {
204 if inner.subscribed.swap(true, Ordering::SeqCst) {
205 return Ok((
206 failed_stream(StreamError::Failed(
207 "source ref has already been materialized".to_owned(),
208 )),
209 NotUsed,
210 ));
211 }
212
213 if let Some(timeout) = &inner.timeout {
214 timeout.cancel();
215 }
216 match inner.direct.claim_input() {
217 Ok(input) => Ok((
218 Box::new(SourceRefDirectStream::new(
219 input,
220 inner.settings,
221 Arc::clone(&inner),
222 )) as BoxStream<T>,
223 NotUsed,
224 )),
225 Err(error) => Ok((failed_stream(error), NotUsed)),
226 }
227 })
228}
229
230pub struct SinkRef<T> {
235 inner: Arc<SinkRefInner<T>>,
236}
237
238struct SinkRefInner<T> {
239 direct: Arc<SinkRefDirectState<T>>,
240 subscribed: AtomicBool,
241}
242
243impl<T> Clone for SinkRef<T> {
244 fn clone(&self) -> Self {
245 Self {
246 inner: Arc::clone(&self.inner),
247 }
248 }
249}
250
251impl<T> fmt::Debug for SinkRef<T> {
252 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253 f.debug_struct("SinkRef").finish_non_exhaustive()
254 }
255}
256
257impl<T> SinkRef<T>
258where
259 T: stream_ref_actor_bound::Bound + Send + 'static,
260{
261 #[must_use]
262 pub fn sink(&self) -> Sink<T, StreamCompletion<NotUsed>> {
263 let inner = Arc::clone(&self.inner);
264 Sink::from_runner(move |input, _materializer| {
265 if inner.subscribed.swap(true, Ordering::SeqCst) {
266 return Ok(StreamCompletion::ready(Err(StreamError::Failed(
267 "sink ref has already been materialized".to_owned(),
268 ))));
269 }
270
271 Ok(inner.direct.attach_input(input))
272 })
273 }
274}
275
276#[allow(dead_code)]
277enum ProducerCommand<T> {
278 Subscribe {
279 consumer: ActorRef<ConsumerCommand<T>>,
280 },
281 Demand {
282 consumer: ActorRef<ConsumerCommand<T>>,
283 cumulative: u64,
284 },
285 Cancel {
286 consumer: ActorRef<ConsumerCommand<T>>,
287 },
288 RemoteFailure {
289 consumer: ActorRef<ConsumerCommand<T>>,
290 error: StreamError,
291 },
292 Ack,
293}
294
295#[allow(dead_code)]
296enum ConsumerCommand<T> {
297 OnSubscribe {
298 producer: ActorRef<ProducerCommand<T>>,
299 },
300 Element {
301 producer: ActorRef<ProducerCommand<T>>,
302 seq: u64,
303 item: T,
304 },
305 Complete {
306 producer: ActorRef<ProducerCommand<T>>,
307 seq: u64,
308 },
309 Failure {
310 producer: ActorRef<ProducerCommand<T>>,
311 error: StreamError,
312 },
313}
314
315#[allow(dead_code)]
316enum ConsumerEndpoint<T> {
317 Actor(ActorRef<ConsumerCommand<T>>),
318 Direct(Weak<ConsumerShared<T>>),
319}
320
321impl<T> Clone for ConsumerEndpoint<T> {
322 fn clone(&self) -> Self {
323 match self {
324 Self::Actor(actor) => Self::Actor(actor.clone()),
325 Self::Direct(shared) => Self::Direct(Weak::clone(shared)),
326 }
327 }
328}
329
330#[allow(dead_code)]
331enum ProducerEndpoint<T> {
332 Actor(ActorRef<ProducerCommand<T>>),
333 Direct(Weak<ProducerShared<T>>),
334}
335
336impl<T> Clone for ProducerEndpoint<T> {
337 fn clone(&self) -> Self {
338 match self {
339 Self::Actor(actor) => Self::Actor(actor.clone()),
340 Self::Direct(shared) => Self::Direct(Weak::clone(shared)),
341 }
342 }
343}
344
345#[allow(dead_code)]
346enum ProducerEndpointHandle<T> {
347 Actor {
348 actor: ActorRef<ProducerCommand<T>>,
349 shared: Arc<ProducerShared<T>>,
350 },
351 Direct(Arc<ProducerShared<T>>),
352}
353
354impl<T> ProducerEndpointHandle<T>
355where
356 T: stream_ref_actor_bound::Bound + Send + 'static,
357{
358 fn stop_actor(&self) {
359 if let Self::Actor { actor, .. } = self {
360 actor.stop(None);
361 }
362 }
363
364 fn actor_ref(&self) -> Option<ActorRef<ProducerCommand<T>>> {
365 match self {
366 Self::Actor { actor, .. } => Some(actor.clone()),
367 Self::Direct(_) => None,
368 }
369 }
370
371 fn direct_shared(&self) -> Option<Arc<ProducerShared<T>>> {
372 match self {
373 Self::Actor { shared, .. } | Self::Direct(shared) => Some(Arc::clone(shared)),
374 }
375 }
376}
377
378#[allow(dead_code)]
379struct LazyProducerStatus<T> {
380 direct: Arc<SourceRefDirectState<T>>,
381}
382
383#[allow(dead_code)]
384impl<T> LazyProducerStatus<T>
385where
386 T: Send + 'static,
387{
388 fn new(direct: Arc<SourceRefDirectState<T>>) -> Self {
389 Self { direct }
390 }
391
392 fn get_status(&self) -> ractor::ActorStatus {
393 self.direct.status()
394 }
395}
396
397struct SourceRefDirectState<T> {
398 inner: Mutex<SourceRefDirectInner<T>>,
399}
400
401struct SourceRefDirectInner<T> {
402 input: Option<BoxStream<T>>,
403 terminal: Option<StreamError>,
404 subscribed: bool,
405}
406
407impl<T> SourceRefDirectState<T>
408where
409 T: Send + 'static,
410{
411 fn new(input: BoxStream<T>) -> Self {
412 Self {
413 inner: Mutex::new(SourceRefDirectInner {
414 input: Some(input),
415 terminal: None,
416 subscribed: false,
417 }),
418 }
419 }
420
421 fn lock(&self) -> MutexGuard<'_, SourceRefDirectInner<T>> {
422 self.inner
423 .lock()
424 .unwrap_or_else(|poison| poison.into_inner())
425 }
426
427 fn claim_input(&self) -> StreamResult<BoxStream<T>> {
428 let mut inner = self.lock();
429 inner.subscribed = true;
430 if let Some(error) = inner.terminal.clone() {
431 return Err(error);
432 }
433 inner.input.take().ok_or(StreamError::ActorTerminated)
434 }
435
436 fn timeout_if_unsubscribed(&self) {
437 let mut inner = self.lock();
438 if !inner.subscribed && inner.terminal.is_none() {
439 let input = inner.input.take();
440 inner.terminal = Some(subscription_timeout_error("stream ref sink"));
441 drop(inner);
442 drop(input);
443 }
444 }
445
446 #[allow(dead_code)]
447 fn status(&self) -> ractor::ActorStatus {
448 let inner = self.lock();
449 if inner.terminal.is_some() || inner.input.is_none() {
450 ractor::ActorStatus::Stopped
451 } else {
452 ractor::ActorStatus::Running
453 }
454 }
455}
456
457struct SourceRefDirectStream<T>
458where
459 T: Send + 'static,
460{
461 input: Option<BoxStream<T>>,
462 queue: VecDeque<T>,
463 terminal: Option<ConsumerTerminal>,
464 settings: StreamRefSettings,
465 terminated: bool,
466 _keep_alive: Arc<SourceRefInner<T>>,
467}
468
469impl<T> SourceRefDirectStream<T>
470where
471 T: Send + 'static,
472{
473 fn new(
474 input: BoxStream<T>,
475 settings: StreamRefSettings,
476 keep_alive: Arc<SourceRefInner<T>>,
477 ) -> Self {
478 Self {
479 input: Some(input),
480 queue: VecDeque::new(),
481 terminal: None,
482 settings,
483 terminated: false,
484 _keep_alive: keep_alive,
485 }
486 }
487
488 fn fill_prefetch(&mut self) {
489 while self.terminal.is_none() && self.queue.len() < self.settings.buffer_capacity {
490 let Some(input) = self.input.as_mut() else {
491 self.terminal = Some(ConsumerTerminal::Error(StreamError::Cancelled));
492 break;
493 };
494 match input.next() {
495 Some(Ok(item)) => self.queue.push_back(item),
496 Some(Err(error)) => {
497 self.input.take();
498 self.terminal = Some(ConsumerTerminal::Error(error));
499 }
500 None => {
501 self.input.take();
502 self.terminal = Some(ConsumerTerminal::Complete);
503 }
504 }
505 }
506 }
507}
508
509impl<T> Iterator for SourceRefDirectStream<T>
510where
511 T: Send + 'static,
512{
513 type Item = StreamResult<T>;
514
515 fn next(&mut self) -> Option<Self::Item> {
516 if self.terminated {
517 return None;
518 }
519
520 self.fill_prefetch();
521 if let Some(item) = self.queue.pop_front() {
522 self.fill_prefetch();
523 return Some(Ok(item));
524 }
525
526 match self.terminal.clone() {
527 Some(ConsumerTerminal::Complete) => {
528 self.terminated = true;
529 None
530 }
531 Some(ConsumerTerminal::Error(error)) => {
532 self.terminated = true;
533 Some(Err(error))
534 }
535 None => None,
536 }
537 }
538}
539
540impl<T> Drop for SourceRefDirectStream<T>
541where
542 T: Send + 'static,
543{
544 fn drop(&mut self) {
545 self.input.take();
546 self.queue.clear();
547 }
548}
549
550struct SinkRefDirectState<T> {
551 inner: Mutex<SinkRefDirectInner<T>>,
552 changed: Condvar,
553}
554
555struct SinkRefDirectInner<T> {
556 input: Option<BoxStream<T>>,
557 completion: Option<oneshot::Sender<StreamResult<NotUsed>>>,
558 completion_cancelled: Option<Arc<AtomicBool>>,
559 terminal: Option<StreamError>,
560 subscribed: bool,
561}
562
563impl<T> SinkRefDirectState<T>
564where
565 T: Send + 'static,
566{
567 fn new() -> Self {
568 Self {
569 inner: Mutex::new(SinkRefDirectInner {
570 input: None,
571 completion: None,
572 completion_cancelled: None,
573 terminal: None,
574 subscribed: false,
575 }),
576 changed: Condvar::new(),
577 }
578 }
579
580 fn lock(&self) -> MutexGuard<'_, SinkRefDirectInner<T>> {
581 self.inner
582 .lock()
583 .unwrap_or_else(|poison| poison.into_inner())
584 }
585
586 fn attach_input(&self, input: BoxStream<T>) -> StreamCompletion<NotUsed> {
587 let mut inner = self.lock();
588 if let Some(error) = inner.terminal.clone() {
589 return StreamCompletion::ready(Err(error));
590 }
591 if inner.subscribed {
592 return StreamCompletion::ready(Err(StreamError::Failed(
593 "stream ref was already subscribed by another endpoint".to_owned(),
594 )));
595 }
596
597 let (sender, receiver) = oneshot::channel();
598 let cancellation = StreamCancellation::for_external_completion();
599 inner.input = Some(input);
600 inner.completion = Some(sender);
601 inner.completion_cancelled = Some(cancellation.cancelled());
602 inner.subscribed = true;
603 drop(inner);
604 self.changed.notify_all();
605 StreamCompletion::from_receiver(receiver, Some(cancellation))
606 }
607
608 fn wait_for_input(&self, settings: StreamRefSettings) -> StreamResult<BoxStream<T>> {
609 let deadline = Instant::now()
610 .checked_add(settings.subscription_timeout)
611 .unwrap_or_else(far_future);
612 let mut inner = self.lock();
613 loop {
614 if let Some(cancelled) = &inner.completion_cancelled
615 && cancelled.load(Ordering::SeqCst)
616 {
617 let input = inner.input.take();
618 inner.terminal = Some(StreamError::Cancelled);
619 inner.completion.take();
620 inner.completion_cancelled = None;
621 drop(inner);
622 drop(input);
623 self.changed.notify_all();
624 return Err(StreamError::Cancelled);
625 }
626 if let Some(input) = inner.input.take() {
627 return Ok(input);
628 }
629 if let Some(error) = inner.terminal.clone() {
630 return Err(error);
631 }
632 let now = Instant::now();
633 if now >= deadline {
634 let error = subscription_timeout_error("stream ref source");
635 inner.terminal = Some(error.clone());
636 if let Some(sender) = inner.completion.take() {
637 let _ = sender.send(Err(error.clone()));
638 }
639 return Err(error);
640 }
641 let remaining = deadline.saturating_duration_since(now);
642 let (next, _) =
643 wait_timeout_unpoison(&self.changed, inner, remaining.min(STREAM_REF_WAIT_POLL));
644 inner = next;
645 }
646 }
647
648 fn settle(&self, result: StreamResult<NotUsed>) {
649 let mut inner = self.lock();
650 if let Some(sender) = inner.completion.take() {
651 let _ = sender.send(result);
652 }
653 inner.completion_cancelled = None;
654 drop(inner);
655 self.changed.notify_all();
656 }
657
658 fn fail_unattached(&self, error: StreamError) {
659 let mut inner = self.lock();
660 if inner.terminal.is_none() {
661 let input = inner.input.take();
662 inner.terminal = Some(error.clone());
663 if let Some(sender) = inner.completion.take() {
664 let _ = sender.send(Err(error));
665 }
666 inner.completion_cancelled = None;
667 drop(inner);
668 drop(input);
669 self.changed.notify_all();
670 }
671 }
672
673 fn is_completion_cancelled(&self) -> bool {
674 self.lock()
675 .completion_cancelled
676 .as_ref()
677 .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
678 }
679}
680
681struct SinkRefDirectStream<T>
682where
683 T: Send + 'static,
684{
685 state: Arc<SinkRefDirectState<T>>,
686 input: Option<BoxStream<T>>,
687 queue: VecDeque<T>,
688 terminal: Option<ConsumerTerminal>,
689 settings: StreamRefSettings,
690 terminated: bool,
691}
692
693impl<T> SinkRefDirectStream<T>
694where
695 T: Send + 'static,
696{
697 fn new(state: Arc<SinkRefDirectState<T>>, settings: StreamRefSettings) -> Self {
698 Self {
699 state,
700 input: None,
701 queue: VecDeque::new(),
702 terminal: None,
703 settings,
704 terminated: false,
705 }
706 }
707
708 fn ensure_input(&mut self) -> StreamResult<()> {
709 if self.input.is_none() && self.terminal.is_none() {
710 self.input = Some(self.state.wait_for_input(self.settings)?);
711 }
712 Ok(())
713 }
714
715 fn fill_prefetch(&mut self) -> StreamResult<()> {
716 self.ensure_input()?;
717 while self.terminal.is_none() && self.queue.len() < self.settings.buffer_capacity {
718 if self.state.is_completion_cancelled() {
719 self.input.take();
720 self.state.settle(Err(StreamError::Cancelled));
721 self.terminal = Some(ConsumerTerminal::Error(StreamError::Cancelled));
722 break;
723 }
724 let Some(input) = self.input.as_mut() else {
725 self.terminal = Some(ConsumerTerminal::Error(StreamError::Cancelled));
726 break;
727 };
728 match input.next() {
729 Some(Ok(item)) => self.queue.push_back(item),
730 Some(Err(error)) => {
731 self.input.take();
732 self.state.settle(Err(error.clone()));
733 self.terminal = Some(ConsumerTerminal::Error(error));
734 }
735 None => {
736 self.input.take();
737 self.state.settle(Ok(NotUsed));
738 self.terminal = Some(ConsumerTerminal::Complete);
739 }
740 }
741 }
742 Ok(())
743 }
744}
745
746impl<T> Iterator for SinkRefDirectStream<T>
747where
748 T: Send + 'static,
749{
750 type Item = StreamResult<T>;
751
752 fn next(&mut self) -> Option<Self::Item> {
753 if self.terminated {
754 return None;
755 }
756
757 if let Err(error) = self.fill_prefetch() {
758 self.terminated = true;
759 return Some(Err(error));
760 }
761 if let Some(item) = self.queue.pop_front() {
762 if let Err(error) = self.fill_prefetch() {
763 self.terminal = Some(ConsumerTerminal::Error(error));
764 }
765 return Some(Ok(item));
766 }
767
768 match self.terminal.clone() {
769 Some(ConsumerTerminal::Complete) => {
770 self.terminated = true;
771 None
772 }
773 Some(ConsumerTerminal::Error(error)) => {
774 self.terminated = true;
775 Some(Err(error))
776 }
777 None => None,
778 }
779 }
780}
781
782impl<T> Drop for SinkRefDirectStream<T>
783where
784 T: Send + 'static,
785{
786 fn drop(&mut self) {
787 if !self.terminated {
788 self.input.take();
789 self.queue.clear();
790 self.state.fail_unattached(StreamError::Cancelled);
791 self.state.settle(Err(StreamError::Cancelled));
792 }
793 }
794}
795
796#[allow(dead_code)]
797struct ProducerShared<T> {
798 inner: Mutex<ProducerInner<T>>,
799 changed: Condvar,
800}
801
802#[allow(dead_code)]
803struct ProducerInner<T> {
804 consumer: Option<ConsumerEndpoint<T>>,
805 cumulative_demand: u64,
806 sent: u64,
807 stopped: Option<StreamError>,
808}
809
810#[allow(dead_code)]
811impl<T> ProducerShared<T>
812where
813 T: stream_ref_actor_bound::Bound + Send + 'static,
814{
815 fn new() -> Self {
816 Self {
817 inner: Mutex::new(ProducerInner {
818 consumer: None,
819 cumulative_demand: 0,
820 sent: 0,
821 stopped: None,
822 }),
823 changed: Condvar::new(),
824 }
825 }
826
827 fn lock(&self) -> MutexGuard<'_, ProducerInner<T>> {
828 self.inner
829 .lock()
830 .unwrap_or_else(|poison| poison.into_inner())
831 }
832
833 fn set_consumer(
834 &self,
835 consumer: ActorRef<ConsumerCommand<T>>,
836 ) -> Result<bool, ActorRef<ConsumerCommand<T>>> {
837 let mut inner = self.lock();
838 match &inner.consumer {
839 Some(ConsumerEndpoint::Actor(existing)) if same_actor(existing, &consumer) => Ok(false),
840 Some(_) => Err(consumer),
841 None if inner.stopped.is_some() => Err(consumer),
842 None => {
843 inner.consumer = Some(ConsumerEndpoint::Actor(consumer));
844 drop(inner);
845 self.changed.notify_all();
846 Ok(true)
847 }
848 }
849 }
850
851 fn set_direct_consumer(&self, consumer: &Arc<ConsumerShared<T>>) -> StreamResult<bool> {
852 let mut inner = self.lock();
853 match &inner.consumer {
854 Some(ConsumerEndpoint::Direct(existing))
855 if existing.ptr_eq(&Arc::downgrade(consumer)) =>
856 {
857 Ok(false)
858 }
859 Some(_) => Err(StreamError::Failed(
860 "stream ref was already subscribed by another endpoint".to_owned(),
861 )),
862 None => {
863 if let Some(error) = inner.stopped.clone() {
864 return Err(error);
865 }
866 inner.consumer = Some(ConsumerEndpoint::Direct(Arc::downgrade(consumer)));
867 drop(inner);
868 self.changed.notify_all();
869 Ok(true)
870 }
871 }
872 }
873
874 fn update_demand(&self, consumer: &ActorRef<ConsumerCommand<T>>, cumulative: u64) {
875 self.update_demand_if(cumulative, |existing| {
876 matches!(existing, ConsumerEndpoint::Actor(actor) if same_actor(actor, consumer))
877 });
878 }
879
880 fn update_direct_demand(&self, consumer: &Arc<ConsumerShared<T>>, cumulative: u64) {
881 let consumer = Arc::downgrade(consumer);
882 self.update_demand_if(cumulative, |existing| {
883 matches!(existing, ConsumerEndpoint::Direct(shared) if shared.ptr_eq(&consumer))
884 });
885 }
886
887 fn update_demand_if(
888 &self,
889 cumulative: u64,
890 matches_consumer: impl FnOnce(&ConsumerEndpoint<T>) -> bool,
891 ) {
892 let mut inner = self.lock();
893 if inner.consumer.as_ref().is_some_and(matches_consumer)
894 && cumulative > inner.cumulative_demand
895 {
896 inner.cumulative_demand = cumulative;
897 drop(inner);
898 self.changed.notify_all();
899 }
900 }
901
902 fn stop_from_consumer(&self, consumer: &ActorRef<ConsumerCommand<T>>, error: StreamError) {
903 self.stop_from_consumer_if(error, |existing| {
904 matches!(existing, ConsumerEndpoint::Actor(actor) if same_actor(actor, consumer))
905 });
906 }
907
908 fn stop_from_direct_consumer(&self, consumer: &Arc<ConsumerShared<T>>, error: StreamError) {
909 let consumer = Arc::downgrade(consumer);
910 self.stop_from_consumer_if(error, |existing| {
911 matches!(existing, ConsumerEndpoint::Direct(shared) if shared.ptr_eq(&consumer))
912 });
913 }
914
915 fn stop_from_consumer_if(
916 &self,
917 error: StreamError,
918 matches_consumer: impl FnOnce(&ConsumerEndpoint<T>) -> bool,
919 ) {
920 let mut inner = self.lock();
921 if inner.consumer.as_ref().is_none_or(matches_consumer) && inner.stopped.is_none() {
922 inner.stopped = Some(error);
923 drop(inner);
924 self.changed.notify_all();
925 }
926 }
927
928 fn stop_unless_finished(&self, error: StreamError) {
929 let mut inner = self.lock();
930 if inner.stopped.is_none() {
931 inner.stopped = Some(error);
932 drop(inner);
933 self.changed.notify_all();
934 }
935 }
936}
937
938#[allow(dead_code)]
939struct ConsumerShared<T> {
940 inner: Mutex<ConsumerInner<T>>,
941 changed: Condvar,
942}
943
944#[allow(dead_code)]
945struct ConsumerInner<T> {
946 producer: Option<ProducerEndpoint<T>>,
947 queue: VecDeque<T>,
948 terminal: Option<ConsumerTerminal>,
949 expected_seq: u64,
950 delivered: u64,
951 cumulative_demand: u64,
952}
953
954#[derive(Clone)]
955enum ConsumerTerminal {
956 Complete,
957 Error(StreamError),
958}
959
960#[allow(dead_code)]
961impl<T> ConsumerShared<T>
962where
963 T: stream_ref_actor_bound::Bound + Send + 'static,
964{
965 fn new(_settings: StreamRefSettings) -> Self {
966 Self {
967 inner: Mutex::new(ConsumerInner {
968 producer: None,
969 queue: VecDeque::new(),
970 terminal: None,
971 expected_seq: 0,
972 delivered: 0,
973 cumulative_demand: 0,
974 }),
975 changed: Condvar::new(),
976 }
977 }
978
979 fn lock(&self) -> MutexGuard<'_, ConsumerInner<T>> {
980 self.inner
981 .lock()
982 .unwrap_or_else(|poison| poison.into_inner())
983 }
984
985 fn set_producer(&self, producer: ActorRef<ProducerCommand<T>>) -> bool {
986 let mut inner = self.lock();
987 match &inner.producer {
988 Some(ProducerEndpoint::Actor(existing)) => same_actor(existing, &producer),
989 Some(ProducerEndpoint::Direct(_)) => false,
990 None => {
991 inner.producer = Some(ProducerEndpoint::Actor(producer));
992 drop(inner);
993 self.changed.notify_all();
994 true
995 }
996 }
997 }
998
999 fn set_direct_producer(&self, producer: &Arc<ProducerShared<T>>) -> StreamResult<bool> {
1000 let mut inner = self.lock();
1001 match &inner.producer {
1002 Some(ProducerEndpoint::Direct(existing))
1003 if existing.ptr_eq(&Arc::downgrade(producer)) =>
1004 {
1005 Ok(false)
1006 }
1007 Some(_) => Err(StreamError::Failed(
1008 "stream ref was already subscribed by another endpoint".to_owned(),
1009 )),
1010 None => {
1011 if let Some(terminal) = inner.terminal.clone() {
1012 return match terminal {
1013 ConsumerTerminal::Complete => Err(StreamError::ActorTerminated),
1014 ConsumerTerminal::Error(error) => Err(error),
1015 };
1016 }
1017 inner.producer = Some(ProducerEndpoint::Direct(Arc::downgrade(producer)));
1018 drop(inner);
1019 self.changed.notify_all();
1020 Ok(true)
1021 }
1022 }
1023 }
1024
1025 fn push(&self, producer: &ActorRef<ProducerCommand<T>>, seq: u64, item: T) {
1026 self.push_if(seq, item, |inner| producer_matches_actor(inner, producer));
1027 }
1028
1029 fn push_direct(&self, producer: &Arc<ProducerShared<T>>, seq: u64, item: T) {
1030 let producer = Arc::downgrade(producer);
1031 self.push_if(seq, item, |inner| producer_matches_direct(inner, &producer));
1032 }
1033
1034 fn push_if(&self, seq: u64, item: T, matches_producer: impl FnOnce(&ConsumerInner<T>) -> bool) {
1035 let mut inner = self.lock();
1036 if inner.terminal.is_some() || !matches_producer(&inner) {
1037 return;
1038 }
1039 let should_notify;
1040 if seq != inner.expected_seq {
1041 inner.queue.clear();
1042 inner.terminal = Some(ConsumerTerminal::Error(invalid_sequence_error(
1043 inner.expected_seq,
1044 seq,
1045 "stream ref element sequence gap",
1046 )));
1047 should_notify = true;
1048 } else {
1049 should_notify = inner.queue.is_empty();
1050 inner.expected_seq += 1;
1051 inner.queue.push_back(item);
1052 }
1053 drop(inner);
1054 if should_notify {
1055 self.changed.notify_all();
1056 }
1057 }
1058
1059 fn complete(&self, producer: &ActorRef<ProducerCommand<T>>, seq: u64) {
1060 self.complete_if(seq, |inner| producer_matches_actor(inner, producer));
1061 }
1062
1063 fn complete_direct(&self, producer: &Arc<ProducerShared<T>>, seq: u64) {
1064 let producer = Arc::downgrade(producer);
1065 self.complete_if(seq, |inner| producer_matches_direct(inner, &producer));
1066 }
1067
1068 fn complete_if(&self, seq: u64, matches_producer: impl FnOnce(&ConsumerInner<T>) -> bool) {
1069 let mut inner = self.lock();
1070 if inner.terminal.is_some() || !matches_producer(&inner) {
1071 return;
1072 }
1073 if seq != inner.expected_seq {
1074 inner.queue.clear();
1075 inner.terminal = Some(ConsumerTerminal::Error(invalid_sequence_error(
1076 inner.expected_seq,
1077 seq,
1078 "stream ref completion sequence gap",
1079 )));
1080 } else {
1081 inner.terminal = Some(ConsumerTerminal::Complete);
1082 }
1083 drop(inner);
1084 self.changed.notify_all();
1085 }
1086
1087 fn fail(&self, producer: &ActorRef<ProducerCommand<T>>, error: StreamError) {
1088 self.fail_if(error, |inner| producer_matches_actor(inner, producer));
1089 }
1090
1091 fn fail_direct(&self, producer: &Arc<ProducerShared<T>>, error: StreamError) {
1092 let producer = Arc::downgrade(producer);
1093 self.fail_if(error, |inner| producer_matches_direct(inner, &producer));
1094 }
1095
1096 fn fail_if(
1097 &self,
1098 error: StreamError,
1099 matches_producer: impl FnOnce(&ConsumerInner<T>) -> bool,
1100 ) {
1101 let mut inner = self.lock();
1102 if inner.terminal.is_some() || !matches_producer(&inner) {
1103 return;
1104 }
1105 inner.queue.clear();
1106 inner.terminal = Some(ConsumerTerminal::Error(error));
1107 drop(inner);
1108 self.changed.notify_all();
1109 }
1110
1111 fn fail_local(&self, error: StreamError) {
1112 let mut inner = self.lock();
1113 if inner.terminal.is_none() {
1114 inner.queue.clear();
1115 inner.terminal = Some(ConsumerTerminal::Error(error));
1116 drop(inner);
1117 self.changed.notify_all();
1118 }
1119 }
1120}
1121
1122#[allow(dead_code)]
1123fn producer_matches_actor<T>(
1124 inner: &ConsumerInner<T>,
1125 producer: &ActorRef<ProducerCommand<T>>,
1126) -> bool
1127where
1128 T: stream_ref_actor_bound::Bound + Send + 'static,
1129{
1130 matches!(
1131 inner.producer.as_ref(),
1132 Some(ProducerEndpoint::Actor(existing)) if same_actor(existing, producer)
1133 )
1134}
1135
1136#[allow(dead_code)]
1137fn producer_matches_direct<T>(inner: &ConsumerInner<T>, producer: &Weak<ProducerShared<T>>) -> bool
1138where
1139 T: stream_ref_actor_bound::Bound + Send + 'static,
1140{
1141 matches!(
1142 inner.producer.as_ref(),
1143 Some(ProducerEndpoint::Direct(existing)) if existing.ptr_eq(producer)
1144 )
1145}
1146
1147#[allow(dead_code)]
1148struct ProducerActor<T> {
1149 shared: Arc<ProducerShared<T>>,
1150 initial_consumer: Option<ActorRef<ConsumerCommand<T>>>,
1151}
1152
1153impl<T> Actor for ProducerActor<T>
1154where
1155 T: stream_ref_actor_bound::Bound + Send + 'static,
1156{
1157 type Msg = ProducerCommand<T>;
1158 type State = ();
1159 type Arguments = ();
1160
1161 async fn pre_start(
1162 &self,
1163 myself: ActorRef<Self::Msg>,
1164 _args: Self::Arguments,
1165 ) -> Result<Self::State, ActorProcessingErr> {
1166 if let Some(consumer) = &self.initial_consumer {
1167 register_producer_consumer(&self.shared, myself, consumer.clone());
1168 }
1169 Ok(())
1170 }
1171
1172 async fn handle(
1173 &self,
1174 myself: ActorRef<Self::Msg>,
1175 message: Self::Msg,
1176 _state: &mut Self::State,
1177 ) -> ActorResult {
1178 match message {
1179 ProducerCommand::Subscribe { consumer } => {
1180 register_producer_consumer(&self.shared, myself, consumer);
1181 }
1182 ProducerCommand::Demand {
1183 consumer,
1184 cumulative,
1185 } => self.shared.update_demand(&consumer, cumulative),
1186 ProducerCommand::Cancel { consumer } => {
1187 self.shared
1188 .stop_from_consumer(&consumer, StreamError::Cancelled);
1189 }
1190 ProducerCommand::RemoteFailure { consumer, error } => {
1191 self.shared.stop_from_consumer(&consumer, error);
1192 }
1193 ProducerCommand::Ack => myself.stop(None),
1194 }
1195 Ok(())
1196 }
1197
1198 async fn post_stop(
1199 &self,
1200 _myself: ActorRef<Self::Msg>,
1201 _state: &mut Self::State,
1202 ) -> ActorResult {
1203 self.shared
1204 .stop_unless_finished(StreamError::ActorTerminated);
1205 Ok(())
1206 }
1207}
1208
1209#[allow(dead_code)]
1210fn register_producer_consumer<T>(
1211 shared: &Arc<ProducerShared<T>>,
1212 producer: ActorRef<ProducerCommand<T>>,
1213 consumer: ActorRef<ConsumerCommand<T>>,
1214) where
1215 T: stream_ref_actor_bound::Bound + Send + 'static,
1216{
1217 match shared.set_consumer(consumer.clone()) {
1218 Ok(true) | Ok(false) => {
1219 let _ = cast_consumer_on_subscribe(&consumer, producer.clone());
1220 }
1221 Err(duplicate) => {
1222 let _ = cast_consumer_failure(
1223 &duplicate,
1224 producer,
1225 StreamError::Failed(
1226 "stream ref was already subscribed by another endpoint".to_owned(),
1227 ),
1228 );
1229 }
1230 }
1231}
1232
1233#[allow(dead_code)]
1234fn register_direct_producer_consumer<T>(
1235 producer: &Arc<ProducerShared<T>>,
1236 consumer: &Arc<ConsumerShared<T>>,
1237) -> StreamResult<()>
1238where
1239 T: stream_ref_actor_bound::Bound + Send + 'static,
1240{
1241 consumer.set_direct_producer(producer)?;
1242 if let Err(error) = producer.set_direct_consumer(consumer) {
1243 consumer.fail_local(error.clone());
1244 return Err(error);
1245 }
1246 Ok(())
1247}
1248
1249#[allow(dead_code)]
1250struct ConsumerActor<T> {
1251 shared: Arc<ConsumerShared<T>>,
1252 initial_producer: Option<ActorRef<ProducerCommand<T>>>,
1253}
1254
1255impl<T> Actor for ConsumerActor<T>
1256where
1257 T: stream_ref_actor_bound::Bound + Send + 'static,
1258{
1259 type Msg = ConsumerCommand<T>;
1260 type State = ();
1261 type Arguments = ();
1262
1263 async fn pre_start(
1264 &self,
1265 myself: ActorRef<Self::Msg>,
1266 _args: Self::Arguments,
1267 ) -> Result<Self::State, ActorProcessingErr> {
1268 if let Some(producer) = &self.initial_producer {
1269 self.shared.set_producer(producer.clone());
1270 if let Err(error) = cast_producer_subscribe(producer, myself.clone()) {
1271 self.shared.fail_local(error);
1272 myself.stop(None);
1273 }
1274 }
1275 Ok(())
1276 }
1277
1278 async fn handle(
1279 &self,
1280 _myself: ActorRef<Self::Msg>,
1281 message: Self::Msg,
1282 _state: &mut Self::State,
1283 ) -> ActorResult {
1284 match message {
1285 ConsumerCommand::OnSubscribe { producer } => {
1286 if !self.shared.set_producer(producer.clone()) {
1287 let _ = cast_producer_remote_failure(
1288 &producer,
1289 _myself.clone(),
1290 StreamError::Failed(
1291 "stream ref was already subscribed by another endpoint".to_owned(),
1292 ),
1293 );
1294 }
1295 }
1296 ConsumerCommand::Element {
1297 producer,
1298 seq,
1299 item,
1300 } => self.shared.push(&producer, seq, item),
1301 ConsumerCommand::Complete { producer, seq } => {
1302 self.shared.complete(&producer, seq);
1303 let _ = cast_actor(&producer, ProducerCommand::Ack);
1304 }
1305 ConsumerCommand::Failure { producer, error } => {
1306 self.shared.fail(&producer, error);
1307 let _ = cast_actor(&producer, ProducerCommand::Ack);
1308 }
1309 }
1310 Ok(())
1311 }
1312
1313 async fn post_stop(
1314 &self,
1315 _myself: ActorRef<Self::Msg>,
1316 _state: &mut Self::State,
1317 ) -> ActorResult {
1318 self.shared.fail_local(StreamError::ActorTerminated);
1319 Ok(())
1320 }
1321}
1322
1323#[allow(dead_code)]
1324struct ConsumerStream<T>
1325where
1326 T: stream_ref_actor_bound::Bound + Send + 'static,
1327{
1328 shared: Arc<ConsumerShared<T>>,
1329 actor_ref: Option<ActorRef<ConsumerCommand<T>>>,
1330 settings: StreamRefSettings,
1331 terminated: bool,
1332 source_ref_keep_alive: Option<Arc<SourceRefInner<T>>>,
1333}
1334
1335impl<T> Iterator for ConsumerStream<T>
1336where
1337 T: stream_ref_actor_bound::Bound + Send + 'static,
1338{
1339 type Item = StreamResult<T>;
1340
1341 fn next(&mut self) -> Option<Self::Item> {
1342 if self.terminated {
1343 return None;
1344 }
1345
1346 if let Err(error) = self.wait_for_subscription() {
1347 self.terminated = true;
1348 return Some(Err(error));
1349 }
1350
1351 if let Err(error) = self.redeliver_or_extend_demand() {
1352 self.terminated = true;
1353 return Some(Err(error));
1354 }
1355
1356 let mut next_redelivery = next_redelivery_deadline(self.settings);
1357 loop {
1358 let demand_after_pop = {
1359 let mut inner = self.shared.lock();
1360 if let Some(item) = inner.queue.pop_front() {
1361 inner.delivered = inner.delivered.saturating_add(1);
1362 let demand = next_demand(&mut inner, self.settings);
1363 drop(inner);
1364 if let Some((producer, cumulative)) = demand
1365 && let Err(error) =
1366 send_demand(&self.shared, &self.actor_ref, &producer, cumulative)
1367 {
1368 self.terminated = true;
1369 return Some(Err(error));
1370 }
1371 return Some(Ok(item));
1372 }
1373
1374 if let Some(terminal) = inner.terminal.clone() {
1375 drop(inner);
1376 match terminal {
1377 ConsumerTerminal::Complete => {
1378 self.terminated = true;
1379 self.stop_actor();
1380 return None;
1381 }
1382 ConsumerTerminal::Error(error) => {
1383 self.terminated = true;
1384 self.stop_actor();
1385 return Some(Err(error));
1386 }
1387 }
1388 }
1389 None::<(ProducerEndpoint<T>, u64)>
1390 };
1391 debug_assert!(demand_after_pop.is_none());
1392
1393 let now = Instant::now();
1394 let timeout = next_redelivery.saturating_duration_since(now);
1395 let mut inner = self.shared.lock();
1396 if !inner.queue.is_empty() || inner.terminal.is_some() {
1397 continue;
1398 }
1399 let (next_inner, result) = wait_timeout_unpoison(
1400 &self.shared.changed,
1401 inner,
1402 timeout.min(STREAM_REF_WAIT_POLL),
1403 );
1404 inner = next_inner;
1405 drop(inner);
1406 if result.timed_out() && Instant::now() >= next_redelivery {
1407 if let Err(error) = self.redeliver_demand() {
1408 self.terminated = true;
1409 return Some(Err(error));
1410 }
1411 next_redelivery = next_redelivery_deadline(self.settings);
1412 }
1413 }
1414 }
1415}
1416
1417#[allow(dead_code)]
1418impl<T> ConsumerStream<T>
1419where
1420 T: stream_ref_actor_bound::Bound + Send + 'static,
1421{
1422 fn wait_for_subscription(&self) -> StreamResult<()> {
1423 let deadline = Instant::now()
1424 .checked_add(self.settings.subscription_timeout)
1425 .unwrap_or_else(far_future);
1426 let mut inner = self.shared.lock();
1427 loop {
1428 if inner.producer.is_some() {
1429 return Ok(());
1430 }
1431 if let Some(terminal) = inner.terminal.clone() {
1432 return match terminal {
1433 ConsumerTerminal::Complete => Ok(()),
1434 ConsumerTerminal::Error(error) => Err(error),
1435 };
1436 }
1437 let now = Instant::now();
1438 if now >= deadline {
1439 drop(inner);
1440 let error = subscription_timeout_error("stream ref source");
1441 self.shared.fail_local(error.clone());
1442 self.stop_actor_ref();
1443 return Err(error);
1444 }
1445 let remaining = deadline.saturating_duration_since(now);
1446 let (next, _) = wait_timeout_unpoison(
1447 &self.shared.changed,
1448 inner,
1449 remaining.min(STREAM_REF_WAIT_POLL),
1450 );
1451 inner = next;
1452 }
1453 }
1454
1455 fn redeliver_or_extend_demand(&self) -> StreamResult<()> {
1456 let demand = {
1457 let mut inner = self.shared.lock();
1458 next_demand(&mut inner, self.settings)
1459 };
1460 if let Some((producer, cumulative)) = demand {
1461 send_demand(&self.shared, &self.actor_ref, &producer, cumulative)?;
1462 }
1463 Ok(())
1464 }
1465
1466 fn redeliver_demand(&self) -> StreamResult<()> {
1467 let (producer, cumulative) = {
1468 let inner = self.shared.lock();
1469 let Some(producer) = inner.producer.clone() else {
1470 return Ok(());
1471 };
1472 if inner.cumulative_demand == 0 {
1473 return Ok(());
1474 }
1475 (producer, inner.cumulative_demand)
1476 };
1477 send_demand(&self.shared, &self.actor_ref, &producer, cumulative)
1478 }
1479
1480 fn stop_actor(&mut self) {
1481 if let Some(actor_ref) = self.actor_ref.take() {
1482 actor_ref.stop(None);
1483 }
1484 }
1485
1486 fn stop_actor_ref(&self) {
1487 if let Some(actor_ref) = &self.actor_ref {
1488 actor_ref.stop(None);
1489 }
1490 }
1491}
1492
1493impl<T> Drop for ConsumerStream<T>
1494where
1495 T: stream_ref_actor_bound::Bound + Send + 'static,
1496{
1497 fn drop(&mut self) {
1498 if !self.terminated {
1499 let producer = self.shared.lock().producer.clone();
1500 match producer {
1501 Some(ProducerEndpoint::Actor(producer)) => {
1502 if let Some(consumer) = &self.actor_ref {
1503 let _ = cast_producer_cancel(&producer, consumer.clone());
1504 }
1505 }
1506 Some(ProducerEndpoint::Direct(producer)) => {
1507 if let Some(producer) = producer.upgrade() {
1508 producer.stop_from_direct_consumer(&self.shared, StreamError::Cancelled);
1509 }
1510 }
1511 None => {}
1512 }
1513 }
1514 self.stop_actor();
1515 drop(self.source_ref_keep_alive.take());
1516 }
1517}
1518
1519#[allow(dead_code)]
1520fn next_demand<T>(
1521 inner: &mut ConsumerInner<T>,
1522 settings: StreamRefSettings,
1523) -> Option<(ProducerEndpoint<T>, u64)> {
1524 if inner.terminal.is_some() {
1525 return None;
1526 }
1527 let remaining_credit = inner.cumulative_demand.saturating_sub(inner.delivered);
1528 if inner.cumulative_demand != 0 && remaining_credit > demand_replenish_threshold(settings) {
1529 return None;
1530 }
1531 let target = inner
1532 .delivered
1533 .saturating_add(settings.buffer_capacity as u64);
1534 if inner.cumulative_demand >= target {
1535 return None;
1536 }
1537 inner.cumulative_demand = target;
1538 inner
1539 .producer
1540 .as_ref()
1541 .map(|producer| (producer.clone(), inner.cumulative_demand))
1542}
1543
1544#[allow(dead_code)]
1545fn send_demand<T>(
1546 consumer_shared: &Arc<ConsumerShared<T>>,
1547 consumer: &Option<ActorRef<ConsumerCommand<T>>>,
1548 producer: &ProducerEndpoint<T>,
1549 cumulative: u64,
1550) -> StreamResult<()>
1551where
1552 T: stream_ref_actor_bound::Bound + Send + 'static,
1553{
1554 match producer {
1555 ProducerEndpoint::Actor(producer) => {
1556 let Some(consumer) = consumer else {
1557 return Err(StreamError::ActorTerminated);
1558 };
1559 cast_producer_demand(producer, consumer.clone(), cumulative)
1560 }
1561 ProducerEndpoint::Direct(producer) => {
1562 let Some(producer) = producer.upgrade() else {
1563 return Err(StreamError::ActorTerminated);
1564 };
1565 producer.update_direct_demand(consumer_shared, cumulative);
1566 Ok(())
1567 }
1568 }
1569}
1570
1571#[allow(dead_code)]
1572fn demand_replenish_threshold(settings: StreamRefSettings) -> u64 {
1573 (settings.buffer_capacity as u64) / 2
1574}
1575
1576fn stream_ref_source_sink<T>(settings: StreamRefSettings) -> Sink<T, SourceRef<T>>
1577where
1578 T: Send + 'static,
1579{
1580 Sink::from_runner(move |input, materializer| {
1581 let direct = Arc::new(SourceRefDirectState::new(input));
1582 let direct_for_timeout = Arc::clone(&direct);
1583 let timeout = materializer.schedule_once(settings.subscription_timeout, move || {
1584 direct_for_timeout.timeout_if_unsubscribed();
1585 });
1586 Ok(SourceRef {
1587 inner: Arc::new(SourceRefInner {
1588 producer: LazyProducerStatus::new(Arc::clone(&direct)),
1589 direct,
1590 settings,
1591 subscribed: AtomicBool::new(false),
1592 timeout: Some(timeout),
1593 }),
1594 })
1595 })
1596}
1597
1598fn stream_ref_sink_source<T>(settings: StreamRefSettings) -> Source<T, SinkRef<T>>
1599where
1600 T: Send + 'static,
1601{
1602 Source::from_materialized_factory(move |_materializer| {
1603 let direct = Arc::new(SinkRefDirectState::new());
1604 let sink_ref = SinkRef {
1605 inner: Arc::new(SinkRefInner {
1606 direct: Arc::clone(&direct),
1607 subscribed: AtomicBool::new(false),
1608 }),
1609 };
1610 Ok((
1611 Box::new(SinkRefDirectStream::new(direct, settings)) as BoxStream<T>,
1612 sink_ref,
1613 ))
1614 })
1615}
1616
1617#[allow(dead_code)]
1618fn send_element_to_consumer<T>(
1619 consumer: &ConsumerEndpoint<T>,
1620 producer: &ProducerEndpointHandle<T>,
1621 seq: u64,
1622 item: T,
1623) -> StreamResult<()>
1624where
1625 T: stream_ref_actor_bound::Bound + Send + 'static,
1626{
1627 match consumer {
1628 ConsumerEndpoint::Actor(consumer) => {
1629 let Some(producer_actor) = producer.actor_ref() else {
1630 return Err(StreamError::ActorTerminated);
1631 };
1632 cast_actor(
1633 consumer,
1634 ConsumerCommand::Element {
1635 producer: producer_actor,
1636 seq,
1637 item,
1638 },
1639 )
1640 }
1641 ConsumerEndpoint::Direct(consumer) => {
1642 let Some(consumer) = consumer.upgrade() else {
1643 return Err(StreamError::Cancelled);
1644 };
1645 let Some(producer) = producer.direct_shared() else {
1646 return Err(StreamError::ActorTerminated);
1647 };
1648 consumer.push_direct(&producer, seq, item);
1649 Ok(())
1650 }
1651 }
1652}
1653
1654#[allow(dead_code)]
1655fn fail_consumer<T>(
1656 consumer: &ConsumerEndpoint<T>,
1657 producer: &ProducerEndpointHandle<T>,
1658 error: StreamError,
1659) -> StreamResult<bool>
1660where
1661 T: stream_ref_actor_bound::Bound + Send + 'static,
1662{
1663 match consumer {
1664 ConsumerEndpoint::Actor(consumer) => {
1665 let Some(producer_actor) = producer.actor_ref() else {
1666 return Err(StreamError::ActorTerminated);
1667 };
1668 cast_actor(
1669 consumer,
1670 ConsumerCommand::Failure {
1671 producer: producer_actor,
1672 error,
1673 },
1674 )?;
1675 Ok(false)
1676 }
1677 ConsumerEndpoint::Direct(consumer) => {
1678 let Some(consumer) = consumer.upgrade() else {
1679 return Err(StreamError::Cancelled);
1680 };
1681 let Some(producer) = producer.direct_shared() else {
1682 return Err(StreamError::ActorTerminated);
1683 };
1684 consumer.fail_direct(&producer, error);
1685 Ok(true)
1686 }
1687 }
1688}
1689
1690#[allow(dead_code)]
1691fn complete_consumer<T>(
1692 consumer: &ConsumerEndpoint<T>,
1693 producer: &ProducerEndpointHandle<T>,
1694 seq: u64,
1695) -> StreamResult<bool>
1696where
1697 T: stream_ref_actor_bound::Bound + Send + 'static,
1698{
1699 match consumer {
1700 ConsumerEndpoint::Actor(consumer) => {
1701 let Some(producer_actor) = producer.actor_ref() else {
1702 return Err(StreamError::ActorTerminated);
1703 };
1704 cast_actor(
1705 consumer,
1706 ConsumerCommand::Complete {
1707 producer: producer_actor,
1708 seq,
1709 },
1710 )?;
1711 Ok(false)
1712 }
1713 ConsumerEndpoint::Direct(consumer) => {
1714 let Some(consumer) = consumer.upgrade() else {
1715 return Err(StreamError::Cancelled);
1716 };
1717 let Some(producer) = producer.direct_shared() else {
1718 return Err(StreamError::ActorTerminated);
1719 };
1720 consumer.complete_direct(&producer, seq);
1721 Ok(true)
1722 }
1723 }
1724}
1725
1726#[allow(dead_code)]
1727fn run_producer_endpoint<T>(
1728 mut input: BoxStream<T>,
1729 shared: Arc<ProducerShared<T>>,
1730 producer_endpoint: ProducerEndpointHandle<T>,
1731 settings: StreamRefSettings,
1732 cancelled: Arc<AtomicBool>,
1733) -> StreamResult<NotUsed>
1734where
1735 T: stream_ref_actor_bound::Bound + Send + 'static,
1736{
1737 let deadline = Instant::now()
1738 .checked_add(settings.subscription_timeout)
1739 .unwrap_or_else(far_future);
1740
1741 loop {
1742 let consumer = match wait_for_remote_demand(&shared, deadline, &cancelled) {
1743 Ok(consumer) => consumer,
1744 Err(error) => {
1745 producer_endpoint.stop_actor();
1746 return Err(error);
1747 }
1748 };
1749 if cancelled.load(Ordering::SeqCst) {
1750 return Err(StreamError::Cancelled);
1751 }
1752
1753 match input.next() {
1754 Some(Ok(item)) => {
1755 let seq = {
1756 let mut inner = shared.lock();
1757 let seq = inner.sent;
1758 inner.sent = inner.sent.saturating_add(1);
1759 seq
1760 };
1761 if let Err(error) =
1762 send_element_to_consumer(&consumer, &producer_endpoint, seq, item)
1763 {
1764 producer_endpoint.stop_actor();
1765 return Err(match error {
1766 StreamError::ActorTerminated => StreamError::Cancelled,
1767 other => other,
1768 });
1769 }
1770 }
1771 Some(Err(error)) => {
1772 if fail_consumer(&consumer, &producer_endpoint, error.clone()).unwrap_or(false) {
1773 producer_endpoint.stop_actor();
1774 }
1775 return Err(error);
1776 }
1777 None => {
1778 let seq = shared.lock().sent;
1779 if complete_consumer(&consumer, &producer_endpoint, seq).unwrap_or(false) {
1780 producer_endpoint.stop_actor();
1781 }
1782 return Ok(NotUsed);
1783 }
1784 }
1785 }
1786}
1787
1788#[allow(dead_code)]
1789fn wait_for_remote_demand<T>(
1790 shared: &Arc<ProducerShared<T>>,
1791 deadline: Instant,
1792 cancelled: &Arc<AtomicBool>,
1793) -> StreamResult<ConsumerEndpoint<T>>
1794where
1795 T: stream_ref_actor_bound::Bound + Send + 'static,
1796{
1797 let mut inner = shared.lock();
1798 loop {
1799 if cancelled.load(Ordering::SeqCst) {
1800 return Err(StreamError::Cancelled);
1801 }
1802 if let Some(error) = inner.stopped.clone() {
1803 return Err(error);
1804 }
1805 if let Some(consumer) = &inner.consumer
1806 && inner.sent < inner.cumulative_demand
1807 {
1808 return Ok(consumer.clone());
1809 }
1810 let now = Instant::now();
1811 if inner.consumer.is_none() && now >= deadline {
1812 return Err(subscription_timeout_error("stream ref sink"));
1813 }
1814 let remaining = deadline.saturating_duration_since(now);
1815 let timeout = if inner.consumer.is_none() {
1816 remaining.min(STREAM_REF_WAIT_POLL)
1817 } else {
1818 STREAM_REF_WAIT_POLL
1819 };
1820 let (next, _) = wait_timeout_unpoison(&shared.changed, inner, timeout);
1821 inner = next;
1822 }
1823}
1824
1825#[allow(dead_code)]
1826fn spawn_producer_actor<T>(
1827 initial_consumer: Option<ActorRef<ConsumerCommand<T>>>,
1828 shared: Arc<ProducerShared<T>>,
1829) -> StreamResult<(
1830 ActorRef<ProducerCommand<T>>,
1831 ractor::concurrency::JoinHandle<()>,
1832)>
1833where
1834 T: stream_ref_actor_bound::Bound + Send + 'static,
1835{
1836 block_on_ractor_runtime(Actor::spawn(
1837 None,
1838 ProducerActor {
1839 shared,
1840 initial_consumer,
1841 },
1842 (),
1843 ))?
1844 .map_err(|error| {
1845 StreamError::Failed(format!(
1846 "stream ref producer actor failed to spawn: {error}"
1847 ))
1848 })
1849}
1850
1851#[allow(dead_code)]
1852fn spawn_consumer_actor<T>(
1853 initial_producer: Option<ActorRef<ProducerCommand<T>>>,
1854 shared: Arc<ConsumerShared<T>>,
1855) -> StreamResult<(
1856 ActorRef<ConsumerCommand<T>>,
1857 ractor::concurrency::JoinHandle<()>,
1858)>
1859where
1860 T: stream_ref_actor_bound::Bound + Send + 'static,
1861{
1862 block_on_ractor_runtime(Actor::spawn(
1863 None,
1864 ConsumerActor {
1865 shared,
1866 initial_producer,
1867 },
1868 (),
1869 ))?
1870 .map_err(|error| {
1871 StreamError::Failed(format!(
1872 "stream ref consumer actor failed to spawn: {error}"
1873 ))
1874 })
1875}
1876
1877#[allow(dead_code)]
1878fn cast_producer_subscribe<T>(
1879 producer: &ActorRef<ProducerCommand<T>>,
1880 consumer: ActorRef<ConsumerCommand<T>>,
1881) -> StreamResult<()>
1882where
1883 T: stream_ref_actor_bound::Bound + Send + 'static,
1884{
1885 cast_actor(producer, ProducerCommand::Subscribe { consumer })
1886}
1887
1888#[allow(dead_code)]
1889fn cast_producer_demand<T>(
1890 producer: &ActorRef<ProducerCommand<T>>,
1891 consumer: ActorRef<ConsumerCommand<T>>,
1892 cumulative: u64,
1893) -> StreamResult<()>
1894where
1895 T: stream_ref_actor_bound::Bound + Send + 'static,
1896{
1897 cast_actor(
1898 producer,
1899 ProducerCommand::Demand {
1900 consumer,
1901 cumulative,
1902 },
1903 )
1904}
1905
1906#[allow(dead_code)]
1907fn cast_producer_cancel<T>(
1908 producer: &ActorRef<ProducerCommand<T>>,
1909 consumer: ActorRef<ConsumerCommand<T>>,
1910) -> StreamResult<()>
1911where
1912 T: stream_ref_actor_bound::Bound + Send + 'static,
1913{
1914 cast_actor(producer, ProducerCommand::Cancel { consumer })
1915}
1916
1917#[allow(dead_code)]
1918fn cast_producer_remote_failure<T>(
1919 producer: &ActorRef<ProducerCommand<T>>,
1920 consumer: ActorRef<ConsumerCommand<T>>,
1921 error: StreamError,
1922) -> StreamResult<()>
1923where
1924 T: stream_ref_actor_bound::Bound + Send + 'static,
1925{
1926 cast_actor(producer, ProducerCommand::RemoteFailure { consumer, error })
1927}
1928
1929#[allow(dead_code)]
1930fn cast_consumer_on_subscribe<T>(
1931 consumer: &ActorRef<ConsumerCommand<T>>,
1932 producer: ActorRef<ProducerCommand<T>>,
1933) -> StreamResult<()>
1934where
1935 T: stream_ref_actor_bound::Bound + Send + 'static,
1936{
1937 cast_actor(consumer, ConsumerCommand::OnSubscribe { producer })
1938}
1939
1940#[allow(dead_code)]
1941fn cast_consumer_failure<T>(
1942 consumer: &ActorRef<ConsumerCommand<T>>,
1943 producer: ActorRef<ProducerCommand<T>>,
1944 error: StreamError,
1945) -> StreamResult<()>
1946where
1947 T: stream_ref_actor_bound::Bound + Send + 'static,
1948{
1949 cast_actor(consumer, ConsumerCommand::Failure { producer, error })
1950}
1951
1952#[allow(dead_code)]
1953fn cast_actor<Msg>(actor_ref: &ActorRef<Msg>, message: Msg) -> StreamResult<()>
1954where
1955 Msg: Message,
1956{
1957 match actor_ref.cast(message) {
1958 Ok(()) => Ok(()),
1959 Err(ractor::MessagingErr::SendErr(_)) | Err(ractor::MessagingErr::ChannelClosed) => {
1960 Err(StreamError::ActorTerminated)
1961 }
1962 Err(error) => Err(StreamError::ActorAskSendFailed {
1963 reason: error.to_string(),
1964 }),
1965 }
1966}
1967
1968#[allow(dead_code)]
1969fn same_actor<MsgA, MsgB>(left: &ActorRef<MsgA>, right: &ActorRef<MsgB>) -> bool
1970where
1971 MsgA: Message,
1972 MsgB: Message,
1973{
1974 left.get_cell().get_id() == right.get_cell().get_id()
1975}
1976
1977fn failed_once<T>(reason: &str) -> BoxStream<T>
1978where
1979 T: stream_ref_actor_bound::Bound + Send + 'static,
1980{
1981 failed_stream(StreamError::Failed(reason.to_owned()))
1982}
1983
1984fn failed_stream<T>(error: StreamError) -> BoxStream<T>
1985where
1986 T: Send + 'static,
1987{
1988 Box::new(std::iter::once(Err(error)))
1989}
1990
1991fn subscription_timeout_error(side: &str) -> StreamError {
1992 StreamError::Failed(format!(
1993 "{side} remote side did not subscribe within subscription timeout"
1994 ))
1995}
1996
1997#[allow(dead_code)]
1998fn invalid_sequence_error(expected: u64, got: u64, context: &str) -> StreamError {
1999 StreamError::Failed(format!(
2000 "{context}: expected sequence {expected}, got {got}"
2001 ))
2002}
2003
2004#[allow(dead_code)]
2005fn next_redelivery_deadline(settings: StreamRefSettings) -> Instant {
2006 Instant::now()
2007 .checked_add(settings.demand_redelivery_interval)
2008 .unwrap_or_else(far_future)
2009}
2010
2011fn far_future() -> Instant {
2012 Instant::now() + Duration::from_secs(60 * 60 * 24 * 365)
2013}
2014
2015fn wait_timeout_unpoison<'a, T>(
2016 condvar: &Condvar,
2017 guard: MutexGuard<'a, T>,
2018 timeout: Duration,
2019) -> (MutexGuard<'a, T>, std::sync::WaitTimeoutResult) {
2020 condvar
2021 .wait_timeout(guard, timeout)
2022 .unwrap_or_else(|poison| poison.into_inner())
2023}
2024
2025#[cfg(test)]
2026mod tests {
2027 use super::*;
2028 use crate::{
2029 stream::{Keep, Source},
2030 testkit::TestSink,
2031 };
2032 use std::sync::{
2033 Arc as StdArc,
2034 atomic::{AtomicBool, AtomicUsize, Ordering},
2035 };
2036
2037 fn wait_until(timeout: Duration, mut condition: impl FnMut() -> bool) -> bool {
2038 let deadline = Instant::now() + timeout;
2039 while Instant::now() < deadline {
2040 if condition() {
2041 return true;
2042 }
2043 std::thread::park_timeout(Duration::from_millis(1));
2044 }
2045 condition()
2046 }
2047
2048 fn assert_condition_holds(timeout: Duration, mut condition: impl FnMut() -> bool) {
2049 let deadline = Instant::now() + timeout;
2050 while Instant::now() < deadline {
2051 assert!(condition());
2052 std::thread::park_timeout(Duration::from_millis(1));
2053 }
2054 assert!(condition());
2055 }
2056
2057 fn short_settings() -> StreamRefSettings {
2058 StreamRefSettings::default()
2059 .with_buffer_capacity(1)
2060 .with_subscription_timeout(Duration::from_millis(50))
2061 .with_demand_redelivery_interval(Duration::from_millis(10))
2062 }
2063
2064 #[test]
2065 fn source_ref_streams_elements_and_completion() {
2066 let source_ref = Source::from_iter(1_u64..=3)
2067 .run_with(StreamRefs::source_ref())
2068 .unwrap();
2069
2070 assert_eq!(source_ref.source().run_collect().unwrap(), vec![1, 2, 3]);
2071 }
2072
2073 #[test]
2074 fn same_process_source_ref_terminal_fold_keeps_direct_splice() {
2075 let closed = StdArc::new(AtomicBool::new(false));
2076 let close_flag = StdArc::clone(&closed);
2077 let source = Source::unfold_resource(
2078 || Ok(0_u64),
2079 |next| {
2080 if *next >= 128 {
2081 Ok(None)
2082 } else {
2083 let value = *next;
2084 *next += 1;
2085 Ok(Some(value))
2086 }
2087 },
2088 move |_state| {
2089 close_flag.store(true, Ordering::SeqCst);
2090 Ok(())
2091 },
2092 );
2093 let source_ref = source.run_with(StreamRefs::source_ref()).unwrap();
2094
2095 let sum = source_ref
2096 .source()
2097 .run_with(Sink::fold(0_u64, |acc, item| acc + item))
2098 .unwrap()
2099 .wait()
2100 .unwrap();
2101
2102 assert_eq!(sum, (0_u64..128).sum::<u64>());
2103 assert!(closed.load(Ordering::SeqCst));
2104 }
2105
2106 #[test]
2107 fn sink_ref_streams_elements_and_completion() {
2108 let (sink_ref, completion) = StreamRefs::sink_ref::<u64>()
2109 .to_mat(Sink::collect(), Keep::both)
2110 .run()
2111 .unwrap();
2112
2113 Source::from_iter(1_u64..=3)
2114 .run_with(sink_ref.sink())
2115 .unwrap()
2116 .wait()
2117 .unwrap();
2118
2119 assert_eq!(completion.wait().unwrap(), vec![1, 2, 3]);
2120 }
2121
2122 #[test]
2123 fn source_ref_propagates_upstream_failure() {
2124 let source_ref = Source::<u64>::failed(StreamError::Failed("boom".to_owned()))
2125 .run_with(StreamRefs::source_ref())
2126 .unwrap();
2127
2128 assert_eq!(
2129 source_ref.source().run_collect(),
2130 Err(StreamError::Failed("boom".to_owned()))
2131 );
2132 }
2133
2134 #[test]
2135 fn sink_ref_propagates_upstream_failure() {
2136 let (sink_ref, completion) = StreamRefs::sink_ref::<u64>()
2137 .to_mat(Sink::collect(), Keep::both)
2138 .run()
2139 .unwrap();
2140
2141 let failure = Source::<u64>::failed(StreamError::Failed("remote boom".to_owned()))
2142 .run_with(sink_ref.sink())
2143 .unwrap()
2144 .wait();
2145
2146 assert_eq!(failure, Err(StreamError::Failed("remote boom".to_owned())));
2147 assert_eq!(
2148 completion.wait(),
2149 Err(StreamError::Failed("remote boom".to_owned()))
2150 );
2151 }
2152
2153 #[test]
2154 fn source_ref_cancellation_reaches_origin() {
2155 let closed = StdArc::new(AtomicBool::new(false));
2156 let close_flag = StdArc::clone(&closed);
2157 let source = Source::unfold_resource(
2158 || Ok(()),
2159 |_state| Ok(Some(1_u64)),
2160 move |_state| {
2161 close_flag.store(true, Ordering::SeqCst);
2162 Ok(())
2163 },
2164 );
2165 let source_ref = source.run_with(StreamRefs::source_ref()).unwrap();
2166
2167 assert_eq!(source_ref.source().take(1).run_collect().unwrap(), vec![1]);
2168 assert!(wait_until(Duration::from_secs(1), || {
2169 closed.load(Ordering::SeqCst)
2170 }));
2171 }
2172
2173 #[test]
2174 fn sink_ref_cancellation_stops_remote_producer() {
2175 let (sink_ref, completion) = StreamRefs::sink_ref::<u64>()
2176 .take(1)
2177 .to_mat(Sink::collect(), Keep::both)
2178 .run()
2179 .unwrap();
2180
2181 let producer = Source::repeat(1_u64)
2182 .run_with(sink_ref.sink())
2183 .unwrap()
2184 .wait();
2185
2186 assert_eq!(completion.wait().unwrap(), vec![1]);
2187 assert_eq!(producer, Err(StreamError::Cancelled));
2188 }
2189
2190 #[test]
2191 fn source_ref_backpressures_across_ref() {
2192 let pulled = StdArc::new(AtomicUsize::new(0));
2193 let pulled_for_source = StdArc::clone(&pulled);
2194 let source = Source::unfold(0_u64, move |next| {
2195 pulled_for_source.fetch_add(1, Ordering::SeqCst);
2196 Some((next + 1, next))
2197 });
2198 let source_ref = source
2199 .run_with(StreamRefs::source_ref_with_settings(short_settings()))
2200 .unwrap();
2201 let mut probe = source_ref.source().run_with(TestSink::probe()).unwrap();
2202
2203 probe.request(1);
2204 probe.assert_next(0);
2205 assert!(wait_until(Duration::from_secs(1), || {
2206 pulled.load(Ordering::SeqCst) >= 2
2207 }));
2208 assert_condition_holds(Duration::from_millis(50), || {
2209 pulled.load(Ordering::SeqCst) <= 2
2210 });
2211
2212 probe.request(1);
2213 probe.assert_next(1);
2214 assert!(wait_until(Duration::from_secs(1), || {
2215 pulled.load(Ordering::SeqCst) >= 3
2216 }));
2217 assert_condition_holds(Duration::from_millis(50), || {
2218 pulled.load(Ordering::SeqCst) <= 3
2219 });
2220
2221 probe.cancel();
2222 }
2223
2224 #[test]
2225 fn source_ref_late_subscription_observes_timeout() {
2226 let source_ref = Source::repeat(1_u64)
2227 .run_with(StreamRefs::source_ref_with_settings(short_settings()))
2228 .unwrap();
2229
2230 assert!(wait_until(Duration::from_secs(1), || {
2231 source_ref.inner.producer.get_status() == ractor::ActorStatus::Stopped
2232 }));
2233
2234 let result = source_ref.source().run_collect();
2235 assert!(matches!(
2236 result,
2237 Err(StreamError::ActorTerminated) | Err(StreamError::Failed(_))
2238 ));
2239 }
2240
2241 #[test]
2242 fn sink_ref_subscription_timeout_fails_local_source() {
2243 let (_sink_ref, probe) = StreamRefs::sink_ref_with_settings::<u64>(short_settings())
2244 .to_mat(TestSink::probe(), Keep::both)
2245 .run()
2246 .unwrap();
2247
2248 probe.request(1);
2249 let error = probe.expect_error();
2250 assert!(
2251 matches!(error, StreamError::Failed(message) if message.contains("did not subscribe"))
2252 );
2253 }
2254
2255 #[test]
2256 fn stream_refs_are_one_shot() {
2257 let source_ref = Source::from_iter([1_u64])
2258 .run_with(StreamRefs::source_ref())
2259 .unwrap();
2260 assert_eq!(source_ref.source().run_collect().unwrap(), vec![1]);
2261 assert!(matches!(
2262 source_ref.source().run_collect(),
2263 Err(StreamError::Failed(message)) if message.contains("already")
2264 ));
2265
2266 let (sink_ref, completion) = StreamRefs::sink_ref::<u64>()
2267 .to_mat(Sink::collect(), Keep::both)
2268 .run()
2269 .unwrap();
2270 Source::single(1_u64)
2271 .run_with(sink_ref.sink())
2272 .unwrap()
2273 .wait()
2274 .unwrap();
2275 assert!(matches!(
2276 Source::single(2_u64).run_with(sink_ref.sink()).unwrap().wait(),
2277 Err(StreamError::Failed(message)) if message.contains("already")
2278 ));
2279 assert_eq!(completion.wait().unwrap(), vec![1]);
2280 }
2281}