1use std::{
13 collections::{BTreeMap, HashMap, VecDeque},
14 fmt,
15 future::Future,
16 hash::Hash,
17 marker::PhantomData,
18 panic::{AssertUnwindSafe, catch_unwind},
19 pin::Pin,
20 sync::{
21 Arc, Condvar, Mutex, OnceLock,
22 atomic::{AtomicBool, AtomicUsize, Ordering},
23 },
24 task::{Context, Poll},
25 thread,
26 time::Duration,
27};
28
29use futures::{channel::oneshot, executor::block_on};
30use thiserror::Error;
31use tokio::{
32 runtime::{Builder as TokioRuntimeBuilder, Runtime as TokioRuntime},
33 task::{JoinError, JoinHandle},
34};
35
36pub(crate) type BoxStream<T> = Box<dyn Iterator<Item = StreamResult<T>> + Send>;
37pub(crate) type PureTransform<In, Out> = Arc<dyn Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync>;
38pub(crate) type RuntimeTransform<In, Out> =
39 Arc<dyn Fn(BoxStream<In>, &Materializer) -> StreamResult<BoxStream<Out>> + Send + Sync>;
40type SinkRunner<In, Mat> = dyn Fn(BoxStream<In>, &Materializer) -> StreamResult<Mat> + Send + Sync;
41type HintedSinkRunner<In, Mat> =
42 dyn Fn(BoxStream<In>, &Materializer, SourceRuntimeHints) -> StreamResult<Mat> + Send + Sync;
43type RunnableGraphRunner<Mat> = dyn Fn(&Materializer) -> StreamResult<Mat> + Send + Sync;
44const STREAM_READY_SPINS: usize = 256;
45const STREAM_SPIN_BACKOFF: usize = 8;
50const STREAM_MAX_PARK: Duration = Duration::from_millis(1);
51
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56struct InlineMicroSourceHint {
57 max_success_items: usize,
58}
59
60#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
61struct SourceHints {
62 inline_head_terminal: bool,
63 inline_micro: Option<InlineMicroSourceHint>,
66 terminal_consumer_batch: bool,
70}
71
72#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
73pub(crate) struct SourceRuntimeHints {
74 pub(crate) inline_micro_max_success_items: Option<usize>,
75 pub(crate) terminal_consumer_batch: bool,
76}
77
78impl SourceHints {
79 const fn with_inline_micro(max_success_items: usize) -> Self {
80 Self {
81 inline_head_terminal: true,
82 inline_micro: Some(InlineMicroSourceHint { max_success_items }),
83 terminal_consumer_batch: true,
84 }
85 }
86
87 const fn with_terminal_consumer_batch() -> Self {
88 Self {
89 inline_head_terminal: false,
90 inline_micro: None,
91 terminal_consumer_batch: true,
92 }
93 }
94
95 fn after_flow(self, flow: FlowHints) -> Self {
96 if flow.preserves_inline_head_terminal {
97 Self {
100 inline_head_terminal: true,
101 inline_micro: None,
102 terminal_consumer_batch: self.terminal_consumer_batch
103 && flow.preserves_terminal_consumer_batch,
104 }
105 } else {
106 Self {
107 inline_head_terminal: false,
108 inline_micro: None,
109 terminal_consumer_batch: self.terminal_consumer_batch
110 && flow.preserves_terminal_consumer_batch,
111 }
112 }
113 }
114
115 fn without_inline_micro(self) -> Self {
116 Self {
117 inline_head_terminal: self.inline_head_terminal,
118 inline_micro: None,
119 terminal_consumer_batch: self.terminal_consumer_batch,
120 }
121 }
122
123 fn runtime(self) -> SourceRuntimeHints {
124 SourceRuntimeHints {
125 inline_micro_max_success_items: self.inline_micro.map(|hint| hint.max_success_items),
126 terminal_consumer_batch: self.terminal_consumer_batch,
127 }
128 }
129}
130
131#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
132struct FlowHints {
133 preserves_inline_head_terminal: bool,
134 preserves_terminal_consumer_batch: bool,
135 scalar_chunk_prefix: bool,
136}
137
138impl FlowHints {
139 const PRESERVES_INLINE_HEAD_TERMINAL: Self = Self {
140 preserves_inline_head_terminal: true,
141 preserves_terminal_consumer_batch: true,
142 scalar_chunk_prefix: false,
143 };
144
145 const PRESERVES_TERMINAL_CONSUMER_BATCH: Self = Self {
146 preserves_inline_head_terminal: false,
147 preserves_terminal_consumer_batch: true,
148 scalar_chunk_prefix: false,
149 };
150
151 fn then(self, next: Self) -> Self {
152 Self {
153 preserves_inline_head_terminal: self.preserves_inline_head_terminal
154 && next.preserves_inline_head_terminal,
155 preserves_terminal_consumer_batch: self.preserves_terminal_consumer_batch
156 && next.preserves_terminal_consumer_batch,
157 scalar_chunk_prefix: self.scalar_chunk_prefix && next.scalar_chunk_prefix,
158 }
159 }
160
161 fn without_scalar_chunk_prefix(mut self) -> Self {
162 self.scalar_chunk_prefix = false;
163 self
164 }
165}
166
167struct PartitionSlot<Key, Out> {
168 key: Option<Key>,
169 active: usize,
170 queued: VecDeque<(usize, Out)>,
171 in_ready_queue: bool,
172}
173
174struct AbortOnDropHandle<T> {
175 handle: JoinHandle<T>,
176}
177
178impl<T> AbortOnDropHandle<T> {
179 fn new(handle: JoinHandle<T>) -> Self {
180 Self { handle }
181 }
182}
183
184impl<T> Drop for AbortOnDropHandle<T> {
185 fn drop(&mut self) {
186 self.handle.abort();
187 }
188}
189
190impl<T> Future for AbortOnDropHandle<T> {
191 type Output = Result<T, JoinError>;
192
193 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
194 Pin::new(&mut self.handle).poll(cx)
195 }
196}
197
198impl<T> Unpin for AbortOnDropHandle<T> {}
199
200pub(crate) fn stream_tokio_runtime() -> &'static TokioRuntime {
201 static RUNTIME: OnceLock<TokioRuntime> = OnceLock::new();
202 RUNTIME.get_or_init(|| {
203 TokioRuntimeBuilder::new_multi_thread()
204 .enable_all()
205 .thread_name("datum-stream-tokio")
206 .build()
207 .expect("stream tokio runtime")
208 })
209}
210
211fn spawn_tokio_task<Fut, T>(future: Fut) -> AbortOnDropHandle<T>
212where
213 Fut: Future<Output = T> + Send + 'static,
214 T: Send + 'static,
215{
216 AbortOnDropHandle::new(stream_tokio_runtime().spawn(future))
217}
218
219pub(crate) fn current_stream_cancelled() -> Option<Arc<AtomicBool>> {
220 runtime::current_stream_cancelled()
221}
222
223pub(super) fn catch_unwind_failed<T, F>(context: &'static str, f: F) -> StreamResult<T>
224where
225 F: FnOnce() -> T,
226{
227 catch_unwind(AssertUnwindSafe(f))
228 .map_err(|_| StreamError::Failed(format!("{context} panicked")))
229}
230
231impl<Key, Out> PartitionSlot<Key, Out> {
232 fn new(key: Key) -> Self {
233 Self {
234 key: Some(key),
235 active: 0,
236 queued: VecDeque::new(),
237 in_ready_queue: false,
238 }
239 }
240}
241
242#[inline(always)]
243fn partition_slot_for<Key, Out>(
244 key: Key,
245 slots_by_key: &mut HashMap<Key, usize>,
246 slots: &mut Vec<PartitionSlot<Key, Out>>,
247 free_slots: &mut Vec<usize>,
248) -> usize
249where
250 Key: Clone + Eq + Hash,
251{
252 if let Some(slot) = slots_by_key.get(&key) {
253 return *slot;
254 }
255
256 let slot = if let Some(slot) = free_slots.pop() {
257 let state = &mut slots[slot];
258 state.key = Some(key.clone());
259 state.active = 0;
260 state.queued.clear();
261 state.in_ready_queue = false;
262 slot
263 } else {
264 slots.push(PartitionSlot::new(key.clone()));
265 slots.len() - 1
266 };
267 slots_by_key.insert(key, slot);
268 slot
269}
270
271#[inline(always)]
272fn retire_partition_slot<Key, Out>(
273 slot: usize,
274 slots_by_key: &mut HashMap<Key, usize>,
275 slots: &mut [PartitionSlot<Key, Out>],
276 free_slots: &mut Vec<usize>,
277) where
278 Key: Eq + Hash,
279{
280 let state = &mut slots[slot];
281 if let Some(key) = state.key.take() {
282 slots_by_key.remove(&key);
283 }
284 state.active = 0;
285 state.queued.clear();
286 state.in_ready_queue = false;
287 free_slots.push(slot);
288}
289
290#[inline(always)]
291fn ready_partition_slot<Key, Out>(
292 slots: &mut [PartitionSlot<Key, Out>],
293 ready_slots: &mut VecDeque<usize>,
294 slot: usize,
295 per_partition: usize,
296) {
297 if let Some(state) = slots.get_mut(slot)
298 && state.key.is_some()
299 && !state.in_ready_queue
300 && state.active < per_partition
301 && !state.queued.is_empty()
302 {
303 state.in_ready_queue = true;
304 ready_slots.push_back(slot);
305 }
306}
307
308#[inline(always)]
309fn pop_ready_partition_slot<Key, Out>(
310 slots: &mut [PartitionSlot<Key, Out>],
311 ready_slots: &mut VecDeque<usize>,
312 per_partition: usize,
313) -> Option<(usize, usize, Out)> {
314 while let Some(slot) = ready_slots.pop_front() {
315 let mut requeue = false;
316 let item = if let Some(state) = slots.get_mut(slot) {
317 state.in_ready_queue = false;
318 if state.key.is_some() && state.active < per_partition {
319 let item = state.queued.pop_front().map(|(index, item)| {
320 state.active += 1;
321 (index, slot, item)
322 });
323 if !state.queued.is_empty() && state.active < per_partition {
324 state.in_ready_queue = true;
325 requeue = true;
326 }
327 item
328 } else {
329 None
330 }
331 } else {
332 None
333 };
334
335 if requeue {
336 ready_slots.push_back(slot);
337 }
338 if item.is_some() {
339 return item;
340 }
341 }
342 None
343}
344
345pub(crate) trait SourceFactory<Out, Mat>: Send + Sync {
346 fn create(self: Arc<Self>, materializer: &Materializer) -> StreamResult<(BoxStream<Out>, Mat)>;
347
348 fn append_scalar_i64(
349 self: Arc<Self>,
350 _steps: &[scalar::ScalarChunkStep<i64>],
351 ) -> Option<Arc<dyn SourceFactory<Out, Mat>>> {
352 None
353 }
354}
355
356struct FnSourceFactory<F>(F);
357
358impl<Out, Mat, F> SourceFactory<Out, Mat> for FnSourceFactory<F>
359where
360 F: Fn(&Materializer) -> StreamResult<(BoxStream<Out>, Mat)> + Send + Sync,
361{
362 fn create(self: Arc<Self>, materializer: &Materializer) -> StreamResult<(BoxStream<Out>, Mat)> {
363 (self.0)(materializer)
364 }
365}
366
367struct MapSourceFactory<In, Out, Mat, F> {
368 source: Arc<dyn SourceFactory<In, Mat>>,
369 stage: F,
370 _marker: PhantomData<fn(In) -> Out>,
371}
372
373impl<In, Out, Mat, F> SourceFactory<Out, Mat> for MapSourceFactory<In, Out, Mat, F>
374where
375 In: Send + 'static,
376 Out: Send + 'static,
377 Mat: Send + 'static,
378 F: Fn(In) -> Out + Send + Sync + 'static,
379{
380 fn create(self: Arc<Self>, materializer: &Materializer) -> StreamResult<(BoxStream<Out>, Mat)> {
381 let (stream, mat) = Arc::clone(&self.source).create(materializer)?;
382 Ok((
383 Box::new(MapSourceStream {
384 input: stream,
385 factory: self,
386 }),
387 mat,
388 ))
389 }
390}
391
392struct MapSourceStream<In, Out, Mat, F> {
393 input: BoxStream<In>,
394 factory: Arc<MapSourceFactory<In, Out, Mat, F>>,
395}
396
397impl<In, Out, Mat, F> Iterator for MapSourceStream<In, Out, Mat, F>
398where
399 F: Fn(In) -> Out,
400{
401 type Item = StreamResult<Out>;
402
403 fn next(&mut self) -> Option<Self::Item> {
404 self.input
405 .next()
406 .map(|item| item.map(|item| (self.factory.stage)(item)))
407 }
408}
409
410fn merge_streams<Out>(streams: Vec<BoxStream<Out>>, eager_complete: bool) -> BoxStream<Out>
411where
412 Out: Send + 'static,
413{
414 let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
415 let mut current = 0usize;
416 Box::new(std::iter::from_fn(move || {
417 loop {
418 let index = next_active_optional_stream(&streams, current, |_| true)?;
419 current = (index + 1) % streams.len().max(1);
420 let Some(stream) = streams[index].as_mut() else {
421 continue;
422 };
423 match stream.next() {
424 Some(item) => return Some(item),
425 None => {
426 streams[index] = None;
427 if eager_complete {
428 return None;
429 }
430 }
431 }
432 }
433 }))
434}
435
436fn merge_prioritized_streams<Out>(
437 streams: Vec<BoxStream<Out>>,
438 priorities: Vec<usize>,
439 eager_complete: bool,
440) -> BoxStream<Out>
441where
442 Out: Send + 'static,
443{
444 let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
445 let schedule: Vec<usize> = priorities
446 .into_iter()
447 .enumerate()
448 .flat_map(|(index, weight)| std::iter::repeat_n(index, weight))
449 .collect();
450 let mut schedule_index = 0usize;
451 Box::new(std::iter::from_fn(move || {
452 loop {
453 if streams.iter().all(Option::is_none) {
454 return None;
455 }
456 let index = next_weighted_stream(&streams, &schedule, &mut schedule_index)?;
457 let Some(stream) = streams[index].as_mut() else {
458 continue;
459 };
460 match stream.next() {
461 Some(item) => return Some(item),
462 None => {
463 streams[index] = None;
464 if eager_complete {
465 return None;
466 }
467 }
468 }
469 }
470 }))
471}
472
473fn merge_sorted_stream<Out>(mut left: BoxStream<Out>, mut right: BoxStream<Out>) -> BoxStream<Out>
474where
475 Out: Ord + Send + 'static,
476{
477 let mut left_next: Option<Out> = None;
478 let mut right_next: Option<Out> = None;
479 let mut left_done = false;
480 let mut right_done = false;
481 Box::new(std::iter::from_fn(move || {
482 loop {
483 if left_next.is_none() && !left_done {
484 match left.next() {
485 Some(Ok(item)) => left_next = Some(item),
486 Some(Err(error)) => return Some(Err(error)),
487 None => left_done = true,
488 }
489 }
490 if right_next.is_none() && !right_done {
491 match right.next() {
492 Some(Ok(item)) => right_next = Some(item),
493 Some(Err(error)) => return Some(Err(error)),
494 None => right_done = true,
495 }
496 }
497
498 let next = match (&left_next, &right_next) {
499 (Some(left_item), Some(right_item)) => {
500 if left_item <= right_item {
501 left_next.take()
502 } else {
503 right_next.take()
504 }
505 }
506 (Some(_), None) if right_done => left_next.take(),
507 (None, Some(_)) if left_done => right_next.take(),
508 (None, None) if left_done && right_done => return None,
509 _ => continue,
510 };
511 if let Some(item) = next {
512 return Some(Ok(item));
513 }
514 }
515 }))
516}
517
518fn merge_latest_streams<Out>(
519 streams: Vec<BoxStream<Out>>,
520 eager_complete: bool,
521) -> BoxStream<Vec<Out>>
522where
523 Out: Clone + Send + 'static,
524{
525 let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
526 let mut latest = vec![None; streams.len()];
527 let mut seen = 0usize;
528 let mut current = 0usize;
529 let mut pending = VecDeque::<Vec<Out>>::new();
530 Box::new(std::iter::from_fn(move || {
531 loop {
532 if let Some(output) = pending.pop_front() {
533 return Some(Ok(output));
534 }
535 if streams.iter().all(Option::is_none) {
536 return None;
537 }
538 let index = next_active_optional_stream(&streams, current, |_| true)?;
539 current = (index + 1) % streams.len().max(1);
540 let Some(stream) = streams[index].as_mut() else {
541 continue;
542 };
543 match stream.next() {
544 Some(Ok(item)) => {
545 if latest[index].is_none() {
546 seen += 1;
547 }
548 latest[index] = Some(item);
549 if seen == latest.len() {
550 pending.push_back(
551 latest
552 .iter()
553 .map(|item| item.clone().expect("merge-latest initialized"))
554 .collect(),
555 );
556 }
557 }
558 Some(Err(error)) => return Some(Err(error)),
559 None => {
560 streams[index] = None;
561 if eager_complete {
562 return None;
563 }
564 }
565 }
566 }
567 }))
568}
569
570fn zip_streams<Left, Right>(
571 mut left: BoxStream<Left>,
572 mut right: BoxStream<Right>,
573) -> BoxStream<(Left, Right)>
574where
575 Left: Send + 'static,
576 Right: Send + 'static,
577{
578 let mut left_next: Option<Left> = None;
579 let mut right_next: Option<Right> = None;
580 let mut left_done = false;
581 let mut right_done = false;
582 Box::new(std::iter::from_fn(move || {
583 loop {
584 if left_next.is_none() && !left_done {
585 match left.next() {
586 Some(Ok(item)) => left_next = Some(item),
587 Some(Err(error)) => return Some(Err(error)),
588 None => left_done = true,
589 }
590 }
591 if right_next.is_none() && !right_done {
592 match right.next() {
593 Some(Ok(item)) => right_next = Some(item),
594 Some(Err(error)) => return Some(Err(error)),
595 None => right_done = true,
596 }
597 }
598 match (left_next.take(), right_next.take()) {
599 (Some(left_item), Some(right_item)) => return Some(Ok((left_item, right_item))),
600 (left_item, right_item) => {
601 left_next = left_item;
602 right_next = right_item;
603 if (left_done && left_next.is_none()) || (right_done && right_next.is_none()) {
604 return None;
605 }
606 }
607 }
608 }
609 }))
610}
611
612fn zip_latest_with_stream<Left, Right, Out, F>(
613 mut left: BoxStream<Left>,
614 mut right: BoxStream<Right>,
615 eager_complete: bool,
616 combine: Arc<F>,
617) -> BoxStream<Out>
618where
619 Left: Clone + Send + 'static,
620 Right: Clone + Send + 'static,
621 Out: Send + 'static,
622 F: Fn(Left, Right) -> Out + Send + Sync + 'static,
623{
624 let mut left_latest: Option<Left> = None;
625 let mut right_latest: Option<Right> = None;
626 let mut left_done = false;
627 let mut right_done = false;
628 let mut turn_left = true;
629 let mut pending = VecDeque::<Out>::new();
630
631 Box::new(std::iter::from_fn(move || {
632 loop {
633 if let Some(output) = pending.pop_front() {
634 return Some(Ok(output));
635 }
636 if eager_complete && (left_done || right_done) {
637 return None;
638 }
639 if left_done && right_done {
640 return None;
641 }
642 if (left_done && left_latest.is_none()) || (right_done && right_latest.is_none()) {
646 return None;
647 }
648
649 let pull_left = if left_done {
650 false
651 } else if right_done {
652 true
653 } else {
654 let value = turn_left;
655 turn_left = !turn_left;
656 value
657 };
658
659 if pull_left {
660 match left.next() {
661 Some(Ok(item)) => {
662 left_latest = Some(item);
663 if let (Some(left_item), Some(right_item)) = (&left_latest, &right_latest) {
664 pending.push_back(combine(left_item.clone(), right_item.clone()));
665 }
666 }
667 Some(Err(error)) => return Some(Err(error)),
668 None => {
669 left_done = true;
670 if eager_complete {
671 return None;
672 }
673 }
674 }
675 } else {
676 match right.next() {
677 Some(Ok(item)) => {
678 right_latest = Some(item);
679 if let (Some(left_item), Some(right_item)) = (&left_latest, &right_latest) {
680 pending.push_back(combine(left_item.clone(), right_item.clone()));
681 }
682 }
683 Some(Err(error)) => return Some(Err(error)),
684 None => {
685 right_done = true;
686 if eager_complete {
687 return None;
688 }
689 }
690 }
691 }
692 }
693 }))
694}
695
696fn zip_all_stream<Left, Right>(
697 mut left: BoxStream<Left>,
698 mut right: BoxStream<Right>,
699 left_fill: Left,
700 right_fill: Right,
701) -> BoxStream<(Left, Right)>
702where
703 Left: Clone + Send + 'static,
704 Right: Clone + Send + 'static,
705{
706 let mut left_done = false;
707 let mut right_done = false;
708 Box::new(std::iter::from_fn(move || {
709 if left_done && right_done {
710 return None;
711 }
712
713 let left_item = if left_done {
714 None
715 } else {
716 match left.next() {
717 Some(Ok(item)) => Some(item),
718 Some(Err(error)) => return Some(Err(error)),
719 None => {
720 left_done = true;
721 None
722 }
723 }
724 };
725 let right_item = if right_done {
726 None
727 } else {
728 match right.next() {
729 Some(Ok(item)) => Some(item),
730 Some(Err(error)) => return Some(Err(error)),
731 None => {
732 right_done = true;
733 None
734 }
735 }
736 };
737
738 match (left_item, right_item) {
739 (None, None) if left_done && right_done => None,
740 (Some(left_value), Some(right_value)) => Some(Ok((left_value, right_value))),
741 (Some(left_value), None) => Some(Ok((left_value, right_fill.clone()))),
742 (None, Some(right_value)) => Some(Ok((left_fill.clone(), right_value))),
743 (None, None) => None,
744 }
745 }))
746}
747
748fn zip_n_streams<Out, Next, F>(streams: Vec<BoxStream<Out>>, zipper: Arc<F>) -> BoxStream<Next>
749where
750 Out: Send + 'static,
751 Next: Send + 'static,
752 F: Fn(Vec<Out>) -> Next + Send + Sync + 'static,
753{
754 let count = streams.len();
755 if count == 0 {
756 return Box::new(std::iter::empty());
757 }
758 let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
759 let mut slots: Vec<Option<Out>> = (0..count).map(|_| None).collect();
760 let mut current = 0usize;
761 Box::new(std::iter::from_fn(move || {
762 loop {
763 if slots.iter().all(Option::is_some) {
764 let values = slots
765 .iter_mut()
766 .map(|slot| slot.take().expect("zip-n slot filled"))
767 .collect();
768 return Some(Ok(zipper(values)));
769 }
770
771 let index = next_active_optional_stream(&streams, current, |idx| slots[idx].is_none())?;
772 current = (index + 1) % count.max(1);
773 let Some(stream) = streams[index].as_mut() else {
774 continue;
775 };
776 match stream.next() {
777 Some(Ok(item)) => slots[index] = Some(item),
778 Some(Err(error)) => return Some(Err(error)),
779 None => {
780 streams[index] = None;
781 slots[index].as_ref()?;
782 }
783 }
784 }
785 }))
786}
787
788fn next_active_optional_stream<T, F>(
789 streams: &[Option<BoxStream<T>>],
790 current: usize,
791 predicate: F,
792) -> Option<usize>
793where
794 T: Send + 'static,
795 F: Fn(usize) -> bool,
796{
797 if streams.is_empty() {
798 return None;
799 }
800 for offset in 0..streams.len() {
801 let index = (current + offset) % streams.len();
802 if streams[index].is_some() && predicate(index) {
803 return Some(index);
804 }
805 }
806 None
807}
808
809fn next_weighted_stream<T>(
810 streams: &[Option<BoxStream<T>>],
811 schedule: &[usize],
812 schedule_index: &mut usize,
813) -> Option<usize>
814where
815 T: Send + 'static,
816{
817 if streams.is_empty() || schedule.is_empty() {
818 return None;
819 }
820 for _ in 0..schedule.len() {
821 let index = schedule[*schedule_index % schedule.len()];
822 *schedule_index = (*schedule_index + 1) % schedule.len();
823 if streams.get(index).is_some_and(Option::is_some) {
824 return Some(index);
825 }
826 }
827 None
828}
829
830pub(crate) mod async_boundary;
831mod completion;
832mod error;
833mod flow;
834mod rate;
835mod restart;
836mod runtime;
837pub(crate) mod scalar;
838mod sink;
839mod source;
840mod time;
841mod timer;
842
843pub(crate) trait SplitSegmentHookDyn: Send + Sync + 'static {
847 fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync>;
848}
849
850pub(crate) trait TerminalSourceHookDyn<In>: Send + Sync + 'static {
856 fn drain_terminal_batch(
857 &self,
858 materializer: &Materializer,
859 cancelled: &Arc<AtomicBool>,
860 batch: &mut Vec<In>,
861 ) -> StreamResult<TerminalSourceStatus>;
862
863 fn supports_direct_terminal(&self) -> bool {
864 false
865 }
866
867 fn try_register_direct_terminal(
868 &self,
869 _consumer: Box<dyn TerminalSinkConsumerDyn<In>>,
870 _cancelled: Arc<AtomicBool>,
871 ) -> Option<StreamResult<()>> {
872 None
873 }
874
875 fn cancel_terminal(&self) {}
876}
877
878#[derive(Clone, Copy, Debug, PartialEq, Eq)]
879pub(crate) enum TerminalSourceStatus {
880 Active,
881 Completed,
882}
883
884pub(crate) trait TerminalSinkConsumerDyn<In>: Send + 'static {
885 fn on_item(&mut self, item: In) -> StreamResult<()>;
886 fn finish(self: Box<Self>, result: StreamResult<()>);
887}
888
889pub(crate) trait FoldFastPathDyn<In: Send + 'static>: Send + Sync + 'static {
892 fn try_register(
896 &self,
897 hook: Arc<dyn SplitSegmentHookDyn>,
898 ) -> Option<StreamResult<Box<dyn std::any::Any + Send>>>;
899
900 fn supports_terminal_drain(&self) -> bool {
901 false
902 }
903
904 fn try_register_direct_terminal(
905 &self,
906 _hook: Arc<dyn TerminalSourceHookDyn<In>>,
907 _materializer: &Materializer,
908 ) -> Option<StreamResult<Box<dyn std::any::Any + Send>>> {
909 None
910 }
911
912 fn try_register_terminal_drain(
913 &self,
914 _hook: Arc<dyn TerminalSourceHookDyn<In>>,
915 _materializer: &Materializer,
916 ) -> Option<StreamResult<Box<dyn std::any::Any + Send>>> {
917 None
918 }
919}
920
921use self::runtime::{runtime_checked_stream, set_current_stream_cancelled};
922
923pub(crate) use self::completion::StreamCancellation;
924
925pub use self::{
926 completion::{Cancellable, StreamCompletion},
927 error::{StreamError, StreamResult, Supervision, SupervisionDecider, SupervisionDirective},
928 flow::{BidiFlow, Flow},
929 rate::{AggregateTimer, OverflowStrategy},
930 restart::{RestartFlow, RestartSettings, RestartSink, RestartSource, RetryFlow},
931 runtime::{Materializer, Runtime},
932 scalar::{
933 ScalarArithmeticOp, ScalarCompareOp, scalar_add_codegen_probe,
934 scalar_compare_codegen_probe, scalar_subtract_codegen_probe,
935 },
936 sink::{RunnableGraph, Sink, SinkCombineStrategy},
937 source::{
938 Demand, IntoSource, Keep, MaybeHandle, NotUsed, PushOutlet, Source, SourceCombineStrategy,
939 },
940 time::{DelayOverflowStrategy, ThrottleMode},
941};
942
943#[cfg(test)]
944mod tests {
945 use super::*;
946 use crate::Attributes;
947 use crate::testkit::TestSink;
948 use std::fs;
949 use std::sync::{
950 Arc as StdArc,
951 atomic::{
952 AtomicBool as StdAtomicBool, AtomicUsize as StdAtomicUsize, Ordering as StdOrdering,
953 },
954 mpsc,
955 };
956 use std::time::Duration as StdDuration;
957 use std::time::Instant;
958
959 fn wait<T>(completion: StreamCompletion<T>) -> T {
960 completion.wait().unwrap()
961 }
962
963 fn wait_until(timeout: Duration, mut condition: impl FnMut() -> bool) -> bool {
964 let deadline = Instant::now() + timeout;
965 while Instant::now() < deadline {
966 if condition() {
967 return true;
968 }
969 thread::sleep(Duration::from_millis(2));
970 }
971 condition()
972 }
973
974 fn linux_thread_count(thread_name: &str) -> usize {
975 fs::read_dir("/proc/self/task")
976 .expect("task directory readable")
977 .filter_map(Result::ok)
978 .filter_map(|entry| fs::read_to_string(entry.path().join("comm")).ok())
979 .filter(|name| name.trim() == thread_name)
980 .count()
981 }
982
983 #[test]
984 fn source_run_terminal_shortcuts_match_explicit_sinks() {
985 let explicit_fold: StreamResult<StreamCompletion<u64>> =
986 Source::from_iter(1_u64..=4).run_with(Sink::fold(0_u64, |acc, item| acc + item));
987 let sugared_fold: StreamResult<StreamCompletion<u64>> =
988 Source::from_iter(1_u64..=4).run_fold(0_u64, |acc, item| acc + item);
989 assert_eq!(wait(explicit_fold.unwrap()), 10);
990 assert_eq!(wait(sugared_fold.unwrap()), 10);
991
992 let explicit_reduce: StreamResult<StreamCompletion<u64>> =
993 Source::from_iter(1_u64..=4).run_with(Sink::reduce(|left: u64, right| left + right));
994 let sugared_reduce: StreamResult<StreamCompletion<u64>> =
995 Source::from_iter(1_u64..=4).run_reduce(|left, right| left + right);
996 assert_eq!(wait(explicit_reduce.unwrap()), 10);
997 assert_eq!(wait(sugared_reduce.unwrap()), 10);
998
999 let explicit_empty_reduce = Source::<u64>::empty()
1000 .run_with(Sink::reduce(|left: u64, right| left + right))
1001 .unwrap()
1002 .wait();
1003 let sugared_empty_reduce = Source::<u64>::empty()
1004 .run_reduce(|left, right| left + right)
1005 .unwrap()
1006 .wait();
1007 assert_eq!(sugared_empty_reduce, explicit_empty_reduce);
1008 assert_eq!(sugared_empty_reduce, Err(StreamError::EmptyStream));
1009
1010 let explicit_sum = StdArc::new(StdAtomicUsize::new(0));
1011 let explicit_sum_sink = StdArc::clone(&explicit_sum);
1012 let explicit_foreach: StreamResult<StreamCompletion<NotUsed>> =
1013 Source::from_iter(1_usize..=4).run_with(Sink::foreach(move |item| {
1014 explicit_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1015 }));
1016 assert_eq!(wait(explicit_foreach.unwrap()), NotUsed);
1017
1018 let sugared_sum = StdArc::new(StdAtomicUsize::new(0));
1019 let sugared_sum_sink = StdArc::clone(&sugared_sum);
1020 let sugared_foreach: StreamResult<StreamCompletion<NotUsed>> =
1021 Source::from_iter(1_usize..=4).run_foreach(move |item| {
1022 sugared_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1023 });
1024 assert_eq!(wait(sugared_foreach.unwrap()), NotUsed);
1025
1026 let alias_sum = StdArc::new(StdAtomicUsize::new(0));
1027 let alias_sum_sink = StdArc::clone(&alias_sum);
1028 let alias_for_each: StreamResult<StreamCompletion<NotUsed>> =
1029 Source::from_iter(1_usize..=4).run_for_each(move |item| {
1030 alias_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1031 });
1032 assert_eq!(wait(alias_for_each.unwrap()), NotUsed);
1033
1034 assert_eq!(explicit_sum.load(StdOrdering::SeqCst), 10);
1035 assert_eq!(sugared_sum.load(StdOrdering::SeqCst), 10);
1036 assert_eq!(alias_sum.load(StdOrdering::SeqCst), 10);
1037 }
1038
1039 #[test]
1040 fn source_constructor_sugar_matches_from_iter() {
1041 let expected = Source::from_iter(1_u64..=4).run_collect().unwrap();
1042
1043 assert_eq!(
1044 Source::from(vec![1_u64, 2, 3, 4]).run_collect().unwrap(),
1045 expected
1046 );
1047 assert_eq!(
1048 Source::from([1_u64, 2, 3, 4]).run_collect().unwrap(),
1049 expected
1050 );
1051 assert_eq!((1_u64..=4).into_source().run_collect().unwrap(), expected);
1052 }
1053
1054 #[test]
1055 fn source_constructor_sugar_keeps_existing_inference_paths() {
1056 let from_vec_into: Source<u64> = vec![1, 2, 3].into();
1057 let from_array_into: Source<u64> = [1, 2, 3].into();
1058 let from_iter = Source::from_iter(1_u64..=3);
1059 let from_iterable = Source::from_iterable(1_u64..=3);
1060 let from_iterator: Source<u64> = (1_u64..=3).collect();
1061 let from_range_into_source: Source<u64> = (1_u64..=3).into_source();
1062
1063 let expected = vec![1, 2, 3];
1064 assert_eq!(from_vec_into.run_collect().unwrap(), expected);
1065 assert_eq!(from_array_into.run_collect().unwrap(), expected);
1066 assert_eq!(from_iter.run_collect().unwrap(), expected);
1067 assert_eq!(from_iterable.run_collect().unwrap(), expected);
1068 assert_eq!(from_iterator.run_collect().unwrap(), expected);
1069 assert_eq!(from_range_into_source.run_collect().unwrap(), expected);
1070 }
1071
1072 #[test]
1073 fn source_async_boundary_preserves_results() {
1074 let expected = Source::from_iter(0_u64..128)
1075 .map(|item| item.wrapping_add(1))
1076 .filter(|item| item % 3 != 0)
1077 .map(|item| item * 2)
1078 .run_collect()
1079 .unwrap();
1080
1081 let actual = Source::from_iter(0_u64..128)
1082 .map(|item| item.wrapping_add(1))
1083 .async_boundary()
1084 .filter(|item| item % 3 != 0)
1085 .map(|item| item * 2)
1086 .run_collect()
1087 .unwrap();
1088
1089 assert_eq!(actual, expected);
1090 }
1091
1092 #[test]
1093 fn flow_async_boundary_preserves_results() {
1094 let expected = Source::from_iter(0_u64..128)
1095 .map(|item| item + 1)
1096 .map(|item| item * 3)
1097 .run_collect()
1098 .unwrap();
1099
1100 let flow = Flow::identity()
1101 .map(|item: u64| item + 1)
1102 .r#async()
1103 .map(|item| item * 3);
1104 let actual = Source::from_iter(0_u64..128)
1105 .via(flow)
1106 .run_collect()
1107 .unwrap();
1108
1109 assert_eq!(actual, expected);
1110 }
1111
1112 #[test]
1113 fn linear_async_boundary_matches_graph_async_boundary_shape() {
1114 use crate::{
1115 AsyncBoundary, AsyncBoundaryExecutionConfig, FusedExecutionConfig, GraphDsl,
1116 GraphFlowShape, MapStage,
1117 };
1118
1119 let graph = GraphDsl::try_create(|builder| {
1120 let first = builder.add(MapStage::new(|item: u64| item + 1));
1121 let boundary = builder.add(AsyncBoundary::<u64>::new());
1122 let second = builder.add(MapStage::new(|item: u64| item * 2));
1123
1124 builder.connect(first.outlet(), boundary.inlet())?;
1125 builder.connect(boundary.outlet(), second.inlet())?;
1126
1127 Ok(GraphFlowShape::new(first.inlet(), second.outlet()))
1128 })
1129 .unwrap();
1130
1131 let linear = Source::from_iter(1_u64..=4)
1132 .map(|item| item + 1)
1133 .async_boundary_with_buffer(4)
1134 .map(|item| item * 2)
1135 .run_collect()
1136 .unwrap();
1137 let graph_output = graph.run_with_input(1_u64..=4).unwrap();
1138 let report = graph
1139 .run_async_boundary_count_with_input_report(
1140 1_u64..=4,
1141 AsyncBoundaryExecutionConfig {
1142 fused: FusedExecutionConfig { event_limit: 1024 },
1143 buffer_size: 4,
1144 },
1145 )
1146 .unwrap();
1147
1148 assert_eq!(linear, graph_output);
1149 assert_eq!(report.result, linear.len());
1150 assert_eq!(report.async_boundary_crossings, linear.len());
1151 }
1152
1153 #[test]
1154 fn async_boundary_regions_run_concurrently() {
1155 let (upstream_tx, upstream_rx) = mpsc::channel::<u64>();
1156 let (downstream_blocked_tx, downstream_blocked_rx) = mpsc::channel::<()>();
1157 let (release_tx, release_rx) = mpsc::channel::<()>();
1158 let release_rx = StdArc::new(Mutex::new(release_rx));
1159
1160 let completion = Source::from_iter(0_u64..3)
1161 .map(move |item| {
1162 upstream_tx.send(item).expect("upstream probe receives");
1163 item
1164 })
1165 .async_boundary_with_buffer(1)
1166 .map({
1167 let release_rx = StdArc::clone(&release_rx);
1168 move |item| {
1169 if item == 0 {
1170 downstream_blocked_tx
1171 .send(())
1172 .expect("downstream probe receives");
1173 release_rx
1174 .lock()
1175 .expect("release receiver lock")
1176 .recv_timeout(StdDuration::from_secs(2))
1177 .expect("downstream release arrives");
1178 }
1179 item
1180 }
1181 })
1182 .run_with(Sink::collect())
1183 .unwrap();
1184
1185 assert_eq!(
1186 downstream_blocked_rx.recv_timeout(StdDuration::from_secs(2)),
1187 Ok(())
1188 );
1189 assert_eq!(upstream_rx.recv_timeout(StdDuration::from_secs(2)), Ok(0));
1190 assert_eq!(upstream_rx.recv_timeout(StdDuration::from_secs(2)), Ok(1));
1191
1192 release_tx.send(()).expect("release downstream");
1193 assert_eq!(completion.wait().unwrap(), vec![0, 1, 2]);
1194 }
1195
1196 #[test]
1197 fn async_boundary_backpressures_slow_downstream() {
1198 let (produced_tx, produced_rx) = mpsc::channel::<u64>();
1199 let (release_tx, release_rx) = mpsc::channel::<()>();
1200 let release_rx = StdArc::new(Mutex::new(release_rx));
1201
1202 let completion = Source::from_iter(0_u64..8)
1203 .map(move |item| {
1204 produced_tx.send(item).expect("producer probe receives");
1205 item
1206 })
1207 .async_boundary_with_buffer(1)
1208 .map({
1209 let release_rx = StdArc::clone(&release_rx);
1210 move |item| {
1211 if item == 0 {
1212 release_rx
1213 .lock()
1214 .expect("release receiver lock")
1215 .recv_timeout(StdDuration::from_secs(2))
1216 .expect("downstream release arrives");
1217 }
1218 item
1219 }
1220 })
1221 .run_with(Sink::collect())
1222 .unwrap();
1223
1224 assert_eq!(produced_rx.recv_timeout(StdDuration::from_secs(2)), Ok(0));
1225 assert_eq!(produced_rx.recv_timeout(StdDuration::from_secs(2)), Ok(1));
1226 if let Ok(item) = produced_rx.recv_timeout(StdDuration::from_millis(100)) {
1227 assert_eq!(item, 2);
1228 }
1229 match produced_rx.recv_timeout(StdDuration::from_millis(100)) {
1230 Err(mpsc::RecvTimeoutError::Timeout) => {}
1231 other => panic!("async boundary handoff was not bounded: {other:?}"),
1232 }
1233
1234 release_tx.send(()).expect("release downstream");
1235 assert_eq!(completion.wait().unwrap(), (0_u64..8).collect::<Vec<_>>());
1236 }
1237
1238 #[test]
1239 fn source_blueprints_are_reusable() {
1240 let source = Source::from_iter(0..5).map(|item| item + 1);
1241
1242 assert_eq!(source.clone().run_collect().unwrap(), vec![1, 2, 3, 4, 5]);
1243 assert_eq!(source.run_collect().unwrap(), vec![1, 2, 3, 4, 5]);
1244 }
1245
1246 #[test]
1247 fn source_map_preserves_materialized_value() {
1248 let graph = Source::single(1)
1249 .map_materialized_value(|_| "source")
1250 .map(|item| item + 1)
1251 .to_mat(Sink::head(), Keep::both);
1252
1253 let materialized = graph.run().unwrap();
1254 assert_eq!(materialized.0, "source");
1255 assert_eq!(wait(materialized.1), 2);
1256 }
1257
1258 #[test]
1259 fn source_and_flow_compose() {
1260 let flow = Flow::identity()
1261 .map(|item: i32| item * 2)
1262 .filter(|item| item % 3 == 0);
1263
1264 let result = Source::from_iter(0..8).via(flow).run_collect().unwrap();
1265
1266 assert_eq!(result, vec![0, 6, 12]);
1267 }
1268
1269 #[test]
1270 fn sink_setup_sees_materializer_defaults_and_local_attributes() {
1271 let observed = StdArc::new(Mutex::new(None));
1272 let observed_in_setup = StdArc::clone(&observed);
1273 let sink = Sink::<i32, StreamCompletion<NotUsed>>::setup(move |_materializer, attrs| {
1274 *observed_in_setup.lock().unwrap() = Some((
1275 attrs.name().map(str::to_owned),
1276 attrs.input_buffer_hint(),
1277 attrs.dispatcher_hint().map(str::to_owned),
1278 ));
1279 Sink::ignore()
1280 })
1281 .add_attributes(Attributes::named("sink-inner"))
1282 .add_attributes(Attributes::input_buffer(4, 4))
1283 .add_attributes(Attributes::dispatcher("bench-dispatcher"));
1284
1285 let materializer = Materializer::new().with_attributes(Attributes::named("mat-outer"));
1286 wait(
1287 Source::from_iter([1, 2, 3])
1288 .run_with_materializer(sink, &materializer)
1289 .unwrap(),
1290 );
1291
1292 assert_eq!(
1293 *observed.lock().unwrap(),
1294 Some((
1295 Some("sink-inner".to_owned()),
1296 Some((4, 4)),
1297 Some("bench-dispatcher".to_owned())
1298 ))
1299 );
1300 }
1301
1302 #[test]
1303 fn sink_pre_materialize_feeds_existing_materialization() {
1304 let materializer = Materializer::new();
1305 let (completion, pre) = Sink::<i32, StreamCompletion<Vec<i32>>>::collect()
1306 .pre_materialize(&materializer)
1307 .unwrap();
1308
1309 Source::from_iter([1, 2, 3])
1310 .run_with_materializer(pre, &materializer)
1311 .unwrap();
1312
1313 assert_eq!(wait(completion), vec![1, 2, 3]);
1314 }
1315
1316 #[test]
1317 fn flow_from_sink_and_source_connects_both_sides() {
1318 assert_eq!(
1319 Source::from_iter([1, 2, 3])
1320 .via(Flow::from_sink_and_source(
1321 Sink::foreach(|_item: i32| {}),
1322 Source::from_iter([10, 20, 30]),
1323 ))
1324 .run_collect()
1325 .unwrap(),
1326 vec![10, 20, 30]
1327 );
1328 }
1329
1330 #[test]
1331 fn from_sink_and_source_keeps_sink_running_after_source_side_completes() {
1332 let completed = StdArc::new(StdAtomicBool::new(false));
1333 let on_complete = StdArc::clone(&completed);
1334 let flow = Flow::from_sink_and_source(
1335 Sink::on_complete(move || {
1336 on_complete.store(true, StdOrdering::SeqCst);
1337 }),
1338 Source::single(10),
1339 );
1340
1341 let result = Source::from_iter([1, 2, 3])
1342 .via(flow)
1343 .run_collect()
1344 .unwrap();
1345
1346 assert_eq!(result, vec![10]);
1347 assert!(wait_until(StdDuration::from_secs(1), || {
1348 completed.load(StdOrdering::SeqCst)
1349 }));
1350 }
1351
1352 #[test]
1353 fn from_sink_and_source_coupled_cancels_source_when_sink_finishes_first() {
1354 let cancellable = StdArc::new(Mutex::new(None));
1355 let observed = StdArc::clone(&cancellable);
1356 let flow = Flow::from_sink_and_source_coupled(
1357 Sink::ignore(),
1358 Source::tick(
1359 StdDuration::from_millis(50),
1360 StdDuration::from_millis(50),
1361 10,
1362 )
1363 .map_materialized_value(move |handle| {
1364 *observed.lock().unwrap() = Some(handle.clone());
1365 handle
1366 }),
1367 );
1368
1369 let completion = Source::from_iter(std::iter::empty::<i32>())
1370 .via(flow)
1371 .run_with(Sink::ignore())
1372 .unwrap();
1373 assert!(wait_until(StdDuration::from_secs(1), || {
1374 cancellable
1375 .lock()
1376 .unwrap()
1377 .as_ref()
1378 .is_some_and(Cancellable::is_cancelled)
1379 }));
1380 assert_eq!(wait(completion), NotUsed);
1381 }
1382
1383 #[test]
1384 fn bidi_flow_join_and_atop_compose() {
1385 let codec = BidiFlow::from_flows(
1386 Flow::identity().map(|item: i32| item + 1),
1387 Flow::identity().map(|item: i32| item * 2),
1388 )
1389 .named("codec");
1390 let framing = BidiFlow::from_flows(
1391 Flow::identity().map(|item: i32| item * 3),
1392 Flow::identity().map(|item: i32| item - 4),
1393 );
1394
1395 let joined = codec
1396 .clone()
1397 .join(Flow::identity().map(|item: i32| item - 5));
1398 let stacked = codec.atop(framing).join(Flow::identity());
1399
1400 assert_eq!(
1401 Source::single(10).via(joined).run_collect().unwrap(),
1402 vec![12]
1403 );
1404 assert_eq!(
1405 Source::single(10).via(stacked).run_collect().unwrap(),
1406 vec![58]
1407 );
1408 }
1409
1410 #[test]
1411 fn flow_buffer_then_map_runs_end_to_end() {
1412 let flow = Flow::identity()
1413 .buffer(8, OverflowStrategy::Backpressure)
1414 .map(|item: i32| item + 1);
1415
1416 let result = Source::from_iter(0..4).via(flow).run_collect().unwrap();
1417
1418 assert_eq!(result, vec![1, 2, 3, 4]);
1419 }
1420
1421 #[test]
1422 fn public_flow_combinators_preserve_runtime_transform_after_buffer() {
1423 fn buffered_flow() -> Flow<i32, i32> {
1424 Flow::identity().buffer(8, OverflowStrategy::Backpressure)
1425 }
1426
1427 assert_eq!(
1428 Source::from_iter(0..4)
1429 .via(buffered_flow().filter(|item| *item % 2 == 0))
1430 .run_collect()
1431 .unwrap(),
1432 vec![0, 2]
1433 );
1434 assert_eq!(
1435 Source::from_iter(0..4)
1436 .via(buffered_flow().filter_not(|item| *item % 2 == 0))
1437 .run_collect()
1438 .unwrap(),
1439 vec![1, 3]
1440 );
1441 assert_eq!(
1442 Source::from_iter(0..4)
1443 .via(buffered_flow().filter_map(|item| (item % 2 == 0).then_some(item + 10)))
1444 .run_collect()
1445 .unwrap(),
1446 vec![10, 12]
1447 );
1448 assert_eq!(
1449 Source::from_iter(0..3)
1450 .via(buffered_flow().map_concat(|item| [item, item + 10]))
1451 .run_collect()
1452 .unwrap(),
1453 vec![0, 10, 1, 11, 2, 12]
1454 );
1455 assert_eq!(
1456 Source::from_iter(0..3)
1457 .via(buffered_flow().stateful_map(5, |state, item| {
1458 *state += item;
1459 *state
1460 }))
1461 .run_collect()
1462 .unwrap(),
1463 vec![5, 6, 8]
1464 );
1465 assert_eq!(
1466 Source::from_iter(0..3)
1467 .via(buffered_flow().stateful_map_concat(0, |state, item| {
1468 *state += item;
1469 [*state, item]
1470 }))
1471 .run_collect()
1472 .unwrap(),
1473 vec![0, 0, 1, 1, 3, 2]
1474 );
1475 assert_eq!(
1476 Source::from_iter(0..4)
1477 .via(buffered_flow().map_async(2, |item| async move { Ok(item + 1) }))
1478 .run_collect()
1479 .unwrap(),
1480 vec![1, 2, 3, 4]
1481 );
1482 assert_eq!(
1483 Source::from_iter(0..4)
1484 .via(buffered_flow().map_async_unordered(2, |item| async move { Ok(item + 1) }))
1485 .run_collect()
1486 .unwrap(),
1487 vec![1, 2, 3, 4]
1488 );
1489 assert_eq!(
1490 Source::from_iter(0..4)
1491 .via(buffered_flow().map_async_partitioned(
1492 2,
1493 1,
1494 |item| item % 2,
1495 |item| async move { Ok(item + 1) },
1496 ))
1497 .run_collect()
1498 .unwrap(),
1499 vec![1, 2, 3, 4]
1500 );
1501 assert_eq!(
1502 Source::from_iter(0..5)
1503 .via(buffered_flow().take(3))
1504 .run_collect()
1505 .unwrap(),
1506 vec![0, 1, 2]
1507 );
1508 assert_eq!(
1509 Source::from_iter(0..5)
1510 .via(buffered_flow().drop(2))
1511 .run_collect()
1512 .unwrap(),
1513 vec![2, 3, 4]
1514 );
1515 assert_eq!(
1516 Source::from_iter(0..5)
1517 .via(buffered_flow().take_while(|item| *item < 3))
1518 .run_collect()
1519 .unwrap(),
1520 vec![0, 1, 2]
1521 );
1522 assert_eq!(
1523 Source::from_iter(0..5)
1524 .via(buffered_flow().drop_while(|item| *item < 3))
1525 .run_collect()
1526 .unwrap(),
1527 vec![3, 4]
1528 );
1529 assert_eq!(
1530 Source::from_iter(0..3)
1531 .via(buffered_flow().limit(5))
1532 .run_collect()
1533 .unwrap(),
1534 vec![0, 1, 2]
1535 );
1536 assert_eq!(
1537 Source::from_iter(0..5)
1538 .via(buffered_flow().grouped(2))
1539 .run_collect()
1540 .unwrap(),
1541 vec![vec![0, 1], vec![2, 3], vec![4]]
1542 );
1543 assert_eq!(
1544 Source::from_iter(1..=3)
1545 .via(buffered_flow().scan(0, |acc, item| acc + item))
1546 .run_collect()
1547 .unwrap(),
1548 vec![0, 1, 3, 6]
1549 );
1550 assert_eq!(
1551 Source::from_iter(1..=4)
1552 .via(buffered_flow().sliding(2, 1))
1553 .run_collect()
1554 .unwrap(),
1555 vec![vec![1, 2], vec![2, 3], vec![3, 4]]
1556 );
1557 assert_eq!(
1558 Source::from_iter(1..=4)
1559 .via(buffered_flow().fold(0, |acc, item| acc + item))
1560 .run_collect()
1561 .unwrap(),
1562 vec![10]
1563 );
1564 assert_eq!(
1565 Source::from_iter(1..=4)
1566 .via(buffered_flow().reduce(|acc, item| acc + item))
1567 .run_collect()
1568 .unwrap(),
1569 vec![10]
1570 );
1571 assert_eq!(
1572 Source::from_factory(|| {
1573 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into())), Ok(2)].into_iter())
1574 })
1575 .via(buffered_flow().map_error(|_| StreamError::Failed("mapped".into())))
1576 .run_collect(),
1577 Err(StreamError::Failed("mapped".into()))
1578 );
1579 assert_eq!(
1580 Source::<i32>::failed(StreamError::Failed("boom".into()))
1581 .via(buffered_flow().recover(|_| Some(42)))
1582 .run_collect()
1583 .unwrap(),
1584 vec![42]
1585 );
1586 assert_eq!(
1587 Source::<i32>::failed(StreamError::Failed("boom".into()))
1588 .via(buffered_flow().recover_with(|_| Some(Source::from_iter([7, 8]))))
1589 .run_collect()
1590 .unwrap(),
1591 vec![7, 8]
1592 );
1593 assert_eq!(
1594 Source::<i32>::failed(StreamError::Failed("boom".into()))
1595 .via(buffered_flow().recover_with_retries(1, |_| Some(Source::from_iter([9]))))
1596 .run_collect()
1597 .unwrap(),
1598 vec![9]
1599 );
1600 assert_eq!(
1601 Source::from_factory(|| {
1602 Box::new(vec![Ok(1), Err(StreamError::Failed("ignored".into())), Ok(2)].into_iter())
1603 })
1604 .via(buffered_flow().on_error_complete())
1605 .run_collect()
1606 .unwrap(),
1607 vec![1]
1608 );
1609
1610 let materialized = Source::from_iter([1, 2, 3])
1611 .run_with(
1612 buffered_flow()
1613 .via(Flow::identity().map(|item| item + 1))
1614 .map_materialized_value(|_| "buffered-flow")
1615 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both),
1616 )
1617 .unwrap();
1618 assert_eq!(materialized.0, "buffered-flow");
1619 assert_eq!(wait(materialized.1), 9);
1620
1621 let kept = Source::from_iter([1, 2, 3])
1622 .run_with(
1623 buffered_flow()
1624 .via_mat_with(Flow::identity().map(|item| item + 1), |_, _| "combined")
1625 .to(Sink::fold(0, |acc, item| acc + item)),
1626 )
1627 .unwrap();
1628 assert_eq!(kept, "combined");
1629 }
1630
1631 #[test]
1632 fn runtime_rate_flows_compose_in_flow_form() {
1633 let conflate = Flow::identity()
1634 .conflate(|left: i32, right| left + right)
1635 .map(|item| item + 1);
1636 assert_eq!(
1637 Source::single(4).via(conflate).run_collect().unwrap(),
1638 vec![5]
1639 );
1640
1641 let batch = Flow::identity()
1642 .batch(4, |item: i32| item, |left, right| left + right)
1643 .map(|item| item + 1);
1644 assert_eq!(Source::single(4).via(batch).run_collect().unwrap(), vec![5]);
1645
1646 let expand = Flow::identity()
1647 .expand(std::iter::once::<i32>)
1648 .map(|item| item + 1);
1649 assert_eq!(
1650 Source::from_iter(0..4).via(expand).run_collect().unwrap(),
1651 vec![1, 2, 3, 4]
1652 );
1653
1654 let aggregate = Flow::identity()
1655 .aggregate_with_boundary(
1656 Vec::<i32>::new,
1657 |mut items, item| {
1658 items.push(item);
1659 let ready = !items.is_empty();
1660 (items, ready)
1661 },
1662 |items| items.into_iter().sum::<i32>(),
1663 None,
1664 )
1665 .map(|item| item + 1);
1666 assert_eq!(
1667 Source::from_iter(0..4)
1668 .via(aggregate)
1669 .run_collect()
1670 .unwrap(),
1671 vec![1, 2, 3, 4]
1672 );
1673
1674 let detached = Flow::identity().detach().map(|item: i32| item + 1);
1675 assert_eq!(
1676 Source::from_iter(0..4).via(detached).run_collect().unwrap(),
1677 vec![1, 2, 3, 4]
1678 );
1679 }
1680
1681 #[test]
1682 fn high_use_source_flow_operators_work() {
1683 let result = Source::from_iter(0..8)
1684 .drop(1)
1685 .take(5)
1686 .filter_not(|item| item % 2 == 0)
1687 .map_concat(|item| [item, item + 10])
1688 .grouped(3)
1689 .run_collect()
1690 .unwrap();
1691
1692 assert_eq!(result, vec![vec![1, 11, 3], vec![13, 5, 15]]);
1693 }
1694
1695 #[test]
1696 fn prefix_and_tail_emits_prefix_and_live_tail() {
1697 let mut outer = Source::from_iter(0..5)
1698 .prefix_and_tail(2)
1699 .run_collect()
1700 .unwrap();
1701 assert_eq!(outer.len(), 1);
1702 let (prefix, tail) = outer.pop().unwrap();
1703 assert_eq!(prefix, vec![0, 1]);
1704 assert_eq!(tail.clone().run_collect().unwrap(), vec![2, 3, 4]);
1705 assert_eq!(
1706 tail.run_collect(),
1707 Err(StreamError::Failed(
1708 "substream source cannot be materialized more than once".into()
1709 ))
1710 );
1711 }
1712
1713 #[test]
1714 fn prefix_and_tail_fails_before_prefix_is_ready() {
1715 let result = Source::from_factory(|| {
1716 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into())), Ok(2)].into_iter())
1717 })
1718 .prefix_and_tail(2)
1719 .run_collect();
1720 assert!(matches!(result, Err(StreamError::Failed(message)) if message == "boom"));
1721 }
1722
1723 #[test]
1724 fn prefix_and_tail_tail_propagates_late_upstream_failure() {
1725 let mut outer = Source::from_factory(|| {
1726 Box::new(vec![Ok(1), Ok(2), Err(StreamError::Failed("boom".into())), Ok(3)].into_iter())
1727 })
1728 .prefix_and_tail(2)
1729 .run_collect()
1730 .unwrap();
1731 let (prefix, tail) = outer.pop().unwrap();
1732 assert_eq!(prefix, vec![1, 2]);
1733 assert_eq!(tail.run_collect(), Err(StreamError::Failed("boom".into())));
1734 }
1735
1736 #[test]
1737 fn prefix_and_tail_accepts_non_clone_elements() {
1738 #[derive(Debug, PartialEq, Eq)]
1739 struct NonClone(u8);
1740
1741 let mut outer = Source::from_factory(|| {
1742 Box::new(vec![Ok(NonClone(1)), Ok(NonClone(2)), Ok(NonClone(3))].into_iter())
1743 })
1744 .prefix_and_tail(2)
1745 .run_collect()
1746 .unwrap();
1747 let (prefix, tail) = outer.pop().unwrap();
1748 assert_eq!(prefix, vec![NonClone(1), NonClone(2)]);
1749 assert_eq!(tail.run_collect().unwrap(), vec![NonClone(3)]);
1750 }
1751
1752 #[test]
1753 fn flat_map_prefix_materializes_on_short_upstream_completion() {
1754 let values = Source::from_iter([1, 2])
1755 .flat_map_prefix(3, |prefix| {
1756 let sum = prefix.into_iter().sum::<i32>();
1757 Flow::identity().prepend(Source::single(sum))
1758 })
1759 .run_collect()
1760 .unwrap();
1761 assert_eq!(values, vec![3]);
1762 }
1763
1764 #[test]
1765 fn flat_map_prefix_does_not_materialize_on_early_upstream_failure() {
1766 let invoked = StdArc::new(StdAtomicBool::new(false));
1767 let invoked_for_stage = StdArc::clone(&invoked);
1768 let result = Source::from_factory(|| {
1769 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into()))].into_iter())
1770 })
1771 .flat_map_prefix(3, move |_prefix| {
1772 invoked_for_stage.store(true, StdOrdering::SeqCst);
1773 Flow::identity()
1774 })
1775 .run_collect();
1776 assert_eq!(result, Err(StreamError::Failed("boom".into())));
1777 assert!(!invoked.load(StdOrdering::SeqCst));
1778 }
1779
1780 #[test]
1781 fn flat_map_concat_flattens_nested_sources_sequentially() {
1782 let values = Source::from_iter([1, 2, 3])
1783 .flat_map_concat(|item| Source::from_iter(0..item))
1784 .run_collect()
1785 .unwrap();
1786 assert_eq!(values, vec![0, 0, 1, 0, 1, 2]);
1787 }
1788
1789 #[test]
1790 fn flat_map_merge_respects_breadth_bound() {
1791 let active = StdArc::new(StdAtomicUsize::new(0));
1792 let max_active = StdArc::new(StdAtomicUsize::new(0));
1793 let active_for_stage = StdArc::clone(&active);
1794 let max_for_stage = StdArc::clone(&max_active);
1795
1796 let mut values = Source::from_iter(0..6)
1797 .flat_map_merge(2, move |item| {
1798 let active = StdArc::clone(&active_for_stage);
1799 let max_active = StdArc::clone(&max_for_stage);
1800 Source::future(move || {
1801 let active = StdArc::clone(&active);
1802 let max_active = StdArc::clone(&max_active);
1803 async move {
1804 let now = active.fetch_add(1, StdOrdering::SeqCst) + 1;
1805 loop {
1806 let seen = max_active.load(StdOrdering::SeqCst);
1807 if now <= seen {
1808 break;
1809 }
1810 if max_active
1811 .compare_exchange(
1812 seen,
1813 now,
1814 StdOrdering::SeqCst,
1815 StdOrdering::SeqCst,
1816 )
1817 .is_ok()
1818 {
1819 break;
1820 }
1821 }
1822 thread::sleep(StdDuration::from_millis(20));
1823 active.fetch_sub(1, StdOrdering::SeqCst);
1824 Ok(item)
1825 }
1826 })
1827 })
1828 .run_collect()
1829 .unwrap();
1830 values.sort_unstable();
1831 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
1832 assert!(max_active.load(StdOrdering::SeqCst) <= 2);
1833 }
1834
1835 #[test]
1836 fn flat_map_merge_propagates_inner_failures() {
1837 let result = Source::from_iter([0, 1, 2])
1838 .flat_map_merge(2, |item| {
1839 if item == 1 {
1840 Source::failed(StreamError::Failed("boom".into()))
1841 } else {
1842 Source::single(item)
1843 }
1844 })
1845 .run_collect();
1846 assert_eq!(result, Err(StreamError::Failed("boom".into())));
1847 }
1848
1849 #[test]
1850 fn flat_map_merge_emits_ready_inner_output_while_upstream_is_blocked() {
1851 let (release_tx, release_rx) = mpsc::channel();
1852 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
1853 let queue = Source::from_factory(move || {
1854 let release_rx = StdArc::clone(&release_rx);
1855 let mut step = 0_u8;
1856 Box::new(std::iter::from_fn(move || {
1857 let item = match step {
1858 0 => Some(Ok(0)),
1859 1 => {
1860 release_rx
1861 .lock()
1862 .unwrap()
1863 .as_ref()
1864 .expect("release receiver available")
1865 .recv_timeout(StdDuration::from_secs(1))
1866 .expect("timed out waiting to release second upstream element");
1867 Some(Ok(1))
1868 }
1869 _ => None,
1870 };
1871 step += 1;
1872 item
1873 }))
1874 })
1875 .flat_map_merge(2, |item| Source::single(item + 10))
1876 .run_with(Sink::queue())
1877 .unwrap();
1878
1879 assert_eq!(queue.pull().unwrap(), Some(10));
1880 release_tx.send(()).unwrap();
1881 assert_eq!(queue.pull().unwrap(), Some(11));
1882 assert!(queue.pull().unwrap().is_none());
1883 }
1884
1885 #[test]
1886 fn group_by_routes_keys_and_drops_closed_keys() {
1887 let outer = Source::from_iter([0, 1, 2, 3, 4])
1888 .group_by(4, |item| item % 2, false)
1889 .run_with(Sink::queue())
1890 .unwrap();
1891
1892 let even = outer.pull().unwrap().unwrap();
1893 let even_completion = even.run_with(Sink::ignore()).unwrap();
1894 let odd = outer.pull().unwrap().unwrap();
1895 drop(even_completion);
1896
1897 assert_eq!(odd.run_collect().unwrap(), vec![1, 3]);
1898 assert!(outer.pull().unwrap().is_none());
1899 }
1900
1901 #[test]
1902 fn group_by_fails_when_distinct_key_limit_is_exceeded() {
1903 let outer = Source::from_iter([0, 1, 2])
1904 .group_by(2, |item| *item, false)
1905 .run_with(Sink::queue())
1906 .unwrap();
1907
1908 let _ = outer.pull().unwrap().unwrap();
1909 let _ = outer.pull().unwrap().unwrap();
1910 assert!(matches!(
1911 outer.pull(),
1912 Err(StreamError::Failed(message)) if message == "group_by reached max_substreams (2)"
1913 ));
1914 }
1915
1916 #[test]
1917 fn group_by_can_recreate_closed_substreams_when_enabled() {
1918 let (release_tx, release_rx) = mpsc::channel();
1919 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
1920 let outer = Source::from_factory(move || {
1921 let release_rx = StdArc::clone(&release_rx);
1922 let mut step = 0_u8;
1923 Box::new(std::iter::from_fn(move || {
1924 let item = match step {
1925 0 => Some(Ok(0)),
1926 1 => Some(Ok(1)),
1927 2 => {
1928 release_rx
1929 .lock()
1930 .unwrap()
1931 .as_ref()
1932 .expect("release receiver available")
1933 .recv_timeout(StdDuration::from_secs(1))
1934 .expect("timed out waiting to release recreated key");
1935 Some(Ok(0))
1936 }
1937 _ => None,
1938 };
1939 step += 1;
1940 item
1941 }))
1942 })
1943 .group_by(4, |item| item % 2, true)
1944 .run_with(Sink::queue())
1945 .unwrap();
1946
1947 let even = outer.pull().unwrap().unwrap();
1948 assert_eq!(wait(even.run_with(Sink::head()).unwrap()), 0);
1949 release_tx.send(()).unwrap();
1950
1951 let odd = outer.pull().unwrap().unwrap();
1952 assert_eq!(odd.run_collect().unwrap(), vec![1]);
1953
1954 let recreated_even = outer.pull().unwrap().unwrap();
1955 assert_eq!(recreated_even.run_collect().unwrap(), vec![0]);
1956 assert!(outer.pull().unwrap().is_none());
1957 }
1958
1959 #[test]
1960 fn group_by_panicking_key_fn_abruptly_terminates_live_substreams() {
1961 let outer = Source::from_iter([0, 1])
1962 .group_by(
1963 4,
1964 |item| {
1965 assert_ne!(*item, 1, "boom");
1966 item % 2
1967 },
1968 false,
1969 )
1970 .run_with(Sink::queue())
1971 .unwrap();
1972
1973 let substream = outer.pull().unwrap().unwrap();
1974 let (result_tx, result_rx) = mpsc::channel();
1975 thread::spawn(move || {
1976 let _ = result_tx.send(substream.run_collect());
1977 });
1978
1979 assert_eq!(
1980 result_rx.recv_timeout(StdDuration::from_secs(1)).unwrap(),
1981 Err(StreamError::AbruptTermination)
1982 );
1983 assert!(matches!(outer.pull(), Err(StreamError::AbruptTermination)));
1984 }
1985
1986 #[test]
1987 fn split_when_starts_new_substream_on_boundary_element() {
1988 let outer = Source::from_iter([1, 2, 0, 3, 0, 4, 5])
1989 .split_when(|item| *item == 0)
1990 .run_with(Sink::queue())
1991 .unwrap();
1992
1993 let first = outer.pull().unwrap().unwrap();
1994 assert_eq!(first.run_collect().unwrap(), vec![1, 2]);
1995 let second = outer.pull().unwrap().unwrap();
1996 assert_eq!(second.run_collect().unwrap(), vec![0, 3]);
1997 let third = outer.pull().unwrap().unwrap();
1998 assert_eq!(third.run_collect().unwrap(), vec![0, 4, 5]);
1999 assert!(outer.pull().unwrap().is_none());
2000 }
2001
2002 #[test]
2003 fn split_after_ends_current_substream_on_boundary_element() {
2004 let outer = Source::from_iter([1, 2, 0, 3, 0, 4, 5])
2005 .split_after(|item| *item == 0)
2006 .run_with(Sink::queue())
2007 .unwrap();
2008
2009 let first = outer.pull().unwrap().unwrap();
2010 assert_eq!(first.run_collect().unwrap(), vec![1, 2, 0]);
2011 let second = outer.pull().unwrap().unwrap();
2012 assert_eq!(second.run_collect().unwrap(), vec![3, 0]);
2013 let third = outer.pull().unwrap().unwrap();
2014 assert_eq!(third.run_collect().unwrap(), vec![4, 5]);
2015 assert!(outer.pull().unwrap().is_none());
2016 }
2017
2018 #[test]
2019 fn split_when_panicking_predicate_abruptly_terminates_live_substreams() {
2020 let outer = Source::from_iter([1, 2])
2021 .split_when(|item| {
2022 assert_ne!(*item, 2, "boom");
2023 false
2024 })
2025 .run_with(Sink::queue())
2026 .unwrap();
2027
2028 let substream = outer.pull().unwrap().unwrap();
2029 let (result_tx, result_rx) = mpsc::channel();
2030 thread::spawn(move || {
2031 let _ = result_tx.send(substream.run_collect());
2032 });
2033
2034 assert_eq!(
2035 result_rx.recv_timeout(StdDuration::from_secs(1)).unwrap(),
2036 Err(StreamError::AbruptTermination)
2037 );
2038 assert!(matches!(outer.pull(), Err(StreamError::AbruptTermination)));
2039 }
2040
2041 #[test]
2042 fn split_when_pre_buffer_segments_match_expected_count() {
2043 let outer = Source::from_iter(0..100)
2044 .split_when(|item| *item != 0 && *item % 10 == 0)
2045 .run_with(Sink::queue())
2046 .unwrap();
2047 let mut segment_count = 0;
2048 while let Some(substream) = outer.pull().unwrap() {
2049 let items: Vec<i32> = substream.run_collect().unwrap();
2050 assert!(!items.is_empty(), "segment should not be empty");
2051 segment_count += 1;
2052 }
2053 assert_eq!(segment_count, 10, "100 elements in segments of 10");
2054 }
2055
2056 #[test]
2057 fn split_after_pre_buffer_segments_match_expected_count() {
2058 let outer = Source::from_iter(0..100)
2059 .split_after(|item| (*item + 1) % 10 == 0)
2060 .run_with(Sink::queue())
2061 .unwrap();
2062 let mut segment_count = 0;
2063 let mut total = 0_i32;
2064 while let Some(substream) = outer.pull().unwrap() {
2065 let items: Vec<i32> = substream.run_collect().unwrap();
2066 assert!(!items.is_empty(), "segment should not be empty");
2067 total += items.len() as i32;
2068 segment_count += 1;
2069 }
2070 assert_eq!(segment_count, 10);
2071 assert_eq!(total, 100);
2072 }
2073
2074 #[test]
2075 fn group_by_single_key_fused_matches_general_path() {
2076 let outer = Source::from_iter(0..1000i64)
2077 .group_by(1, |_| 0u8, false)
2078 .run_with(Sink::queue())
2079 .unwrap();
2080 let substream = outer.pull().unwrap().unwrap();
2081 let items: Vec<i64> = substream.run_collect().unwrap();
2082 assert_eq!(items.len(), 1000);
2083 assert_eq!(items[0], 0);
2084 assert_eq!(items[999], 999);
2085 assert!(outer.pull().unwrap().is_none());
2086 }
2087
2088 #[test]
2089 fn group_by_single_key_fused_handles_key_change_with_substream_limit() {
2090 let outer = Source::from_iter([0, 1, 0])
2091 .group_by(2, |item| *item, false)
2092 .run_with(Sink::queue())
2093 .unwrap();
2094 let mut sources = vec![];
2095 while let Some(source) = outer.pull().unwrap() {
2096 sources.push(source);
2097 }
2098 assert_eq!(sources.len(), 2);
2099 assert_eq!(sources[0].clone().run_collect().unwrap(), vec![0, 0]);
2100 assert_eq!(sources[1].clone().run_collect().unwrap(), vec![1]);
2101 }
2102
2103 #[test]
2104 fn flat_map_merge_lock_lighter_matches_expected_count() {
2105 let items = Source::from_iter(0..20)
2106 .flat_map_merge(2, |item| Source::single(item + 100))
2107 .run_with(Sink::queue())
2108 .unwrap();
2109 let mut count = 0;
2110 while items.pull().unwrap().is_some() {
2111 count += 1;
2112 }
2113 assert_eq!(count, 20);
2114 }
2115
2116 #[test]
2126 fn group_by_single_key_emits_substream_before_upstream_completes() {
2127 let (tx, rx) = mpsc::sync_channel::<i32>(0);
2130 let rx = StdArc::new(std::sync::Mutex::new(rx));
2131
2132 let outer = Source::from_factory({
2133 let rx = StdArc::clone(&rx);
2134 move || {
2135 let rx = StdArc::clone(&rx);
2136 Box::new(std::iter::from_fn(move || {
2137 rx.lock().unwrap().recv().ok().map(Ok)
2138 })) as BoxStream<i32>
2139 }
2140 })
2141 .group_by(1, |_| 0u8, false)
2142 .run_with(Sink::queue())
2143 .unwrap();
2144
2145 let (sub_tx, sub_rx) = mpsc::channel::<Source<i32>>();
2148 let outer_thread = thread::spawn(move || {
2149 let substream = outer.pull().unwrap().expect("expected a substream");
2150 sub_tx.send(substream).unwrap();
2151 });
2152
2153 tx.send(0).unwrap();
2156
2157 let substream = sub_rx
2159 .recv_timeout(StdDuration::from_secs(5))
2160 .expect("timed out — group_by buffered first element before emitting substream");
2161
2162 for i in 1..100_i32 {
2164 tx.send(i).unwrap();
2165 }
2166 drop(tx);
2167
2168 let items: Vec<i32> = substream.run_collect().unwrap();
2169 assert_eq!(items.len(), 100);
2170 outer_thread.join().unwrap();
2171 }
2172
2173 #[test]
2174 fn group_by_concurrent_live_substreams_do_not_hold_ready_item_stress() {
2175 const STREAMS: usize = 32;
2176 const ROUNDS: usize = 8;
2177 const ITEMS: i64 = 8;
2178
2179 for _ in 0..ROUNDS {
2180 let barrier = StdArc::new(std::sync::Barrier::new(STREAMS));
2181 let mut handles = Vec::with_capacity(STREAMS);
2182
2183 for _ in 0..STREAMS {
2184 let barrier = StdArc::clone(&barrier);
2185 handles.push(thread::spawn(move || {
2186 let (tx, rx) = mpsc::sync_channel::<i64>(0);
2187 let rx = StdArc::new(std::sync::Mutex::new(rx));
2188
2189 let outer = Source::from_factory({
2190 let rx = StdArc::clone(&rx);
2191 move || {
2192 let rx = StdArc::clone(&rx);
2193 Box::new(std::iter::from_fn(move || {
2194 rx.lock().unwrap().recv().ok().map(Ok)
2195 })) as BoxStream<i64>
2196 }
2197 })
2198 .group_by(1, |_| 0_u8, false)
2199 .run_with(Sink::queue())
2200 .unwrap();
2201
2202 barrier.wait();
2203
2204 tx.send(0).unwrap();
2205 let substream = outer.pull().unwrap().expect("expected group_by substream");
2206 let subqueue = substream.run_with(Sink::queue()).unwrap();
2207 assert_eq!(subqueue.pull().unwrap(), Some(0));
2208
2209 for item in 1..ITEMS {
2210 tx.send(item).unwrap();
2211 assert_eq!(subqueue.pull().unwrap(), Some(item));
2212 }
2213 drop(tx);
2214
2215 assert!(subqueue.pull().unwrap().is_none());
2216 assert!(outer.pull().unwrap().is_none());
2217 }));
2218 }
2219
2220 for handle in handles {
2221 handle.join().expect("group_by stress worker panicked");
2222 }
2223 }
2224 }
2225
2226 #[test]
2233 fn split_when_emits_substream_before_segment_ends() {
2234 const SEGMENT_LEN: usize = 300;
2238
2239 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT_LEN * 2 + 4);
2240 for i in 0..SEGMENT_LEN as i32 {
2241 tx.send(i).unwrap();
2242 }
2243 tx.send(-1).unwrap(); tx.send(99).unwrap(); drop(tx);
2246
2247 let outer = Source::from_iter(rx)
2248 .split_when(|item| *item == -1)
2249 .run_with(Sink::queue())
2250 .unwrap();
2251
2252 let (result_tx, result_rx) = mpsc::channel();
2253 thread::spawn(move || {
2254 let first = outer.pull().unwrap().expect("expected first substream");
2255 let items: Vec<i32> = first.run_collect().unwrap();
2256 let second = outer.pull().unwrap().expect("expected second substream");
2257 let items2: Vec<i32> = second.run_collect().unwrap();
2258 let done = outer.pull().unwrap().is_none();
2259 let _ = result_tx.send((items, items2, done));
2260 });
2261
2262 let (items, items2, done) = result_rx
2263 .recv_timeout(StdDuration::from_secs(5))
2264 .expect("timed out — split_when is buffering the whole segment");
2265 assert_eq!(items.len(), SEGMENT_LEN);
2266 assert_eq!(items2, vec![-1, 99]);
2267 assert!(done);
2268 }
2269
2270 #[test]
2271 fn split_after_emits_substream_before_segment_ends() {
2272 const SEGMENT_LEN: usize = 300;
2273
2274 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT_LEN * 2 + 4);
2275 for i in 0..SEGMENT_LEN as i32 {
2276 tx.send(i).unwrap();
2277 }
2278 tx.send(-1).unwrap(); tx.send(99).unwrap();
2280 drop(tx);
2281
2282 let outer = Source::from_iter(rx)
2283 .split_after(|item| *item == -1)
2284 .run_with(Sink::queue())
2285 .unwrap();
2286
2287 let (result_tx, result_rx) = mpsc::channel();
2288 thread::spawn(move || {
2289 let first = outer.pull().unwrap().expect("expected first substream");
2290 let items: Vec<i32> = first.run_collect().unwrap();
2291 let second = outer.pull().unwrap().expect("expected second substream");
2292 let items2: Vec<i32> = second.run_collect().unwrap();
2293 let done = outer.pull().unwrap().is_none();
2294 let _ = result_tx.send((items, items2, done));
2295 });
2296
2297 let (items, items2, done) = result_rx
2298 .recv_timeout(StdDuration::from_secs(5))
2299 .expect("timed out — split_after is buffering the whole segment");
2300 assert_eq!(items.len(), SEGMENT_LEN + 1);
2302 assert_eq!(items2, vec![99]);
2303 assert!(done);
2304 }
2305
2306 #[test]
2310 fn flat_map_merge_coordinator_no_lost_wakeup_stress() {
2311 for _ in 0..20 {
2312 let result = Source::from_iter(0..50_i32)
2313 .flat_map_merge(8, |item| Source::from_iter(item..item + 3))
2314 .run_with(Sink::fold(0i64, |acc, item| acc + item as i64))
2315 .unwrap()
2316 .wait();
2317 assert_eq!(result, Ok(3825), "flat_map_merge produced wrong sum");
2320 }
2321 }
2322
2323 #[test]
2328 fn flat_map_merge_single_mutex_race_stress() {
2329 for _ in 0..20 {
2330 let result = Source::from_iter(0..100_i64)
2331 .flat_map_merge(16, |item| Source::from_iter([item, item + 1000]))
2332 .run_with(Sink::fold(0i64, |acc, v| acc + v))
2333 .unwrap()
2334 .wait();
2335 assert_eq!(result, Ok(109_900), "flat_map_merge single-mutex stress");
2339 }
2340 }
2341
2342 #[test]
2349 fn split_when_bounded_memory_rendezvous() {
2350 const SEGMENT: usize = 100;
2353 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT * 4);
2354 for i in 0..SEGMENT as i32 {
2355 tx.send(i).unwrap();
2356 }
2357 tx.send(-1).unwrap(); for i in 0..10_i32 {
2360 tx.send(i).unwrap();
2361 }
2362 drop(tx);
2363
2364 let outer = Source::from_iter(rx)
2365 .split_when(|item| *item == -1)
2366 .run_with(Sink::queue())
2367 .unwrap();
2368
2369 let (result_tx, result_rx) = mpsc::channel();
2370 thread::spawn(move || {
2371 let first = outer.pull().unwrap().expect("first segment");
2372 let seg1: Vec<i32> = first.run_collect().unwrap();
2373 let second = outer.pull().unwrap().expect("second segment");
2374 let seg2: Vec<i32> = second.run_collect().unwrap();
2375 let done = outer.pull().unwrap().is_none();
2376 result_tx.send((seg1, seg2, done)).unwrap();
2377 });
2378
2379 let (seg1, seg2, done) = result_rx
2380 .recv_timeout(StdDuration::from_secs(5))
2381 .expect("timed out — split_when writer held items past LIVE_SUBSTREAM_BATCH");
2382 assert_eq!(seg1.len(), SEGMENT, "first segment length");
2383 assert_eq!(seg2[0], -1, "boundary element starts second segment");
2384 assert_eq!(seg2.len(), 11, "second segment: boundary + 10 items");
2385 assert!(done);
2386 }
2387
2388 #[test]
2392 fn group_by_single_key_bounded_memory_rendezvous() {
2393 const N: usize = 200;
2396 let outer = Source::from_iter(0..N as i64)
2397 .group_by(1, |_| 0u8, false)
2398 .run_with(Sink::queue())
2399 .unwrap();
2400
2401 let (result_tx, result_rx) = mpsc::channel();
2402 thread::spawn(move || {
2403 let substream = outer.pull().unwrap().expect("substream");
2404 let items: Vec<i64> = substream.run_collect().unwrap();
2405 let done = outer.pull().unwrap().is_none();
2406 result_tx.send((items, done)).unwrap();
2407 });
2408
2409 let (items, done) = result_rx
2410 .recv_timeout(StdDuration::from_secs(5))
2411 .expect("timed out — group_by write batch held items beyond LIVE_SUBSTREAM_BATCH");
2412 assert_eq!(items.len(), N, "all items delivered");
2413 assert_eq!(items[0], 0);
2414 assert_eq!(items[N - 1], (N - 1) as i64);
2415 assert!(done);
2416 }
2417
2418 #[test]
2419 fn scan_emits_seed_and_accumulated_values() {
2420 let result = Source::from_iter(1..=3)
2421 .scan(0, |acc, item| acc + item)
2422 .run_collect()
2423 .unwrap();
2424
2425 assert_eq!(result, vec![0, 1, 3, 6]);
2426 }
2427
2428 #[test]
2429 fn limit_fails_after_max_elements() {
2430 let result = Source::from_iter(0..3).limit(2).run_collect();
2431
2432 assert_eq!(result, Err(StreamError::LimitExceeded { max: 2 }));
2433 }
2434
2435 #[test]
2436 fn limit_weighted_fails_with_limit_error_like_akka() {
2437 let result = Source::from_iter(["this", "is", "some", "string"])
2438 .via(Flow::identity().limit_weighted(15, |item: &&str| item.len()))
2439 .run_collect();
2440
2441 assert_eq!(result, Err(StreamError::LimitExceeded { max: 15 }));
2442 }
2443
2444 #[test]
2445 fn grouped_weighted_allows_oversized_first_element_like_akka() {
2446 let result = Source::from_iter([10_usize, 1, 2])
2447 .via(Flow::identity().grouped_weighted(5, |item: &usize| *item))
2448 .run_collect()
2449 .unwrap();
2450
2451 assert_eq!(result, vec![vec![10], vec![1, 2]]);
2452 }
2453
2454 #[test]
2455 fn grouped_weighted_keeps_oversized_later_element_in_current_group_like_akka() {
2456 let result = Source::from_iter([1_usize, 10, 2])
2457 .via(Flow::identity().grouped_weighted(5, |item: &usize| *item))
2458 .run_collect()
2459 .unwrap();
2460
2461 assert_eq!(result, vec![vec![1, 10], vec![2]]);
2462 }
2463
2464 #[test]
2465 fn sink_terminals_materialize_results() {
2466 let sum = Source::from_iter(1..=4)
2467 .run_with(Sink::fold(0, |acc, item| acc + item))
2468 .unwrap();
2469
2470 assert_eq!(wait(sum), 10);
2471 assert_eq!(
2472 wait(Source::from_iter(1..=4).run_with(Sink::head()).unwrap()),
2473 1
2474 );
2475 assert_eq!(
2476 wait(Source::from_iter(1..=4).run_with(Sink::last()).unwrap()),
2477 4
2478 );
2479 }
2480
2481 #[test]
2482 fn all_terminal_sink_variants_complete() {
2483 assert_eq!(
2484 wait(
2485 Source::from_iter([1, 2, 3])
2486 .run_with(Sink::collect())
2487 .unwrap()
2488 ),
2489 vec![1, 2, 3]
2490 );
2491 assert_eq!(
2492 wait(
2493 Source::<i32>::empty()
2494 .run_with(Sink::head_option())
2495 .unwrap()
2496 ),
2497 None
2498 );
2499 assert_eq!(
2500 wait(
2501 Source::from_iter([1, 2, 3])
2502 .run_with(Sink::last_option())
2503 .unwrap()
2504 ),
2505 Some(3)
2506 );
2507 assert_eq!(
2508 wait(
2509 Source::from_iter([1, 2, 3])
2510 .run_with(Sink::reduce(|acc, item| acc + item))
2511 .unwrap()
2512 ),
2513 6
2514 );
2515
2516 let seen = StdArc::new(StdAtomicUsize::new(0));
2517 let seen_by_sink = StdArc::clone(&seen);
2518 assert_eq!(
2519 wait(
2520 Source::from_iter([1_usize, 2, 3])
2521 .run_with(Sink::foreach(move |item| {
2522 seen_by_sink.fetch_add(item, StdOrdering::SeqCst);
2523 }))
2524 .unwrap()
2525 ),
2526 NotUsed
2527 );
2528 assert_eq!(seen.load(StdOrdering::SeqCst), 6);
2529 }
2530
2531 #[test]
2532 fn take_last_zero_returns_empty_vector() {
2533 let result = Source::from_iter([1, 2, 3])
2534 .run_with(Sink::take_last(0))
2535 .unwrap();
2536
2537 assert_eq!(wait(result), Vec::<i32>::new());
2538 }
2539
2540 #[test]
2541 fn bounded_head_terminals_complete_inline() {
2542 let materializer = Materializer::new();
2543
2544 let mut head = Source::from_iter(0_u64..1_000)
2545 .run_with_materializer(Sink::head(), &materializer)
2546 .unwrap();
2547 assert_eq!(materializer.active_streams(), 0);
2548 assert_eq!(head.try_wait(), Some(Ok(0)));
2549
2550 let mut filtered_head = Source::from_iter(0_u64..1_000)
2551 .filter(|item| *item >= 10)
2552 .run_with_materializer(Sink::head(), &materializer)
2553 .unwrap();
2554 assert_eq!(materializer.active_streams(), 0);
2555 assert_eq!(filtered_head.try_wait(), Some(Ok(10)));
2556
2557 let mut head_option = Source::<u64>::empty()
2558 .run_with_materializer(Sink::head_option(), &materializer)
2559 .unwrap();
2560 assert_eq!(materializer.active_streams(), 0);
2561 assert_eq!(head_option.try_wait(), Some(Ok(None)));
2562 }
2563
2564 #[test]
2565 fn bounded_head_fast_path_preserves_terminal_errors() {
2566 let materializer = Materializer::new();
2567
2568 let mut empty = Source::<u64>::empty()
2569 .run_with_materializer(Sink::head(), &materializer)
2570 .unwrap();
2571 assert_eq!(empty.try_wait(), Some(Err(StreamError::EmptyStream)));
2572
2573 let mut failed = Source::<u64>::failed(StreamError::Failed("boom".into()))
2574 .run_with_materializer(Sink::head(), &materializer)
2575 .unwrap();
2576 assert_eq!(
2577 failed.try_wait(),
2578 Some(Err(StreamError::Failed("boom".into())))
2579 );
2580 assert_eq!(materializer.active_streams(), 0);
2581 }
2582
2583 #[test]
2584 fn runnable_graph_composes_source_and_sink() {
2585 let graph = Source::from_iter(1..=4)
2586 .map(|item| item * 2)
2587 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::right);
2588
2589 assert_eq!(wait(graph.run().unwrap()), 20);
2590
2591 let graph = Source::single(1)
2592 .map_materialized_value(|_| 20)
2593 .to(Sink::ignore())
2594 .map_materialized_value(|value| value + 1);
2595 assert_eq!(graph.run().unwrap(), 21);
2596
2597 let ignored = Source::single(1).to(Sink::ignore()).run().unwrap();
2598 assert_eq!(ignored, NotUsed);
2599 }
2600
2601 #[test]
2602 fn materialized_values_follow_keep_defaults() {
2603 let source = Source::single(1).map_materialized_value(|_| "source");
2604 let flow = Flow::identity().map_materialized_value(|_| "flow");
2605
2606 let source_mat = source.clone().via(flow.clone()).to(Sink::ignore()).run();
2607 assert_eq!(source_mat.unwrap(), "source");
2608
2609 let combined = source
2610 .via_mat(flow, Keep::both)
2611 .to_mat(Sink::ignore(), Keep::both)
2612 .run()
2613 .unwrap();
2614 assert_eq!(combined.0, ("source", "flow"));
2615 assert_eq!(wait(combined.1), NotUsed);
2616
2617 let sink_mat = Source::single(41)
2618 .map_materialized_value(|_| "ignored source")
2619 .run_with(Sink::fold(1, |acc, item| acc + item))
2620 .unwrap();
2621 assert_eq!(wait(sink_mat), 42);
2622 }
2623
2624 #[test]
2625 fn flow_to_sink_preserves_flow_materialized_value_by_default() {
2626 let sink = Flow::identity()
2627 .map(|item: i32| item + 1)
2628 .map_materialized_value(|_| "flow")
2629 .to(Sink::fold(0, |acc, item| acc + item));
2630
2631 let materialized = Source::from_iter([1, 2, 3]).run_with(sink).unwrap();
2632
2633 assert_eq!(materialized, "flow");
2634 let explicit = Flow::identity()
2635 .map(|item: i32| item + 1)
2636 .map_materialized_value(|_| "flow")
2637 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both)
2638 .run_with(Source::from_iter([1, 2, 3]))
2639 .unwrap();
2640 assert_eq!(explicit, NotUsed);
2641
2642 let explicit = Source::from_iter([1, 2, 3])
2643 .run_with(
2644 Flow::identity()
2645 .map(|item: i32| item + 1)
2646 .map_materialized_value(|_| "flow")
2647 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both),
2648 )
2649 .unwrap();
2650 assert_eq!(explicit.0, "flow");
2651 assert_eq!(wait(explicit.1), 9);
2652 }
2653
2654 #[test]
2655 fn materializer_shutdown_fails_materialization() {
2656 let materializer = Materializer::new();
2657 let named = materializer.with_name_prefix("test-stream");
2658 materializer.shutdown();
2659
2660 let graph = Source::single(1).to(Sink::ignore());
2661
2662 assert_eq!(named.name_prefix(), "test-stream");
2663 assert_eq!(
2664 graph.run_with_materializer(&named),
2665 Err(StreamError::AbruptTermination)
2666 );
2667 }
2668
2669 #[test]
2670 fn materializer_shutdown_fails_running_stream_completion() {
2671 let materializer = Materializer::new();
2672 let completion = Source::repeat(1)
2673 .run_with_materializer(Sink::ignore(), &materializer)
2674 .unwrap();
2675
2676 assert_eq!(materializer.active_streams(), 1);
2677 materializer.shutdown();
2678 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
2679 assert_eq!(materializer.active_streams(), 0);
2680 }
2681
2682 #[test]
2683 fn dropped_stream_completion_cancels_running_stream() {
2684 let materializer = Materializer::new();
2685 let completion = Source::repeat(1)
2686 .run_with_materializer(Sink::ignore(), &materializer)
2687 .unwrap();
2688
2689 assert_eq!(materializer.active_streams(), 1);
2690 drop(completion);
2691 for _ in 0..50 {
2692 if materializer.active_streams() == 0 {
2693 break;
2694 }
2695 thread::sleep(Duration::from_millis(5));
2696 }
2697 assert_eq!(materializer.active_streams(), 0);
2698 }
2699
2700 #[test]
2701 fn runtime_timers_fire_cancel_and_stop_on_shutdown() {
2702 let materializer = Materializer::new();
2703 let (once_tx, once_rx) = mpsc::channel();
2704 let once = materializer.schedule_once(Duration::from_millis(5), move || {
2705 once_tx.send(()).unwrap();
2706 });
2707 once_rx.recv_timeout(Duration::from_millis(250)).unwrap();
2708 assert!(!once.is_cancelled());
2709
2710 let (cancelled_tx, cancelled_rx) = mpsc::channel();
2711 let cancelled = materializer.schedule_once(Duration::from_millis(25), move || {
2712 cancelled_tx.send(()).unwrap();
2713 });
2714 assert!(cancelled.cancel());
2715 assert!(!cancelled.cancel());
2716 assert!(cancelled.is_cancelled());
2717 assert!(
2718 cancelled_rx
2719 .recv_timeout(Duration::from_millis(75))
2720 .is_err()
2721 );
2722
2723 let fixed_delay_count = StdArc::new(StdAtomicUsize::new(0));
2724 let fixed_delay_task_count = StdArc::clone(&fixed_delay_count);
2725 let fixed_delay = materializer.schedule_with_fixed_delay(
2726 Duration::from_millis(1),
2727 Duration::from_millis(5),
2728 move || {
2729 fixed_delay_task_count.fetch_add(1, StdOrdering::SeqCst);
2730 },
2731 );
2732 thread::sleep(Duration::from_millis(25));
2733 assert!(fixed_delay_count.load(StdOrdering::SeqCst) > 0);
2734 fixed_delay.cancel();
2735
2736 let fixed_rate_count = StdArc::new(StdAtomicUsize::new(0));
2737 let fixed_rate_task_count = StdArc::clone(&fixed_rate_count);
2738 let fixed_rate = materializer.schedule_at_fixed_rate(
2739 Duration::from_millis(1),
2740 Duration::from_millis(5),
2741 move || {
2742 fixed_rate_task_count.fetch_add(1, StdOrdering::SeqCst);
2743 },
2744 );
2745 thread::sleep(Duration::from_millis(25));
2746 assert!(fixed_rate_count.load(StdOrdering::SeqCst) > 0);
2747 fixed_rate.cancel();
2748
2749 let shutdown_materializer = Materializer::new();
2750 let (shutdown_tx, shutdown_rx) = mpsc::channel();
2751 shutdown_materializer.schedule_once(Duration::from_millis(25), move || {
2752 shutdown_tx.send(()).unwrap();
2753 });
2754 shutdown_materializer.shutdown();
2755 assert!(shutdown_rx.recv_timeout(Duration::from_millis(75)).is_err());
2756 }
2757
2758 #[test]
2759 fn runtime_timer_driver_preserves_fixed_rate_cadence_under_slow_tasks() {
2760 use std::sync::{Condvar, Mutex};
2761
2762 #[derive(Debug)]
2763 enum TimerEvent {
2764 Started(usize, Instant),
2765 Completed(usize, Instant),
2766 }
2767
2768 let recv_event = |rx: &mpsc::Receiver<TimerEvent>, label: &str| {
2769 rx.recv_timeout(Duration::from_secs(20))
2770 .unwrap_or_else(|err| panic!("{label}: expected timer event within 20 s: {err}"))
2771 };
2772 let release = |gate: &StdArc<(Mutex<bool>, Condvar)>| {
2773 let (released, condvar) = &**gate;
2774 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2775 *released = true;
2776 condvar.notify_all();
2777 };
2778
2779 let interval = Duration::from_secs(2);
2780 let overrun = interval + Duration::from_millis(250);
2781
2782 let rate_materializer = Materializer::new();
2783 let (rate_tx, rate_rx) = mpsc::channel();
2784 let rate_runs = StdArc::new(StdAtomicUsize::new(0));
2785 let rate_task_runs = StdArc::clone(&rate_runs);
2786 let rate_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2787 let rate_task_gate = StdArc::clone(&rate_gate);
2788 let fixed_rate =
2789 rate_materializer.schedule_at_fixed_rate(Duration::ZERO, interval, move || {
2790 let run = rate_task_runs.fetch_add(1, StdOrdering::SeqCst) + 1;
2791 rate_tx
2792 .send(TimerEvent::Started(run, Instant::now()))
2793 .unwrap();
2794 if run == 1 {
2795 let (released, condvar) = &*rate_task_gate;
2796 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2797 while !*released {
2798 released = condvar
2799 .wait(released)
2800 .unwrap_or_else(|poison| poison.into_inner());
2801 }
2802 rate_tx
2803 .send(TimerEvent::Completed(run, Instant::now()))
2804 .unwrap();
2805 }
2806 });
2807 let rate_first_started = match recv_event(&rate_rx, "fixed-rate first task") {
2808 TimerEvent::Started(1, at) => at,
2809 other => panic!("fixed-rate first task: unexpected event {other:?}"),
2810 };
2811 assert!(wait_until(Duration::from_secs(20), || {
2812 rate_first_started.elapsed() >= overrun
2813 }));
2814 release(&rate_gate);
2815 let rate_first_completed = match recv_event(&rate_rx, "fixed-rate first completion") {
2816 TimerEvent::Completed(1, at) => at,
2817 other => panic!("fixed-rate first completion: unexpected event {other:?}"),
2818 };
2819 let rate_second_started = match recv_event(&rate_rx, "fixed-rate second task") {
2820 TimerEvent::Started(2, at) => at,
2821 other => panic!("fixed-rate second task: unexpected event {other:?}"),
2822 };
2823 fixed_rate.cancel();
2824 rate_materializer.shutdown();
2825
2826 let delay_materializer = Materializer::new();
2827 let (delay_tx, delay_rx) = mpsc::channel();
2828 let delay_runs = StdArc::new(StdAtomicUsize::new(0));
2829 let delay_task_runs = StdArc::clone(&delay_runs);
2830 let delay_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2831 let delay_task_gate = StdArc::clone(&delay_gate);
2832 let fixed_delay =
2833 delay_materializer.schedule_with_fixed_delay(Duration::ZERO, interval, move || {
2834 let run = delay_task_runs.fetch_add(1, StdOrdering::SeqCst) + 1;
2835 delay_tx
2836 .send(TimerEvent::Started(run, Instant::now()))
2837 .unwrap();
2838 if run == 1 {
2839 let (released, condvar) = &*delay_task_gate;
2840 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2841 while !*released {
2842 released = condvar
2843 .wait(released)
2844 .unwrap_or_else(|poison| poison.into_inner());
2845 }
2846 delay_tx
2847 .send(TimerEvent::Completed(run, Instant::now()))
2848 .unwrap();
2849 }
2850 });
2851 let delay_first_started = match recv_event(&delay_rx, "fixed-delay first task") {
2852 TimerEvent::Started(1, at) => at,
2853 other => panic!("fixed-delay first task: unexpected event {other:?}"),
2854 };
2855 assert!(wait_until(Duration::from_secs(20), || {
2856 delay_first_started.elapsed() >= overrun
2857 }));
2858 release(&delay_gate);
2859 let delay_first_completed = match recv_event(&delay_rx, "fixed-delay first completion") {
2860 TimerEvent::Completed(1, at) => at,
2861 other => panic!("fixed-delay first completion: unexpected event {other:?}"),
2862 };
2863 let delay_second_started = match recv_event(&delay_rx, "fixed-delay second task") {
2864 TimerEvent::Started(2, at) => at,
2865 other => panic!("fixed-delay second task: unexpected event {other:?}"),
2866 };
2867 fixed_delay.cancel();
2868 delay_materializer.shutdown();
2869
2870 let rate_task_time = rate_first_completed.duration_since(rate_first_started);
2871 let rate_catch_up = rate_second_started.duration_since(rate_first_completed);
2872 let delay_task_time = delay_first_completed.duration_since(delay_first_started);
2873 let delay_gap = delay_second_started.duration_since(delay_first_completed);
2874 assert!(
2875 rate_task_time >= interval,
2876 "fixed-rate first task should overrun its interval; ran for {rate_task_time:?}"
2877 );
2878 assert!(
2879 rate_catch_up < interval,
2880 "fixed-rate second task should catch up after an overrun; waited {rate_catch_up:?}"
2881 );
2882 assert!(
2883 delay_task_time >= interval,
2884 "fixed-delay first task should overrun its interval; ran for {delay_task_time:?}"
2885 );
2886 assert!(
2887 delay_gap >= interval,
2888 "fixed-delay second task fired before one full delay elapsed after completion: {delay_gap:?}",
2889 );
2890 }
2891
2892 #[test]
2893 fn runtime_repeating_timer_cancellation_stops_future_fires() {
2894 let materializer = Materializer::new();
2895 let (tx, rx) = mpsc::channel();
2896 let timer = materializer.schedule_at_fixed_rate(
2897 Duration::from_millis(1),
2898 Duration::from_millis(30),
2899 move || {
2900 tx.send(()).unwrap();
2901 },
2902 );
2903
2904 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2905 assert!(timer.cancel());
2906 assert!(rx.recv_timeout(Duration::from_millis(90)).is_err());
2907 materializer.shutdown();
2908 }
2909
2910 #[test]
2911 fn runtime_panicking_once_timer_does_not_kill_driver_or_later_timers() {
2912 let materializer = Materializer::new();
2913 materializer.schedule_once(Duration::from_millis(1), || {
2914 panic!("timer boom");
2915 });
2916
2917 let (tx, rx) = mpsc::channel();
2918 materializer.schedule_once(Duration::from_millis(20), move || {
2919 tx.send(()).unwrap();
2920 });
2921
2922 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2923 materializer.shutdown();
2924 }
2925
2926 #[test]
2927 fn runtime_panicking_fixed_rate_timer_stops_itself_and_leaves_driver_alive() {
2928 let materializer = Materializer::new();
2929 let panic_count = StdArc::new(StdAtomicUsize::new(0));
2930 let panic_count_task = StdArc::clone(&panic_count);
2931 materializer.schedule_at_fixed_rate(Duration::ZERO, Duration::from_millis(20), move || {
2932 panic_count_task.fetch_add(1, StdOrdering::SeqCst);
2933 panic!("fixed-rate boom");
2934 });
2935
2936 assert!(wait_until(Duration::from_millis(150), || {
2937 panic_count.load(StdOrdering::SeqCst) == 1
2938 }));
2939
2940 let (tx, rx) = mpsc::channel();
2941 materializer.schedule_once(Duration::from_millis(30), move || {
2942 tx.send(()).unwrap();
2943 });
2944 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2945
2946 thread::sleep(Duration::from_millis(90));
2947 assert_eq!(panic_count.load(StdOrdering::SeqCst), 1);
2948 materializer.shutdown();
2949 }
2950
2951 #[test]
2952 fn runtime_slow_timer_task_does_not_delay_unrelated_timers() {
2953 let materializer = Materializer::new();
2954 let (started_tx, started_rx) = mpsc::channel();
2955 let release_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2956 let release_task = StdArc::clone(&release_gate);
2957 let runs = StdArc::new(StdAtomicUsize::new(0));
2958 let task_runs = StdArc::clone(&runs);
2959 let slow_timer = materializer.schedule_at_fixed_rate(
2960 Duration::ZERO,
2961 Duration::from_millis(250),
2962 move || {
2963 if task_runs.fetch_add(1, StdOrdering::SeqCst) == 0 {
2964 started_tx.send(()).unwrap();
2965 let (released, condvar) = &*release_task;
2966 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2967 while !*released {
2968 released = condvar
2969 .wait(released)
2970 .unwrap_or_else(|poison| poison.into_inner());
2971 }
2972 }
2973 },
2974 );
2975
2976 started_rx
2977 .recv_timeout(Duration::from_secs(10))
2978 .expect("slow timer task should start");
2979
2980 let (tx, rx) = mpsc::channel();
2981 materializer.schedule_once(Duration::from_millis(10), move || {
2982 tx.send(()).unwrap();
2983 });
2984 let fired = rx.recv_timeout(Duration::from_secs(10));
2985
2986 let (released, condvar) = &*release_gate;
2987 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2988 *released = true;
2989 condvar.notify_all();
2990 drop(released);
2991
2992 slow_timer.cancel();
2993 materializer.shutdown();
2994 fired.expect("unrelated timer should fire while slow timer task is still blocked");
2995 }
2996
2997 #[test]
2998 fn runtime_shutdown_stops_timer_driver_thread() {
2999 let materializer = Materializer::new();
3000 assert!(wait_until(Duration::from_secs(1), || materializer
3001 .timer_driver_is_live()));
3002
3003 materializer.shutdown();
3004 assert!(wait_until(Duration::from_secs(2), || !materializer
3005 .timer_driver_is_live()));
3006 }
3007
3008 #[test]
3009 fn runtime_timer_driver_orders_many_timers_by_deadline() {
3010 let materializer = Materializer::new();
3011 let (tx, rx) = mpsc::channel();
3012 let schedule = [(450_u64, 4_u8), (50, 1), (350, 3), (150, 2), (550, 5)];
3013
3014 for (delay_ms, value) in schedule {
3015 let tx = tx.clone();
3016 materializer.schedule_once(Duration::from_millis(delay_ms), move || {
3017 tx.send(value).unwrap();
3018 });
3019 }
3020 drop(tx);
3021
3022 let mut received = Vec::new();
3023 for _ in 0..schedule.len() {
3024 received.push(rx.recv_timeout(Duration::from_secs(10)).unwrap());
3025 }
3026 materializer.shutdown();
3027
3028 assert_eq!(received, vec![1, 2, 3, 4, 5]);
3029 }
3030
3031 #[test]
3032 fn runtime_timer_driver_uses_one_thread_per_runtime_regardless_of_timer_count() {
3033 let materializer = Materializer::new();
3034 let thread_name = materializer.timer_thread_name().to_owned();
3035 let linux_thread_name = thread_name.chars().take(15).collect::<String>();
3041 assert!(wait_until(Duration::from_secs(5), || {
3042 materializer.timer_driver_is_live() && linux_thread_count(&linux_thread_name) >= 1
3043 }));
3044 let live_timer_threads = linux_thread_count(&linux_thread_name);
3045
3046 for _ in 0..128 {
3047 materializer.schedule_once(Duration::from_secs(60), || {});
3048 }
3049
3050 assert!(
3051 wait_until(Duration::from_secs(5), || {
3052 materializer.timer_driver_is_live()
3053 && linux_thread_count(&linux_thread_name) == live_timer_threads
3054 }),
3055 "scheduling timers should not create extra timer threads for a runtime",
3056 );
3057 materializer.shutdown();
3058 assert!(wait_until(Duration::from_secs(5), || {
3059 !materializer.timer_driver_is_live()
3060 && linux_thread_count(&linux_thread_name) < live_timer_threads
3061 }));
3062 }
3063
3064 #[test]
3065 fn cancelled_and_never_sinks_have_distinct_materialization_results() {
3066 assert_eq!(
3067 Source::repeat(1)
3068 .run_with(Sink::cancelled())
3069 .expect("cancelled sink materializes"),
3070 NotUsed
3071 );
3072 assert_eq!(
3073 Source::single(1)
3074 .run_with(Sink::never())
3075 .expect("never sink materializes")
3076 .try_wait(),
3077 None
3078 );
3079 }
3080
3081 #[test]
3082 fn never_sink_finishes_on_materializer_shutdown() {
3083 let materializer = Materializer::new();
3084 let completion = Source::single(1)
3085 .run_with_materializer(Sink::never(), &materializer)
3086 .unwrap();
3087
3088 materializer.shutdown();
3089 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
3090 }
3091
3092 #[test]
3093 fn dropping_source_never_completion_releases_parked_worker() {
3094 let materializer = Materializer::new();
3095 let completion = Source::<i32>::never()
3096 .run_with_materializer(Sink::ignore(), &materializer)
3097 .unwrap();
3098
3099 assert!(wait_until(StdDuration::from_secs(1), || {
3100 materializer.active_streams() == 1
3101 }));
3102 assert_eq!(materializer.active_streams(), 1);
3103
3104 drop(completion);
3105
3106 assert!(wait_until(StdDuration::from_secs(15), || {
3107 materializer.active_streams() == 0
3108 }));
3109 assert_eq!(materializer.active_streams(), 0);
3110 }
3111
3112 #[test]
3113 fn future_and_maybe_sources_emit_values() {
3114 let future_value = Source::future(|| async { Ok(7) }).run_collect().unwrap();
3115 assert_eq!(future_value, vec![7]);
3116
3117 let future_source = Source::future_source(|| async { Ok(Source::from_iter([1, 2, 3])) })
3118 .run_collect()
3119 .unwrap();
3120 assert_eq!(future_source, vec![1, 2, 3]);
3121
3122 let (handle, source) = Source::maybe();
3123 assert_eq!(
3124 source.clone().run_collect(),
3125 Err(StreamError::MaybeIncomplete)
3126 );
3127 handle.complete(9).unwrap();
3128 assert_eq!(source.run_collect().unwrap(), vec![9]);
3129 }
3130
3131 #[test]
3132 fn wp6b_source_generators_emit_and_fail_like_stream_errors() {
3133 assert_eq!(
3134 Source::cycle(|| [1, 2, 3].into_iter())
3135 .take(8)
3136 .run_collect()
3137 .unwrap(),
3138 vec![1, 2, 3, 1, 2, 3, 1, 2]
3139 );
3140 assert_eq!(
3141 Source::<i32>::cycle(std::iter::empty::<i32>).run_collect(),
3142 Err(StreamError::Failed("empty iterator".into()))
3143 );
3144 assert_eq!(
3145 Source::unfold(0, |state| (state < 4).then_some((state + 1, state)))
3146 .run_collect()
3147 .unwrap(),
3148 vec![0, 1, 2, 3]
3149 );
3150 assert_eq!(
3151 Source::unfold_async(0, |state| async move {
3152 Ok((state < 4).then_some((state + 1, state * 2)))
3153 })
3154 .run_collect()
3155 .unwrap(),
3156 vec![0, 2, 4, 6]
3157 );
3158 assert!(matches!(
3159 Source::<i32>::lazy_single(|| panic!("boom")).run_collect(),
3160 Err(StreamError::Failed(message)) if message == "lazy_single factory panicked"
3161 ));
3162 }
3163
3164 #[test]
3165 fn wp6b_lazy_sources_defer_until_first_pull_and_complete_deferred_mat() {
3166 let created = StdArc::new(StdAtomicUsize::new(0));
3167 let created_for_source = StdArc::clone(&created);
3168 let source = Source::<i32>::lazy_source(move || {
3169 created_for_source.fetch_add(1, StdOrdering::SeqCst);
3170 Source::from_iter([7, 8]).map_materialized_value(|_| 99)
3171 });
3172 let materializer = Materializer::new();
3173 let (mut stream, mut mat) = StdArc::clone(&source.factory)
3174 .create(&materializer)
3175 .unwrap();
3176
3177 assert_eq!(created.load(StdOrdering::SeqCst), 0);
3178 assert!(mat.try_wait().is_none());
3179 assert_eq!(stream.next().unwrap().unwrap(), 7);
3180 assert_eq!(mat.wait().unwrap(), 99);
3181 assert_eq!(created.load(StdOrdering::SeqCst), 1);
3182 assert_eq!(stream.next().unwrap().unwrap(), 8);
3183
3184 let never_created = StdArc::new(StdAtomicUsize::new(0));
3185 let never_created_for_source = StdArc::clone(&never_created);
3186 let mat = Source::<i32>::lazy_future_source(move || {
3187 never_created_for_source.fetch_add(1, StdOrdering::SeqCst);
3188 async { Ok(Source::single(1)) }
3189 })
3190 .to(Sink::cancelled())
3191 .run()
3192 .unwrap();
3193 assert!(matches!(mat.wait(), Err(StreamError::Failed(_))));
3194 assert_eq!(never_created.load(StdOrdering::SeqCst), 0);
3195
3196 let lazy_future = StdArc::new(StdAtomicUsize::new(0));
3197 let lazy_future_for_source = StdArc::clone(&lazy_future);
3198 let source = Source::lazy_future(move || {
3199 lazy_future_for_source.fetch_add(1, StdOrdering::SeqCst);
3200 async { Ok(42) }
3201 });
3202 let (mut stream, _) = StdArc::clone(&source.factory)
3203 .create(&Materializer::new())
3204 .unwrap();
3205 assert_eq!(lazy_future.load(StdOrdering::SeqCst), 0);
3206 assert_eq!(stream.next().unwrap().unwrap(), 42);
3207 assert_eq!(lazy_future.load(StdOrdering::SeqCst), 1);
3208 }
3209
3210 #[test]
3211 fn wp6b_unfold_resource_closes_on_completion_failure_and_cancellation() {
3212 let closed = StdArc::new(StdAtomicUsize::new(0));
3213 let closed_on_complete = StdArc::clone(&closed);
3214 let values = Source::unfold_resource(
3215 || Ok(std::collections::VecDeque::from([1, 2, 3])),
3216 |items| Ok(items.pop_front()),
3217 move |_items| {
3218 closed_on_complete.fetch_add(1, StdOrdering::SeqCst);
3219 Ok(())
3220 },
3221 )
3222 .run_collect()
3223 .unwrap();
3224 assert_eq!(values, vec![1, 2, 3]);
3225 assert_eq!(closed.load(StdOrdering::SeqCst), 1);
3226
3227 let closed_on_failure = StdArc::new(StdAtomicUsize::new(0));
3228 let closed_on_failure_for_close = StdArc::clone(&closed_on_failure);
3229 let failed = Source::<i32>::unfold_resource(
3230 || Ok(()),
3231 |_| Err(StreamError::Failed("read".into())),
3232 move |_| {
3233 closed_on_failure_for_close.fetch_add(1, StdOrdering::SeqCst);
3234 Err(StreamError::Failed("close".into()))
3235 },
3236 )
3237 .run_collect();
3238 assert_eq!(failed, Err(StreamError::Failed("read".into())));
3239 assert_eq!(closed_on_failure.load(StdOrdering::SeqCst), 1);
3240
3241 let closed_on_cancel = StdArc::new(StdAtomicUsize::new(0));
3242 let closed_on_cancel_for_close = StdArc::clone(&closed_on_cancel);
3243 let first = Source::unfold_resource(
3244 || Ok(0_usize),
3245 |next| {
3246 let item = *next;
3247 *next += 1;
3248 Ok(Some(item))
3249 },
3250 move |_| {
3251 closed_on_cancel_for_close.fetch_add(1, StdOrdering::SeqCst);
3252 Ok(())
3253 },
3254 )
3255 .run_with(Sink::head())
3256 .unwrap();
3257 assert_eq!(first.wait().unwrap(), 0);
3258 assert!(wait_until(Duration::from_millis(250), || {
3259 closed_on_cancel.load(StdOrdering::SeqCst) == 1
3260 }));
3261 }
3262
3263 #[test]
3264 fn wp6b_async_resource_and_async_accumulators_are_sequential() {
3265 let closed = StdArc::new(StdAtomicUsize::new(0));
3266 let closed_for_close = StdArc::clone(&closed);
3267 let values = Source::unfold_resource_async(
3268 || async { Ok(std::collections::VecDeque::from([1, 2, 3])) },
3269 |items| {
3270 let item = items.pop_front();
3271 async move { Ok(item) }
3272 },
3273 move |_items| {
3274 let closed = StdArc::clone(&closed_for_close);
3275 async move {
3276 closed.fetch_add(1, StdOrdering::SeqCst);
3277 Ok(())
3278 }
3279 },
3280 )
3281 .run_collect()
3282 .unwrap();
3283 assert_eq!(values, vec![1, 2, 3]);
3284 assert_eq!(closed.load(StdOrdering::SeqCst), 1);
3285
3286 let closed_on_failure = StdArc::new(StdAtomicUsize::new(0));
3287 let closed_on_failure_for_close = StdArc::clone(&closed_on_failure);
3288 let failed = Source::<i32>::unfold_resource_async(
3289 || async { Ok(()) },
3290 |_resource| async { Err(StreamError::Failed("read".into())) },
3291 move |_resource| {
3292 let closed_on_failure = StdArc::clone(&closed_on_failure_for_close);
3293 async move {
3294 closed_on_failure.fetch_add(1, StdOrdering::SeqCst);
3295 Err(StreamError::Failed("close".into()))
3296 }
3297 },
3298 )
3299 .run_collect();
3300 assert_eq!(failed, Err(StreamError::Failed("read".into())));
3301 assert_eq!(closed_on_failure.load(StdOrdering::SeqCst), 1);
3302
3303 let active = StdArc::new(StdAtomicUsize::new(0));
3304 let max_active = StdArc::new(StdAtomicUsize::new(0));
3305 let active_for_stage = StdArc::clone(&active);
3306 let max_for_stage = StdArc::clone(&max_active);
3307 let scanned = Source::from_iter(1..=4)
3308 .scan_async(0, move |acc, item| {
3309 let active = StdArc::clone(&active_for_stage);
3310 let max_active = StdArc::clone(&max_for_stage);
3311 async move {
3312 let now = active.fetch_add(1, StdOrdering::SeqCst) + 1;
3313 max_active.fetch_max(now, StdOrdering::SeqCst);
3314 tokio::time::sleep(Duration::from_millis(1)).await;
3315 active.fetch_sub(1, StdOrdering::SeqCst);
3316 Ok(acc + item)
3317 }
3318 })
3319 .run_collect()
3320 .unwrap();
3321 assert_eq!(scanned, vec![0, 1, 3, 6, 10]);
3322 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
3323
3324 let folded = Source::from_iter(1..=4)
3325 .fold_async(0, |acc, item| async move { Ok(acc + item) })
3326 .run_collect()
3327 .unwrap();
3328 assert_eq!(folded, vec![10]);
3329 }
3330
3331 #[test]
3332 fn wp6b_fold_async_materialization_does_not_drain_upstream() {
3333 let release = StdArc::new((std::sync::Mutex::new(false), std::sync::Condvar::new()));
3334 let started = StdArc::new(StdAtomicBool::new(false));
3335 let source = {
3336 let release = StdArc::clone(&release);
3337 let started = StdArc::clone(&started);
3338 Source::from_factory(move || {
3339 let release = StdArc::clone(&release);
3340 let started = StdArc::clone(&started);
3341 let mut emitted = false;
3342 Box::new(std::iter::from_fn(move || {
3343 if emitted {
3344 return None;
3345 }
3346 emitted = true;
3347 started.store(true, StdOrdering::SeqCst);
3348 let (released, available) = &*release;
3349 let mut released = released.lock().unwrap();
3350 while !*released {
3351 released = available.wait(released).unwrap();
3352 }
3353 Some(Ok(1))
3354 }))
3355 })
3356 };
3357
3358 let (materialized_tx, materialized_rx) = mpsc::channel();
3359 let join = thread::spawn(move || {
3360 let queue = source
3361 .fold_async(0, |acc, item| async move { Ok(acc + item) })
3362 .run_with(Sink::queue())
3363 .unwrap();
3364 materialized_tx.send(queue).unwrap();
3365 });
3366
3367 let queue = match materialized_rx.recv_timeout(StdDuration::from_secs(1)) {
3368 Ok(queue) => queue,
3369 Err(error) => {
3370 let (released, available) = &*release;
3371 *released.lock().unwrap() = true;
3372 available.notify_all();
3373 let _ = join.join();
3374 panic!("fold_async materialization did not return before first pull: {error}");
3375 }
3376 };
3377 let (released, _) = &*release;
3378 assert!(
3379 !*released.lock().unwrap(),
3380 "test source was released before materialization returned"
3381 );
3382
3383 let (released, available) = &*release;
3384 *released.lock().unwrap() = true;
3385 available.notify_all();
3386 assert_eq!(queue.pull().unwrap(), Some(1));
3387 assert_eq!(queue.pull().unwrap(), None);
3388 join.join().unwrap();
3389 assert!(started.load(StdOrdering::SeqCst));
3390 }
3391
3392 #[test]
3393 fn wp6b_lazy_sink_and_flow_wait_for_first_element() {
3394 let lazy_sink_created = StdArc::new(StdAtomicUsize::new(0));
3395 let lazy_sink_created_for_factory = StdArc::clone(&lazy_sink_created);
3396 let empty_sink = Source::<i32>::empty()
3397 .run_with(Sink::lazy_sink(move || {
3398 lazy_sink_created_for_factory.fetch_add(1, StdOrdering::SeqCst);
3399 Sink::ignore()
3400 }))
3401 .unwrap();
3402 assert!(matches!(empty_sink.wait(), Err(StreamError::Failed(_))));
3403 assert_eq!(lazy_sink_created.load(StdOrdering::SeqCst), 0);
3404
3405 let foreach_sum = StdArc::new(StdAtomicUsize::new(0));
3406 let foreach_sum_for_sink = StdArc::clone(&foreach_sum);
3407 Source::from_iter([1_usize, 2, 3])
3408 .run_with(Sink::foreach_async(2, move |item| {
3409 let foreach_sum = StdArc::clone(&foreach_sum_for_sink);
3410 async move {
3411 foreach_sum.fetch_add(item, StdOrdering::SeqCst);
3412 Ok(())
3413 }
3414 }))
3415 .unwrap()
3416 .wait()
3417 .unwrap();
3418 assert_eq!(foreach_sum.load(StdOrdering::SeqCst), 6);
3419
3420 let lazy_flow_created = StdArc::new(StdAtomicUsize::new(0));
3421 let lazy_flow_created_for_factory = StdArc::clone(&lazy_flow_created);
3422 let lazy_flow = Flow::<i32, i32>::lazy_flow(move || {
3423 lazy_flow_created_for_factory.fetch_add(1, StdOrdering::SeqCst);
3424 Flow::identity()
3425 .map(|item: i32| item + 10)
3426 .map_materialized_value(|_| 123)
3427 });
3428 let mat = (lazy_flow.materialize)().unwrap();
3429 let mut stream = match lazy_flow.transform {
3430 flow::FlowTransform::Runtime(transform) => {
3431 transform(Box::new([Ok(1), Ok(2)].into_iter()), &Materializer::new()).unwrap()
3432 }
3433 flow::FlowTransform::Pure(_) => panic!("lazy flow must be runtime-backed"),
3434 };
3435 assert_eq!(lazy_flow_created.load(StdOrdering::SeqCst), 0);
3436 assert_eq!(stream.next().unwrap().unwrap(), 11);
3437 assert_eq!(mat.wait().unwrap(), 123);
3438 assert_eq!(lazy_flow_created.load(StdOrdering::SeqCst), 1);
3439 assert_eq!(stream.next().unwrap().unwrap(), 12);
3440
3441 let future_flow = Source::from_iter([1, 2])
3442 .via_mat(
3443 Flow::future_flow(|| async {
3444 Ok(Flow::identity()
3445 .map(|item: i32| item * 2)
3446 .map_materialized_value(|_| 77))
3447 }),
3448 Keep::right,
3449 )
3450 .to_mat(Sink::collect(), Keep::both)
3451 .run()
3452 .unwrap();
3453 assert_eq!(future_flow.0.wait().unwrap(), 77);
3454 assert_eq!(future_flow.1.wait().unwrap(), vec![2, 4]);
3455 }
3456
3457 #[test]
3458 fn wp6b_lazy_flow_double_use_in_one_chain_pairs_instances_in_order() {
3459 for round in 0..50 {
3464 let counter = StdArc::new(StdAtomicUsize::new(1));
3465 let factory_counter = StdArc::clone(&counter);
3466 let lazy: Flow<usize, usize, _> = Flow::lazy_flow(move || {
3467 let id = factory_counter.fetch_add(1, StdOrdering::SeqCst);
3468 Flow::identity()
3469 .map(move |x: usize| x * 100 + id)
3470 .map_materialized_value(move |_| id)
3471 });
3472 let lazy_again = lazy.clone();
3473
3474 let ((first_mat, second_mat), out) = Source::from_iter([0usize])
3475 .via_mat(lazy, Keep::right)
3476 .via_mat(lazy_again, Keep::both)
3477 .to_mat(Sink::collect(), Keep::both)
3478 .run()
3479 .unwrap();
3480
3481 let first_id = first_mat.wait().unwrap();
3482 let second_id = second_mat.wait().unwrap();
3483 let element = out.wait().unwrap()[0];
3484 assert_eq!(
3485 element,
3486 first_id * 100 + second_id,
3487 "round {round}: mats ({first_id},{second_id}) cross-wired with transform order"
3488 );
3489 assert_ne!(
3490 first_id, second_id,
3491 "round {round}: same factory instance paired twice"
3492 );
3493 }
3494 }
3495
3496 #[test]
3497 fn wp6b_lazy_flow_clones_materialize_concurrently_without_cross_wiring() {
3498 for _ in 0..20 {
3499 let next_id = StdArc::new(StdAtomicUsize::new(0));
3500 let next_id_for_factory = StdArc::clone(&next_id);
3501 let flow = Flow::<i32, i32>::lazy_flow(move || {
3502 let id = next_id_for_factory.fetch_add(1, StdOrdering::SeqCst) + 1;
3503 Flow::identity()
3504 .map(move |item: i32| item + (id as i32 * 100))
3505 .map_materialized_value(move |_| id)
3506 });
3507 let barrier = StdArc::new(std::sync::Barrier::new(3));
3508
3509 let spawn_materialization = |input: i32| {
3510 let flow = flow.clone();
3511 let barrier = StdArc::clone(&barrier);
3512 thread::spawn(move || {
3513 barrier.wait();
3514 let (mat, values) = Source::single(input)
3515 .via_mat(flow, Keep::right)
3516 .to_mat(Sink::collect(), Keep::both)
3517 .run()
3518 .unwrap();
3519 (input, mat.wait().unwrap(), values.wait().unwrap())
3520 })
3521 };
3522
3523 let first = spawn_materialization(1);
3524 let second = spawn_materialization(2);
3525 barrier.wait();
3526
3527 for result in [first.join().unwrap(), second.join().unwrap()] {
3528 let (input, mat_id, values) = result;
3529 assert_eq!(values, vec![input + (mat_id as i32 * 100)]);
3530 }
3531 assert_eq!(next_id.load(StdOrdering::SeqCst), 2);
3532 }
3533 }
3534
3535 #[test]
3536 fn wp6b_map_with_resource_emits_close_item_before_terminal_error() {
3537 let queue = Source::from_factory(|| {
3538 Box::new(vec![Ok(1), Err(StreamError::Failed("upstream".into()))].into_iter())
3539 })
3540 .map_with_resource(
3541 || Ok(()),
3542 |_resource, item| Ok(item + 10),
3543 |_resource| Ok(Some(99)),
3544 )
3545 .run_with(Sink::queue())
3546 .unwrap();
3547
3548 assert_eq!(queue.pull().unwrap(), Some(11));
3549 assert_eq!(queue.pull().unwrap(), Some(99));
3550 assert_eq!(queue.pull(), Err(StreamError::Failed("upstream".into())));
3551
3552 let failed: StreamResult<Vec<i32>> = Source::single(1)
3553 .map_with_resource(
3554 || Ok(()),
3555 |_resource, _item| -> StreamResult<i32> { Err(StreamError::Failed("map".into())) },
3556 |_resource| -> StreamResult<Option<i32>> {
3557 Err(StreamError::Failed("close".into()))
3558 },
3559 )
3560 .run_collect();
3561 assert_eq!(failed, Err(StreamError::Failed("map".into())));
3562 }
3563
3564 #[test]
3565 fn stateful_and_terminal_source_operators_work() {
3566 let stateful = Source::from_iter([1, 2, 3])
3567 .stateful_map(0, |sum, item| {
3568 *sum += item;
3569 *sum
3570 })
3571 .run_collect()
3572 .unwrap();
3573 assert_eq!(stateful, vec![1, 3, 6]);
3574
3575 let concat = Source::from_iter([1, 2, 3])
3576 .stateful_map_concat(0, |sum, item| {
3577 *sum += item;
3578 [item, *sum]
3579 })
3580 .run_collect()
3581 .unwrap();
3582 assert_eq!(concat, vec![1, 1, 2, 3, 3, 6]);
3583
3584 assert_eq!(
3585 Source::from_iter([1, 2, 3])
3586 .fold(10, |acc, item| acc + item)
3587 .run_collect()
3588 .unwrap(),
3589 vec![16]
3590 );
3591 assert_eq!(
3592 Source::from_iter([1, 2, 3])
3593 .reduce(|acc, item| acc + item)
3594 .run_collect()
3595 .unwrap(),
3596 vec![6]
3597 );
3598 }
3599
3600 #[test]
3601 fn concat_and_sliding_emit_before_unbounded_upstream_finishes() {
3602 let concat = Source::single(())
3603 .map_concat(|_| 0_u64..)
3604 .take(1)
3605 .run_collect()
3606 .unwrap();
3607 assert_eq!(concat, vec![0]);
3608
3609 let sliding = Source::repeat(1_u64)
3610 .sliding(2, 1)
3611 .take(1)
3612 .run_collect()
3613 .unwrap();
3614 assert_eq!(sliding, vec![vec![1, 1]]);
3615 }
3616
3617 #[test]
3618 fn fan_in_source_operators_follow_ordering_rules() {
3619 assert_eq!(
3620 Source::from_iter([1, 2])
3621 .concat(Source::from_iter([3, 4]))
3622 .run_collect()
3623 .unwrap(),
3624 vec![1, 2, 3, 4]
3625 );
3626 assert_eq!(
3627 Source::from_iter([3, 4])
3628 .prepend(Source::from_iter([1, 2]))
3629 .run_collect()
3630 .unwrap(),
3631 vec![1, 2, 3, 4]
3632 );
3633 assert_eq!(
3634 Source::empty()
3635 .or_else(Source::from_iter([10, 20]))
3636 .run_collect()
3637 .unwrap(),
3638 vec![10, 20]
3639 );
3640 assert_eq!(
3641 Source::from_iter([1, 2])
3642 .or_else(Source::from_iter([10, 20]))
3643 .run_collect()
3644 .unwrap(),
3645 vec![1, 2]
3646 );
3647 assert_eq!(
3648 Source::from_iter([1, 2, 3])
3649 .interleave(Source::from_iter([10, 11, 12]), 2)
3650 .run_collect()
3651 .unwrap(),
3652 vec![1, 2, 10, 11, 3, 12]
3653 );
3654 }
3655
3656 #[test]
3657 fn fan_in_flow_operators_compose_with_primary_stream() {
3658 let concat = Source::from_iter([1, 2])
3659 .via(Flow::identity().concat(Source::from_iter([3, 4])))
3660 .run_collect()
3661 .unwrap();
3662 assert_eq!(concat, vec![1, 2, 3, 4]);
3663
3664 let prepend = Source::from_iter([3, 4])
3665 .via(Flow::identity().prepend(Source::from_iter([1, 2])))
3666 .run_collect()
3667 .unwrap();
3668 assert_eq!(prepend, vec![1, 2, 3, 4]);
3669
3670 let interleave = Source::from_iter([1, 2, 3])
3671 .via(Flow::identity().interleave(Source::from_iter([10, 11, 12]), 1))
3672 .run_collect()
3673 .unwrap();
3674 assert_eq!(interleave, vec![1, 10, 2, 11, 3, 12]);
3675
3676 let merge_sorted = Source::from_iter([1, 4])
3677 .via(Flow::identity().merge_sorted(Source::from_iter([2, 3, 5])))
3678 .run_collect()
3679 .unwrap();
3680 assert_eq!(merge_sorted, vec![1, 2, 3, 4, 5]);
3681
3682 let zip_latest = Source::from_iter([1, 2])
3683 .via(Flow::identity().zip_latest(Source::single(10)))
3684 .run_collect()
3685 .unwrap();
3686 assert_eq!(zip_latest, vec![(1, 10), (2, 10)]);
3687
3688 let zip_latest_with = Source::from_iter([1, 2])
3689 .via(
3690 Flow::identity()
3691 .zip_latest_with(Source::single(10), false, |left, right| left + right),
3692 )
3693 .run_collect()
3694 .unwrap();
3695 assert_eq!(zip_latest_with, vec![11, 12]);
3696 }
3697
3698 #[test]
3699 fn fan_in_operators_propagate_errors_and_eager_close() {
3700 assert!(matches!(
3701 Source::failed(StreamError::Failed("boom".into()))
3702 .or_else(Source::from_iter([1, 2]))
3703 .run_collect(),
3704 Err(StreamError::Failed(_))
3705 ));
3706 assert!(matches!(
3707 Source::from_iter([1, 2])
3708 .prepend(Source::failed(StreamError::Failed("boom".into())))
3709 .run_collect(),
3710 Err(StreamError::Failed(_))
3711 ));
3712 assert_eq!(
3713 Source::from_iter([1, 2])
3714 .interleave_all([Source::empty()], 1, true)
3715 .run_collect()
3716 .unwrap(),
3717 vec![1]
3718 );
3719 }
3720
3721 #[test]
3722 fn interleave_lazy_pulls_only_inputs_needed_for_first_segment() {
3723 use std::sync::{Arc, atomic::AtomicUsize, atomic::Ordering};
3724
3725 let pulls: Arc<[AtomicUsize; 3]> = Arc::new([
3726 AtomicUsize::new(0),
3727 AtomicUsize::new(0),
3728 AtomicUsize::new(0),
3729 ]);
3730
3731 let make_source = |idx: usize| {
3732 let pulls = Arc::clone(&pulls);
3733 Source::from_materialized_factory(move |_| {
3734 let pulls = Arc::clone(&pulls);
3735 let mut emitted = false;
3736 Ok((
3737 Box::new(std::iter::from_fn(move || {
3738 pulls[idx].fetch_add(1, Ordering::SeqCst);
3739 if !emitted && idx == 0 {
3740 emitted = true;
3741 Some(Ok(42))
3742 } else {
3743 None
3744 }
3745 })) as BoxStream<i32>,
3746 NotUsed,
3747 ))
3748 })
3749 };
3750
3751 let result = make_source(0)
3752 .interleave_all([make_source(1), make_source(2)], 1, false)
3753 .run_with(Sink::head());
3754
3755 assert_eq!(wait(result.unwrap()), 42);
3756 assert_eq!(pulls[0].load(Ordering::SeqCst), 1);
3757 assert_eq!(
3758 pulls[1].load(Ordering::SeqCst),
3759 0,
3760 "second input should not be pulled when downstream cancels after first element"
3761 );
3762 assert_eq!(
3763 pulls[2].load(Ordering::SeqCst),
3764 0,
3765 "third input should not be pulled before its turn"
3766 );
3767 }
3768
3769 #[test]
3770 fn interleave_non_eager_drains_remaining_when_one_input_completes() {
3771 assert_eq!(
3772 Source::from_iter([1, 2, 3, 4])
3773 .interleave_all(
3774 [Source::from_iter([10]), Source::from_iter([20, 21, 22])],
3775 1,
3776 false
3777 )
3778 .run_collect()
3779 .unwrap(),
3780 vec![1, 10, 20, 2, 21, 3, 22, 4]
3781 );
3782 }
3783
3784 #[test]
3785 fn remaining_merge_and_zip_family_matches_expected_ordering() {
3786 assert_eq!(
3787 Source::from_iter([1, 4])
3788 .merge_sorted(Source::from_iter([2, 3, 5]))
3789 .run_collect()
3790 .unwrap(),
3791 vec![1, 2, 3, 4, 5]
3792 );
3793
3794 assert_eq!(
3795 Source::from_iter([1, 2])
3796 .merge_latest(Source::single(10), false)
3797 .run_collect()
3798 .unwrap(),
3799 vec![vec![1, 10], vec![2, 10]]
3800 );
3801
3802 assert_eq!(
3803 Source::from_iter([1, 2, 3])
3804 .merge_all([Source::from_iter([10, 11])], false)
3805 .run_collect()
3806 .unwrap(),
3807 vec![1, 10, 2, 11, 3]
3808 );
3809
3810 assert_eq!(
3811 Source::from_iter([1, 2, 3])
3812 .zip_with(Source::from_iter([10, 11, 12]), |left, right| left + right)
3813 .run_collect()
3814 .unwrap(),
3815 vec![11, 13, 15]
3816 );
3817
3818 assert_eq!(
3819 Source::from_iter([1, 2])
3820 .zip_latest(Source::single(10))
3821 .run_collect()
3822 .unwrap(),
3823 vec![(1, 10), (2, 10)]
3824 );
3825
3826 assert_eq!(
3827 Source::from_iter([1, 2, 3])
3828 .zip_latest_with(Source::from_iter([10]), false, |left, right| left + right)
3829 .run_collect()
3830 .unwrap(),
3831 vec![11, 12, 13]
3832 );
3833
3834 assert_eq!(
3835 Source::from_iter([1, 2])
3836 .zip_all(Source::from_iter([10, 11, 12]), -1, -2)
3837 .run_collect()
3838 .unwrap(),
3839 vec![(1, 10), (2, 11), (-1, 12)]
3840 );
3841
3842 assert_eq!(
3843 Source::from_iter([5, 6, 7])
3844 .zip_with_index()
3845 .run_collect()
3846 .unwrap(),
3847 vec![(5, 0), (6, 1), (7, 2)]
3848 );
3849
3850 assert_eq!(
3851 Source::zip_n([Source::from_iter([1, 2]), Source::from_iter([10, 20])])
3852 .run_collect()
3853 .unwrap(),
3854 vec![vec![1, 10], vec![2, 20]]
3855 );
3856
3857 assert_eq!(
3858 Source::zip_with_n(
3859 [
3860 Source::from_iter([1, 2]),
3861 Source::from_iter([10, 20]),
3862 Source::from_iter([100, 200]),
3863 ],
3864 |values| values.into_iter().sum::<i32>(),
3865 )
3866 .run_collect()
3867 .unwrap(),
3868 vec![111, 222]
3869 );
3870
3871 assert_eq!(
3872 Source::merge_prioritized_n(
3873 [
3874 (Source::from_iter([1, 2, 3, 4]), 2),
3875 (Source::from_iter([10, 11]), 1),
3876 ],
3877 false,
3878 )
3879 .run_collect()
3880 .unwrap(),
3881 vec![1, 2, 10, 3, 4, 11]
3882 );
3883
3884 assert_eq!(
3885 Source::combine(
3886 Source::from_iter([1, 2, 3]),
3887 Source::from_iter([10, 11]),
3888 std::iter::empty::<Source<i32, NotUsed>>(),
3889 SourceCombineStrategy::Merge {
3890 eager_complete: false,
3891 },
3892 )
3893 .run_collect()
3894 .unwrap(),
3895 vec![1, 10, 2, 11, 3]
3896 );
3897
3898 let combined_sink = Sink::combine(
3899 Sink::ignore(),
3900 Sink::ignore(),
3901 std::iter::empty::<Sink<i32, NotUsed>>(),
3902 SinkCombineStrategy::Broadcast,
3903 );
3904 assert_eq!(
3905 Source::from_iter([1, 2, 3])
3906 .run_with(combined_sink)
3907 .unwrap(),
3908 NotUsed
3909 );
3910 }
3911
3912 #[test]
3913 fn sink_combine_broadcast_delivers_every_element_to_every_child() {
3914 let first_count = StdArc::new(StdAtomicUsize::new(0));
3918 let second_count = StdArc::new(StdAtomicUsize::new(0));
3919 let first_counter = StdArc::clone(&first_count);
3920 let second_counter = StdArc::clone(&second_count);
3921 let combined = Sink::combine(
3922 Sink::foreach(move |_: i32| {
3923 first_counter.fetch_add(1, StdOrdering::SeqCst);
3924 }),
3925 Sink::foreach(move |_: i32| {
3926 second_counter.fetch_add(1, StdOrdering::SeqCst);
3927 }),
3928 std::iter::empty::<Sink<i32, NotUsed>>(),
3929 SinkCombineStrategy::Broadcast,
3930 );
3931 assert_eq!(
3932 Source::from_iter(0..100).run_with(combined).unwrap(),
3933 NotUsed
3934 );
3935 assert!(wait_until(StdDuration::from_secs(1), || {
3940 first_count.load(StdOrdering::SeqCst) == 100
3941 && second_count.load(StdOrdering::SeqCst) == 100
3942 }));
3943 }
3944
3945 #[test]
3946 fn zip_latest_completes_when_one_side_finishes_without_emitting() {
3947 assert_eq!(
3951 Source::from_iter(std::iter::empty::<i32>())
3952 .zip_latest_with(Source::repeat(10), false, |left, right| left + right)
3953 .run_collect()
3954 .unwrap(),
3955 Vec::<i32>::new()
3956 );
3957 assert_eq!(
3958 Source::repeat(10)
3959 .zip_latest_with(
3960 Source::from_iter(std::iter::empty::<i32>()),
3961 false,
3962 |left, right| left + right,
3963 )
3964 .run_collect()
3965 .unwrap(),
3966 Vec::<i32>::new()
3967 );
3968 }
3969
3970 #[test]
3971 fn zip_family_completion_boundaries_match_expected_results() {
3972 assert_eq!(
3973 Source::from_iter([1, 2, 3])
3974 .zip_with(Source::from_iter([10]), |left, right| left + right)
3975 .run_collect()
3976 .unwrap(),
3977 vec![11]
3978 );
3979
3980 assert_eq!(
3981 Source::from_iter([1, 2, 3])
3982 .zip_latest_with(Source::from_iter([10]), true, |left, right| left + right)
3983 .run_collect()
3984 .unwrap(),
3985 vec![11, 12]
3986 );
3987
3988 assert_eq!(
3989 Source::zip_n([
3990 Source::from_iter([1, 2, 3]),
3991 Source::from_iter([10]),
3992 Source::from_iter([100, 200, 300]),
3993 ])
3994 .run_collect()
3995 .unwrap(),
3996 vec![vec![1, 10, 100]]
3997 );
3998 }
3999
4000 #[test]
4001 fn combine_strategies_follow_merge_concat_and_priority_rules() {
4002 assert_eq!(
4003 Source::combine(
4004 Source::from_iter([1, 2]),
4005 Source::from_iter([10, 11]),
4006 [Source::from_iter([100])],
4007 SourceCombineStrategy::Concat,
4008 )
4009 .run_collect()
4010 .unwrap(),
4011 vec![1, 2, 10, 11, 100]
4012 );
4013
4014 assert_eq!(
4015 Source::combine(
4016 Source::from_iter([1, 2, 3, 4]),
4017 Source::from_iter([10, 11]),
4018 std::iter::empty::<Source<i32, NotUsed>>(),
4019 SourceCombineStrategy::Prioritized {
4020 priorities: vec![2, 1],
4021 eager_complete: false,
4022 },
4023 )
4024 .run_collect()
4025 .unwrap(),
4026 vec![1, 2, 10, 3, 4, 11]
4027 );
4028 }
4029
4030 #[test]
4031 fn concat_lazy_defers_follow_on_source_until_needed() {
4032 let source_counter = StdArc::new(StdAtomicUsize::new(0));
4033 let source_counter_clone = StdArc::clone(&source_counter);
4034 let lazy_source = Source::from_materialized_factory(move |_| {
4035 source_counter_clone.fetch_add(1, StdOrdering::SeqCst);
4036 Ok((Box::new(std::iter::once(Ok(99))), NotUsed))
4037 });
4038 let source_head = Source::single(1)
4039 .concat_lazy(lazy_source)
4040 .run_with(Sink::head());
4041 assert_eq!(wait(source_head.unwrap()), 1);
4042 assert_eq!(source_counter.load(StdOrdering::SeqCst), 0);
4043
4044 let flow_counter = StdArc::new(StdAtomicUsize::new(0));
4045 let flow_counter_clone = StdArc::clone(&flow_counter);
4046 let lazy_flow_source = Source::from_materialized_factory(move |_| {
4047 flow_counter_clone.fetch_add(1, StdOrdering::SeqCst);
4048 Ok((Box::new(std::iter::once(Ok(99))), NotUsed))
4049 });
4050 let flow_head = Source::single(1)
4051 .via(Flow::identity().concat_lazy(lazy_flow_source))
4052 .run_with(Sink::head());
4053 assert_eq!(wait(flow_head.unwrap()), 1);
4054 assert_eq!(flow_counter.load(StdOrdering::SeqCst), 0);
4055 }
4056
4057 #[test]
4058 fn also_to_completes_when_side_sink_cancels() {
4059 assert_eq!(
4060 Source::from_iter([1, 2, 3])
4061 .also_to(Sink::cancelled())
4062 .run_collect()
4063 .unwrap(),
4064 Vec::<i32>::new()
4065 );
4066 assert_eq!(
4067 Source::from_iter([1, 2, 3])
4068 .also_to_all([Sink::cancelled(), Sink::cancelled()])
4069 .run_collect()
4070 .unwrap(),
4071 Vec::<i32>::new()
4072 );
4073 }
4074
4075 #[test]
4076 fn also_to_completes_gracefully_when_side_sink_disconnects() {
4077 let result = Source::from_iter(0..100)
4078 .also_to(Sink::head())
4079 .run_collect()
4080 .unwrap();
4081 assert!(!result.is_empty(), "main should emit at least one element");
4082 assert!(
4083 result.len() < 100,
4084 "main should complete early when side disconnects"
4085 );
4086 }
4087
4088 #[test]
4089 fn also_to_propagates_original_error_when_side_is_disconnected() {
4090 let err = StreamError::Failed("distinctive-boom".into());
4091 assert!(matches!(
4092 Source::<i32>::failed(err.clone())
4093 .also_to(Sink::cancelled())
4094 .run_collect(),
4095 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4096 ));
4097 assert!(matches!(
4098 Source::<i32>::failed(err.clone())
4099 .also_to_all([Sink::cancelled()])
4100 .run_collect(),
4101 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4102 ));
4103 assert!(matches!(
4104 Source::<i32>::failed(err)
4105 .divert_to(Sink::cancelled(), |_: &i32| true)
4106 .run_collect(),
4107 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4108 ));
4109 }
4110
4111 #[test]
4112 fn divert_to_routes_matching_elements_to_side_sink() {
4113 let diverted = Source::from_iter([1, 2, 3, 4])
4114 .divert_to(Sink::ignore(), |item| item % 2 == 0)
4115 .run_collect()
4116 .unwrap();
4117 assert_eq!(diverted, vec![1, 3]);
4118 }
4119
4120 #[test]
4121 fn wire_tap_drops_when_side_sink_backpressures() {
4122 let tapped = Source::from_iter([1, 2, 3])
4123 .wire_tap(Sink::head())
4124 .run_collect()
4125 .unwrap();
4126 assert_eq!(tapped, vec![1, 2, 3]);
4127
4128 let tapped_via_flow = Source::from_iter([1, 2, 3])
4129 .via(Flow::identity().wire_tap(Sink::head()))
4130 .run_collect()
4131 .unwrap();
4132 assert_eq!(tapped_via_flow, vec![1, 2, 3]);
4133 }
4134
4135 #[test]
4136 fn async_mapping_variants_complete() {
4137 let ordered = Source::from_iter(0..4)
4138 .map_async(2, |item| async move { Ok(item * 2) })
4139 .run_collect()
4140 .unwrap();
4141 assert_eq!(ordered, vec![0, 2, 4, 6]);
4142
4143 let unordered = Source::from_iter(0..4)
4144 .map_async_unordered(2, |item| async move { Ok(item * 2) })
4145 .run_collect()
4146 .unwrap();
4147 assert_eq!(unordered, vec![0, 2, 4, 6]);
4148
4149 let partitioned = Source::from_iter(0..4)
4150 .map_async_partitioned(4, 1, |item| item % 2, |item| async move { Ok(item + 1) })
4151 .run_collect()
4152 .unwrap();
4153 assert_eq!(partitioned, vec![1, 2, 3, 4]);
4154 }
4155
4156 #[test]
4157 fn map_async_ordered_bounds_pulls_behind_stuck_head() {
4158 let pulls = StdArc::new(StdAtomicUsize::new(0));
4159 let pulls_for_source = StdArc::clone(&pulls);
4160 let probe = Source::from_fn_iter(move || {
4161 let pulls = StdArc::clone(&pulls_for_source);
4162 std::iter::from_fn(move || {
4163 let next = pulls.fetch_add(1, StdOrdering::SeqCst);
4164 Some(next)
4165 })
4166 })
4167 .map_async(2, |item| async move {
4168 if item == 0 {
4169 tokio::time::sleep(StdDuration::from_millis(300)).await;
4170 }
4171 Ok(item)
4172 })
4173 .run_with(TestSink::probe())
4174 .unwrap();
4175
4176 probe.request(16);
4177 thread::sleep(StdDuration::from_millis(100));
4178 assert!(
4179 pulls.load(StdOrdering::SeqCst) <= 3,
4180 "pulled {} elements with parallelism=2 behind a stuck ordered head",
4181 pulls.load(StdOrdering::SeqCst)
4182 );
4183 }
4184
4185 #[test]
4186 fn async_mapping_parks_until_woken_future_completes() {
4187 struct WakeOnceFuture {
4188 value: Option<u64>,
4189 ready: StdArc<StdAtomicBool>,
4190 polls: StdArc<StdAtomicUsize>,
4191 poll_tx: mpsc::Sender<usize>,
4192 latest_waker: StdArc<Mutex<Option<std::task::Waker>>>,
4193 }
4194
4195 impl std::future::Future for WakeOnceFuture {
4196 type Output = StreamResult<u64>;
4197
4198 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
4199 let this = self.as_mut().get_mut();
4200 let poll = this.polls.fetch_add(1, StdOrdering::SeqCst) + 1;
4201 if this.ready.load(StdOrdering::SeqCst) {
4202 let _ = this.poll_tx.send(poll);
4203 return Poll::Ready(Ok(this.value.take().unwrap()));
4204 }
4205
4206 *this.latest_waker.lock().expect("latest wake slot mutex") =
4207 Some(cx.waker().clone());
4208 let _ = this.poll_tx.send(poll);
4209 Poll::Pending
4210 }
4211 }
4212
4213 let ready = StdArc::new(StdAtomicBool::new(false));
4214 let polls = StdArc::new(StdAtomicUsize::new(0));
4215 let latest_waker = StdArc::new(Mutex::new(None));
4216 let (poll_tx, poll_rx) = mpsc::channel();
4217 let (result_tx, result_rx) = mpsc::channel();
4218 let ready_for_stage = StdArc::clone(&ready);
4219 let polls_for_stage = StdArc::clone(&polls);
4220 let poll_tx_for_stage = poll_tx.clone();
4221 let latest_waker_for_stage = StdArc::clone(&latest_waker);
4222 let worker = thread::spawn(move || {
4223 let result = Source::single(41)
4224 .map_async(1, move |item| WakeOnceFuture {
4225 value: Some(item + 1),
4226 ready: StdArc::clone(&ready_for_stage),
4227 polls: StdArc::clone(&polls_for_stage),
4228 poll_tx: poll_tx_for_stage.clone(),
4229 latest_waker: StdArc::clone(&latest_waker_for_stage),
4230 })
4231 .run_collect();
4232 let _ = result_tx.send(result);
4233 });
4234
4235 assert_eq!(poll_rx.recv_timeout(StdDuration::from_secs(10)).unwrap(), 1);
4236 assert!(
4237 matches!(
4238 result_rx.try_recv(),
4239 Err(std::sync::mpsc::TryRecvError::Empty)
4240 ),
4241 "stream completed before the future was woken"
4242 );
4243 assert_eq!(poll_rx.recv_timeout(StdDuration::from_secs(10)).unwrap(), 2);
4244 assert!(
4245 matches!(
4246 result_rx.try_recv(),
4247 Err(std::sync::mpsc::TryRecvError::Empty)
4248 ),
4249 "stream completed before the future was marked ready"
4250 );
4251
4252 ready.store(true, StdOrdering::SeqCst);
4253 latest_waker
4254 .lock()
4255 .expect("latest wake slot mutex")
4256 .take()
4257 .expect("pending future should have registered a waker")
4258 .wake();
4259
4260 let values = result_rx
4261 .recv_timeout(StdDuration::from_secs(10))
4262 .expect("stream should complete after the pending future is woken")
4263 .unwrap();
4264 worker.join().expect("stream worker should not panic");
4265 assert_eq!(values, vec![42]);
4266 assert_eq!(
4267 polls.load(StdOrdering::SeqCst),
4268 3,
4269 "pending future should be polled once inline, once to register the waker, and once after wake"
4270 );
4271 }
4272
4273 #[test]
4274 fn async_mapping_emits_before_unbounded_upstream_finishes() {
4275 let ordered = Source::repeat(1)
4276 .map_async(2, |item| async move { Ok(item + 1) })
4277 .take(1)
4278 .run_collect()
4279 .unwrap();
4280 assert_eq!(ordered, vec![2]);
4281
4282 let unordered = Source::repeat(1)
4283 .map_async_unordered(2, |item| async move { Ok(item + 1) })
4284 .take(1)
4285 .run_collect()
4286 .unwrap();
4287 assert_eq!(unordered, vec![2]);
4288
4289 let partitioned = Source::repeat(1)
4290 .map_async_partitioned(2, 1, |_| 0_u8, |item| async move { Ok(item + 1) })
4291 .take(1)
4292 .run_collect()
4293 .unwrap();
4294 assert_eq!(partitioned, vec![2]);
4295 }
4296
4297 #[test]
4298 fn partitioned_async_mapping_limits_same_key_concurrency() {
4299 let active = StdArc::new(StdAtomicUsize::new(0));
4300 let max_active = StdArc::new(StdAtomicUsize::new(0));
4301 let active_for_stage = StdArc::clone(&active);
4302 let max_for_stage = StdArc::clone(&max_active);
4303
4304 let values = Source::from_iter(0..6)
4305 .map_async_partitioned(
4306 4,
4307 1,
4308 |_| 0_u8,
4309 move |item| {
4310 let active = StdArc::clone(&active_for_stage);
4311 let max_active = StdArc::clone(&max_for_stage);
4312 let current = active.fetch_add(1, StdOrdering::SeqCst) + 1;
4313 max_active.fetch_max(current, StdOrdering::SeqCst);
4314 async move {
4315 thread::sleep(Duration::from_millis(1));
4316 active.fetch_sub(1, StdOrdering::SeqCst);
4317 Ok(item)
4318 }
4319 },
4320 )
4321 .run_collect()
4322 .unwrap();
4323
4324 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
4325 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
4326 }
4327
4328 #[test]
4329 fn partitioned_async_mapping_scans_past_blocked_pending_key() {
4330 let active = StdArc::new(StdAtomicUsize::new(0));
4331 let max_active = StdArc::new(StdAtomicUsize::new(0));
4332 let active_for_stage = StdArc::clone(&active);
4333 let max_for_stage = StdArc::clone(&max_active);
4334 let (release_tx, release_rx) = oneshot::channel::<()>();
4335 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
4336 let release_rx_for_stage = StdArc::clone(&release_rx);
4337 let max_for_release = StdArc::clone(&max_active);
4338
4339 let releaser = thread::spawn(move || {
4340 let deadline = Instant::now() + StdDuration::from_secs(1);
4341 while max_for_release.load(StdOrdering::SeqCst) < 2 && Instant::now() < deadline {
4342 thread::yield_now();
4343 }
4344 let _ = release_tx.send(());
4345 });
4346
4347 let values = Source::from_iter([0, 2, 1])
4348 .map_async_partitioned(
4349 2,
4350 1,
4351 |item| item % 2,
4352 move |item| {
4353 let active = StdArc::clone(&active_for_stage);
4354 let max_active = StdArc::clone(&max_for_stage);
4355 let release_rx = StdArc::clone(&release_rx_for_stage);
4356 let current = active.fetch_add(1, StdOrdering::SeqCst) + 1;
4357 max_active.fetch_max(current, StdOrdering::SeqCst);
4358 async move {
4359 if item == 0 {
4360 let receiver = release_rx
4361 .lock()
4362 .expect("release receiver mutex")
4363 .take()
4364 .expect("release receiver present");
4365 let _ = receiver.await;
4366 }
4367 active.fetch_sub(1, StdOrdering::SeqCst);
4368 Ok(item)
4369 }
4370 },
4371 )
4372 .run_collect()
4373 .unwrap();
4374 releaser.join().unwrap();
4375
4376 assert_eq!(values, vec![0, 2, 1]);
4377 assert_eq!(max_active.load(StdOrdering::SeqCst), 2);
4378 }
4379
4380 #[test]
4381 fn partitioned_async_mapping_p1_still_evaluates_partition() {
4382 let partitions = StdArc::new(StdAtomicUsize::new(0));
4383 let partitions_for_stage = StdArc::clone(&partitions);
4384
4385 let values = Source::from_iter(0..8)
4386 .map_async_partitioned(
4387 1,
4388 1,
4389 move |item| {
4390 partitions_for_stage.fetch_add(1, StdOrdering::SeqCst);
4391 item % 2
4392 },
4393 |item| async move { Ok(item + 1) },
4394 )
4395 .run_collect()
4396 .unwrap();
4397
4398 assert_eq!(values, (1..9).collect::<Vec<_>>());
4399 assert_eq!(partitions.load(StdOrdering::SeqCst), 8);
4400 }
4401
4402 #[test]
4403 fn partitioned_async_mapping_handles_many_keys_high_parallelism() {
4404 let active_by_key =
4405 StdArc::new((0..16).map(|_| StdAtomicUsize::new(0)).collect::<Vec<_>>());
4406 let max_by_key = StdArc::new((0..16).map(|_| StdAtomicUsize::new(0)).collect::<Vec<_>>());
4407 let active_for_stage = StdArc::clone(&active_by_key);
4408 let max_for_stage = StdArc::clone(&max_by_key);
4409
4410 let values = Source::from_iter(0..512_usize)
4411 .map_async_partitioned(
4412 32,
4413 1,
4414 |item| item % 16,
4415 move |item| {
4416 let active = StdArc::clone(&active_for_stage);
4417 let max_active = StdArc::clone(&max_for_stage);
4418 let key = item % 16;
4419 let current = active[key].fetch_add(1, StdOrdering::SeqCst) + 1;
4420 max_active[key].fetch_max(current, StdOrdering::SeqCst);
4421 async move {
4422 active[key].fetch_sub(1, StdOrdering::SeqCst);
4423 Ok(item)
4424 }
4425 },
4426 )
4427 .run_collect()
4428 .unwrap();
4429
4430 assert_eq!(values, (0..512).collect::<Vec<_>>());
4431 for max_active in max_by_key.iter() {
4432 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
4433 }
4434 }
4435
4436 #[test]
4437 fn error_operators_map_recover_and_complete() {
4438 let mapped = Source::<i32>::failed(StreamError::Failed("boom".into()))
4439 .map_error(|_| StreamError::Failed("mapped".into()))
4440 .run_collect();
4441 assert_eq!(mapped, Err(StreamError::Failed("mapped".into())));
4442
4443 let recovered = Source::<i32>::failed(StreamError::Failed("boom".into()))
4444 .recover(|error| match error {
4445 StreamError::Failed(_) => Some(42),
4446 _ => None,
4447 })
4448 .run_collect()
4449 .unwrap();
4450 assert_eq!(recovered, vec![42]);
4451
4452 let unrecovered = Source::<i32>::failed(StreamError::Failed("original".into()))
4453 .recover(|_| None)
4454 .run_collect();
4455 assert_eq!(unrecovered, Err(StreamError::Failed("original".into())));
4456
4457 let recovered_with = Source::<i32>::failed(StreamError::Failed("boom".into()))
4458 .recover_with_retries(1, |_| Some(Source::from_iter([1, 2])))
4459 .run_collect()
4460 .unwrap();
4461 assert_eq!(recovered_with, vec![1, 2]);
4462
4463 let declined_recover_with = Source::<i32>::failed(StreamError::Failed("declined".into()))
4464 .recover_with_retries(1, |_| None)
4465 .run_collect();
4466 assert_eq!(
4467 declined_recover_with,
4468 Err(StreamError::Failed("declined".into()))
4469 );
4470
4471 let completed = Source::from_factory(|| {
4472 Box::new(vec![Ok(1), Err(StreamError::Failed("ignored".into())), Ok(2)].into_iter())
4473 })
4474 .on_error_complete()
4475 .run_collect()
4476 .unwrap();
4477 assert_eq!(completed, vec![1]);
4478 }
4479
4480 #[test]
4481 fn sliding_matches_akka_window_semantics() {
4482 assert_eq!(
4484 Source::from_iter(1..=4)
4485 .sliding(3, 1)
4486 .run_collect()
4487 .unwrap(),
4488 vec![vec![1, 2, 3], vec![2, 3, 4]]
4489 );
4490 assert_eq!(
4491 Source::from_iter(1..=4)
4492 .sliding(2, 1)
4493 .run_collect()
4494 .unwrap(),
4495 vec![vec![1, 2], vec![2, 3], vec![3, 4]]
4496 );
4497 assert_eq!(
4499 Source::from_iter(1..=3)
4500 .sliding(3, 1)
4501 .run_collect()
4502 .unwrap(),
4503 vec![vec![1, 2, 3]]
4504 );
4505 assert_eq!(
4507 Source::from_iter(1..=2)
4508 .sliding(3, 1)
4509 .run_collect()
4510 .unwrap(),
4511 vec![vec![1, 2]]
4512 );
4513 assert_eq!(
4515 Source::from_iter(1..=3)
4516 .sliding(1, 1)
4517 .run_collect()
4518 .unwrap(),
4519 vec![vec![1], vec![2], vec![3]]
4520 );
4521 assert_eq!(
4523 Source::from_iter(1..=6)
4524 .sliding(2, 3)
4525 .run_collect()
4526 .unwrap(),
4527 vec![vec![1, 2], vec![4, 5]]
4528 );
4529 assert_eq!(
4531 Source::from_iter(1..=3)
4532 .sliding(2, 4)
4533 .run_collect()
4534 .unwrap(),
4535 vec![vec![1, 2]]
4536 );
4537 }
4538
4539 #[test]
4540 fn recover_with_retries_indefinitely_like_akka() {
4541 let attempts = StdArc::new(StdAtomicUsize::new(0));
4542 let attempts_in_stage = StdArc::clone(&attempts);
4543 let recovered = Source::<i32>::failed(StreamError::Failed("boom".into()))
4546 .recover_with(move |_error| {
4547 if attempts_in_stage.fetch_add(1, StdOrdering::SeqCst) < 5 {
4548 Some(Source::<i32>::failed(StreamError::Failed("again".into())))
4549 } else {
4550 Some(Source::from_iter([42]))
4551 }
4552 })
4553 .run_collect()
4554 .unwrap();
4555 assert_eq!(recovered, vec![42]);
4556 assert_eq!(attempts.load(StdOrdering::SeqCst), 6);
4557 }
4558
4559 #[test]
4560 fn many_concurrent_streams_do_not_starve_the_pool() {
4561 let materializer = Materializer::new();
4570 let busy = 6_usize;
4571
4572 let mut held = Vec::with_capacity(busy);
4573 for _ in 0..busy {
4574 held.push(
4575 Source::single(1_u64)
4576 .run_with_materializer(Sink::never(), &materializer)
4577 .unwrap(),
4578 );
4579 }
4580
4581 for _ in 0..400 {
4582 if materializer.active_streams() >= busy {
4583 break;
4584 }
4585 thread::sleep(Duration::from_millis(5));
4586 }
4587 assert_eq!(materializer.active_streams(), busy);
4588
4589 let sum = Source::from_iter(0_u64..5)
4592 .run_with_materializer(Sink::fold(0_u64, |acc, item| acc + item), &materializer)
4593 .unwrap();
4594 assert_eq!(sum.wait().unwrap(), 10);
4595
4596 materializer.shutdown();
4597 for completion in held {
4598 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
4599 }
4600 }
4601}