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::{ScalarArithmeticOp, ScalarCompareOp},
933 sink::{RunnableGraph, Sink, SinkCombineStrategy},
934 source::{
935 Demand, IntoSource, Keep, MaybeHandle, NotUsed, PushOutlet, Source, SourceCombineStrategy,
936 },
937 time::{DelayOverflowStrategy, ThrottleMode},
938};
939
940#[cfg(test)]
941mod tests {
942 use super::*;
943 use crate::Attributes;
944 use crate::testkit::TestSink;
945 use std::fs;
946 use std::sync::{
947 Arc as StdArc,
948 atomic::{
949 AtomicBool as StdAtomicBool, AtomicUsize as StdAtomicUsize, Ordering as StdOrdering,
950 },
951 mpsc,
952 };
953 use std::time::Duration as StdDuration;
954 use std::time::Instant;
955
956 fn wait<T>(completion: StreamCompletion<T>) -> T {
957 completion.wait().unwrap()
958 }
959
960 fn wait_until(timeout: Duration, mut condition: impl FnMut() -> bool) -> bool {
961 let deadline = Instant::now() + timeout;
962 while Instant::now() < deadline {
963 if condition() {
964 return true;
965 }
966 thread::sleep(Duration::from_millis(2));
967 }
968 condition()
969 }
970
971 fn linux_thread_count(thread_name: &str) -> usize {
972 fs::read_dir("/proc/self/task")
973 .expect("task directory readable")
974 .filter_map(Result::ok)
975 .filter_map(|entry| fs::read_to_string(entry.path().join("comm")).ok())
976 .filter(|name| name.trim() == thread_name)
977 .count()
978 }
979
980 #[test]
981 fn source_run_terminal_shortcuts_match_explicit_sinks() {
982 let explicit_fold: StreamResult<StreamCompletion<u64>> =
983 Source::from_iter(1_u64..=4).run_with(Sink::fold(0_u64, |acc, item| acc + item));
984 let sugared_fold: StreamResult<StreamCompletion<u64>> =
985 Source::from_iter(1_u64..=4).run_fold(0_u64, |acc, item| acc + item);
986 assert_eq!(wait(explicit_fold.unwrap()), 10);
987 assert_eq!(wait(sugared_fold.unwrap()), 10);
988
989 let explicit_reduce: StreamResult<StreamCompletion<u64>> =
990 Source::from_iter(1_u64..=4).run_with(Sink::reduce(|left: u64, right| left + right));
991 let sugared_reduce: StreamResult<StreamCompletion<u64>> =
992 Source::from_iter(1_u64..=4).run_reduce(|left, right| left + right);
993 assert_eq!(wait(explicit_reduce.unwrap()), 10);
994 assert_eq!(wait(sugared_reduce.unwrap()), 10);
995
996 let explicit_empty_reduce = Source::<u64>::empty()
997 .run_with(Sink::reduce(|left: u64, right| left + right))
998 .unwrap()
999 .wait();
1000 let sugared_empty_reduce = Source::<u64>::empty()
1001 .run_reduce(|left, right| left + right)
1002 .unwrap()
1003 .wait();
1004 assert_eq!(sugared_empty_reduce, explicit_empty_reduce);
1005 assert_eq!(sugared_empty_reduce, Err(StreamError::EmptyStream));
1006
1007 let explicit_sum = StdArc::new(StdAtomicUsize::new(0));
1008 let explicit_sum_sink = StdArc::clone(&explicit_sum);
1009 let explicit_foreach: StreamResult<StreamCompletion<NotUsed>> =
1010 Source::from_iter(1_usize..=4).run_with(Sink::foreach(move |item| {
1011 explicit_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1012 }));
1013 assert_eq!(wait(explicit_foreach.unwrap()), NotUsed);
1014
1015 let sugared_sum = StdArc::new(StdAtomicUsize::new(0));
1016 let sugared_sum_sink = StdArc::clone(&sugared_sum);
1017 let sugared_foreach: StreamResult<StreamCompletion<NotUsed>> =
1018 Source::from_iter(1_usize..=4).run_foreach(move |item| {
1019 sugared_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1020 });
1021 assert_eq!(wait(sugared_foreach.unwrap()), NotUsed);
1022
1023 let alias_sum = StdArc::new(StdAtomicUsize::new(0));
1024 let alias_sum_sink = StdArc::clone(&alias_sum);
1025 let alias_for_each: StreamResult<StreamCompletion<NotUsed>> =
1026 Source::from_iter(1_usize..=4).run_for_each(move |item| {
1027 alias_sum_sink.fetch_add(item, StdOrdering::SeqCst);
1028 });
1029 assert_eq!(wait(alias_for_each.unwrap()), NotUsed);
1030
1031 assert_eq!(explicit_sum.load(StdOrdering::SeqCst), 10);
1032 assert_eq!(sugared_sum.load(StdOrdering::SeqCst), 10);
1033 assert_eq!(alias_sum.load(StdOrdering::SeqCst), 10);
1034 }
1035
1036 #[test]
1037 fn source_constructor_sugar_matches_from_iter() {
1038 let expected = Source::from_iter(1_u64..=4).run_collect().unwrap();
1039
1040 assert_eq!(
1041 Source::from(vec![1_u64, 2, 3, 4]).run_collect().unwrap(),
1042 expected
1043 );
1044 assert_eq!(
1045 Source::from([1_u64, 2, 3, 4]).run_collect().unwrap(),
1046 expected
1047 );
1048 assert_eq!((1_u64..=4).into_source().run_collect().unwrap(), expected);
1049 }
1050
1051 #[test]
1052 fn source_constructor_sugar_keeps_existing_inference_paths() {
1053 let from_vec_into: Source<u64> = vec![1, 2, 3].into();
1054 let from_array_into: Source<u64> = [1, 2, 3].into();
1055 let from_iter = Source::from_iter(1_u64..=3);
1056 let from_iterable = Source::from_iterable(1_u64..=3);
1057 let from_iterator: Source<u64> = (1_u64..=3).collect();
1058 let from_range_into_source: Source<u64> = (1_u64..=3).into_source();
1059
1060 let expected = vec![1, 2, 3];
1061 assert_eq!(from_vec_into.run_collect().unwrap(), expected);
1062 assert_eq!(from_array_into.run_collect().unwrap(), expected);
1063 assert_eq!(from_iter.run_collect().unwrap(), expected);
1064 assert_eq!(from_iterable.run_collect().unwrap(), expected);
1065 assert_eq!(from_iterator.run_collect().unwrap(), expected);
1066 assert_eq!(from_range_into_source.run_collect().unwrap(), expected);
1067 }
1068
1069 #[test]
1070 fn source_async_boundary_preserves_results() {
1071 let expected = Source::from_iter(0_u64..128)
1072 .map(|item| item.wrapping_add(1))
1073 .filter(|item| item % 3 != 0)
1074 .map(|item| item * 2)
1075 .run_collect()
1076 .unwrap();
1077
1078 let actual = Source::from_iter(0_u64..128)
1079 .map(|item| item.wrapping_add(1))
1080 .async_boundary()
1081 .filter(|item| item % 3 != 0)
1082 .map(|item| item * 2)
1083 .run_collect()
1084 .unwrap();
1085
1086 assert_eq!(actual, expected);
1087 }
1088
1089 #[test]
1090 fn flow_async_boundary_preserves_results() {
1091 let expected = Source::from_iter(0_u64..128)
1092 .map(|item| item + 1)
1093 .map(|item| item * 3)
1094 .run_collect()
1095 .unwrap();
1096
1097 let flow = Flow::identity()
1098 .map(|item: u64| item + 1)
1099 .r#async()
1100 .map(|item| item * 3);
1101 let actual = Source::from_iter(0_u64..128)
1102 .via(flow)
1103 .run_collect()
1104 .unwrap();
1105
1106 assert_eq!(actual, expected);
1107 }
1108
1109 #[test]
1110 fn linear_async_boundary_matches_graph_async_boundary_shape() {
1111 use crate::{
1112 AsyncBoundary, AsyncBoundaryExecutionConfig, FusedExecutionConfig, GraphDsl,
1113 GraphFlowShape, MapStage,
1114 };
1115
1116 let graph = GraphDsl::try_create(|builder| {
1117 let first = builder.add(MapStage::new(|item: u64| item + 1));
1118 let boundary = builder.add(AsyncBoundary::<u64>::new());
1119 let second = builder.add(MapStage::new(|item: u64| item * 2));
1120
1121 builder.connect(first.outlet(), boundary.inlet())?;
1122 builder.connect(boundary.outlet(), second.inlet())?;
1123
1124 Ok(GraphFlowShape::new(first.inlet(), second.outlet()))
1125 })
1126 .unwrap();
1127
1128 let linear = Source::from_iter(1_u64..=4)
1129 .map(|item| item + 1)
1130 .async_boundary_with_buffer(4)
1131 .map(|item| item * 2)
1132 .run_collect()
1133 .unwrap();
1134 let graph_output = graph.run_with_input(1_u64..=4).unwrap();
1135 let report = graph
1136 .run_async_boundary_count_with_input_report(
1137 1_u64..=4,
1138 AsyncBoundaryExecutionConfig {
1139 fused: FusedExecutionConfig { event_limit: 1024 },
1140 buffer_size: 4,
1141 },
1142 )
1143 .unwrap();
1144
1145 assert_eq!(linear, graph_output);
1146 assert_eq!(report.result, linear.len());
1147 assert_eq!(report.async_boundary_crossings, linear.len());
1148 }
1149
1150 #[test]
1151 fn async_boundary_regions_run_concurrently() {
1152 let (upstream_tx, upstream_rx) = mpsc::channel::<u64>();
1153 let (downstream_blocked_tx, downstream_blocked_rx) = mpsc::channel::<()>();
1154 let (release_tx, release_rx) = mpsc::channel::<()>();
1155 let release_rx = StdArc::new(Mutex::new(release_rx));
1156
1157 let completion = Source::from_iter(0_u64..3)
1158 .map(move |item| {
1159 upstream_tx.send(item).expect("upstream probe receives");
1160 item
1161 })
1162 .async_boundary_with_buffer(1)
1163 .map({
1164 let release_rx = StdArc::clone(&release_rx);
1165 move |item| {
1166 if item == 0 {
1167 downstream_blocked_tx
1168 .send(())
1169 .expect("downstream probe receives");
1170 release_rx
1171 .lock()
1172 .expect("release receiver lock")
1173 .recv_timeout(StdDuration::from_secs(2))
1174 .expect("downstream release arrives");
1175 }
1176 item
1177 }
1178 })
1179 .run_with(Sink::collect())
1180 .unwrap();
1181
1182 assert_eq!(
1183 downstream_blocked_rx.recv_timeout(StdDuration::from_secs(2)),
1184 Ok(())
1185 );
1186 assert_eq!(upstream_rx.recv_timeout(StdDuration::from_secs(2)), Ok(0));
1187 assert_eq!(upstream_rx.recv_timeout(StdDuration::from_secs(2)), Ok(1));
1188
1189 release_tx.send(()).expect("release downstream");
1190 assert_eq!(completion.wait().unwrap(), vec![0, 1, 2]);
1191 }
1192
1193 #[test]
1194 fn async_boundary_backpressures_slow_downstream() {
1195 let (produced_tx, produced_rx) = mpsc::channel::<u64>();
1196 let (release_tx, release_rx) = mpsc::channel::<()>();
1197 let release_rx = StdArc::new(Mutex::new(release_rx));
1198
1199 let completion = Source::from_iter(0_u64..8)
1200 .map(move |item| {
1201 produced_tx.send(item).expect("producer probe receives");
1202 item
1203 })
1204 .async_boundary_with_buffer(1)
1205 .map({
1206 let release_rx = StdArc::clone(&release_rx);
1207 move |item| {
1208 if item == 0 {
1209 release_rx
1210 .lock()
1211 .expect("release receiver lock")
1212 .recv_timeout(StdDuration::from_secs(2))
1213 .expect("downstream release arrives");
1214 }
1215 item
1216 }
1217 })
1218 .run_with(Sink::collect())
1219 .unwrap();
1220
1221 assert_eq!(produced_rx.recv_timeout(StdDuration::from_secs(2)), Ok(0));
1222 assert_eq!(produced_rx.recv_timeout(StdDuration::from_secs(2)), Ok(1));
1223 if let Ok(item) = produced_rx.recv_timeout(StdDuration::from_millis(100)) {
1224 assert_eq!(item, 2);
1225 }
1226 match produced_rx.recv_timeout(StdDuration::from_millis(100)) {
1227 Err(mpsc::RecvTimeoutError::Timeout) => {}
1228 other => panic!("async boundary handoff was not bounded: {other:?}"),
1229 }
1230
1231 release_tx.send(()).expect("release downstream");
1232 assert_eq!(completion.wait().unwrap(), (0_u64..8).collect::<Vec<_>>());
1233 }
1234
1235 #[test]
1236 fn source_blueprints_are_reusable() {
1237 let source = Source::from_iter(0..5).map(|item| item + 1);
1238
1239 assert_eq!(source.clone().run_collect().unwrap(), vec![1, 2, 3, 4, 5]);
1240 assert_eq!(source.run_collect().unwrap(), vec![1, 2, 3, 4, 5]);
1241 }
1242
1243 #[test]
1244 fn source_map_preserves_materialized_value() {
1245 let graph = Source::single(1)
1246 .map_materialized_value(|_| "source")
1247 .map(|item| item + 1)
1248 .to_mat(Sink::head(), Keep::both);
1249
1250 let materialized = graph.run().unwrap();
1251 assert_eq!(materialized.0, "source");
1252 assert_eq!(wait(materialized.1), 2);
1253 }
1254
1255 #[test]
1256 fn source_and_flow_compose() {
1257 let flow = Flow::identity()
1258 .map(|item: i32| item * 2)
1259 .filter(|item| item % 3 == 0);
1260
1261 let result = Source::from_iter(0..8).via(flow).run_collect().unwrap();
1262
1263 assert_eq!(result, vec![0, 6, 12]);
1264 }
1265
1266 #[test]
1267 fn sink_setup_sees_materializer_defaults_and_local_attributes() {
1268 let observed = StdArc::new(Mutex::new(None));
1269 let observed_in_setup = StdArc::clone(&observed);
1270 let sink = Sink::<i32, StreamCompletion<NotUsed>>::setup(move |_materializer, attrs| {
1271 *observed_in_setup.lock().unwrap() = Some((
1272 attrs.name().map(str::to_owned),
1273 attrs.input_buffer_hint(),
1274 attrs.dispatcher_hint().map(str::to_owned),
1275 ));
1276 Sink::ignore()
1277 })
1278 .add_attributes(Attributes::named("sink-inner"))
1279 .add_attributes(Attributes::input_buffer(4, 4))
1280 .add_attributes(Attributes::dispatcher("bench-dispatcher"));
1281
1282 let materializer = Materializer::new().with_attributes(Attributes::named("mat-outer"));
1283 wait(
1284 Source::from_iter([1, 2, 3])
1285 .run_with_materializer(sink, &materializer)
1286 .unwrap(),
1287 );
1288
1289 assert_eq!(
1290 *observed.lock().unwrap(),
1291 Some((
1292 Some("sink-inner".to_owned()),
1293 Some((4, 4)),
1294 Some("bench-dispatcher".to_owned())
1295 ))
1296 );
1297 }
1298
1299 #[test]
1300 fn sink_pre_materialize_feeds_existing_materialization() {
1301 let materializer = Materializer::new();
1302 let (completion, pre) = Sink::<i32, StreamCompletion<Vec<i32>>>::collect()
1303 .pre_materialize(&materializer)
1304 .unwrap();
1305
1306 Source::from_iter([1, 2, 3])
1307 .run_with_materializer(pre, &materializer)
1308 .unwrap();
1309
1310 assert_eq!(wait(completion), vec![1, 2, 3]);
1311 }
1312
1313 #[test]
1314 fn flow_from_sink_and_source_connects_both_sides() {
1315 assert_eq!(
1316 Source::from_iter([1, 2, 3])
1317 .via(Flow::from_sink_and_source(
1318 Sink::foreach(|_item: i32| {}),
1319 Source::from_iter([10, 20, 30]),
1320 ))
1321 .run_collect()
1322 .unwrap(),
1323 vec![10, 20, 30]
1324 );
1325 }
1326
1327 #[test]
1328 fn from_sink_and_source_keeps_sink_running_after_source_side_completes() {
1329 let completed = StdArc::new(StdAtomicBool::new(false));
1330 let on_complete = StdArc::clone(&completed);
1331 let flow = Flow::from_sink_and_source(
1332 Sink::on_complete(move || {
1333 on_complete.store(true, StdOrdering::SeqCst);
1334 }),
1335 Source::single(10),
1336 );
1337
1338 let result = Source::from_iter([1, 2, 3])
1339 .via(flow)
1340 .run_collect()
1341 .unwrap();
1342
1343 assert_eq!(result, vec![10]);
1344 assert!(wait_until(StdDuration::from_secs(1), || {
1345 completed.load(StdOrdering::SeqCst)
1346 }));
1347 }
1348
1349 #[test]
1350 fn from_sink_and_source_coupled_cancels_source_when_sink_finishes_first() {
1351 let cancellable = StdArc::new(Mutex::new(None));
1352 let observed = StdArc::clone(&cancellable);
1353 let flow = Flow::from_sink_and_source_coupled(
1354 Sink::ignore(),
1355 Source::tick(
1356 StdDuration::from_millis(50),
1357 StdDuration::from_millis(50),
1358 10,
1359 )
1360 .map_materialized_value(move |handle| {
1361 *observed.lock().unwrap() = Some(handle.clone());
1362 handle
1363 }),
1364 );
1365
1366 let completion = Source::from_iter(std::iter::empty::<i32>())
1367 .via(flow)
1368 .run_with(Sink::ignore())
1369 .unwrap();
1370 assert!(wait_until(StdDuration::from_secs(1), || {
1371 cancellable
1372 .lock()
1373 .unwrap()
1374 .as_ref()
1375 .is_some_and(Cancellable::is_cancelled)
1376 }));
1377 assert_eq!(wait(completion), NotUsed);
1378 }
1379
1380 #[test]
1381 fn bidi_flow_join_and_atop_compose() {
1382 let codec = BidiFlow::from_flows(
1383 Flow::identity().map(|item: i32| item + 1),
1384 Flow::identity().map(|item: i32| item * 2),
1385 )
1386 .named("codec");
1387 let framing = BidiFlow::from_flows(
1388 Flow::identity().map(|item: i32| item * 3),
1389 Flow::identity().map(|item: i32| item - 4),
1390 );
1391
1392 let joined = codec
1393 .clone()
1394 .join(Flow::identity().map(|item: i32| item - 5));
1395 let stacked = codec.atop(framing).join(Flow::identity());
1396
1397 assert_eq!(
1398 Source::single(10).via(joined).run_collect().unwrap(),
1399 vec![12]
1400 );
1401 assert_eq!(
1402 Source::single(10).via(stacked).run_collect().unwrap(),
1403 vec![58]
1404 );
1405 }
1406
1407 #[test]
1408 fn flow_buffer_then_map_runs_end_to_end() {
1409 let flow = Flow::identity()
1410 .buffer(8, OverflowStrategy::Backpressure)
1411 .map(|item: i32| item + 1);
1412
1413 let result = Source::from_iter(0..4).via(flow).run_collect().unwrap();
1414
1415 assert_eq!(result, vec![1, 2, 3, 4]);
1416 }
1417
1418 #[test]
1419 fn public_flow_combinators_preserve_runtime_transform_after_buffer() {
1420 fn buffered_flow() -> Flow<i32, i32> {
1421 Flow::identity().buffer(8, OverflowStrategy::Backpressure)
1422 }
1423
1424 assert_eq!(
1425 Source::from_iter(0..4)
1426 .via(buffered_flow().filter(|item| *item % 2 == 0))
1427 .run_collect()
1428 .unwrap(),
1429 vec![0, 2]
1430 );
1431 assert_eq!(
1432 Source::from_iter(0..4)
1433 .via(buffered_flow().filter_not(|item| *item % 2 == 0))
1434 .run_collect()
1435 .unwrap(),
1436 vec![1, 3]
1437 );
1438 assert_eq!(
1439 Source::from_iter(0..4)
1440 .via(buffered_flow().filter_map(|item| (item % 2 == 0).then_some(item + 10)))
1441 .run_collect()
1442 .unwrap(),
1443 vec![10, 12]
1444 );
1445 assert_eq!(
1446 Source::from_iter(0..3)
1447 .via(buffered_flow().map_concat(|item| [item, item + 10]))
1448 .run_collect()
1449 .unwrap(),
1450 vec![0, 10, 1, 11, 2, 12]
1451 );
1452 assert_eq!(
1453 Source::from_iter(0..3)
1454 .via(buffered_flow().stateful_map(5, |state, item| {
1455 *state += item;
1456 *state
1457 }))
1458 .run_collect()
1459 .unwrap(),
1460 vec![5, 6, 8]
1461 );
1462 assert_eq!(
1463 Source::from_iter(0..3)
1464 .via(buffered_flow().stateful_map_concat(0, |state, item| {
1465 *state += item;
1466 [*state, item]
1467 }))
1468 .run_collect()
1469 .unwrap(),
1470 vec![0, 0, 1, 1, 3, 2]
1471 );
1472 assert_eq!(
1473 Source::from_iter(0..4)
1474 .via(buffered_flow().map_async(2, |item| async move { Ok(item + 1) }))
1475 .run_collect()
1476 .unwrap(),
1477 vec![1, 2, 3, 4]
1478 );
1479 assert_eq!(
1480 Source::from_iter(0..4)
1481 .via(buffered_flow().map_async_unordered(2, |item| async move { Ok(item + 1) }))
1482 .run_collect()
1483 .unwrap(),
1484 vec![1, 2, 3, 4]
1485 );
1486 assert_eq!(
1487 Source::from_iter(0..4)
1488 .via(buffered_flow().map_async_partitioned(
1489 2,
1490 1,
1491 |item| item % 2,
1492 |item| async move { Ok(item + 1) },
1493 ))
1494 .run_collect()
1495 .unwrap(),
1496 vec![1, 2, 3, 4]
1497 );
1498 assert_eq!(
1499 Source::from_iter(0..5)
1500 .via(buffered_flow().take(3))
1501 .run_collect()
1502 .unwrap(),
1503 vec![0, 1, 2]
1504 );
1505 assert_eq!(
1506 Source::from_iter(0..5)
1507 .via(buffered_flow().drop(2))
1508 .run_collect()
1509 .unwrap(),
1510 vec![2, 3, 4]
1511 );
1512 assert_eq!(
1513 Source::from_iter(0..5)
1514 .via(buffered_flow().take_while(|item| *item < 3))
1515 .run_collect()
1516 .unwrap(),
1517 vec![0, 1, 2]
1518 );
1519 assert_eq!(
1520 Source::from_iter(0..5)
1521 .via(buffered_flow().drop_while(|item| *item < 3))
1522 .run_collect()
1523 .unwrap(),
1524 vec![3, 4]
1525 );
1526 assert_eq!(
1527 Source::from_iter(0..3)
1528 .via(buffered_flow().limit(5))
1529 .run_collect()
1530 .unwrap(),
1531 vec![0, 1, 2]
1532 );
1533 assert_eq!(
1534 Source::from_iter(0..5)
1535 .via(buffered_flow().grouped(2))
1536 .run_collect()
1537 .unwrap(),
1538 vec![vec![0, 1], vec![2, 3], vec![4]]
1539 );
1540 assert_eq!(
1541 Source::from_iter(1..=3)
1542 .via(buffered_flow().scan(0, |acc, item| acc + item))
1543 .run_collect()
1544 .unwrap(),
1545 vec![0, 1, 3, 6]
1546 );
1547 assert_eq!(
1548 Source::from_iter(1..=4)
1549 .via(buffered_flow().sliding(2, 1))
1550 .run_collect()
1551 .unwrap(),
1552 vec![vec![1, 2], vec![2, 3], vec![3, 4]]
1553 );
1554 assert_eq!(
1555 Source::from_iter(1..=4)
1556 .via(buffered_flow().fold(0, |acc, item| acc + item))
1557 .run_collect()
1558 .unwrap(),
1559 vec![10]
1560 );
1561 assert_eq!(
1562 Source::from_iter(1..=4)
1563 .via(buffered_flow().reduce(|acc, item| acc + item))
1564 .run_collect()
1565 .unwrap(),
1566 vec![10]
1567 );
1568 assert_eq!(
1569 Source::from_factory(|| {
1570 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into())), Ok(2)].into_iter())
1571 })
1572 .via(buffered_flow().map_error(|_| StreamError::Failed("mapped".into())))
1573 .run_collect(),
1574 Err(StreamError::Failed("mapped".into()))
1575 );
1576 assert_eq!(
1577 Source::<i32>::failed(StreamError::Failed("boom".into()))
1578 .via(buffered_flow().recover(|_| Some(42)))
1579 .run_collect()
1580 .unwrap(),
1581 vec![42]
1582 );
1583 assert_eq!(
1584 Source::<i32>::failed(StreamError::Failed("boom".into()))
1585 .via(buffered_flow().recover_with(|_| Some(Source::from_iter([7, 8]))))
1586 .run_collect()
1587 .unwrap(),
1588 vec![7, 8]
1589 );
1590 assert_eq!(
1591 Source::<i32>::failed(StreamError::Failed("boom".into()))
1592 .via(buffered_flow().recover_with_retries(1, |_| Some(Source::from_iter([9]))))
1593 .run_collect()
1594 .unwrap(),
1595 vec![9]
1596 );
1597 assert_eq!(
1598 Source::from_factory(|| {
1599 Box::new(vec![Ok(1), Err(StreamError::Failed("ignored".into())), Ok(2)].into_iter())
1600 })
1601 .via(buffered_flow().on_error_complete())
1602 .run_collect()
1603 .unwrap(),
1604 vec![1]
1605 );
1606
1607 let materialized = Source::from_iter([1, 2, 3])
1608 .run_with(
1609 buffered_flow()
1610 .via(Flow::identity().map(|item| item + 1))
1611 .map_materialized_value(|_| "buffered-flow")
1612 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both),
1613 )
1614 .unwrap();
1615 assert_eq!(materialized.0, "buffered-flow");
1616 assert_eq!(wait(materialized.1), 9);
1617
1618 let kept = Source::from_iter([1, 2, 3])
1619 .run_with(
1620 buffered_flow()
1621 .via_mat_with(Flow::identity().map(|item| item + 1), |_, _| "combined")
1622 .to(Sink::fold(0, |acc, item| acc + item)),
1623 )
1624 .unwrap();
1625 assert_eq!(kept, "combined");
1626 }
1627
1628 #[test]
1629 fn runtime_rate_flows_compose_in_flow_form() {
1630 let conflate = Flow::identity()
1631 .conflate(|left: i32, right| left + right)
1632 .map(|item| item + 1);
1633 assert_eq!(
1634 Source::single(4).via(conflate).run_collect().unwrap(),
1635 vec![5]
1636 );
1637
1638 let batch = Flow::identity()
1639 .batch(4, |item: i32| item, |left, right| left + right)
1640 .map(|item| item + 1);
1641 assert_eq!(Source::single(4).via(batch).run_collect().unwrap(), vec![5]);
1642
1643 let expand = Flow::identity()
1644 .expand(std::iter::once::<i32>)
1645 .map(|item| item + 1);
1646 assert_eq!(
1647 Source::from_iter(0..4).via(expand).run_collect().unwrap(),
1648 vec![1, 2, 3, 4]
1649 );
1650
1651 let aggregate = Flow::identity()
1652 .aggregate_with_boundary(
1653 Vec::<i32>::new,
1654 |mut items, item| {
1655 items.push(item);
1656 let ready = !items.is_empty();
1657 (items, ready)
1658 },
1659 |items| items.into_iter().sum::<i32>(),
1660 None,
1661 )
1662 .map(|item| item + 1);
1663 assert_eq!(
1664 Source::from_iter(0..4)
1665 .via(aggregate)
1666 .run_collect()
1667 .unwrap(),
1668 vec![1, 2, 3, 4]
1669 );
1670
1671 let detached = Flow::identity().detach().map(|item: i32| item + 1);
1672 assert_eq!(
1673 Source::from_iter(0..4).via(detached).run_collect().unwrap(),
1674 vec![1, 2, 3, 4]
1675 );
1676 }
1677
1678 #[test]
1679 fn high_use_source_flow_operators_work() {
1680 let result = Source::from_iter(0..8)
1681 .drop(1)
1682 .take(5)
1683 .filter_not(|item| item % 2 == 0)
1684 .map_concat(|item| [item, item + 10])
1685 .grouped(3)
1686 .run_collect()
1687 .unwrap();
1688
1689 assert_eq!(result, vec![vec![1, 11, 3], vec![13, 5, 15]]);
1690 }
1691
1692 #[test]
1693 fn prefix_and_tail_emits_prefix_and_live_tail() {
1694 let mut outer = Source::from_iter(0..5)
1695 .prefix_and_tail(2)
1696 .run_collect()
1697 .unwrap();
1698 assert_eq!(outer.len(), 1);
1699 let (prefix, tail) = outer.pop().unwrap();
1700 assert_eq!(prefix, vec![0, 1]);
1701 assert_eq!(tail.clone().run_collect().unwrap(), vec![2, 3, 4]);
1702 assert_eq!(
1703 tail.run_collect(),
1704 Err(StreamError::Failed(
1705 "substream source cannot be materialized more than once".into()
1706 ))
1707 );
1708 }
1709
1710 #[test]
1711 fn prefix_and_tail_fails_before_prefix_is_ready() {
1712 let result = Source::from_factory(|| {
1713 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into())), Ok(2)].into_iter())
1714 })
1715 .prefix_and_tail(2)
1716 .run_collect();
1717 assert!(matches!(result, Err(StreamError::Failed(message)) if message == "boom"));
1718 }
1719
1720 #[test]
1721 fn prefix_and_tail_tail_propagates_late_upstream_failure() {
1722 let mut outer = Source::from_factory(|| {
1723 Box::new(vec![Ok(1), Ok(2), Err(StreamError::Failed("boom".into())), Ok(3)].into_iter())
1724 })
1725 .prefix_and_tail(2)
1726 .run_collect()
1727 .unwrap();
1728 let (prefix, tail) = outer.pop().unwrap();
1729 assert_eq!(prefix, vec![1, 2]);
1730 assert_eq!(tail.run_collect(), Err(StreamError::Failed("boom".into())));
1731 }
1732
1733 #[test]
1734 fn prefix_and_tail_accepts_non_clone_elements() {
1735 #[derive(Debug, PartialEq, Eq)]
1736 struct NonClone(u8);
1737
1738 let mut outer = Source::from_factory(|| {
1739 Box::new(vec![Ok(NonClone(1)), Ok(NonClone(2)), Ok(NonClone(3))].into_iter())
1740 })
1741 .prefix_and_tail(2)
1742 .run_collect()
1743 .unwrap();
1744 let (prefix, tail) = outer.pop().unwrap();
1745 assert_eq!(prefix, vec![NonClone(1), NonClone(2)]);
1746 assert_eq!(tail.run_collect().unwrap(), vec![NonClone(3)]);
1747 }
1748
1749 #[test]
1750 fn flat_map_prefix_materializes_on_short_upstream_completion() {
1751 let values = Source::from_iter([1, 2])
1752 .flat_map_prefix(3, |prefix| {
1753 let sum = prefix.into_iter().sum::<i32>();
1754 Flow::identity().prepend(Source::single(sum))
1755 })
1756 .run_collect()
1757 .unwrap();
1758 assert_eq!(values, vec![3]);
1759 }
1760
1761 #[test]
1762 fn flat_map_prefix_does_not_materialize_on_early_upstream_failure() {
1763 let invoked = StdArc::new(StdAtomicBool::new(false));
1764 let invoked_for_stage = StdArc::clone(&invoked);
1765 let result = Source::from_factory(|| {
1766 Box::new(vec![Ok(1), Err(StreamError::Failed("boom".into()))].into_iter())
1767 })
1768 .flat_map_prefix(3, move |_prefix| {
1769 invoked_for_stage.store(true, StdOrdering::SeqCst);
1770 Flow::identity()
1771 })
1772 .run_collect();
1773 assert_eq!(result, Err(StreamError::Failed("boom".into())));
1774 assert!(!invoked.load(StdOrdering::SeqCst));
1775 }
1776
1777 #[test]
1778 fn flat_map_concat_flattens_nested_sources_sequentially() {
1779 let values = Source::from_iter([1, 2, 3])
1780 .flat_map_concat(|item| Source::from_iter(0..item))
1781 .run_collect()
1782 .unwrap();
1783 assert_eq!(values, vec![0, 0, 1, 0, 1, 2]);
1784 }
1785
1786 #[test]
1787 fn flat_map_merge_respects_breadth_bound() {
1788 let active = StdArc::new(StdAtomicUsize::new(0));
1789 let max_active = StdArc::new(StdAtomicUsize::new(0));
1790 let active_for_stage = StdArc::clone(&active);
1791 let max_for_stage = StdArc::clone(&max_active);
1792
1793 let mut values = Source::from_iter(0..6)
1794 .flat_map_merge(2, move |item| {
1795 let active = StdArc::clone(&active_for_stage);
1796 let max_active = StdArc::clone(&max_for_stage);
1797 Source::future(move || {
1798 let active = StdArc::clone(&active);
1799 let max_active = StdArc::clone(&max_active);
1800 async move {
1801 let now = active.fetch_add(1, StdOrdering::SeqCst) + 1;
1802 loop {
1803 let seen = max_active.load(StdOrdering::SeqCst);
1804 if now <= seen {
1805 break;
1806 }
1807 if max_active
1808 .compare_exchange(
1809 seen,
1810 now,
1811 StdOrdering::SeqCst,
1812 StdOrdering::SeqCst,
1813 )
1814 .is_ok()
1815 {
1816 break;
1817 }
1818 }
1819 thread::sleep(StdDuration::from_millis(20));
1820 active.fetch_sub(1, StdOrdering::SeqCst);
1821 Ok(item)
1822 }
1823 })
1824 })
1825 .run_collect()
1826 .unwrap();
1827 values.sort_unstable();
1828 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
1829 assert!(max_active.load(StdOrdering::SeqCst) <= 2);
1830 }
1831
1832 #[test]
1833 fn flat_map_merge_propagates_inner_failures() {
1834 let result = Source::from_iter([0, 1, 2])
1835 .flat_map_merge(2, |item| {
1836 if item == 1 {
1837 Source::failed(StreamError::Failed("boom".into()))
1838 } else {
1839 Source::single(item)
1840 }
1841 })
1842 .run_collect();
1843 assert_eq!(result, Err(StreamError::Failed("boom".into())));
1844 }
1845
1846 #[test]
1847 fn flat_map_merge_emits_ready_inner_output_while_upstream_is_blocked() {
1848 let (release_tx, release_rx) = mpsc::channel();
1849 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
1850 let queue = Source::from_factory(move || {
1851 let release_rx = StdArc::clone(&release_rx);
1852 let mut step = 0_u8;
1853 Box::new(std::iter::from_fn(move || {
1854 let item = match step {
1855 0 => Some(Ok(0)),
1856 1 => {
1857 release_rx
1858 .lock()
1859 .unwrap()
1860 .as_ref()
1861 .expect("release receiver available")
1862 .recv_timeout(StdDuration::from_secs(1))
1863 .expect("timed out waiting to release second upstream element");
1864 Some(Ok(1))
1865 }
1866 _ => None,
1867 };
1868 step += 1;
1869 item
1870 }))
1871 })
1872 .flat_map_merge(2, |item| Source::single(item + 10))
1873 .run_with(Sink::queue())
1874 .unwrap();
1875
1876 assert_eq!(queue.pull().unwrap(), Some(10));
1877 release_tx.send(()).unwrap();
1878 assert_eq!(queue.pull().unwrap(), Some(11));
1879 assert!(queue.pull().unwrap().is_none());
1880 }
1881
1882 #[test]
1883 fn group_by_routes_keys_and_drops_closed_keys() {
1884 let outer = Source::from_iter([0, 1, 2, 3, 4])
1885 .group_by(4, |item| item % 2, false)
1886 .run_with(Sink::queue())
1887 .unwrap();
1888
1889 let even = outer.pull().unwrap().unwrap();
1890 let even_completion = even.run_with(Sink::ignore()).unwrap();
1891 let odd = outer.pull().unwrap().unwrap();
1892 drop(even_completion);
1893
1894 assert_eq!(odd.run_collect().unwrap(), vec![1, 3]);
1895 assert!(outer.pull().unwrap().is_none());
1896 }
1897
1898 #[test]
1899 fn group_by_fails_when_distinct_key_limit_is_exceeded() {
1900 let outer = Source::from_iter([0, 1, 2])
1901 .group_by(2, |item| *item, false)
1902 .run_with(Sink::queue())
1903 .unwrap();
1904
1905 let _ = outer.pull().unwrap().unwrap();
1906 let _ = outer.pull().unwrap().unwrap();
1907 assert!(matches!(
1908 outer.pull(),
1909 Err(StreamError::Failed(message)) if message == "group_by reached max_substreams (2)"
1910 ));
1911 }
1912
1913 #[test]
1914 fn group_by_can_recreate_closed_substreams_when_enabled() {
1915 let (release_tx, release_rx) = mpsc::channel();
1916 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
1917 let outer = Source::from_factory(move || {
1918 let release_rx = StdArc::clone(&release_rx);
1919 let mut step = 0_u8;
1920 Box::new(std::iter::from_fn(move || {
1921 let item = match step {
1922 0 => Some(Ok(0)),
1923 1 => Some(Ok(1)),
1924 2 => {
1925 release_rx
1926 .lock()
1927 .unwrap()
1928 .as_ref()
1929 .expect("release receiver available")
1930 .recv_timeout(StdDuration::from_secs(1))
1931 .expect("timed out waiting to release recreated key");
1932 Some(Ok(0))
1933 }
1934 _ => None,
1935 };
1936 step += 1;
1937 item
1938 }))
1939 })
1940 .group_by(4, |item| item % 2, true)
1941 .run_with(Sink::queue())
1942 .unwrap();
1943
1944 let even = outer.pull().unwrap().unwrap();
1945 assert_eq!(wait(even.run_with(Sink::head()).unwrap()), 0);
1946 release_tx.send(()).unwrap();
1947
1948 let odd = outer.pull().unwrap().unwrap();
1949 assert_eq!(odd.run_collect().unwrap(), vec![1]);
1950
1951 let recreated_even = outer.pull().unwrap().unwrap();
1952 assert_eq!(recreated_even.run_collect().unwrap(), vec![0]);
1953 assert!(outer.pull().unwrap().is_none());
1954 }
1955
1956 #[test]
1957 fn group_by_panicking_key_fn_abruptly_terminates_live_substreams() {
1958 let outer = Source::from_iter([0, 1])
1959 .group_by(
1960 4,
1961 |item| {
1962 assert_ne!(*item, 1, "boom");
1963 item % 2
1964 },
1965 false,
1966 )
1967 .run_with(Sink::queue())
1968 .unwrap();
1969
1970 let substream = outer.pull().unwrap().unwrap();
1971 let (result_tx, result_rx) = mpsc::channel();
1972 thread::spawn(move || {
1973 let _ = result_tx.send(substream.run_collect());
1974 });
1975
1976 assert_eq!(
1977 result_rx.recv_timeout(StdDuration::from_secs(1)).unwrap(),
1978 Err(StreamError::AbruptTermination)
1979 );
1980 assert!(matches!(outer.pull(), Err(StreamError::AbruptTermination)));
1981 }
1982
1983 #[test]
1984 fn split_when_starts_new_substream_on_boundary_element() {
1985 let outer = Source::from_iter([1, 2, 0, 3, 0, 4, 5])
1986 .split_when(|item| *item == 0)
1987 .run_with(Sink::queue())
1988 .unwrap();
1989
1990 let first = outer.pull().unwrap().unwrap();
1991 assert_eq!(first.run_collect().unwrap(), vec![1, 2]);
1992 let second = outer.pull().unwrap().unwrap();
1993 assert_eq!(second.run_collect().unwrap(), vec![0, 3]);
1994 let third = outer.pull().unwrap().unwrap();
1995 assert_eq!(third.run_collect().unwrap(), vec![0, 4, 5]);
1996 assert!(outer.pull().unwrap().is_none());
1997 }
1998
1999 #[test]
2000 fn split_after_ends_current_substream_on_boundary_element() {
2001 let outer = Source::from_iter([1, 2, 0, 3, 0, 4, 5])
2002 .split_after(|item| *item == 0)
2003 .run_with(Sink::queue())
2004 .unwrap();
2005
2006 let first = outer.pull().unwrap().unwrap();
2007 assert_eq!(first.run_collect().unwrap(), vec![1, 2, 0]);
2008 let second = outer.pull().unwrap().unwrap();
2009 assert_eq!(second.run_collect().unwrap(), vec![3, 0]);
2010 let third = outer.pull().unwrap().unwrap();
2011 assert_eq!(third.run_collect().unwrap(), vec![4, 5]);
2012 assert!(outer.pull().unwrap().is_none());
2013 }
2014
2015 #[test]
2016 fn split_when_panicking_predicate_abruptly_terminates_live_substreams() {
2017 let outer = Source::from_iter([1, 2])
2018 .split_when(|item| {
2019 assert_ne!(*item, 2, "boom");
2020 false
2021 })
2022 .run_with(Sink::queue())
2023 .unwrap();
2024
2025 let substream = outer.pull().unwrap().unwrap();
2026 let (result_tx, result_rx) = mpsc::channel();
2027 thread::spawn(move || {
2028 let _ = result_tx.send(substream.run_collect());
2029 });
2030
2031 assert_eq!(
2032 result_rx.recv_timeout(StdDuration::from_secs(1)).unwrap(),
2033 Err(StreamError::AbruptTermination)
2034 );
2035 assert!(matches!(outer.pull(), Err(StreamError::AbruptTermination)));
2036 }
2037
2038 #[test]
2039 fn split_when_pre_buffer_segments_match_expected_count() {
2040 let outer = Source::from_iter(0..100)
2041 .split_when(|item| *item != 0 && *item % 10 == 0)
2042 .run_with(Sink::queue())
2043 .unwrap();
2044 let mut segment_count = 0;
2045 while let Some(substream) = outer.pull().unwrap() {
2046 let items: Vec<i32> = substream.run_collect().unwrap();
2047 assert!(!items.is_empty(), "segment should not be empty");
2048 segment_count += 1;
2049 }
2050 assert_eq!(segment_count, 10, "100 elements in segments of 10");
2051 }
2052
2053 #[test]
2054 fn split_after_pre_buffer_segments_match_expected_count() {
2055 let outer = Source::from_iter(0..100)
2056 .split_after(|item| (*item + 1) % 10 == 0)
2057 .run_with(Sink::queue())
2058 .unwrap();
2059 let mut segment_count = 0;
2060 let mut total = 0_i32;
2061 while let Some(substream) = outer.pull().unwrap() {
2062 let items: Vec<i32> = substream.run_collect().unwrap();
2063 assert!(!items.is_empty(), "segment should not be empty");
2064 total += items.len() as i32;
2065 segment_count += 1;
2066 }
2067 assert_eq!(segment_count, 10);
2068 assert_eq!(total, 100);
2069 }
2070
2071 #[test]
2072 fn group_by_single_key_fused_matches_general_path() {
2073 let outer = Source::from_iter(0..1000i64)
2074 .group_by(1, |_| 0u8, false)
2075 .run_with(Sink::queue())
2076 .unwrap();
2077 let substream = outer.pull().unwrap().unwrap();
2078 let items: Vec<i64> = substream.run_collect().unwrap();
2079 assert_eq!(items.len(), 1000);
2080 assert_eq!(items[0], 0);
2081 assert_eq!(items[999], 999);
2082 assert!(outer.pull().unwrap().is_none());
2083 }
2084
2085 #[test]
2086 fn group_by_single_key_fused_handles_key_change_with_substream_limit() {
2087 let outer = Source::from_iter([0, 1, 0])
2088 .group_by(2, |item| *item, false)
2089 .run_with(Sink::queue())
2090 .unwrap();
2091 let mut sources = vec![];
2092 while let Some(source) = outer.pull().unwrap() {
2093 sources.push(source);
2094 }
2095 assert_eq!(sources.len(), 2);
2096 assert_eq!(sources[0].clone().run_collect().unwrap(), vec![0, 0]);
2097 assert_eq!(sources[1].clone().run_collect().unwrap(), vec![1]);
2098 }
2099
2100 #[test]
2101 fn flat_map_merge_lock_lighter_matches_expected_count() {
2102 let items = Source::from_iter(0..20)
2103 .flat_map_merge(2, |item| Source::single(item + 100))
2104 .run_with(Sink::queue())
2105 .unwrap();
2106 let mut count = 0;
2107 while items.pull().unwrap().is_some() {
2108 count += 1;
2109 }
2110 assert_eq!(count, 20);
2111 }
2112
2113 #[test]
2123 fn group_by_single_key_emits_substream_before_upstream_completes() {
2124 let (tx, rx) = mpsc::sync_channel::<i32>(0);
2127 let rx = StdArc::new(std::sync::Mutex::new(rx));
2128
2129 let outer = Source::from_factory({
2130 let rx = StdArc::clone(&rx);
2131 move || {
2132 let rx = StdArc::clone(&rx);
2133 Box::new(std::iter::from_fn(move || {
2134 rx.lock().unwrap().recv().ok().map(Ok)
2135 })) as BoxStream<i32>
2136 }
2137 })
2138 .group_by(1, |_| 0u8, false)
2139 .run_with(Sink::queue())
2140 .unwrap();
2141
2142 let (sub_tx, sub_rx) = mpsc::channel::<Source<i32>>();
2145 let outer_thread = thread::spawn(move || {
2146 let substream = outer.pull().unwrap().expect("expected a substream");
2147 sub_tx.send(substream).unwrap();
2148 });
2149
2150 tx.send(0).unwrap();
2153
2154 let substream = sub_rx
2156 .recv_timeout(StdDuration::from_secs(5))
2157 .expect("timed out — group_by buffered first element before emitting substream");
2158
2159 for i in 1..100_i32 {
2161 tx.send(i).unwrap();
2162 }
2163 drop(tx);
2164
2165 let items: Vec<i32> = substream.run_collect().unwrap();
2166 assert_eq!(items.len(), 100);
2167 outer_thread.join().unwrap();
2168 }
2169
2170 #[test]
2171 fn group_by_concurrent_live_substreams_do_not_hold_ready_item_stress() {
2172 const STREAMS: usize = 32;
2173 const ROUNDS: usize = 8;
2174 const ITEMS: i64 = 8;
2175
2176 for _ in 0..ROUNDS {
2177 let barrier = StdArc::new(std::sync::Barrier::new(STREAMS));
2178 let mut handles = Vec::with_capacity(STREAMS);
2179
2180 for _ in 0..STREAMS {
2181 let barrier = StdArc::clone(&barrier);
2182 handles.push(thread::spawn(move || {
2183 let (tx, rx) = mpsc::sync_channel::<i64>(0);
2184 let rx = StdArc::new(std::sync::Mutex::new(rx));
2185
2186 let outer = Source::from_factory({
2187 let rx = StdArc::clone(&rx);
2188 move || {
2189 let rx = StdArc::clone(&rx);
2190 Box::new(std::iter::from_fn(move || {
2191 rx.lock().unwrap().recv().ok().map(Ok)
2192 })) as BoxStream<i64>
2193 }
2194 })
2195 .group_by(1, |_| 0_u8, false)
2196 .run_with(Sink::queue())
2197 .unwrap();
2198
2199 barrier.wait();
2200
2201 tx.send(0).unwrap();
2202 let substream = outer.pull().unwrap().expect("expected group_by substream");
2203 let subqueue = substream.run_with(Sink::queue()).unwrap();
2204 assert_eq!(subqueue.pull().unwrap(), Some(0));
2205
2206 for item in 1..ITEMS {
2207 tx.send(item).unwrap();
2208 assert_eq!(subqueue.pull().unwrap(), Some(item));
2209 }
2210 drop(tx);
2211
2212 assert!(subqueue.pull().unwrap().is_none());
2213 assert!(outer.pull().unwrap().is_none());
2214 }));
2215 }
2216
2217 for handle in handles {
2218 handle.join().expect("group_by stress worker panicked");
2219 }
2220 }
2221 }
2222
2223 #[test]
2230 fn split_when_emits_substream_before_segment_ends() {
2231 const SEGMENT_LEN: usize = 300;
2235
2236 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT_LEN * 2 + 4);
2237 for i in 0..SEGMENT_LEN as i32 {
2238 tx.send(i).unwrap();
2239 }
2240 tx.send(-1).unwrap(); tx.send(99).unwrap(); drop(tx);
2243
2244 let outer = Source::from_iter(rx)
2245 .split_when(|item| *item == -1)
2246 .run_with(Sink::queue())
2247 .unwrap();
2248
2249 let (result_tx, result_rx) = mpsc::channel();
2250 thread::spawn(move || {
2251 let first = outer.pull().unwrap().expect("expected first substream");
2252 let items: Vec<i32> = first.run_collect().unwrap();
2253 let second = outer.pull().unwrap().expect("expected second substream");
2254 let items2: Vec<i32> = second.run_collect().unwrap();
2255 let done = outer.pull().unwrap().is_none();
2256 let _ = result_tx.send((items, items2, done));
2257 });
2258
2259 let (items, items2, done) = result_rx
2260 .recv_timeout(StdDuration::from_secs(5))
2261 .expect("timed out — split_when is buffering the whole segment");
2262 assert_eq!(items.len(), SEGMENT_LEN);
2263 assert_eq!(items2, vec![-1, 99]);
2264 assert!(done);
2265 }
2266
2267 #[test]
2268 fn split_after_emits_substream_before_segment_ends() {
2269 const SEGMENT_LEN: usize = 300;
2270
2271 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT_LEN * 2 + 4);
2272 for i in 0..SEGMENT_LEN as i32 {
2273 tx.send(i).unwrap();
2274 }
2275 tx.send(-1).unwrap(); tx.send(99).unwrap();
2277 drop(tx);
2278
2279 let outer = Source::from_iter(rx)
2280 .split_after(|item| *item == -1)
2281 .run_with(Sink::queue())
2282 .unwrap();
2283
2284 let (result_tx, result_rx) = mpsc::channel();
2285 thread::spawn(move || {
2286 let first = outer.pull().unwrap().expect("expected first substream");
2287 let items: Vec<i32> = first.run_collect().unwrap();
2288 let second = outer.pull().unwrap().expect("expected second substream");
2289 let items2: Vec<i32> = second.run_collect().unwrap();
2290 let done = outer.pull().unwrap().is_none();
2291 let _ = result_tx.send((items, items2, done));
2292 });
2293
2294 let (items, items2, done) = result_rx
2295 .recv_timeout(StdDuration::from_secs(5))
2296 .expect("timed out — split_after is buffering the whole segment");
2297 assert_eq!(items.len(), SEGMENT_LEN + 1);
2299 assert_eq!(items2, vec![99]);
2300 assert!(done);
2301 }
2302
2303 #[test]
2307 fn flat_map_merge_coordinator_no_lost_wakeup_stress() {
2308 for _ in 0..20 {
2309 let result = Source::from_iter(0..50_i32)
2310 .flat_map_merge(8, |item| Source::from_iter(item..item + 3))
2311 .run_with(Sink::fold(0i64, |acc, item| acc + item as i64))
2312 .unwrap()
2313 .wait();
2314 assert_eq!(result, Ok(3825), "flat_map_merge produced wrong sum");
2317 }
2318 }
2319
2320 #[test]
2325 fn flat_map_merge_single_mutex_race_stress() {
2326 for _ in 0..20 {
2327 let result = Source::from_iter(0..100_i64)
2328 .flat_map_merge(16, |item| Source::from_iter([item, item + 1000]))
2329 .run_with(Sink::fold(0i64, |acc, v| acc + v))
2330 .unwrap()
2331 .wait();
2332 assert_eq!(result, Ok(109_900), "flat_map_merge single-mutex stress");
2336 }
2337 }
2338
2339 #[test]
2346 fn split_when_bounded_memory_rendezvous() {
2347 const SEGMENT: usize = 100;
2350 let (tx, rx) = mpsc::sync_channel::<i32>(SEGMENT * 4);
2351 for i in 0..SEGMENT as i32 {
2352 tx.send(i).unwrap();
2353 }
2354 tx.send(-1).unwrap(); for i in 0..10_i32 {
2357 tx.send(i).unwrap();
2358 }
2359 drop(tx);
2360
2361 let outer = Source::from_iter(rx)
2362 .split_when(|item| *item == -1)
2363 .run_with(Sink::queue())
2364 .unwrap();
2365
2366 let (result_tx, result_rx) = mpsc::channel();
2367 thread::spawn(move || {
2368 let first = outer.pull().unwrap().expect("first segment");
2369 let seg1: Vec<i32> = first.run_collect().unwrap();
2370 let second = outer.pull().unwrap().expect("second segment");
2371 let seg2: Vec<i32> = second.run_collect().unwrap();
2372 let done = outer.pull().unwrap().is_none();
2373 result_tx.send((seg1, seg2, done)).unwrap();
2374 });
2375
2376 let (seg1, seg2, done) = result_rx
2377 .recv_timeout(StdDuration::from_secs(5))
2378 .expect("timed out — split_when writer held items past LIVE_SUBSTREAM_BATCH");
2379 assert_eq!(seg1.len(), SEGMENT, "first segment length");
2380 assert_eq!(seg2[0], -1, "boundary element starts second segment");
2381 assert_eq!(seg2.len(), 11, "second segment: boundary + 10 items");
2382 assert!(done);
2383 }
2384
2385 #[test]
2389 fn group_by_single_key_bounded_memory_rendezvous() {
2390 const N: usize = 200;
2393 let outer = Source::from_iter(0..N as i64)
2394 .group_by(1, |_| 0u8, false)
2395 .run_with(Sink::queue())
2396 .unwrap();
2397
2398 let (result_tx, result_rx) = mpsc::channel();
2399 thread::spawn(move || {
2400 let substream = outer.pull().unwrap().expect("substream");
2401 let items: Vec<i64> = substream.run_collect().unwrap();
2402 let done = outer.pull().unwrap().is_none();
2403 result_tx.send((items, done)).unwrap();
2404 });
2405
2406 let (items, done) = result_rx
2407 .recv_timeout(StdDuration::from_secs(5))
2408 .expect("timed out — group_by write batch held items beyond LIVE_SUBSTREAM_BATCH");
2409 assert_eq!(items.len(), N, "all items delivered");
2410 assert_eq!(items[0], 0);
2411 assert_eq!(items[N - 1], (N - 1) as i64);
2412 assert!(done);
2413 }
2414
2415 #[test]
2416 fn scan_emits_seed_and_accumulated_values() {
2417 let result = Source::from_iter(1..=3)
2418 .scan(0, |acc, item| acc + item)
2419 .run_collect()
2420 .unwrap();
2421
2422 assert_eq!(result, vec![0, 1, 3, 6]);
2423 }
2424
2425 #[test]
2426 fn limit_fails_after_max_elements() {
2427 let result = Source::from_iter(0..3).limit(2).run_collect();
2428
2429 assert_eq!(result, Err(StreamError::LimitExceeded { max: 2 }));
2430 }
2431
2432 #[test]
2433 fn limit_weighted_fails_with_limit_error_like_akka() {
2434 let result = Source::from_iter(["this", "is", "some", "string"])
2435 .via(Flow::identity().limit_weighted(15, |item: &&str| item.len()))
2436 .run_collect();
2437
2438 assert_eq!(result, Err(StreamError::LimitExceeded { max: 15 }));
2439 }
2440
2441 #[test]
2442 fn grouped_weighted_allows_oversized_first_element_like_akka() {
2443 let result = Source::from_iter([10_usize, 1, 2])
2444 .via(Flow::identity().grouped_weighted(5, |item: &usize| *item))
2445 .run_collect()
2446 .unwrap();
2447
2448 assert_eq!(result, vec![vec![10], vec![1, 2]]);
2449 }
2450
2451 #[test]
2452 fn grouped_weighted_keeps_oversized_later_element_in_current_group_like_akka() {
2453 let result = Source::from_iter([1_usize, 10, 2])
2454 .via(Flow::identity().grouped_weighted(5, |item: &usize| *item))
2455 .run_collect()
2456 .unwrap();
2457
2458 assert_eq!(result, vec![vec![1, 10], vec![2]]);
2459 }
2460
2461 #[test]
2462 fn sink_terminals_materialize_results() {
2463 let sum = Source::from_iter(1..=4)
2464 .run_with(Sink::fold(0, |acc, item| acc + item))
2465 .unwrap();
2466
2467 assert_eq!(wait(sum), 10);
2468 assert_eq!(
2469 wait(Source::from_iter(1..=4).run_with(Sink::head()).unwrap()),
2470 1
2471 );
2472 assert_eq!(
2473 wait(Source::from_iter(1..=4).run_with(Sink::last()).unwrap()),
2474 4
2475 );
2476 }
2477
2478 #[test]
2479 fn all_terminal_sink_variants_complete() {
2480 assert_eq!(
2481 wait(
2482 Source::from_iter([1, 2, 3])
2483 .run_with(Sink::collect())
2484 .unwrap()
2485 ),
2486 vec![1, 2, 3]
2487 );
2488 assert_eq!(
2489 wait(
2490 Source::<i32>::empty()
2491 .run_with(Sink::head_option())
2492 .unwrap()
2493 ),
2494 None
2495 );
2496 assert_eq!(
2497 wait(
2498 Source::from_iter([1, 2, 3])
2499 .run_with(Sink::last_option())
2500 .unwrap()
2501 ),
2502 Some(3)
2503 );
2504 assert_eq!(
2505 wait(
2506 Source::from_iter([1, 2, 3])
2507 .run_with(Sink::reduce(|acc, item| acc + item))
2508 .unwrap()
2509 ),
2510 6
2511 );
2512
2513 let seen = StdArc::new(StdAtomicUsize::new(0));
2514 let seen_by_sink = StdArc::clone(&seen);
2515 assert_eq!(
2516 wait(
2517 Source::from_iter([1_usize, 2, 3])
2518 .run_with(Sink::foreach(move |item| {
2519 seen_by_sink.fetch_add(item, StdOrdering::SeqCst);
2520 }))
2521 .unwrap()
2522 ),
2523 NotUsed
2524 );
2525 assert_eq!(seen.load(StdOrdering::SeqCst), 6);
2526 }
2527
2528 #[test]
2529 fn take_last_zero_returns_empty_vector() {
2530 let result = Source::from_iter([1, 2, 3])
2531 .run_with(Sink::take_last(0))
2532 .unwrap();
2533
2534 assert_eq!(wait(result), Vec::<i32>::new());
2535 }
2536
2537 #[test]
2538 fn bounded_head_terminals_complete_inline() {
2539 let materializer = Materializer::new();
2540
2541 let mut head = Source::from_iter(0_u64..1_000)
2542 .run_with_materializer(Sink::head(), &materializer)
2543 .unwrap();
2544 assert_eq!(materializer.active_streams(), 0);
2545 assert_eq!(head.try_wait(), Some(Ok(0)));
2546
2547 let mut filtered_head = Source::from_iter(0_u64..1_000)
2548 .filter(|item| *item >= 10)
2549 .run_with_materializer(Sink::head(), &materializer)
2550 .unwrap();
2551 assert_eq!(materializer.active_streams(), 0);
2552 assert_eq!(filtered_head.try_wait(), Some(Ok(10)));
2553
2554 let mut head_option = Source::<u64>::empty()
2555 .run_with_materializer(Sink::head_option(), &materializer)
2556 .unwrap();
2557 assert_eq!(materializer.active_streams(), 0);
2558 assert_eq!(head_option.try_wait(), Some(Ok(None)));
2559 }
2560
2561 #[test]
2562 fn bounded_head_fast_path_preserves_terminal_errors() {
2563 let materializer = Materializer::new();
2564
2565 let mut empty = Source::<u64>::empty()
2566 .run_with_materializer(Sink::head(), &materializer)
2567 .unwrap();
2568 assert_eq!(empty.try_wait(), Some(Err(StreamError::EmptyStream)));
2569
2570 let mut failed = Source::<u64>::failed(StreamError::Failed("boom".into()))
2571 .run_with_materializer(Sink::head(), &materializer)
2572 .unwrap();
2573 assert_eq!(
2574 failed.try_wait(),
2575 Some(Err(StreamError::Failed("boom".into())))
2576 );
2577 assert_eq!(materializer.active_streams(), 0);
2578 }
2579
2580 #[test]
2581 fn runnable_graph_composes_source_and_sink() {
2582 let graph = Source::from_iter(1..=4)
2583 .map(|item| item * 2)
2584 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::right);
2585
2586 assert_eq!(wait(graph.run().unwrap()), 20);
2587
2588 let graph = Source::single(1)
2589 .map_materialized_value(|_| 20)
2590 .to(Sink::ignore())
2591 .map_materialized_value(|value| value + 1);
2592 assert_eq!(graph.run().unwrap(), 21);
2593
2594 let ignored = Source::single(1).to(Sink::ignore()).run().unwrap();
2595 assert_eq!(ignored, NotUsed);
2596 }
2597
2598 #[test]
2599 fn materialized_values_follow_keep_defaults() {
2600 let source = Source::single(1).map_materialized_value(|_| "source");
2601 let flow = Flow::identity().map_materialized_value(|_| "flow");
2602
2603 let source_mat = source.clone().via(flow.clone()).to(Sink::ignore()).run();
2604 assert_eq!(source_mat.unwrap(), "source");
2605
2606 let combined = source
2607 .via_mat(flow, Keep::both)
2608 .to_mat(Sink::ignore(), Keep::both)
2609 .run()
2610 .unwrap();
2611 assert_eq!(combined.0, ("source", "flow"));
2612 assert_eq!(wait(combined.1), NotUsed);
2613
2614 let sink_mat = Source::single(41)
2615 .map_materialized_value(|_| "ignored source")
2616 .run_with(Sink::fold(1, |acc, item| acc + item))
2617 .unwrap();
2618 assert_eq!(wait(sink_mat), 42);
2619 }
2620
2621 #[test]
2622 fn flow_to_sink_preserves_flow_materialized_value_by_default() {
2623 let sink = Flow::identity()
2624 .map(|item: i32| item + 1)
2625 .map_materialized_value(|_| "flow")
2626 .to(Sink::fold(0, |acc, item| acc + item));
2627
2628 let materialized = Source::from_iter([1, 2, 3]).run_with(sink).unwrap();
2629
2630 assert_eq!(materialized, "flow");
2631 let explicit = Flow::identity()
2632 .map(|item: i32| item + 1)
2633 .map_materialized_value(|_| "flow")
2634 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both)
2635 .run_with(Source::from_iter([1, 2, 3]))
2636 .unwrap();
2637 assert_eq!(explicit, NotUsed);
2638
2639 let explicit = Source::from_iter([1, 2, 3])
2640 .run_with(
2641 Flow::identity()
2642 .map(|item: i32| item + 1)
2643 .map_materialized_value(|_| "flow")
2644 .to_mat(Sink::fold(0, |acc, item| acc + item), Keep::both),
2645 )
2646 .unwrap();
2647 assert_eq!(explicit.0, "flow");
2648 assert_eq!(wait(explicit.1), 9);
2649 }
2650
2651 #[test]
2652 fn materializer_shutdown_fails_materialization() {
2653 let materializer = Materializer::new();
2654 let named = materializer.with_name_prefix("test-stream");
2655 materializer.shutdown();
2656
2657 let graph = Source::single(1).to(Sink::ignore());
2658
2659 assert_eq!(named.name_prefix(), "test-stream");
2660 assert_eq!(
2661 graph.run_with_materializer(&named),
2662 Err(StreamError::AbruptTermination)
2663 );
2664 }
2665
2666 #[test]
2667 fn materializer_shutdown_fails_running_stream_completion() {
2668 let materializer = Materializer::new();
2669 let completion = Source::repeat(1)
2670 .run_with_materializer(Sink::ignore(), &materializer)
2671 .unwrap();
2672
2673 assert_eq!(materializer.active_streams(), 1);
2674 materializer.shutdown();
2675 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
2676 assert_eq!(materializer.active_streams(), 0);
2677 }
2678
2679 #[test]
2680 fn dropped_stream_completion_cancels_running_stream() {
2681 let materializer = Materializer::new();
2682 let completion = Source::repeat(1)
2683 .run_with_materializer(Sink::ignore(), &materializer)
2684 .unwrap();
2685
2686 assert_eq!(materializer.active_streams(), 1);
2687 drop(completion);
2688 for _ in 0..50 {
2689 if materializer.active_streams() == 0 {
2690 break;
2691 }
2692 thread::sleep(Duration::from_millis(5));
2693 }
2694 assert_eq!(materializer.active_streams(), 0);
2695 }
2696
2697 #[test]
2698 fn runtime_timers_fire_cancel_and_stop_on_shutdown() {
2699 let materializer = Materializer::new();
2700 let (once_tx, once_rx) = mpsc::channel();
2701 let once = materializer.schedule_once(Duration::from_millis(5), move || {
2702 once_tx.send(()).unwrap();
2703 });
2704 once_rx.recv_timeout(Duration::from_secs(10)).unwrap();
2705 assert!(!once.is_cancelled());
2706
2707 let (cancelled_tx, cancelled_rx) = mpsc::channel();
2711 let cancelled = materializer.schedule_once(Duration::from_millis(500), move || {
2712 cancelled_tx.send(()).unwrap();
2713 });
2714 assert!(cancelled.cancel());
2715 assert!(!cancelled.cancel());
2716 assert!(cancelled.is_cancelled());
2717 let (sentinel_tx, sentinel_rx) = mpsc::channel();
2718 materializer.schedule_once(Duration::from_millis(750), move || {
2719 sentinel_tx.send(()).unwrap();
2720 });
2721 sentinel_rx.recv_timeout(Duration::from_secs(30)).unwrap();
2722 assert!(cancelled_rx.try_recv().is_err());
2723
2724 let fixed_delay_count = StdArc::new(StdAtomicUsize::new(0));
2725 let fixed_delay_task_count = StdArc::clone(&fixed_delay_count);
2726 let fixed_delay = materializer.schedule_with_fixed_delay(
2727 Duration::from_millis(1),
2728 Duration::from_millis(5),
2729 move || {
2730 fixed_delay_task_count.fetch_add(1, StdOrdering::SeqCst);
2731 },
2732 );
2733 for _ in 0..2_000 {
2734 if fixed_delay_count.load(StdOrdering::SeqCst) > 0 {
2735 break;
2736 }
2737 thread::sleep(Duration::from_millis(5));
2738 }
2739 assert!(fixed_delay_count.load(StdOrdering::SeqCst) > 0);
2740 fixed_delay.cancel();
2741
2742 let fixed_rate_count = StdArc::new(StdAtomicUsize::new(0));
2743 let fixed_rate_task_count = StdArc::clone(&fixed_rate_count);
2744 let fixed_rate = materializer.schedule_at_fixed_rate(
2745 Duration::from_millis(1),
2746 Duration::from_millis(5),
2747 move || {
2748 fixed_rate_task_count.fetch_add(1, StdOrdering::SeqCst);
2749 },
2750 );
2751 for _ in 0..2_000 {
2752 if fixed_rate_count.load(StdOrdering::SeqCst) > 0 {
2753 break;
2754 }
2755 thread::sleep(Duration::from_millis(5));
2756 }
2757 assert!(fixed_rate_count.load(StdOrdering::SeqCst) > 0);
2758 fixed_rate.cancel();
2759
2760 let shutdown_materializer = Materializer::new();
2764 let (shutdown_tx, shutdown_rx) = mpsc::channel();
2765 shutdown_materializer.schedule_once(Duration::from_millis(500), move || {
2766 shutdown_tx.send(()).unwrap();
2767 });
2768 shutdown_materializer.shutdown();
2769 let (shutdown_sentinel_tx, shutdown_sentinel_rx) = mpsc::channel();
2770 materializer.schedule_once(Duration::from_millis(750), move || {
2771 shutdown_sentinel_tx.send(()).unwrap();
2772 });
2773 shutdown_sentinel_rx
2774 .recv_timeout(Duration::from_secs(30))
2775 .unwrap();
2776 assert!(shutdown_rx.try_recv().is_err());
2777 }
2778
2779 #[test]
2780 fn runtime_timer_driver_preserves_fixed_rate_cadence_under_slow_tasks() {
2781 use std::sync::{Condvar, Mutex};
2782
2783 #[derive(Debug)]
2784 enum TimerEvent {
2785 Started(usize, Instant),
2786 Completed(usize, Instant),
2787 }
2788
2789 let recv_event = |rx: &mpsc::Receiver<TimerEvent>, label: &str| {
2790 rx.recv_timeout(Duration::from_secs(20))
2791 .unwrap_or_else(|err| panic!("{label}: expected timer event within 20 s: {err}"))
2792 };
2793 let release = |gate: &StdArc<(Mutex<bool>, Condvar)>| {
2794 let (released, condvar) = &**gate;
2795 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2796 *released = true;
2797 condvar.notify_all();
2798 };
2799
2800 let interval = Duration::from_secs(2);
2801 let overrun = interval + Duration::from_millis(250);
2802
2803 let rate_materializer = Materializer::new();
2804 let (rate_tx, rate_rx) = mpsc::channel();
2805 let rate_runs = StdArc::new(StdAtomicUsize::new(0));
2806 let rate_task_runs = StdArc::clone(&rate_runs);
2807 let rate_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2808 let rate_task_gate = StdArc::clone(&rate_gate);
2809 let fixed_rate =
2810 rate_materializer.schedule_at_fixed_rate(Duration::ZERO, interval, move || {
2811 let run = rate_task_runs.fetch_add(1, StdOrdering::SeqCst) + 1;
2812 rate_tx
2813 .send(TimerEvent::Started(run, Instant::now()))
2814 .unwrap();
2815 if run == 1 {
2816 let (released, condvar) = &*rate_task_gate;
2817 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2818 while !*released {
2819 released = condvar
2820 .wait(released)
2821 .unwrap_or_else(|poison| poison.into_inner());
2822 }
2823 rate_tx
2824 .send(TimerEvent::Completed(run, Instant::now()))
2825 .unwrap();
2826 }
2827 });
2828 let rate_first_started = match recv_event(&rate_rx, "fixed-rate first task") {
2829 TimerEvent::Started(1, at) => at,
2830 other => panic!("fixed-rate first task: unexpected event {other:?}"),
2831 };
2832 assert!(wait_until(Duration::from_secs(20), || {
2833 rate_first_started.elapsed() >= overrun
2834 }));
2835 release(&rate_gate);
2836 let rate_first_completed = match recv_event(&rate_rx, "fixed-rate first completion") {
2837 TimerEvent::Completed(1, at) => at,
2838 other => panic!("fixed-rate first completion: unexpected event {other:?}"),
2839 };
2840 let rate_second_started = match recv_event(&rate_rx, "fixed-rate second task") {
2841 TimerEvent::Started(2, at) => at,
2842 other => panic!("fixed-rate second task: unexpected event {other:?}"),
2843 };
2844 fixed_rate.cancel();
2845 rate_materializer.shutdown();
2846
2847 let delay_materializer = Materializer::new();
2848 let (delay_tx, delay_rx) = mpsc::channel();
2849 let delay_runs = StdArc::new(StdAtomicUsize::new(0));
2850 let delay_task_runs = StdArc::clone(&delay_runs);
2851 let delay_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2852 let delay_task_gate = StdArc::clone(&delay_gate);
2853 let fixed_delay =
2854 delay_materializer.schedule_with_fixed_delay(Duration::ZERO, interval, move || {
2855 let run = delay_task_runs.fetch_add(1, StdOrdering::SeqCst) + 1;
2856 delay_tx
2857 .send(TimerEvent::Started(run, Instant::now()))
2858 .unwrap();
2859 if run == 1 {
2860 let (released, condvar) = &*delay_task_gate;
2861 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2862 while !*released {
2863 released = condvar
2864 .wait(released)
2865 .unwrap_or_else(|poison| poison.into_inner());
2866 }
2867 delay_tx
2868 .send(TimerEvent::Completed(run, Instant::now()))
2869 .unwrap();
2870 }
2871 });
2872 let delay_first_started = match recv_event(&delay_rx, "fixed-delay first task") {
2873 TimerEvent::Started(1, at) => at,
2874 other => panic!("fixed-delay first task: unexpected event {other:?}"),
2875 };
2876 assert!(wait_until(Duration::from_secs(20), || {
2877 delay_first_started.elapsed() >= overrun
2878 }));
2879 release(&delay_gate);
2880 let delay_first_completed = match recv_event(&delay_rx, "fixed-delay first completion") {
2881 TimerEvent::Completed(1, at) => at,
2882 other => panic!("fixed-delay first completion: unexpected event {other:?}"),
2883 };
2884 let delay_second_started = match recv_event(&delay_rx, "fixed-delay second task") {
2885 TimerEvent::Started(2, at) => at,
2886 other => panic!("fixed-delay second task: unexpected event {other:?}"),
2887 };
2888 fixed_delay.cancel();
2889 delay_materializer.shutdown();
2890
2891 let rate_task_time = rate_first_completed.duration_since(rate_first_started);
2892 let rate_catch_up = rate_second_started.duration_since(rate_first_completed);
2893 let delay_task_time = delay_first_completed.duration_since(delay_first_started);
2894 let delay_gap = delay_second_started.duration_since(delay_first_completed);
2895 assert!(
2896 rate_task_time >= interval,
2897 "fixed-rate first task should overrun its interval; ran for {rate_task_time:?}"
2898 );
2899 assert!(
2900 rate_catch_up < interval,
2901 "fixed-rate second task should catch up after an overrun; waited {rate_catch_up:?}"
2902 );
2903 assert!(
2904 delay_task_time >= interval,
2905 "fixed-delay first task should overrun its interval; ran for {delay_task_time:?}"
2906 );
2907 assert!(
2908 delay_gap >= interval,
2909 "fixed-delay second task fired before one full delay elapsed after completion: {delay_gap:?}",
2910 );
2911 }
2912
2913 #[test]
2914 fn runtime_repeating_timer_cancellation_stops_future_fires() {
2915 let materializer = Materializer::new();
2916 let (tx, rx) = mpsc::channel();
2917 let timer = materializer.schedule_at_fixed_rate(
2918 Duration::from_millis(1),
2919 Duration::from_millis(30),
2920 move || {
2921 tx.send(()).unwrap();
2922 },
2923 );
2924
2925 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2926 assert!(timer.cancel());
2927 assert!(rx.recv_timeout(Duration::from_millis(90)).is_err());
2928 materializer.shutdown();
2929 }
2930
2931 #[test]
2932 fn runtime_panicking_once_timer_does_not_kill_driver_or_later_timers() {
2933 let materializer = Materializer::new();
2934 materializer.schedule_once(Duration::from_millis(1), || {
2935 panic!("timer boom");
2936 });
2937
2938 let (tx, rx) = mpsc::channel();
2939 materializer.schedule_once(Duration::from_millis(20), move || {
2940 tx.send(()).unwrap();
2941 });
2942
2943 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2944 materializer.shutdown();
2945 }
2946
2947 #[test]
2948 fn runtime_panicking_fixed_rate_timer_stops_itself_and_leaves_driver_alive() {
2949 let materializer = Materializer::new();
2950 let panic_count = StdArc::new(StdAtomicUsize::new(0));
2951 let panic_count_task = StdArc::clone(&panic_count);
2952 materializer.schedule_at_fixed_rate(Duration::ZERO, Duration::from_millis(20), move || {
2953 panic_count_task.fetch_add(1, StdOrdering::SeqCst);
2954 panic!("fixed-rate boom");
2955 });
2956
2957 assert!(wait_until(Duration::from_millis(150), || {
2958 panic_count.load(StdOrdering::SeqCst) == 1
2959 }));
2960
2961 let (tx, rx) = mpsc::channel();
2962 materializer.schedule_once(Duration::from_millis(30), move || {
2963 tx.send(()).unwrap();
2964 });
2965 rx.recv_timeout(Duration::from_millis(250)).unwrap();
2966
2967 thread::sleep(Duration::from_millis(90));
2968 assert_eq!(panic_count.load(StdOrdering::SeqCst), 1);
2969 materializer.shutdown();
2970 }
2971
2972 #[test]
2973 fn runtime_slow_timer_task_does_not_delay_unrelated_timers() {
2974 let materializer = Materializer::new();
2975 let (started_tx, started_rx) = mpsc::channel();
2976 let release_gate = StdArc::new((Mutex::new(false), Condvar::new()));
2977 let release_task = StdArc::clone(&release_gate);
2978 let runs = StdArc::new(StdAtomicUsize::new(0));
2979 let task_runs = StdArc::clone(&runs);
2980 let slow_timer = materializer.schedule_at_fixed_rate(
2981 Duration::ZERO,
2982 Duration::from_millis(250),
2983 move || {
2984 if task_runs.fetch_add(1, StdOrdering::SeqCst) == 0 {
2985 started_tx.send(()).unwrap();
2986 let (released, condvar) = &*release_task;
2987 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
2988 while !*released {
2989 released = condvar
2990 .wait(released)
2991 .unwrap_or_else(|poison| poison.into_inner());
2992 }
2993 }
2994 },
2995 );
2996
2997 started_rx
2998 .recv_timeout(Duration::from_secs(10))
2999 .expect("slow timer task should start");
3000
3001 let (tx, rx) = mpsc::channel();
3002 materializer.schedule_once(Duration::from_millis(10), move || {
3003 tx.send(()).unwrap();
3004 });
3005 let fired = rx.recv_timeout(Duration::from_secs(10));
3006
3007 let (released, condvar) = &*release_gate;
3008 let mut released = released.lock().unwrap_or_else(|poison| poison.into_inner());
3009 *released = true;
3010 condvar.notify_all();
3011 drop(released);
3012
3013 slow_timer.cancel();
3014 materializer.shutdown();
3015 fired.expect("unrelated timer should fire while slow timer task is still blocked");
3016 }
3017
3018 #[test]
3019 fn runtime_shutdown_stops_timer_driver_thread() {
3020 let materializer = Materializer::new();
3021 assert!(wait_until(Duration::from_secs(1), || materializer
3022 .timer_driver_is_live()));
3023
3024 materializer.shutdown();
3025 assert!(wait_until(Duration::from_secs(2), || !materializer
3026 .timer_driver_is_live()));
3027 }
3028
3029 #[test]
3030 fn runtime_timer_driver_orders_many_timers_by_deadline() {
3031 let materializer = Materializer::new();
3032 let (tx, rx) = mpsc::channel();
3033 let schedule = [(450_u64, 4_u8), (50, 1), (350, 3), (150, 2), (550, 5)];
3034
3035 for (delay_ms, value) in schedule {
3036 let tx = tx.clone();
3037 materializer.schedule_once(Duration::from_millis(delay_ms), move || {
3038 tx.send(value).unwrap();
3039 });
3040 }
3041 drop(tx);
3042
3043 let mut received = Vec::new();
3044 for _ in 0..schedule.len() {
3045 received.push(rx.recv_timeout(Duration::from_secs(10)).unwrap());
3046 }
3047 materializer.shutdown();
3048
3049 assert_eq!(received, vec![1, 2, 3, 4, 5]);
3050 }
3051
3052 #[test]
3053 fn runtime_timer_driver_uses_one_thread_per_runtime_regardless_of_timer_count() {
3054 let materializer = Materializer::new();
3055 let thread_name = materializer.timer_thread_name().to_owned();
3056 let linux_thread_name = thread_name.chars().take(15).collect::<String>();
3062 assert!(wait_until(Duration::from_secs(5), || {
3063 materializer.timer_driver_is_live() && linux_thread_count(&linux_thread_name) >= 1
3064 }));
3065 let live_timer_threads = linux_thread_count(&linux_thread_name);
3066
3067 for _ in 0..128 {
3068 materializer.schedule_once(Duration::from_secs(60), || {});
3069 }
3070
3071 assert!(
3072 wait_until(Duration::from_secs(5), || {
3073 materializer.timer_driver_is_live()
3074 && linux_thread_count(&linux_thread_name) == live_timer_threads
3075 }),
3076 "scheduling timers should not create extra timer threads for a runtime",
3077 );
3078 materializer.shutdown();
3079 assert!(wait_until(Duration::from_secs(5), || {
3080 !materializer.timer_driver_is_live()
3081 && linux_thread_count(&linux_thread_name) < live_timer_threads
3082 }));
3083 }
3084
3085 #[test]
3086 fn cancelled_and_never_sinks_have_distinct_materialization_results() {
3087 assert_eq!(
3088 Source::repeat(1)
3089 .run_with(Sink::cancelled())
3090 .expect("cancelled sink materializes"),
3091 NotUsed
3092 );
3093 assert_eq!(
3094 Source::single(1)
3095 .run_with(Sink::never())
3096 .expect("never sink materializes")
3097 .try_wait(),
3098 None
3099 );
3100 }
3101
3102 #[test]
3103 fn never_sink_finishes_on_materializer_shutdown() {
3104 let materializer = Materializer::new();
3105 let completion = Source::single(1)
3106 .run_with_materializer(Sink::never(), &materializer)
3107 .unwrap();
3108
3109 materializer.shutdown();
3110 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
3111 }
3112
3113 #[test]
3114 fn dropping_source_never_completion_releases_parked_worker() {
3115 let materializer = Materializer::new();
3116 let completion = Source::<i32>::never()
3117 .run_with_materializer(Sink::ignore(), &materializer)
3118 .unwrap();
3119
3120 assert!(wait_until(StdDuration::from_secs(1), || {
3121 materializer.active_streams() == 1
3122 }));
3123 assert_eq!(materializer.active_streams(), 1);
3124
3125 drop(completion);
3126
3127 assert!(wait_until(StdDuration::from_secs(15), || {
3128 materializer.active_streams() == 0
3129 }));
3130 assert_eq!(materializer.active_streams(), 0);
3131 }
3132
3133 #[test]
3134 fn future_and_maybe_sources_emit_values() {
3135 let future_value = Source::future(|| async { Ok(7) }).run_collect().unwrap();
3136 assert_eq!(future_value, vec![7]);
3137
3138 let future_source = Source::future_source(|| async { Ok(Source::from_iter([1, 2, 3])) })
3139 .run_collect()
3140 .unwrap();
3141 assert_eq!(future_source, vec![1, 2, 3]);
3142
3143 let (handle, source) = Source::maybe();
3144 assert_eq!(
3145 source.clone().run_collect(),
3146 Err(StreamError::MaybeIncomplete)
3147 );
3148 handle.complete(9).unwrap();
3149 assert_eq!(source.run_collect().unwrap(), vec![9]);
3150 }
3151
3152 #[test]
3153 fn wp6b_source_generators_emit_and_fail_like_stream_errors() {
3154 assert_eq!(
3155 Source::cycle(|| [1, 2, 3].into_iter())
3156 .take(8)
3157 .run_collect()
3158 .unwrap(),
3159 vec![1, 2, 3, 1, 2, 3, 1, 2]
3160 );
3161 assert_eq!(
3162 Source::<i32>::cycle(std::iter::empty::<i32>).run_collect(),
3163 Err(StreamError::Failed("empty iterator".into()))
3164 );
3165 assert_eq!(
3166 Source::unfold(0, |state| (state < 4).then_some((state + 1, state)))
3167 .run_collect()
3168 .unwrap(),
3169 vec![0, 1, 2, 3]
3170 );
3171 assert_eq!(
3172 Source::unfold_async(0, |state| async move {
3173 Ok((state < 4).then_some((state + 1, state * 2)))
3174 })
3175 .run_collect()
3176 .unwrap(),
3177 vec![0, 2, 4, 6]
3178 );
3179 assert!(matches!(
3180 Source::<i32>::lazy_single(|| panic!("boom")).run_collect(),
3181 Err(StreamError::Failed(message)) if message == "lazy_single factory panicked"
3182 ));
3183 }
3184
3185 #[test]
3186 fn wp6b_lazy_sources_defer_until_first_pull_and_complete_deferred_mat() {
3187 let created = StdArc::new(StdAtomicUsize::new(0));
3188 let created_for_source = StdArc::clone(&created);
3189 let source = Source::<i32>::lazy_source(move || {
3190 created_for_source.fetch_add(1, StdOrdering::SeqCst);
3191 Source::from_iter([7, 8]).map_materialized_value(|_| 99)
3192 });
3193 let materializer = Materializer::new();
3194 let (mut stream, mut mat) = StdArc::clone(&source.factory)
3195 .create(&materializer)
3196 .unwrap();
3197
3198 assert_eq!(created.load(StdOrdering::SeqCst), 0);
3199 assert!(mat.try_wait().is_none());
3200 assert_eq!(stream.next().unwrap().unwrap(), 7);
3201 assert_eq!(mat.wait().unwrap(), 99);
3202 assert_eq!(created.load(StdOrdering::SeqCst), 1);
3203 assert_eq!(stream.next().unwrap().unwrap(), 8);
3204
3205 let never_created = StdArc::new(StdAtomicUsize::new(0));
3206 let never_created_for_source = StdArc::clone(&never_created);
3207 let mat = Source::<i32>::lazy_future_source(move || {
3208 never_created_for_source.fetch_add(1, StdOrdering::SeqCst);
3209 async { Ok(Source::single(1)) }
3210 })
3211 .to(Sink::cancelled())
3212 .run()
3213 .unwrap();
3214 assert!(matches!(mat.wait(), Err(StreamError::Failed(_))));
3215 assert_eq!(never_created.load(StdOrdering::SeqCst), 0);
3216
3217 let lazy_future = StdArc::new(StdAtomicUsize::new(0));
3218 let lazy_future_for_source = StdArc::clone(&lazy_future);
3219 let source = Source::lazy_future(move || {
3220 lazy_future_for_source.fetch_add(1, StdOrdering::SeqCst);
3221 async { Ok(42) }
3222 });
3223 let (mut stream, _) = StdArc::clone(&source.factory)
3224 .create(&Materializer::new())
3225 .unwrap();
3226 assert_eq!(lazy_future.load(StdOrdering::SeqCst), 0);
3227 assert_eq!(stream.next().unwrap().unwrap(), 42);
3228 assert_eq!(lazy_future.load(StdOrdering::SeqCst), 1);
3229 }
3230
3231 #[test]
3232 fn wp6b_unfold_resource_closes_on_completion_failure_and_cancellation() {
3233 let closed = StdArc::new(StdAtomicUsize::new(0));
3234 let closed_on_complete = StdArc::clone(&closed);
3235 let values = Source::unfold_resource(
3236 || Ok(std::collections::VecDeque::from([1, 2, 3])),
3237 |items| Ok(items.pop_front()),
3238 move |_items| {
3239 closed_on_complete.fetch_add(1, StdOrdering::SeqCst);
3240 Ok(())
3241 },
3242 )
3243 .run_collect()
3244 .unwrap();
3245 assert_eq!(values, vec![1, 2, 3]);
3246 assert_eq!(closed.load(StdOrdering::SeqCst), 1);
3247
3248 let closed_on_failure = StdArc::new(StdAtomicUsize::new(0));
3249 let closed_on_failure_for_close = StdArc::clone(&closed_on_failure);
3250 let failed = Source::<i32>::unfold_resource(
3251 || Ok(()),
3252 |_| Err(StreamError::Failed("read".into())),
3253 move |_| {
3254 closed_on_failure_for_close.fetch_add(1, StdOrdering::SeqCst);
3255 Err(StreamError::Failed("close".into()))
3256 },
3257 )
3258 .run_collect();
3259 assert_eq!(failed, Err(StreamError::Failed("read".into())));
3260 assert_eq!(closed_on_failure.load(StdOrdering::SeqCst), 1);
3261
3262 let closed_on_cancel = StdArc::new(StdAtomicUsize::new(0));
3263 let closed_on_cancel_for_close = StdArc::clone(&closed_on_cancel);
3264 let first = Source::unfold_resource(
3265 || Ok(0_usize),
3266 |next| {
3267 let item = *next;
3268 *next += 1;
3269 Ok(Some(item))
3270 },
3271 move |_| {
3272 closed_on_cancel_for_close.fetch_add(1, StdOrdering::SeqCst);
3273 Ok(())
3274 },
3275 )
3276 .run_with(Sink::head())
3277 .unwrap();
3278 assert_eq!(first.wait().unwrap(), 0);
3279 assert!(wait_until(Duration::from_millis(250), || {
3280 closed_on_cancel.load(StdOrdering::SeqCst) == 1
3281 }));
3282 }
3283
3284 #[test]
3285 fn wp6b_async_resource_and_async_accumulators_are_sequential() {
3286 let closed = StdArc::new(StdAtomicUsize::new(0));
3287 let closed_for_close = StdArc::clone(&closed);
3288 let values = Source::unfold_resource_async(
3289 || async { Ok(std::collections::VecDeque::from([1, 2, 3])) },
3290 |items| {
3291 let item = items.pop_front();
3292 async move { Ok(item) }
3293 },
3294 move |_items| {
3295 let closed = StdArc::clone(&closed_for_close);
3296 async move {
3297 closed.fetch_add(1, StdOrdering::SeqCst);
3298 Ok(())
3299 }
3300 },
3301 )
3302 .run_collect()
3303 .unwrap();
3304 assert_eq!(values, vec![1, 2, 3]);
3305 assert_eq!(closed.load(StdOrdering::SeqCst), 1);
3306
3307 let closed_on_failure = StdArc::new(StdAtomicUsize::new(0));
3308 let closed_on_failure_for_close = StdArc::clone(&closed_on_failure);
3309 let failed = Source::<i32>::unfold_resource_async(
3310 || async { Ok(()) },
3311 |_resource| async { Err(StreamError::Failed("read".into())) },
3312 move |_resource| {
3313 let closed_on_failure = StdArc::clone(&closed_on_failure_for_close);
3314 async move {
3315 closed_on_failure.fetch_add(1, StdOrdering::SeqCst);
3316 Err(StreamError::Failed("close".into()))
3317 }
3318 },
3319 )
3320 .run_collect();
3321 assert_eq!(failed, Err(StreamError::Failed("read".into())));
3322 assert_eq!(closed_on_failure.load(StdOrdering::SeqCst), 1);
3323
3324 let active = StdArc::new(StdAtomicUsize::new(0));
3325 let max_active = StdArc::new(StdAtomicUsize::new(0));
3326 let active_for_stage = StdArc::clone(&active);
3327 let max_for_stage = StdArc::clone(&max_active);
3328 let scanned = Source::from_iter(1..=4)
3329 .scan_async(0, move |acc, item| {
3330 let active = StdArc::clone(&active_for_stage);
3331 let max_active = StdArc::clone(&max_for_stage);
3332 async move {
3333 let now = active.fetch_add(1, StdOrdering::SeqCst) + 1;
3334 max_active.fetch_max(now, StdOrdering::SeqCst);
3335 tokio::time::sleep(Duration::from_millis(1)).await;
3336 active.fetch_sub(1, StdOrdering::SeqCst);
3337 Ok(acc + item)
3338 }
3339 })
3340 .run_collect()
3341 .unwrap();
3342 assert_eq!(scanned, vec![0, 1, 3, 6, 10]);
3343 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
3344
3345 let folded = Source::from_iter(1..=4)
3346 .fold_async(0, |acc, item| async move { Ok(acc + item) })
3347 .run_collect()
3348 .unwrap();
3349 assert_eq!(folded, vec![10]);
3350 }
3351
3352 #[test]
3353 fn wp6b_fold_async_materialization_does_not_drain_upstream() {
3354 let release = StdArc::new((std::sync::Mutex::new(false), std::sync::Condvar::new()));
3355 let started = StdArc::new(StdAtomicBool::new(false));
3356 let source = {
3357 let release = StdArc::clone(&release);
3358 let started = StdArc::clone(&started);
3359 Source::from_factory(move || {
3360 let release = StdArc::clone(&release);
3361 let started = StdArc::clone(&started);
3362 let mut emitted = false;
3363 Box::new(std::iter::from_fn(move || {
3364 if emitted {
3365 return None;
3366 }
3367 emitted = true;
3368 started.store(true, StdOrdering::SeqCst);
3369 let (released, available) = &*release;
3370 let mut released = released.lock().unwrap();
3371 while !*released {
3372 released = available.wait(released).unwrap();
3373 }
3374 Some(Ok(1))
3375 }))
3376 })
3377 };
3378
3379 let (materialized_tx, materialized_rx) = mpsc::channel();
3380 let join = thread::spawn(move || {
3381 let queue = source
3382 .fold_async(0, |acc, item| async move { Ok(acc + item) })
3383 .run_with(Sink::queue())
3384 .unwrap();
3385 materialized_tx.send(queue).unwrap();
3386 });
3387
3388 let queue = match materialized_rx.recv_timeout(StdDuration::from_secs(1)) {
3389 Ok(queue) => queue,
3390 Err(error) => {
3391 let (released, available) = &*release;
3392 *released.lock().unwrap() = true;
3393 available.notify_all();
3394 let _ = join.join();
3395 panic!("fold_async materialization did not return before first pull: {error}");
3396 }
3397 };
3398 let (released, _) = &*release;
3399 assert!(
3400 !*released.lock().unwrap(),
3401 "test source was released before materialization returned"
3402 );
3403
3404 let (released, available) = &*release;
3405 *released.lock().unwrap() = true;
3406 available.notify_all();
3407 assert_eq!(queue.pull().unwrap(), Some(1));
3408 assert_eq!(queue.pull().unwrap(), None);
3409 join.join().unwrap();
3410 assert!(started.load(StdOrdering::SeqCst));
3411 }
3412
3413 #[test]
3414 fn wp6b_lazy_sink_and_flow_wait_for_first_element() {
3415 let lazy_sink_created = StdArc::new(StdAtomicUsize::new(0));
3416 let lazy_sink_created_for_factory = StdArc::clone(&lazy_sink_created);
3417 let empty_sink = Source::<i32>::empty()
3418 .run_with(Sink::lazy_sink(move || {
3419 lazy_sink_created_for_factory.fetch_add(1, StdOrdering::SeqCst);
3420 Sink::ignore()
3421 }))
3422 .unwrap();
3423 assert!(matches!(empty_sink.wait(), Err(StreamError::Failed(_))));
3424 assert_eq!(lazy_sink_created.load(StdOrdering::SeqCst), 0);
3425
3426 let foreach_sum = StdArc::new(StdAtomicUsize::new(0));
3427 let foreach_sum_for_sink = StdArc::clone(&foreach_sum);
3428 Source::from_iter([1_usize, 2, 3])
3429 .run_with(Sink::foreach_async(2, move |item| {
3430 let foreach_sum = StdArc::clone(&foreach_sum_for_sink);
3431 async move {
3432 foreach_sum.fetch_add(item, StdOrdering::SeqCst);
3433 Ok(())
3434 }
3435 }))
3436 .unwrap()
3437 .wait()
3438 .unwrap();
3439 assert_eq!(foreach_sum.load(StdOrdering::SeqCst), 6);
3440
3441 let lazy_flow_created = StdArc::new(StdAtomicUsize::new(0));
3442 let lazy_flow_created_for_factory = StdArc::clone(&lazy_flow_created);
3443 let lazy_flow = Flow::<i32, i32>::lazy_flow(move || {
3444 lazy_flow_created_for_factory.fetch_add(1, StdOrdering::SeqCst);
3445 Flow::identity()
3446 .map(|item: i32| item + 10)
3447 .map_materialized_value(|_| 123)
3448 });
3449 let mat = (lazy_flow.materialize)().unwrap();
3450 let mut stream = match lazy_flow.transform {
3451 flow::FlowTransform::Runtime(transform) => {
3452 transform(Box::new([Ok(1), Ok(2)].into_iter()), &Materializer::new()).unwrap()
3453 }
3454 flow::FlowTransform::Pure(_) => panic!("lazy flow must be runtime-backed"),
3455 };
3456 assert_eq!(lazy_flow_created.load(StdOrdering::SeqCst), 0);
3457 assert_eq!(stream.next().unwrap().unwrap(), 11);
3458 assert_eq!(mat.wait().unwrap(), 123);
3459 assert_eq!(lazy_flow_created.load(StdOrdering::SeqCst), 1);
3460 assert_eq!(stream.next().unwrap().unwrap(), 12);
3461
3462 let future_flow = Source::from_iter([1, 2])
3463 .via_mat(
3464 Flow::future_flow(|| async {
3465 Ok(Flow::identity()
3466 .map(|item: i32| item * 2)
3467 .map_materialized_value(|_| 77))
3468 }),
3469 Keep::right,
3470 )
3471 .to_mat(Sink::collect(), Keep::both)
3472 .run()
3473 .unwrap();
3474 assert_eq!(future_flow.0.wait().unwrap(), 77);
3475 assert_eq!(future_flow.1.wait().unwrap(), vec![2, 4]);
3476 }
3477
3478 #[test]
3479 fn wp6b_lazy_flow_double_use_in_one_chain_pairs_instances_in_order() {
3480 for round in 0..50 {
3485 let counter = StdArc::new(StdAtomicUsize::new(1));
3486 let factory_counter = StdArc::clone(&counter);
3487 let lazy: Flow<usize, usize, _> = Flow::lazy_flow(move || {
3488 let id = factory_counter.fetch_add(1, StdOrdering::SeqCst);
3489 Flow::identity()
3490 .map(move |x: usize| x * 100 + id)
3491 .map_materialized_value(move |_| id)
3492 });
3493 let lazy_again = lazy.clone();
3494
3495 let ((first_mat, second_mat), out) = Source::from_iter([0usize])
3496 .via_mat(lazy, Keep::right)
3497 .via_mat(lazy_again, Keep::both)
3498 .to_mat(Sink::collect(), Keep::both)
3499 .run()
3500 .unwrap();
3501
3502 let first_id = first_mat.wait().unwrap();
3503 let second_id = second_mat.wait().unwrap();
3504 let element = out.wait().unwrap()[0];
3505 assert_eq!(
3506 element,
3507 first_id * 100 + second_id,
3508 "round {round}: mats ({first_id},{second_id}) cross-wired with transform order"
3509 );
3510 assert_ne!(
3511 first_id, second_id,
3512 "round {round}: same factory instance paired twice"
3513 );
3514 }
3515 }
3516
3517 #[test]
3518 fn wp6b_lazy_flow_clones_materialize_concurrently_without_cross_wiring() {
3519 for _ in 0..20 {
3520 let next_id = StdArc::new(StdAtomicUsize::new(0));
3521 let next_id_for_factory = StdArc::clone(&next_id);
3522 let flow = Flow::<i32, i32>::lazy_flow(move || {
3523 let id = next_id_for_factory.fetch_add(1, StdOrdering::SeqCst) + 1;
3524 Flow::identity()
3525 .map(move |item: i32| item + (id as i32 * 100))
3526 .map_materialized_value(move |_| id)
3527 });
3528 let barrier = StdArc::new(std::sync::Barrier::new(3));
3529
3530 let spawn_materialization = |input: i32| {
3531 let flow = flow.clone();
3532 let barrier = StdArc::clone(&barrier);
3533 thread::spawn(move || {
3534 barrier.wait();
3535 let (mat, values) = Source::single(input)
3536 .via_mat(flow, Keep::right)
3537 .to_mat(Sink::collect(), Keep::both)
3538 .run()
3539 .unwrap();
3540 (input, mat.wait().unwrap(), values.wait().unwrap())
3541 })
3542 };
3543
3544 let first = spawn_materialization(1);
3545 let second = spawn_materialization(2);
3546 barrier.wait();
3547
3548 for result in [first.join().unwrap(), second.join().unwrap()] {
3549 let (input, mat_id, values) = result;
3550 assert_eq!(values, vec![input + (mat_id as i32 * 100)]);
3551 }
3552 assert_eq!(next_id.load(StdOrdering::SeqCst), 2);
3553 }
3554 }
3555
3556 #[test]
3557 fn wp6b_map_with_resource_emits_close_item_before_terminal_error() {
3558 let queue = Source::from_factory(|| {
3559 Box::new(vec![Ok(1), Err(StreamError::Failed("upstream".into()))].into_iter())
3560 })
3561 .map_with_resource(
3562 || Ok(()),
3563 |_resource, item| Ok(item + 10),
3564 |_resource| Ok(Some(99)),
3565 )
3566 .run_with(Sink::queue())
3567 .unwrap();
3568
3569 assert_eq!(queue.pull().unwrap(), Some(11));
3570 assert_eq!(queue.pull().unwrap(), Some(99));
3571 assert_eq!(queue.pull(), Err(StreamError::Failed("upstream".into())));
3572
3573 let failed: StreamResult<Vec<i32>> = Source::single(1)
3574 .map_with_resource(
3575 || Ok(()),
3576 |_resource, _item| -> StreamResult<i32> { Err(StreamError::Failed("map".into())) },
3577 |_resource| -> StreamResult<Option<i32>> {
3578 Err(StreamError::Failed("close".into()))
3579 },
3580 )
3581 .run_collect();
3582 assert_eq!(failed, Err(StreamError::Failed("map".into())));
3583 }
3584
3585 #[test]
3586 fn stateful_and_terminal_source_operators_work() {
3587 let stateful = Source::from_iter([1, 2, 3])
3588 .stateful_map(0, |sum, item| {
3589 *sum += item;
3590 *sum
3591 })
3592 .run_collect()
3593 .unwrap();
3594 assert_eq!(stateful, vec![1, 3, 6]);
3595
3596 let concat = Source::from_iter([1, 2, 3])
3597 .stateful_map_concat(0, |sum, item| {
3598 *sum += item;
3599 [item, *sum]
3600 })
3601 .run_collect()
3602 .unwrap();
3603 assert_eq!(concat, vec![1, 1, 2, 3, 3, 6]);
3604
3605 assert_eq!(
3606 Source::from_iter([1, 2, 3])
3607 .fold(10, |acc, item| acc + item)
3608 .run_collect()
3609 .unwrap(),
3610 vec![16]
3611 );
3612 assert_eq!(
3613 Source::from_iter([1, 2, 3])
3614 .reduce(|acc, item| acc + item)
3615 .run_collect()
3616 .unwrap(),
3617 vec![6]
3618 );
3619 }
3620
3621 #[test]
3622 fn concat_and_sliding_emit_before_unbounded_upstream_finishes() {
3623 let concat = Source::single(())
3624 .map_concat(|_| 0_u64..)
3625 .take(1)
3626 .run_collect()
3627 .unwrap();
3628 assert_eq!(concat, vec![0]);
3629
3630 let sliding = Source::repeat(1_u64)
3631 .sliding(2, 1)
3632 .take(1)
3633 .run_collect()
3634 .unwrap();
3635 assert_eq!(sliding, vec![vec![1, 1]]);
3636 }
3637
3638 #[test]
3639 fn fan_in_source_operators_follow_ordering_rules() {
3640 assert_eq!(
3641 Source::from_iter([1, 2])
3642 .concat(Source::from_iter([3, 4]))
3643 .run_collect()
3644 .unwrap(),
3645 vec![1, 2, 3, 4]
3646 );
3647 assert_eq!(
3648 Source::from_iter([3, 4])
3649 .prepend(Source::from_iter([1, 2]))
3650 .run_collect()
3651 .unwrap(),
3652 vec![1, 2, 3, 4]
3653 );
3654 assert_eq!(
3655 Source::empty()
3656 .or_else(Source::from_iter([10, 20]))
3657 .run_collect()
3658 .unwrap(),
3659 vec![10, 20]
3660 );
3661 assert_eq!(
3662 Source::from_iter([1, 2])
3663 .or_else(Source::from_iter([10, 20]))
3664 .run_collect()
3665 .unwrap(),
3666 vec![1, 2]
3667 );
3668 assert_eq!(
3669 Source::from_iter([1, 2, 3])
3670 .interleave(Source::from_iter([10, 11, 12]), 2)
3671 .run_collect()
3672 .unwrap(),
3673 vec![1, 2, 10, 11, 3, 12]
3674 );
3675 }
3676
3677 #[test]
3678 fn fan_in_flow_operators_compose_with_primary_stream() {
3679 let concat = Source::from_iter([1, 2])
3680 .via(Flow::identity().concat(Source::from_iter([3, 4])))
3681 .run_collect()
3682 .unwrap();
3683 assert_eq!(concat, vec![1, 2, 3, 4]);
3684
3685 let prepend = Source::from_iter([3, 4])
3686 .via(Flow::identity().prepend(Source::from_iter([1, 2])))
3687 .run_collect()
3688 .unwrap();
3689 assert_eq!(prepend, vec![1, 2, 3, 4]);
3690
3691 let interleave = Source::from_iter([1, 2, 3])
3692 .via(Flow::identity().interleave(Source::from_iter([10, 11, 12]), 1))
3693 .run_collect()
3694 .unwrap();
3695 assert_eq!(interleave, vec![1, 10, 2, 11, 3, 12]);
3696
3697 let merge_sorted = Source::from_iter([1, 4])
3698 .via(Flow::identity().merge_sorted(Source::from_iter([2, 3, 5])))
3699 .run_collect()
3700 .unwrap();
3701 assert_eq!(merge_sorted, vec![1, 2, 3, 4, 5]);
3702
3703 let zip_latest = Source::from_iter([1, 2])
3704 .via(Flow::identity().zip_latest(Source::single(10)))
3705 .run_collect()
3706 .unwrap();
3707 assert_eq!(zip_latest, vec![(1, 10), (2, 10)]);
3708
3709 let zip_latest_with = Source::from_iter([1, 2])
3710 .via(
3711 Flow::identity()
3712 .zip_latest_with(Source::single(10), false, |left, right| left + right),
3713 )
3714 .run_collect()
3715 .unwrap();
3716 assert_eq!(zip_latest_with, vec![11, 12]);
3717 }
3718
3719 #[test]
3720 fn fan_in_operators_propagate_errors_and_eager_close() {
3721 assert!(matches!(
3722 Source::failed(StreamError::Failed("boom".into()))
3723 .or_else(Source::from_iter([1, 2]))
3724 .run_collect(),
3725 Err(StreamError::Failed(_))
3726 ));
3727 assert!(matches!(
3728 Source::from_iter([1, 2])
3729 .prepend(Source::failed(StreamError::Failed("boom".into())))
3730 .run_collect(),
3731 Err(StreamError::Failed(_))
3732 ));
3733 assert_eq!(
3734 Source::from_iter([1, 2])
3735 .interleave_all([Source::empty()], 1, true)
3736 .run_collect()
3737 .unwrap(),
3738 vec![1]
3739 );
3740 }
3741
3742 #[test]
3743 fn interleave_lazy_pulls_only_inputs_needed_for_first_segment() {
3744 use std::sync::{Arc, atomic::AtomicUsize, atomic::Ordering};
3745
3746 let pulls: Arc<[AtomicUsize; 3]> = Arc::new([
3747 AtomicUsize::new(0),
3748 AtomicUsize::new(0),
3749 AtomicUsize::new(0),
3750 ]);
3751
3752 let make_source = |idx: usize| {
3753 let pulls = Arc::clone(&pulls);
3754 Source::from_materialized_factory(move |_| {
3755 let pulls = Arc::clone(&pulls);
3756 let mut emitted = false;
3757 Ok((
3758 Box::new(std::iter::from_fn(move || {
3759 pulls[idx].fetch_add(1, Ordering::SeqCst);
3760 if !emitted && idx == 0 {
3761 emitted = true;
3762 Some(Ok(42))
3763 } else {
3764 None
3765 }
3766 })) as BoxStream<i32>,
3767 NotUsed,
3768 ))
3769 })
3770 };
3771
3772 let result = make_source(0)
3773 .interleave_all([make_source(1), make_source(2)], 1, false)
3774 .run_with(Sink::head());
3775
3776 assert_eq!(wait(result.unwrap()), 42);
3777 assert_eq!(pulls[0].load(Ordering::SeqCst), 1);
3778 assert_eq!(
3779 pulls[1].load(Ordering::SeqCst),
3780 0,
3781 "second input should not be pulled when downstream cancels after first element"
3782 );
3783 assert_eq!(
3784 pulls[2].load(Ordering::SeqCst),
3785 0,
3786 "third input should not be pulled before its turn"
3787 );
3788 }
3789
3790 #[test]
3791 fn interleave_non_eager_drains_remaining_when_one_input_completes() {
3792 assert_eq!(
3793 Source::from_iter([1, 2, 3, 4])
3794 .interleave_all(
3795 [Source::from_iter([10]), Source::from_iter([20, 21, 22])],
3796 1,
3797 false
3798 )
3799 .run_collect()
3800 .unwrap(),
3801 vec![1, 10, 20, 2, 21, 3, 22, 4]
3802 );
3803 }
3804
3805 #[test]
3806 fn remaining_merge_and_zip_family_matches_expected_ordering() {
3807 assert_eq!(
3808 Source::from_iter([1, 4])
3809 .merge_sorted(Source::from_iter([2, 3, 5]))
3810 .run_collect()
3811 .unwrap(),
3812 vec![1, 2, 3, 4, 5]
3813 );
3814
3815 assert_eq!(
3816 Source::from_iter([1, 2])
3817 .merge_latest(Source::single(10), false)
3818 .run_collect()
3819 .unwrap(),
3820 vec![vec![1, 10], vec![2, 10]]
3821 );
3822
3823 assert_eq!(
3824 Source::from_iter([1, 2, 3])
3825 .merge_all([Source::from_iter([10, 11])], false)
3826 .run_collect()
3827 .unwrap(),
3828 vec![1, 10, 2, 11, 3]
3829 );
3830
3831 assert_eq!(
3832 Source::from_iter([1, 2, 3])
3833 .zip_with(Source::from_iter([10, 11, 12]), |left, right| left + right)
3834 .run_collect()
3835 .unwrap(),
3836 vec![11, 13, 15]
3837 );
3838
3839 assert_eq!(
3840 Source::from_iter([1, 2])
3841 .zip_latest(Source::single(10))
3842 .run_collect()
3843 .unwrap(),
3844 vec![(1, 10), (2, 10)]
3845 );
3846
3847 assert_eq!(
3848 Source::from_iter([1, 2, 3])
3849 .zip_latest_with(Source::from_iter([10]), false, |left, right| left + right)
3850 .run_collect()
3851 .unwrap(),
3852 vec![11, 12, 13]
3853 );
3854
3855 assert_eq!(
3856 Source::from_iter([1, 2])
3857 .zip_all(Source::from_iter([10, 11, 12]), -1, -2)
3858 .run_collect()
3859 .unwrap(),
3860 vec![(1, 10), (2, 11), (-1, 12)]
3861 );
3862
3863 assert_eq!(
3864 Source::from_iter([5, 6, 7])
3865 .zip_with_index()
3866 .run_collect()
3867 .unwrap(),
3868 vec![(5, 0), (6, 1), (7, 2)]
3869 );
3870
3871 assert_eq!(
3872 Source::zip_n([Source::from_iter([1, 2]), Source::from_iter([10, 20])])
3873 .run_collect()
3874 .unwrap(),
3875 vec![vec![1, 10], vec![2, 20]]
3876 );
3877
3878 assert_eq!(
3879 Source::zip_with_n(
3880 [
3881 Source::from_iter([1, 2]),
3882 Source::from_iter([10, 20]),
3883 Source::from_iter([100, 200]),
3884 ],
3885 |values| values.into_iter().sum::<i32>(),
3886 )
3887 .run_collect()
3888 .unwrap(),
3889 vec![111, 222]
3890 );
3891
3892 assert_eq!(
3893 Source::merge_prioritized_n(
3894 [
3895 (Source::from_iter([1, 2, 3, 4]), 2),
3896 (Source::from_iter([10, 11]), 1),
3897 ],
3898 false,
3899 )
3900 .run_collect()
3901 .unwrap(),
3902 vec![1, 2, 10, 3, 4, 11]
3903 );
3904
3905 assert_eq!(
3906 Source::combine(
3907 Source::from_iter([1, 2, 3]),
3908 Source::from_iter([10, 11]),
3909 std::iter::empty::<Source<i32, NotUsed>>(),
3910 SourceCombineStrategy::Merge {
3911 eager_complete: false,
3912 },
3913 )
3914 .run_collect()
3915 .unwrap(),
3916 vec![1, 10, 2, 11, 3]
3917 );
3918
3919 let combined_sink = Sink::combine(
3920 Sink::ignore(),
3921 Sink::ignore(),
3922 std::iter::empty::<Sink<i32, NotUsed>>(),
3923 SinkCombineStrategy::Broadcast,
3924 );
3925 assert_eq!(
3926 Source::from_iter([1, 2, 3])
3927 .run_with(combined_sink)
3928 .unwrap(),
3929 NotUsed
3930 );
3931 }
3932
3933 #[test]
3934 fn sink_combine_broadcast_delivers_every_element_to_every_child() {
3935 let first_count = StdArc::new(StdAtomicUsize::new(0));
3939 let second_count = StdArc::new(StdAtomicUsize::new(0));
3940 let first_counter = StdArc::clone(&first_count);
3941 let second_counter = StdArc::clone(&second_count);
3942 let combined = Sink::combine(
3943 Sink::foreach(move |_: i32| {
3944 first_counter.fetch_add(1, StdOrdering::SeqCst);
3945 }),
3946 Sink::foreach(move |_: i32| {
3947 second_counter.fetch_add(1, StdOrdering::SeqCst);
3948 }),
3949 std::iter::empty::<Sink<i32, NotUsed>>(),
3950 SinkCombineStrategy::Broadcast,
3951 );
3952 assert_eq!(
3953 Source::from_iter(0..100).run_with(combined).unwrap(),
3954 NotUsed
3955 );
3956 assert!(wait_until(StdDuration::from_secs(1), || {
3961 first_count.load(StdOrdering::SeqCst) == 100
3962 && second_count.load(StdOrdering::SeqCst) == 100
3963 }));
3964 }
3965
3966 #[test]
3967 fn zip_latest_completes_when_one_side_finishes_without_emitting() {
3968 assert_eq!(
3972 Source::from_iter(std::iter::empty::<i32>())
3973 .zip_latest_with(Source::repeat(10), false, |left, right| left + right)
3974 .run_collect()
3975 .unwrap(),
3976 Vec::<i32>::new()
3977 );
3978 assert_eq!(
3979 Source::repeat(10)
3980 .zip_latest_with(
3981 Source::from_iter(std::iter::empty::<i32>()),
3982 false,
3983 |left, right| left + right,
3984 )
3985 .run_collect()
3986 .unwrap(),
3987 Vec::<i32>::new()
3988 );
3989 }
3990
3991 #[test]
3992 fn zip_family_completion_boundaries_match_expected_results() {
3993 assert_eq!(
3994 Source::from_iter([1, 2, 3])
3995 .zip_with(Source::from_iter([10]), |left, right| left + right)
3996 .run_collect()
3997 .unwrap(),
3998 vec![11]
3999 );
4000
4001 assert_eq!(
4002 Source::from_iter([1, 2, 3])
4003 .zip_latest_with(Source::from_iter([10]), true, |left, right| left + right)
4004 .run_collect()
4005 .unwrap(),
4006 vec![11, 12]
4007 );
4008
4009 assert_eq!(
4010 Source::zip_n([
4011 Source::from_iter([1, 2, 3]),
4012 Source::from_iter([10]),
4013 Source::from_iter([100, 200, 300]),
4014 ])
4015 .run_collect()
4016 .unwrap(),
4017 vec![vec![1, 10, 100]]
4018 );
4019 }
4020
4021 #[test]
4022 fn combine_strategies_follow_merge_concat_and_priority_rules() {
4023 assert_eq!(
4024 Source::combine(
4025 Source::from_iter([1, 2]),
4026 Source::from_iter([10, 11]),
4027 [Source::from_iter([100])],
4028 SourceCombineStrategy::Concat,
4029 )
4030 .run_collect()
4031 .unwrap(),
4032 vec![1, 2, 10, 11, 100]
4033 );
4034
4035 assert_eq!(
4036 Source::combine(
4037 Source::from_iter([1, 2, 3, 4]),
4038 Source::from_iter([10, 11]),
4039 std::iter::empty::<Source<i32, NotUsed>>(),
4040 SourceCombineStrategy::Prioritized {
4041 priorities: vec![2, 1],
4042 eager_complete: false,
4043 },
4044 )
4045 .run_collect()
4046 .unwrap(),
4047 vec![1, 2, 10, 3, 4, 11]
4048 );
4049 }
4050
4051 #[test]
4052 fn concat_lazy_defers_follow_on_source_until_needed() {
4053 let source_counter = StdArc::new(StdAtomicUsize::new(0));
4054 let source_counter_clone = StdArc::clone(&source_counter);
4055 let lazy_source = Source::from_materialized_factory(move |_| {
4056 source_counter_clone.fetch_add(1, StdOrdering::SeqCst);
4057 Ok((Box::new(std::iter::once(Ok(99))), NotUsed))
4058 });
4059 let source_head = Source::single(1)
4060 .concat_lazy(lazy_source)
4061 .run_with(Sink::head());
4062 assert_eq!(wait(source_head.unwrap()), 1);
4063 assert_eq!(source_counter.load(StdOrdering::SeqCst), 0);
4064
4065 let flow_counter = StdArc::new(StdAtomicUsize::new(0));
4066 let flow_counter_clone = StdArc::clone(&flow_counter);
4067 let lazy_flow_source = Source::from_materialized_factory(move |_| {
4068 flow_counter_clone.fetch_add(1, StdOrdering::SeqCst);
4069 Ok((Box::new(std::iter::once(Ok(99))), NotUsed))
4070 });
4071 let flow_head = Source::single(1)
4072 .via(Flow::identity().concat_lazy(lazy_flow_source))
4073 .run_with(Sink::head());
4074 assert_eq!(wait(flow_head.unwrap()), 1);
4075 assert_eq!(flow_counter.load(StdOrdering::SeqCst), 0);
4076 }
4077
4078 #[test]
4079 fn also_to_completes_when_side_sink_cancels() {
4080 assert_eq!(
4081 Source::from_iter([1, 2, 3])
4082 .also_to(Sink::cancelled())
4083 .run_collect()
4084 .unwrap(),
4085 Vec::<i32>::new()
4086 );
4087 assert_eq!(
4088 Source::from_iter([1, 2, 3])
4089 .also_to_all([Sink::cancelled(), Sink::cancelled()])
4090 .run_collect()
4091 .unwrap(),
4092 Vec::<i32>::new()
4093 );
4094 }
4095
4096 #[test]
4097 fn also_to_completes_gracefully_when_side_sink_disconnects() {
4098 let result = Source::from_iter(0..100)
4099 .also_to(Sink::head())
4100 .run_collect()
4101 .unwrap();
4102 assert!(!result.is_empty(), "main should emit at least one element");
4103 assert!(
4104 result.len() < 100,
4105 "main should complete early when side disconnects"
4106 );
4107 }
4108
4109 #[test]
4110 fn also_to_propagates_original_error_when_side_is_disconnected() {
4111 let err = StreamError::Failed("distinctive-boom".into());
4112 assert!(matches!(
4113 Source::<i32>::failed(err.clone())
4114 .also_to(Sink::cancelled())
4115 .run_collect(),
4116 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4117 ));
4118 assert!(matches!(
4119 Source::<i32>::failed(err.clone())
4120 .also_to_all([Sink::cancelled()])
4121 .run_collect(),
4122 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4123 ));
4124 assert!(matches!(
4125 Source::<i32>::failed(err)
4126 .divert_to(Sink::cancelled(), |_: &i32| true)
4127 .run_collect(),
4128 Err(StreamError::Failed(msg)) if msg == "distinctive-boom"
4129 ));
4130 }
4131
4132 #[test]
4133 fn divert_to_routes_matching_elements_to_side_sink() {
4134 let diverted = Source::from_iter([1, 2, 3, 4])
4135 .divert_to(Sink::ignore(), |item| item % 2 == 0)
4136 .run_collect()
4137 .unwrap();
4138 assert_eq!(diverted, vec![1, 3]);
4139 }
4140
4141 #[test]
4142 fn wire_tap_drops_when_side_sink_backpressures() {
4143 let tapped = Source::from_iter([1, 2, 3])
4144 .wire_tap(Sink::head())
4145 .run_collect()
4146 .unwrap();
4147 assert_eq!(tapped, vec![1, 2, 3]);
4148
4149 let tapped_via_flow = Source::from_iter([1, 2, 3])
4150 .via(Flow::identity().wire_tap(Sink::head()))
4151 .run_collect()
4152 .unwrap();
4153 assert_eq!(tapped_via_flow, vec![1, 2, 3]);
4154 }
4155
4156 #[test]
4157 fn async_mapping_variants_complete() {
4158 let ordered = Source::from_iter(0..4)
4159 .map_async(2, |item| async move { Ok(item * 2) })
4160 .run_collect()
4161 .unwrap();
4162 assert_eq!(ordered, vec![0, 2, 4, 6]);
4163
4164 let unordered = Source::from_iter(0..4)
4165 .map_async_unordered(2, |item| async move { Ok(item * 2) })
4166 .run_collect()
4167 .unwrap();
4168 assert_eq!(unordered, vec![0, 2, 4, 6]);
4169
4170 let partitioned = Source::from_iter(0..4)
4171 .map_async_partitioned(4, 1, |item| item % 2, |item| async move { Ok(item + 1) })
4172 .run_collect()
4173 .unwrap();
4174 assert_eq!(partitioned, vec![1, 2, 3, 4]);
4175 }
4176
4177 #[test]
4178 fn map_async_ordered_bounds_pulls_behind_stuck_head() {
4179 let pulls = StdArc::new(StdAtomicUsize::new(0));
4180 let pulls_for_source = StdArc::clone(&pulls);
4181 let probe = Source::from_fn_iter(move || {
4182 let pulls = StdArc::clone(&pulls_for_source);
4183 std::iter::from_fn(move || {
4184 let next = pulls.fetch_add(1, StdOrdering::SeqCst);
4185 Some(next)
4186 })
4187 })
4188 .map_async(2, |item| async move {
4189 if item == 0 {
4190 tokio::time::sleep(StdDuration::from_millis(300)).await;
4191 }
4192 Ok(item)
4193 })
4194 .run_with(TestSink::probe())
4195 .unwrap();
4196
4197 probe.request(16);
4198 thread::sleep(StdDuration::from_millis(100));
4199 assert!(
4200 pulls.load(StdOrdering::SeqCst) <= 3,
4201 "pulled {} elements with parallelism=2 behind a stuck ordered head",
4202 pulls.load(StdOrdering::SeqCst)
4203 );
4204 }
4205
4206 #[test]
4207 fn async_mapping_parks_until_woken_future_completes() {
4208 struct WakeOnceFuture {
4209 value: Option<u64>,
4210 ready: StdArc<StdAtomicBool>,
4211 polls: StdArc<StdAtomicUsize>,
4212 poll_tx: mpsc::Sender<usize>,
4213 latest_waker: StdArc<Mutex<Option<std::task::Waker>>>,
4214 }
4215
4216 impl std::future::Future for WakeOnceFuture {
4217 type Output = StreamResult<u64>;
4218
4219 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
4220 let this = self.as_mut().get_mut();
4221 let poll = this.polls.fetch_add(1, StdOrdering::SeqCst) + 1;
4222 if this.ready.load(StdOrdering::SeqCst) {
4223 let _ = this.poll_tx.send(poll);
4224 return Poll::Ready(Ok(this.value.take().unwrap()));
4225 }
4226
4227 *this.latest_waker.lock().expect("latest wake slot mutex") =
4228 Some(cx.waker().clone());
4229 let _ = this.poll_tx.send(poll);
4230 Poll::Pending
4231 }
4232 }
4233
4234 let ready = StdArc::new(StdAtomicBool::new(false));
4235 let polls = StdArc::new(StdAtomicUsize::new(0));
4236 let latest_waker = StdArc::new(Mutex::new(None));
4237 let (poll_tx, poll_rx) = mpsc::channel();
4238 let (result_tx, result_rx) = mpsc::channel();
4239 let ready_for_stage = StdArc::clone(&ready);
4240 let polls_for_stage = StdArc::clone(&polls);
4241 let poll_tx_for_stage = poll_tx.clone();
4242 let latest_waker_for_stage = StdArc::clone(&latest_waker);
4243 let worker = thread::spawn(move || {
4244 let result = Source::single(41)
4245 .map_async(1, move |item| WakeOnceFuture {
4246 value: Some(item + 1),
4247 ready: StdArc::clone(&ready_for_stage),
4248 polls: StdArc::clone(&polls_for_stage),
4249 poll_tx: poll_tx_for_stage.clone(),
4250 latest_waker: StdArc::clone(&latest_waker_for_stage),
4251 })
4252 .run_collect();
4253 let _ = result_tx.send(result);
4254 });
4255
4256 assert_eq!(poll_rx.recv_timeout(StdDuration::from_secs(10)).unwrap(), 1);
4257 assert!(
4258 matches!(
4259 result_rx.try_recv(),
4260 Err(std::sync::mpsc::TryRecvError::Empty)
4261 ),
4262 "stream completed before the future was woken"
4263 );
4264 assert_eq!(poll_rx.recv_timeout(StdDuration::from_secs(10)).unwrap(), 2);
4265 assert!(
4266 matches!(
4267 result_rx.try_recv(),
4268 Err(std::sync::mpsc::TryRecvError::Empty)
4269 ),
4270 "stream completed before the future was marked ready"
4271 );
4272
4273 ready.store(true, StdOrdering::SeqCst);
4274 latest_waker
4275 .lock()
4276 .expect("latest wake slot mutex")
4277 .take()
4278 .expect("pending future should have registered a waker")
4279 .wake();
4280
4281 let values = result_rx
4282 .recv_timeout(StdDuration::from_secs(10))
4283 .expect("stream should complete after the pending future is woken")
4284 .unwrap();
4285 worker.join().expect("stream worker should not panic");
4286 assert_eq!(values, vec![42]);
4287 assert_eq!(
4288 polls.load(StdOrdering::SeqCst),
4289 3,
4290 "pending future should be polled once inline, once to register the waker, and once after wake"
4291 );
4292 }
4293
4294 #[test]
4295 fn async_mapping_emits_before_unbounded_upstream_finishes() {
4296 let ordered = Source::repeat(1)
4297 .map_async(2, |item| async move { Ok(item + 1) })
4298 .take(1)
4299 .run_collect()
4300 .unwrap();
4301 assert_eq!(ordered, vec![2]);
4302
4303 let unordered = Source::repeat(1)
4304 .map_async_unordered(2, |item| async move { Ok(item + 1) })
4305 .take(1)
4306 .run_collect()
4307 .unwrap();
4308 assert_eq!(unordered, vec![2]);
4309
4310 let partitioned = Source::repeat(1)
4311 .map_async_partitioned(2, 1, |_| 0_u8, |item| async move { Ok(item + 1) })
4312 .take(1)
4313 .run_collect()
4314 .unwrap();
4315 assert_eq!(partitioned, vec![2]);
4316 }
4317
4318 #[test]
4319 fn partitioned_async_mapping_limits_same_key_concurrency() {
4320 let active = StdArc::new(StdAtomicUsize::new(0));
4321 let max_active = StdArc::new(StdAtomicUsize::new(0));
4322 let active_for_stage = StdArc::clone(&active);
4323 let max_for_stage = StdArc::clone(&max_active);
4324
4325 let values = Source::from_iter(0..6)
4326 .map_async_partitioned(
4327 4,
4328 1,
4329 |_| 0_u8,
4330 move |item| {
4331 let active = StdArc::clone(&active_for_stage);
4332 let max_active = StdArc::clone(&max_for_stage);
4333 let current = active.fetch_add(1, StdOrdering::SeqCst) + 1;
4334 max_active.fetch_max(current, StdOrdering::SeqCst);
4335 async move {
4336 thread::sleep(Duration::from_millis(1));
4337 active.fetch_sub(1, StdOrdering::SeqCst);
4338 Ok(item)
4339 }
4340 },
4341 )
4342 .run_collect()
4343 .unwrap();
4344
4345 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
4346 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
4347 }
4348
4349 #[test]
4350 fn partitioned_async_mapping_scans_past_blocked_pending_key() {
4351 let active = StdArc::new(StdAtomicUsize::new(0));
4352 let max_active = StdArc::new(StdAtomicUsize::new(0));
4353 let active_for_stage = StdArc::clone(&active);
4354 let max_for_stage = StdArc::clone(&max_active);
4355 let (release_tx, release_rx) = oneshot::channel::<()>();
4356 let release_rx = StdArc::new(std::sync::Mutex::new(Some(release_rx)));
4357 let release_rx_for_stage = StdArc::clone(&release_rx);
4358 let max_for_release = StdArc::clone(&max_active);
4359
4360 let releaser = thread::spawn(move || {
4361 let deadline = Instant::now() + StdDuration::from_secs(1);
4362 while max_for_release.load(StdOrdering::SeqCst) < 2 && Instant::now() < deadline {
4363 thread::yield_now();
4364 }
4365 let _ = release_tx.send(());
4366 });
4367
4368 let values = Source::from_iter([0, 2, 1])
4369 .map_async_partitioned(
4370 2,
4371 1,
4372 |item| item % 2,
4373 move |item| {
4374 let active = StdArc::clone(&active_for_stage);
4375 let max_active = StdArc::clone(&max_for_stage);
4376 let release_rx = StdArc::clone(&release_rx_for_stage);
4377 let current = active.fetch_add(1, StdOrdering::SeqCst) + 1;
4378 max_active.fetch_max(current, StdOrdering::SeqCst);
4379 async move {
4380 if item == 0 {
4381 let receiver = release_rx
4382 .lock()
4383 .expect("release receiver mutex")
4384 .take()
4385 .expect("release receiver present");
4386 let _ = receiver.await;
4387 }
4388 active.fetch_sub(1, StdOrdering::SeqCst);
4389 Ok(item)
4390 }
4391 },
4392 )
4393 .run_collect()
4394 .unwrap();
4395 releaser.join().unwrap();
4396
4397 assert_eq!(values, vec![0, 2, 1]);
4398 assert_eq!(max_active.load(StdOrdering::SeqCst), 2);
4399 }
4400
4401 #[test]
4402 fn partitioned_async_mapping_p1_still_evaluates_partition() {
4403 let partitions = StdArc::new(StdAtomicUsize::new(0));
4404 let partitions_for_stage = StdArc::clone(&partitions);
4405
4406 let values = Source::from_iter(0..8)
4407 .map_async_partitioned(
4408 1,
4409 1,
4410 move |item| {
4411 partitions_for_stage.fetch_add(1, StdOrdering::SeqCst);
4412 item % 2
4413 },
4414 |item| async move { Ok(item + 1) },
4415 )
4416 .run_collect()
4417 .unwrap();
4418
4419 assert_eq!(values, (1..9).collect::<Vec<_>>());
4420 assert_eq!(partitions.load(StdOrdering::SeqCst), 8);
4421 }
4422
4423 #[test]
4424 fn partitioned_async_mapping_handles_many_keys_high_parallelism() {
4425 let active_by_key =
4426 StdArc::new((0..16).map(|_| StdAtomicUsize::new(0)).collect::<Vec<_>>());
4427 let max_by_key = StdArc::new((0..16).map(|_| StdAtomicUsize::new(0)).collect::<Vec<_>>());
4428 let active_for_stage = StdArc::clone(&active_by_key);
4429 let max_for_stage = StdArc::clone(&max_by_key);
4430
4431 let values = Source::from_iter(0..512_usize)
4432 .map_async_partitioned(
4433 32,
4434 1,
4435 |item| item % 16,
4436 move |item| {
4437 let active = StdArc::clone(&active_for_stage);
4438 let max_active = StdArc::clone(&max_for_stage);
4439 let key = item % 16;
4440 let current = active[key].fetch_add(1, StdOrdering::SeqCst) + 1;
4441 max_active[key].fetch_max(current, StdOrdering::SeqCst);
4442 async move {
4443 active[key].fetch_sub(1, StdOrdering::SeqCst);
4444 Ok(item)
4445 }
4446 },
4447 )
4448 .run_collect()
4449 .unwrap();
4450
4451 assert_eq!(values, (0..512).collect::<Vec<_>>());
4452 for max_active in max_by_key.iter() {
4453 assert_eq!(max_active.load(StdOrdering::SeqCst), 1);
4454 }
4455 }
4456
4457 #[test]
4458 fn error_operators_map_recover_and_complete() {
4459 let mapped = Source::<i32>::failed(StreamError::Failed("boom".into()))
4460 .map_error(|_| StreamError::Failed("mapped".into()))
4461 .run_collect();
4462 assert_eq!(mapped, Err(StreamError::Failed("mapped".into())));
4463
4464 let recovered = Source::<i32>::failed(StreamError::Failed("boom".into()))
4465 .recover(|error| match error {
4466 StreamError::Failed(_) => Some(42),
4467 _ => None,
4468 })
4469 .run_collect()
4470 .unwrap();
4471 assert_eq!(recovered, vec![42]);
4472
4473 let unrecovered = Source::<i32>::failed(StreamError::Failed("original".into()))
4474 .recover(|_| None)
4475 .run_collect();
4476 assert_eq!(unrecovered, Err(StreamError::Failed("original".into())));
4477
4478 let recovered_with = Source::<i32>::failed(StreamError::Failed("boom".into()))
4479 .recover_with_retries(1, |_| Some(Source::from_iter([1, 2])))
4480 .run_collect()
4481 .unwrap();
4482 assert_eq!(recovered_with, vec![1, 2]);
4483
4484 let declined_recover_with = Source::<i32>::failed(StreamError::Failed("declined".into()))
4485 .recover_with_retries(1, |_| None)
4486 .run_collect();
4487 assert_eq!(
4488 declined_recover_with,
4489 Err(StreamError::Failed("declined".into()))
4490 );
4491
4492 let completed = Source::from_factory(|| {
4493 Box::new(vec![Ok(1), Err(StreamError::Failed("ignored".into())), Ok(2)].into_iter())
4494 })
4495 .on_error_complete()
4496 .run_collect()
4497 .unwrap();
4498 assert_eq!(completed, vec![1]);
4499 }
4500
4501 #[test]
4502 fn sliding_matches_akka_window_semantics() {
4503 assert_eq!(
4505 Source::from_iter(1..=4)
4506 .sliding(3, 1)
4507 .run_collect()
4508 .unwrap(),
4509 vec![vec![1, 2, 3], vec![2, 3, 4]]
4510 );
4511 assert_eq!(
4512 Source::from_iter(1..=4)
4513 .sliding(2, 1)
4514 .run_collect()
4515 .unwrap(),
4516 vec![vec![1, 2], vec![2, 3], vec![3, 4]]
4517 );
4518 assert_eq!(
4520 Source::from_iter(1..=3)
4521 .sliding(3, 1)
4522 .run_collect()
4523 .unwrap(),
4524 vec![vec![1, 2, 3]]
4525 );
4526 assert_eq!(
4528 Source::from_iter(1..=2)
4529 .sliding(3, 1)
4530 .run_collect()
4531 .unwrap(),
4532 vec![vec![1, 2]]
4533 );
4534 assert_eq!(
4536 Source::from_iter(1..=3)
4537 .sliding(1, 1)
4538 .run_collect()
4539 .unwrap(),
4540 vec![vec![1], vec![2], vec![3]]
4541 );
4542 assert_eq!(
4544 Source::from_iter(1..=6)
4545 .sliding(2, 3)
4546 .run_collect()
4547 .unwrap(),
4548 vec![vec![1, 2], vec![4, 5]]
4549 );
4550 assert_eq!(
4552 Source::from_iter(1..=3)
4553 .sliding(2, 4)
4554 .run_collect()
4555 .unwrap(),
4556 vec![vec![1, 2]]
4557 );
4558 }
4559
4560 #[test]
4561 fn recover_with_retries_indefinitely_like_akka() {
4562 let attempts = StdArc::new(StdAtomicUsize::new(0));
4563 let attempts_in_stage = StdArc::clone(&attempts);
4564 let recovered = Source::<i32>::failed(StreamError::Failed("boom".into()))
4567 .recover_with(move |_error| {
4568 if attempts_in_stage.fetch_add(1, StdOrdering::SeqCst) < 5 {
4569 Some(Source::<i32>::failed(StreamError::Failed("again".into())))
4570 } else {
4571 Some(Source::from_iter([42]))
4572 }
4573 })
4574 .run_collect()
4575 .unwrap();
4576 assert_eq!(recovered, vec![42]);
4577 assert_eq!(attempts.load(StdOrdering::SeqCst), 6);
4578 }
4579
4580 #[test]
4581 fn many_concurrent_streams_do_not_starve_the_pool() {
4582 let materializer = Materializer::new();
4591 let busy = 6_usize;
4592
4593 let mut held = Vec::with_capacity(busy);
4594 for _ in 0..busy {
4595 held.push(
4596 Source::single(1_u64)
4597 .run_with_materializer(Sink::never(), &materializer)
4598 .unwrap(),
4599 );
4600 }
4601
4602 for _ in 0..400 {
4603 if materializer.active_streams() >= busy {
4604 break;
4605 }
4606 thread::sleep(Duration::from_millis(5));
4607 }
4608 assert_eq!(materializer.active_streams(), busy);
4609
4610 let sum = Source::from_iter(0_u64..5)
4613 .run_with_materializer(Sink::fold(0_u64, |acc, item| acc + item), &materializer)
4614 .unwrap();
4615 assert_eq!(sum.wait().unwrap(), 10);
4616
4617 materializer.shutdown();
4618 for completion in held {
4619 assert_eq!(completion.wait(), Err(StreamError::AbruptTermination));
4620 }
4621 }
4622}