1use super::*;
10use crate::Attributes;
11use crate::context::FlowWithContext;
12use crate::stream::error::{decide_supervision, panic_stream_error};
13use futures::{FutureExt, task::noop_waker};
14use std::any::Any;
15use std::task::Context;
16use std::{
17 collections::{HashMap, HashSet},
18 thread,
19};
20
21pub(super) enum FlowTransform<In, Out> {
22 Pure(PureTransform<In, Out>),
23 Runtime(RuntimeTransform<In, Out>),
24}
25
26pub struct Flow<In, Out, Mat = NotUsed> {
27 pub(super) transform: FlowTransform<In, Out>,
28 pub(super) materialize: Arc<dyn Fn() -> StreamResult<Mat> + Send + Sync>,
29 pub(super) hints: FlowHints,
30 pub(super) attributes: Attributes,
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
34pub(super) enum GroupByBatchMode {
35 Immediate,
36 FiniteEagerNoRecreate,
37}
38
39#[derive(Clone)]
40pub struct BidiFlow<I1, O1, I2, O2> {
41 top: Flow<I1, O1, NotUsed>,
42 bottom: Flow<I2, O2, NotUsed>,
43 attributes: Attributes,
44}
45
46impl<In, Out> Clone for FlowTransform<In, Out> {
47 fn clone(&self) -> Self {
48 match self {
49 Self::Pure(transform) => Self::Pure(Arc::clone(transform)),
50 Self::Runtime(transform) => Self::Runtime(Arc::clone(transform)),
51 }
52 }
53}
54
55impl<In, Out, Mat> Clone for Flow<In, Out, Mat> {
56 fn clone(&self) -> Self {
57 Self {
58 transform: self.transform.clone(),
59 materialize: Arc::clone(&self.materialize),
60 hints: self.hints,
61 attributes: self.attributes.clone(),
62 }
63 }
64}
65
66fn call_supervised<T, F>(context: &str, f: F) -> StreamResult<T>
67where
68 F: FnOnce() -> StreamResult<T>,
69{
70 catch_unwind(AssertUnwindSafe(f)).unwrap_or_else(|_| Err(panic_stream_error(context)))
71}
72
73impl<T: Send + 'static> Flow<T, T, NotUsed> {
74 #[must_use]
75 pub fn identity() -> Self {
76 Self::from_preserving_transform(|input| input)
77 }
78}
79
80impl<In: Send + 'static, Out: Send + 'static> Flow<In, Out, NotUsed> {
81 pub(crate) fn from_transform<F>(transform: F) -> Self
82 where
83 F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
84 {
85 Self::from_parts_with_hints(transform, || Ok(NotUsed), FlowHints::default())
86 }
87
88 pub(crate) fn from_preserving_transform<F>(transform: F) -> Self
89 where
90 F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
91 {
92 Self::from_parts_with_hints(
93 transform,
94 || Ok(NotUsed),
95 FlowHints::PRESERVES_INLINE_HEAD_TERMINAL,
96 )
97 }
98
99 pub(crate) fn from_runtime_transform<F>(transform: F) -> Self
100 where
101 F: Fn(BoxStream<In>, &Materializer) -> StreamResult<BoxStream<Out>> + Send + Sync + 'static,
102 {
103 Self::from_runtime_transform_with_hints(transform, FlowHints::default())
104 }
105
106 fn from_runtime_transform_with_hints<F>(transform: F, hints: FlowHints) -> Self
107 where
108 F: Fn(BoxStream<In>, &Materializer) -> StreamResult<BoxStream<Out>> + Send + Sync + 'static,
109 {
110 Self {
111 transform: FlowTransform::Runtime(Arc::new(transform)),
112 materialize: Arc::new(|| Ok(NotUsed)),
113 hints,
114 attributes: Attributes::default(),
115 }
116 }
117
118 #[must_use]
119 pub fn from_sink_and_source<InMat, OutMat>(
120 sink: Sink<In, InMat>,
121 source: Source<Out, OutMat>,
122 ) -> Self
123 where
124 InMat: Send + 'static,
125 OutMat: Send + 'static,
126 {
127 Self::from_runtime_transform(move |input, materializer| {
128 let sink_keepalive = Arc::new(MaterializedKeepalive::default());
129 let sink_input = Box::new(InputKeepaliveStream {
130 inner: input,
131 keepalive: Arc::clone(&sink_keepalive),
132 peer_keepalive: None,
133 });
134 let sink_mat = sink.clone().run(sink_input, materializer)?;
135 sink_keepalive.store(Box::new(sink_mat));
136 let (output, source_mat) = Arc::clone(&source.factory).create(materializer)?;
137 let source_keepalive = Arc::new(MaterializedKeepalive::default());
138 source_keepalive.store(Box::new(source_mat));
139 Ok(Box::new(CoupledStream {
140 inner: output,
141 source_keepalive,
142 sink_keepalive: None,
143 coupled: false,
144 }))
145 })
146 }
147
148 #[must_use]
149 pub fn from_sink_and_source_coupled<InMat, OutMat>(
150 sink: Sink<In, InMat>,
151 source: Source<Out, OutMat>,
152 ) -> Self
153 where
154 InMat: Send + 'static,
155 OutMat: Send + 'static,
156 {
157 Self::from_runtime_transform(move |input, materializer| {
158 let source_keepalive = Arc::new(MaterializedKeepalive::default());
159 let sink_keepalive = Arc::new(MaterializedKeepalive::default());
160 let sink_input = Box::new(InputKeepaliveStream {
161 inner: input,
162 keepalive: Arc::clone(&sink_keepalive),
163 peer_keepalive: Some(Arc::clone(&source_keepalive)),
164 });
165 let sink_mat = sink.clone().run(sink_input, materializer)?;
166 sink_keepalive.store(Box::new(sink_mat));
167 let (output, source_mat) = Arc::clone(&source.factory).create(materializer)?;
168 source_keepalive.store(Box::new(source_mat));
169 Ok(Box::new(CoupledStream {
170 inner: output,
171 source_keepalive,
172 sink_keepalive: Some(sink_keepalive),
173 coupled: true,
174 }))
175 })
176 }
177 #[must_use]
178 pub fn future_flow<InnerMat, F, Fut>(future: F) -> Flow<In, Out, StreamCompletion<InnerMat>>
179 where
180 InnerMat: Send + 'static,
181 F: Fn() -> Fut + Send + Sync + 'static,
182 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
183 {
184 let future = Arc::new(future);
185 Flow::from_runtime_materialized_factory(move || {
186 let (sender, receiver) = oneshot::channel();
187 let sender = Arc::new(Mutex::new(Some(sender)));
188 let future = Arc::clone(&future);
189 let transform: RuntimeTransform<In, Out> =
190 Arc::new(move |input, materializer: &Materializer| {
191 let mat_sender = sender
192 .lock()
193 .expect("future_flow materialized sender poisoned")
194 .take()
195 .ok_or_else(|| {
196 StreamError::Failed("future_flow transform already materialized".into())
197 })?;
198 Ok(Box::new(FutureFlowStream {
199 future: Arc::clone(&future),
200 materializer: materializer
201 .with_name_prefix(materializer.name_prefix().to_owned()),
202 input: Some(input),
203 current: None,
204 mat_sender: Some(mat_sender),
205 initialized: false,
206 terminated: false,
207 _marker: PhantomData,
208 }) as BoxStream<Out>)
209 });
210 (transform, StreamCompletion::from_receiver(receiver, None))
211 })
212 }
213
214 #[must_use]
215 pub fn lazy_flow<InnerMat, F>(create: F) -> Flow<In, Out, StreamCompletion<InnerMat>>
216 where
217 InnerMat: Send + 'static,
218 F: Fn() -> Flow<In, Out, InnerMat> + Send + Sync + 'static,
219 {
220 let create = Arc::new(create);
221 Self::lazy_future_flow(move || {
222 let create = Arc::clone(&create);
223 async move { catch_unwind_failed("lazy_flow factory", || create()) }
224 })
225 }
226
227 #[must_use]
228 pub fn lazy_future_flow<InnerMat, F, Fut>(
229 create: F,
230 ) -> Flow<In, Out, StreamCompletion<InnerMat>>
231 where
232 InnerMat: Send + 'static,
233 F: Fn() -> Fut + Send + Sync + 'static,
234 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
235 {
236 let create = Arc::new(create);
237 Flow::from_runtime_materialized_factory(move || {
238 let (sender, receiver) = oneshot::channel();
239 let sender = Arc::new(Mutex::new(Some(sender)));
240 let create = Arc::clone(&create);
241 let transform: RuntimeTransform<In, Out> =
242 Arc::new(move |input, materializer: &Materializer| {
243 let mat_sender = sender
244 .lock()
245 .expect("lazy_future_flow materialized sender poisoned")
246 .take()
247 .ok_or_else(|| {
248 StreamError::Failed(
249 "lazy_future_flow transform already materialized".into(),
250 )
251 })?;
252 Ok(Box::new(LazyFutureFlowStream {
253 create: Arc::clone(&create),
254 materializer: materializer
255 .with_name_prefix(materializer.name_prefix().to_owned()),
256 input: Some(input),
257 current: None,
258 mat_sender: Some(mat_sender),
259 initialized: false,
260 terminated: false,
261 _marker: PhantomData,
262 }) as BoxStream<Out>)
263 });
264 (transform, StreamCompletion::from_receiver(receiver, None))
265 })
266 }
267}
268
269#[derive(Default)]
270struct MaterializedKeepalive {
271 released: AtomicBool,
272 value: Mutex<Option<Box<dyn Any + Send>>>,
273}
274
275impl MaterializedKeepalive {
276 fn store(&self, value: Box<dyn Any + Send>) {
277 let mut slot = self.value.lock().expect("materialized keepalive poisoned");
278 if !self.released.load(Ordering::SeqCst) {
279 *slot = Some(value);
280 return;
281 }
282 drop(slot);
287 release_materialized_value(value);
288 }
289
290 fn release(&self) {
291 self.released.store(true, Ordering::SeqCst);
292 if let Some(value) = self
293 .value
294 .lock()
295 .expect("materialized keepalive poisoned")
296 .take()
297 {
298 release_materialized_value(value);
299 }
300 }
301
302 fn is_released(&self) -> bool {
303 self.released.load(Ordering::SeqCst)
304 }
305}
306
307fn release_materialized_value(value: Box<dyn Any + Send>) {
308 match value.downcast::<Cancellable>() {
309 Ok(cancellable) => {
310 cancellable.cancel();
311 }
312 Err(value) => drop(value),
313 }
314}
315
316struct InputKeepaliveStream<In> {
317 inner: BoxStream<In>,
318 keepalive: Arc<MaterializedKeepalive>,
319 peer_keepalive: Option<Arc<MaterializedKeepalive>>,
320}
321
322impl<In> Iterator for InputKeepaliveStream<In> {
323 type Item = StreamResult<In>;
324
325 fn next(&mut self) -> Option<Self::Item> {
326 self.inner.next()
327 }
328}
329
330impl<In> Drop for InputKeepaliveStream<In> {
331 fn drop(&mut self) {
332 self.keepalive.release();
333 if let Some(peer_keepalive) = &self.peer_keepalive {
334 peer_keepalive.release();
335 }
336 }
337}
338
339struct CoupledStream<Out> {
340 inner: BoxStream<Out>,
341 source_keepalive: Arc<MaterializedKeepalive>,
342 sink_keepalive: Option<Arc<MaterializedKeepalive>>,
343 coupled: bool,
344}
345
346impl<Out> Iterator for CoupledStream<Out> {
347 type Item = StreamResult<Out>;
348
349 fn next(&mut self) -> Option<Self::Item> {
350 if self.coupled && self.source_keepalive.is_released() {
351 return None;
352 }
353 let next = self.inner.next();
354 if next.is_none() || next.as_ref().is_some_and(|item| item.is_err()) {
355 self.source_keepalive.release();
356 if self.coupled
357 && let Some(sink_keepalive) = &self.sink_keepalive
358 {
359 sink_keepalive.release();
360 }
361 }
362 next
363 }
364}
365
366impl<Out> Drop for CoupledStream<Out> {
367 fn drop(&mut self) {
368 self.source_keepalive.release();
369 if self.coupled
370 && let Some(sink_keepalive) = &self.sink_keepalive
371 {
372 sink_keepalive.release();
373 }
374 }
375}
376
377impl<In: Send + 'static, Out: Send + 'static, Mat: Send + 'static> Flow<In, Out, Mat> {
378 pub fn as_flow_with_context<U, CtxIn, CtxOut, Collapse, Extract>(
379 self,
380 collapse_context: Collapse,
381 extract_context: Extract,
382 ) -> FlowWithContext<U, CtxIn, Out, CtxOut, Mat>
383 where
384 U: Send + 'static,
385 CtxIn: Send + 'static,
386 CtxOut: Send + 'static,
387 Collapse: Fn(U, CtxIn) -> In + Send + Sync + 'static,
388 Extract: Fn(&Out) -> CtxOut + Send + Sync + 'static,
389 {
390 FlowWithContext::from_flow(
391 Flow::identity()
392 .map(move |(value, context)| collapse_context(value, context))
393 .via_mat(self, |_, flow_mat| flow_mat)
394 .map(move |value| {
395 let context = extract_context(&value);
396 (value, context)
397 }),
398 )
399 }
400
401 pub(crate) fn from_parts<F, M>(transform: F, materialize: M) -> Self
402 where
403 F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
404 M: Fn() -> StreamResult<Mat> + Send + Sync + 'static,
405 {
406 Self::from_parts_with_hints(transform, materialize, FlowHints::default())
407 }
408
409 pub(crate) fn from_materialized_factory<F>(factory: F) -> Self
410 where
411 F: Fn() -> (PureTransform<In, Out>, Mat) + Send + Sync + 'static,
412 {
413 struct PendingSlot<In, Out, Mat> {
414 transform: Option<PureTransform<In, Out>>,
415 mat: Option<Mat>,
416 transform_taken: bool,
417 mat_taken: bool,
418 }
419
420 let pending = Arc::new(Mutex::new(HashMap::<
421 thread::ThreadId,
422 Vec<PendingSlot<In, Out, Mat>>,
423 >::new()));
424 let factory = Arc::new(factory);
425
426 let transform = {
427 let pending = Arc::clone(&pending);
428 let factory = Arc::clone(&factory);
429 move |input| {
430 let transform = {
431 let mut pending = pending.lock().expect("flow materialized factory poisoned");
432 let thread_id = thread::current().id();
433 let slots = pending.entry(thread_id).or_default();
434 let index = slots
435 .iter()
436 .position(|slot| !slot.transform_taken && slot.mat_taken)
437 .unwrap_or_else(|| {
438 let (transform, mat) = factory();
439 slots.push(PendingSlot {
440 transform: Some(transform),
441 mat: Some(mat),
442 transform_taken: false,
443 mat_taken: false,
444 });
445 slots.len() - 1
446 });
447 let slot = slots
448 .get_mut(index)
449 .expect("pending flow materialization slot exists");
450 slot.transform_taken = true;
451 let transform = slot
452 .transform
453 .take()
454 .expect("pending flow transform present");
455 if slot.transform_taken && slot.mat_taken {
456 slots.remove(index);
457 }
458 if slots.is_empty() {
459 pending.remove(&thread_id);
460 }
461 transform
462 };
463 transform(input)
464 }
465 };
466
467 let materialize = {
468 let pending = Arc::clone(&pending);
469 let factory = Arc::clone(&factory);
470 move || {
471 let mat = {
472 let mut pending = pending.lock().expect("flow materialized factory poisoned");
473 let thread_id = thread::current().id();
474 let slots = pending.entry(thread_id).or_default();
475 let index = slots
476 .iter()
477 .position(|slot| !slot.mat_taken && slot.transform_taken)
478 .unwrap_or_else(|| {
479 let (transform, mat) = factory();
480 slots.push(PendingSlot {
481 transform: Some(transform),
482 mat: Some(mat),
483 transform_taken: false,
484 mat_taken: false,
485 });
486 slots.len() - 1
487 });
488 let slot = slots
489 .get_mut(index)
490 .expect("pending flow materialization slot exists");
491 slot.mat_taken = true;
492 let mat = slot
493 .mat
494 .take()
495 .expect("pending flow materialized value present");
496 if slot.transform_taken && slot.mat_taken {
497 slots.remove(index);
498 }
499 if slots.is_empty() {
500 pending.remove(&thread_id);
501 }
502 mat
503 };
504 Ok(mat)
505 }
506 };
507
508 Self::from_parts_with_hints(transform, materialize, FlowHints::default())
509 }
510
511 pub(crate) fn from_runtime_materialized_factory<F>(factory: F) -> Self
512 where
513 F: Fn() -> (RuntimeTransform<In, Out>, Mat) + Send + Sync + 'static,
514 {
515 struct PendingSlot<In, Out, Mat> {
516 transform: Option<RuntimeTransform<In, Out>>,
517 mat: Option<Mat>,
518 transform_taken: bool,
519 mat_taken: bool,
520 }
521
522 let pending = Arc::new(Mutex::new(HashMap::<
523 thread::ThreadId,
524 Vec<PendingSlot<In, Out, Mat>>,
525 >::new()));
526 let factory = Arc::new(factory);
527
528 let transform = {
529 let pending = Arc::clone(&pending);
530 let factory = Arc::clone(&factory);
531 move |input, materializer: &Materializer| {
532 let thread_id = thread::current().id();
533 let transform = {
534 let mut pending = pending.lock().expect("flow materialized factory poisoned");
535 let slots = pending.entry(thread_id).or_default();
536 if let Some(index) = slots
537 .iter()
538 .position(|slot| !slot.transform_taken && slot.mat_taken)
539 {
540 let slot = slots
541 .get_mut(index)
542 .expect("pending flow materialization slot exists");
543 slot.transform_taken = true;
544 let transform = slot
545 .transform
546 .take()
547 .expect("pending flow transform present");
548 if slot.transform_taken && slot.mat_taken {
549 slots.remove(index);
550 }
551 if slots.is_empty() {
552 pending.remove(&thread_id);
553 }
554 Some(transform)
555 } else {
556 None
557 }
558 };
559
560 let transform = match transform {
561 Some(transform) => transform,
562 None => {
563 let (transform, mat) = factory();
564 let mut pending =
565 pending.lock().expect("flow materialized factory poisoned");
566 let slots = pending.entry(thread_id).or_default();
567 slots.push(PendingSlot {
568 transform: None,
569 mat: Some(mat),
570 transform_taken: true,
571 mat_taken: false,
572 });
573 transform
574 }
575 };
576
577 transform(input, materializer)
578 }
579 };
580
581 let materialize = {
582 let pending = Arc::clone(&pending);
583 let factory = Arc::clone(&factory);
584 move || {
585 let thread_id = thread::current().id();
586 let mat = {
587 let mut pending = pending.lock().expect("flow materialized factory poisoned");
588 let slots = pending.entry(thread_id).or_default();
589 if let Some(index) = slots
590 .iter()
591 .position(|slot| !slot.mat_taken && slot.transform_taken)
592 {
593 let slot = slots
594 .get_mut(index)
595 .expect("pending flow materialization slot exists");
596 slot.mat_taken = true;
597 let mat = slot
598 .mat
599 .take()
600 .expect("pending flow materialized value present");
601 if slot.transform_taken && slot.mat_taken {
602 slots.remove(index);
603 }
604 if slots.is_empty() {
605 pending.remove(&thread_id);
606 }
607 Some(mat)
608 } else {
609 None
610 }
611 };
612
613 match mat {
614 Some(mat) => Ok(mat),
615 None => {
616 let (transform, mat) = factory();
617 let mut pending =
618 pending.lock().expect("flow materialized factory poisoned");
619 let slots = pending.entry(thread_id).or_default();
620 slots.push(PendingSlot {
621 transform: Some(transform),
622 mat: None,
623 transform_taken: false,
624 mat_taken: true,
625 });
626 Ok(mat)
627 }
628 }
629 }
630 };
631
632 Flow {
633 transform: FlowTransform::Runtime(Arc::new(transform)),
634 materialize: Arc::new(materialize),
635 hints: FlowHints::default(),
636 attributes: Attributes::default(),
637 }
638 }
639
640 fn from_parts_with_hints<F, M>(transform: F, materialize: M, hints: FlowHints) -> Self
641 where
642 F: Fn(BoxStream<In>) -> BoxStream<Out> + Send + Sync + 'static,
643 M: Fn() -> StreamResult<Mat> + Send + Sync + 'static,
644 {
645 Self {
646 transform: FlowTransform::Pure(Arc::new(transform)),
647 materialize: Arc::new(materialize),
648 hints,
649 attributes: Attributes::default(),
650 }
651 }
652
653 #[must_use]
654 pub fn attributes(&self) -> &Attributes {
655 &self.attributes
656 }
657
658 #[must_use]
659 pub fn with_attributes(mut self, attributes: Attributes) -> Self {
660 self.attributes = attributes;
661 self
662 }
663
664 #[must_use]
665 pub fn add_attributes(mut self, attributes: Attributes) -> Self {
666 self.attributes = self.attributes.and(attributes);
667 self
668 }
669
670 #[must_use]
671 pub fn named(self, name: impl Into<String>) -> Self {
672 self.add_attributes(Attributes::named(name))
673 }
674
675 #[must_use]
682 pub fn async_boundary(self) -> Self {
683 self.async_boundary_with_config(crate::graph::AsyncBoundaryExecutionConfig::default())
684 }
685
686 #[must_use]
688 pub fn r#async(self) -> Self {
689 self.async_boundary()
690 }
691
692 #[must_use]
698 pub fn async_boundary_with_config(
699 self,
700 config: crate::graph::AsyncBoundaryExecutionConfig,
701 ) -> Self {
702 self.via(Flow::from_runtime_transform(move |input, materializer| {
703 super::async_boundary::linear_async_boundary_stream(input, materializer, config)
704 }))
705 }
706
707 #[must_use]
709 pub fn async_boundary_with_buffer(self, buffer_size: usize) -> Self {
710 self.async_boundary_with_config(crate::graph::AsyncBoundaryExecutionConfig {
711 buffer_size,
712 ..crate::graph::AsyncBoundaryExecutionConfig::default()
713 })
714 }
715
716 #[must_use]
717 pub fn via<Next, NextMat>(self, next: Flow<Out, Next, NextMat>) -> Flow<In, Next, Mat>
718 where
719 Next: Send + 'static,
720 NextMat: Send + 'static,
721 {
722 self.via_mat(next, Keep::left)
723 }
724
725 #[must_use]
726 pub fn via_mat<Next, NextMat, Combined, F>(
727 self,
728 next: Flow<Out, Next, NextMat>,
729 combine: F,
730 ) -> Flow<In, Next, Combined>
731 where
732 Next: Send + 'static,
733 NextMat: Send + 'static,
734 Combined: Send + 'static,
735 F: Fn(Mat, NextMat) -> Combined + Send + Sync + 'static,
736 {
737 let Flow {
738 transform: first,
739 materialize: materialize_first,
740 hints: first_hints,
741 attributes: first_attributes,
742 } = self;
743 let Flow {
744 transform: second,
745 materialize: materialize_second,
746 hints: second_hints,
747 attributes: second_attributes,
748 } = next;
749 let combine = Arc::new(combine);
750 match (first, second) {
751 (FlowTransform::Pure(first), FlowTransform::Pure(second)) => {
752 let hints = first_hints.then(second_hints);
753 Flow::from_parts_with_hints(
754 move |input| second(first(input)),
755 move || {
756 let left = materialize_first()?;
757 let right = materialize_second()?;
758 Ok(combine(left, right))
759 },
760 hints,
761 )
762 .with_attributes(first_attributes.and(second_attributes))
763 }
764 (first, second) => {
765 let hints = first_hints.then(second_hints);
766 Flow {
767 transform: FlowTransform::Runtime(Arc::new(move |input, materializer| {
768 let stream = match &first {
769 FlowTransform::Pure(first) => first(input),
770 FlowTransform::Runtime(first) => first(input, materializer)?,
771 };
772 match &second {
773 FlowTransform::Pure(second) => Ok(second(stream)),
774 FlowTransform::Runtime(second) => second(stream, materializer),
775 }
776 })),
777 materialize: Arc::new(move || {
778 let left = materialize_first()?;
779 let right = materialize_second()?;
780 Ok(combine(left, right))
781 }),
782 hints,
783 attributes: first_attributes.and(second_attributes),
784 }
785 }
786 }
787 }
788
789 #[must_use]
790 pub fn via_mat_with<Next, NextMat, Combined, F>(
791 self,
792 next: Flow<Out, Next, NextMat>,
793 combine: F,
794 ) -> Flow<In, Next, Combined>
795 where
796 Next: Send + 'static,
797 NextMat: Send + 'static,
798 Combined: Send + 'static,
799 F: Fn(Mat, NextMat) -> Combined + Send + Sync + 'static,
800 {
801 self.via_mat(next, combine)
802 }
803
804 #[must_use]
805 pub fn map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
806 where
807 Next: Send + 'static,
808 F: Fn(Out) -> Next + Send + Sync + 'static,
809 {
810 let stage = Arc::new(f);
811 match &self.transform {
812 FlowTransform::Pure(_) => {
813 let Flow {
814 transform,
815 materialize,
816 hints,
817 attributes,
818 } = self;
819 let FlowTransform::Pure(transform) = transform else {
820 unreachable!("pure transform checked above");
821 };
822 Flow::from_parts_with_hints(
823 move |input| {
824 let stage = Arc::clone(&stage);
825 Box::new(transform(input).map(move |item| item.map(|item| stage(item))))
826 },
827 move || materialize(),
828 hints,
829 )
830 .with_attributes(attributes)
831 }
832 FlowTransform::Runtime(_) => self.via(Flow::from_transform(move |input| {
833 let stage = Arc::clone(&stage);
834 Box::new(input.map(move |item| item.map(|item| stage(item))))
835 })),
836 }
837 }
838
839 #[must_use]
842 pub fn try_map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
843 where
844 Next: Send + 'static,
845 F: Fn(Out) -> StreamResult<Next> + Send + Sync + 'static,
846 {
847 let stage = Arc::new(f);
848 self.via(Flow::from_transform(move |input| {
849 let stage = Arc::clone(&stage);
850 Box::new(input.map(move |item| item.and_then(|item| stage(item))))
851 }))
852 }
853
854 #[must_use]
856 #[deprecated(
857 since = "0.9.0",
858 note = "renamed to `try_map` (idiomatic); the `_result` name still works"
859 )]
860 pub fn map_result<Next, F>(self, f: F) -> Flow<In, Next, Mat>
861 where
862 Next: Send + 'static,
863 F: Fn(Out) -> StreamResult<Next> + Send + Sync + 'static,
864 {
865 self.try_map(f)
866 }
867
868 #[must_use]
873 pub fn map_result_with_supervision<Next, F>(
874 self,
875 f: F,
876 decider: SupervisionDecider,
877 ) -> Flow<In, Next, Mat>
878 where
879 Next: Send + 'static,
880 F: Fn(Out) -> StreamResult<Next> + Send + Sync + 'static,
881 {
882 let stage = Arc::new(f);
883 self.via(Flow::from_transform(move |input| {
884 let stage = Arc::clone(&stage);
885 let decider = Arc::clone(&decider);
886 Box::new(input.filter_map(move |item| match item {
887 Ok(item) => match call_supervised("map_result callback", || stage(item)) {
888 Ok(next) => Some(Ok(next)),
889 Err(error) => match decide_supervision(&decider, &error) {
890 SupervisionDirective::Stop => Some(Err(error)),
891 SupervisionDirective::Resume | SupervisionDirective::Restart => None,
892 },
893 },
894 Err(error) => Some(Err(error)),
895 }))
896 }))
897 }
898
899 #[must_use]
900 pub fn filter<F>(self, predicate: F) -> Flow<In, Out, Mat>
901 where
902 F: Fn(&Out) -> bool + Send + Sync + 'static,
903 {
904 let predicate = Arc::new(predicate);
905 self.via(Flow::from_preserving_transform(move |input| {
906 let predicate = Arc::clone(&predicate);
907 Box::new(input.filter_map(move |item| match item {
908 Ok(item) if predicate(&item) => Some(Ok(item)),
909 Ok(_) => None,
910 Err(error) => Some(Err(error)),
911 }))
912 }))
913 }
914
915 #[must_use]
916 pub fn try_filter<F>(self, predicate: F) -> Flow<In, Out, Mat>
917 where
918 F: Fn(&Out) -> StreamResult<bool> + Send + Sync + 'static,
919 {
920 let predicate = Arc::new(predicate);
921 self.via(Flow::from_transform(move |input| {
922 let predicate = Arc::clone(&predicate);
923 Box::new(input.filter_map(move |item| match item {
924 Ok(item) => match predicate(&item) {
925 Ok(true) => Some(Ok(item)),
926 Ok(false) => None,
927 Err(error) => Some(Err(error)),
928 },
929 Err(error) => Some(Err(error)),
930 }))
931 }))
932 }
933
934 #[must_use]
936 #[deprecated(
937 since = "0.9.0",
938 note = "renamed to `try_filter` (idiomatic); the `_result` name still works"
939 )]
940 pub fn filter_result<F>(self, predicate: F) -> Flow<In, Out, Mat>
941 where
942 F: Fn(&Out) -> StreamResult<bool> + Send + Sync + 'static,
943 {
944 self.try_filter(predicate)
945 }
946
947 #[must_use]
948 pub fn filter_result_with_supervision<F>(
949 self,
950 predicate: F,
951 decider: SupervisionDecider,
952 ) -> Flow<In, Out, Mat>
953 where
954 F: Fn(&Out) -> StreamResult<bool> + Send + Sync + 'static,
955 {
956 let predicate = Arc::new(predicate);
957 self.via(Flow::from_transform(move |input| {
958 let predicate = Arc::clone(&predicate);
959 let decider = Arc::clone(&decider);
960 Box::new(input.filter_map(move |item| match item {
961 Ok(item) => match call_supervised("filter_result callback", || predicate(&item)) {
962 Ok(true) => Some(Ok(item)),
963 Ok(false) => None,
964 Err(error) => match decide_supervision(&decider, &error) {
965 SupervisionDirective::Stop => Some(Err(error)),
966 SupervisionDirective::Resume | SupervisionDirective::Restart => None,
967 },
968 },
969 Err(error) => Some(Err(error)),
970 }))
971 }))
972 }
973
974 #[must_use]
975 pub fn filter_not<F>(self, predicate: F) -> Flow<In, Out, Mat>
976 where
977 F: Fn(&Out) -> bool + Send + Sync + 'static,
978 {
979 let predicate = Arc::new(predicate);
980 self.filter(move |item| !predicate(item))
981 }
982
983 #[must_use]
984 pub fn filter_map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
985 where
986 Next: Send + 'static,
987 F: Fn(Out) -> Option<Next> + Send + Sync + 'static,
988 {
989 let stage = Arc::new(f);
990 self.via(Flow::from_transform(move |input| {
991 let stage = Arc::clone(&stage);
992 Box::new(input.filter_map(move |item| match item {
993 Ok(item) => stage(item).map(Ok),
994 Err(error) => Some(Err(error)),
995 }))
996 }))
997 }
998
999 #[must_use]
1000 pub fn try_filter_map<Next, F>(self, f: F) -> Flow<In, Next, Mat>
1001 where
1002 Next: Send + 'static,
1003 F: Fn(Out) -> StreamResult<Option<Next>> + Send + Sync + 'static,
1004 {
1005 let stage = Arc::new(f);
1006 self.via(Flow::from_transform(move |input| {
1007 let stage = Arc::clone(&stage);
1008 Box::new(input.filter_map(move |item| match item {
1009 Ok(item) => match stage(item) {
1010 Ok(Some(next)) => Some(Ok(next)),
1011 Ok(None) => None,
1012 Err(error) => Some(Err(error)),
1013 },
1014 Err(error) => Some(Err(error)),
1015 }))
1016 }))
1017 }
1018
1019 #[must_use]
1021 #[deprecated(
1022 since = "0.9.0",
1023 note = "renamed to `try_filter_map` (idiomatic); the `_result` name still works"
1024 )]
1025 pub fn filter_map_result<Next, F>(self, f: F) -> Flow<In, Next, Mat>
1026 where
1027 Next: Send + 'static,
1028 F: Fn(Out) -> StreamResult<Option<Next>> + Send + Sync + 'static,
1029 {
1030 self.try_filter_map(f)
1031 }
1032
1033 #[must_use]
1034 pub fn filter_map_result_with_supervision<Next, F>(
1035 self,
1036 f: F,
1037 decider: SupervisionDecider,
1038 ) -> Flow<In, Next, Mat>
1039 where
1040 Next: Send + 'static,
1041 F: Fn(Out) -> StreamResult<Option<Next>> + Send + Sync + 'static,
1042 {
1043 let stage = Arc::new(f);
1044 self.via(Flow::from_transform(move |input| {
1045 let stage = Arc::clone(&stage);
1046 let decider = Arc::clone(&decider);
1047 Box::new(input.filter_map(move |item| match item {
1048 Ok(item) => match call_supervised("filter_map_result callback", || stage(item)) {
1049 Ok(Some(next)) => Some(Ok(next)),
1050 Ok(None) => None,
1051 Err(error) => match decide_supervision(&decider, &error) {
1052 SupervisionDirective::Stop => Some(Err(error)),
1053 SupervisionDirective::Resume | SupervisionDirective::Restart => None,
1054 },
1055 },
1056 Err(error) => Some(Err(error)),
1057 }))
1058 }))
1059 }
1060
1061 #[must_use]
1062 pub fn map_concat<Next, F, I>(self, f: F) -> Flow<In, Next, Mat>
1063 where
1064 Next: Send + 'static,
1065 F: Fn(Out) -> I + Send + Sync + 'static,
1066 I: IntoIterator<Item = Next>,
1067 I::IntoIter: Send + 'static,
1068 {
1069 let stage = Arc::new(f);
1070 self.via(Flow::from_transform(move |mut input| {
1071 let stage = Arc::clone(&stage);
1072 let mut current = None::<I::IntoIter>;
1073 let mut done = false;
1074 Box::new(std::iter::from_fn(move || {
1075 loop {
1076 if let Some(iter) = &mut current {
1077 if let Some(item) = iter.next() {
1078 return Some(Ok(item));
1079 }
1080 current = None;
1081 }
1082
1083 if done {
1084 return None;
1085 }
1086
1087 match input.next()? {
1088 Ok(item) => current = Some(stage(item).into_iter()),
1089 Err(error) => {
1090 done = true;
1091 return Some(Err(error));
1092 }
1093 }
1094 }
1095 }))
1096 }))
1097 }
1098
1099 #[must_use]
1100 pub fn try_map_concat<Next, F, I>(self, f: F) -> Flow<In, Next, Mat>
1101 where
1102 Next: Send + 'static,
1103 F: Fn(Out) -> StreamResult<I> + Send + Sync + 'static,
1104 I: IntoIterator<Item = Next>,
1105 I::IntoIter: Send + 'static,
1106 {
1107 let stage = Arc::new(f);
1108 self.via(Flow::from_transform(move |mut input| {
1109 let stage = Arc::clone(&stage);
1110 let mut current = None::<I::IntoIter>;
1111 let mut done = false;
1112 Box::new(std::iter::from_fn(move || {
1113 loop {
1114 if let Some(iter) = &mut current {
1115 if let Some(item) = iter.next() {
1116 return Some(Ok(item));
1117 }
1118 current = None;
1119 }
1120
1121 if done {
1122 return None;
1123 }
1124
1125 match input.next()? {
1126 Ok(item) => match stage(item) {
1127 Ok(items) => current = Some(items.into_iter()),
1128 Err(error) => {
1129 done = true;
1130 return Some(Err(error));
1131 }
1132 },
1133 Err(error) => {
1134 done = true;
1135 return Some(Err(error));
1136 }
1137 }
1138 }
1139 }))
1140 }))
1141 }
1142
1143 #[must_use]
1145 #[deprecated(
1146 since = "0.9.0",
1147 note = "renamed to `try_map_concat` (idiomatic); the `_result` name still works"
1148 )]
1149 pub fn map_concat_result<Next, F, I>(self, f: F) -> Flow<In, Next, Mat>
1150 where
1151 Next: Send + 'static,
1152 F: Fn(Out) -> StreamResult<I> + Send + Sync + 'static,
1153 I: IntoIterator<Item = Next>,
1154 I::IntoIter: Send + 'static,
1155 {
1156 self.try_map_concat(f)
1157 }
1158
1159 #[must_use]
1160 pub fn map_concat_result_with_supervision<Next, F, I>(
1161 self,
1162 f: F,
1163 decider: SupervisionDecider,
1164 ) -> Flow<In, Next, Mat>
1165 where
1166 Next: Send + 'static,
1167 F: Fn(Out) -> StreamResult<I> + Send + Sync + 'static,
1168 I: IntoIterator<Item = Next>,
1169 I::IntoIter: Send + 'static,
1170 {
1171 let stage = Arc::new(f);
1172 self.via(Flow::from_transform(move |mut input| {
1173 let stage = Arc::clone(&stage);
1174 let decider = Arc::clone(&decider);
1175 let mut current = None::<I::IntoIter>;
1176 let mut done = false;
1177 Box::new(std::iter::from_fn(move || {
1178 loop {
1179 if let Some(iter) = &mut current {
1180 if let Some(item) = iter.next() {
1181 return Some(Ok(item));
1182 }
1183 current = None;
1184 }
1185
1186 if done {
1187 return None;
1188 }
1189
1190 match input.next()? {
1191 Ok(item) => {
1192 match call_supervised("map_concat_result callback", || stage(item)) {
1193 Ok(items) => current = Some(items.into_iter()),
1194 Err(error) => match decide_supervision(&decider, &error) {
1195 SupervisionDirective::Stop => {
1196 done = true;
1197 return Some(Err(error));
1198 }
1199 SupervisionDirective::Resume
1200 | SupervisionDirective::Restart => {}
1201 },
1202 }
1203 }
1204 Err(error) => {
1205 done = true;
1206 return Some(Err(error));
1207 }
1208 }
1209 }
1210 }))
1211 }))
1212 }
1213
1214 #[must_use]
1215 pub fn stateful_map<State, Next, F>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1216 where
1217 State: Clone + Send + Sync + 'static,
1218 Next: Send + 'static,
1219 F: Fn(&mut State, Out) -> Next + Send + Sync + 'static,
1220 {
1221 let stage = Arc::new(f);
1222 self.via(Flow::from_transform(move |input| {
1223 let stage = Arc::clone(&stage);
1224 let mut state = seed.clone();
1225 Box::new(input.map(move |item| item.map(|item| stage(&mut state, item))))
1226 }))
1227 }
1228
1229 #[must_use]
1230 pub fn try_stateful_map<State, Next, F>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1231 where
1232 State: Clone + Send + Sync + 'static,
1233 Next: Send + 'static,
1234 F: Fn(&mut State, Out) -> StreamResult<Next> + Send + Sync + 'static,
1235 {
1236 let stage = Arc::new(f);
1237 self.via(Flow::from_transform(move |input| {
1238 let stage = Arc::clone(&stage);
1239 let mut state = seed.clone();
1240 Box::new(input.map(move |item| match item {
1241 Ok(item) => stage(&mut state, item),
1242 Err(error) => Err(error),
1243 }))
1244 }))
1245 }
1246
1247 #[must_use]
1249 #[deprecated(
1250 since = "0.9.0",
1251 note = "renamed to `try_stateful_map` (idiomatic); the `_result` name still works"
1252 )]
1253 pub fn stateful_map_result<State, Next, F>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1254 where
1255 State: Clone + Send + Sync + 'static,
1256 Next: Send + 'static,
1257 F: Fn(&mut State, Out) -> StreamResult<Next> + Send + Sync + 'static,
1258 {
1259 self.try_stateful_map(seed, f)
1260 }
1261
1262 #[must_use]
1263 pub fn stateful_map_result_with_supervision<State, Next, F>(
1264 self,
1265 seed: State,
1266 f: F,
1267 decider: SupervisionDecider,
1268 ) -> Flow<In, Next, Mat>
1269 where
1270 State: Clone + Send + Sync + 'static,
1271 Next: Send + 'static,
1272 F: Fn(&mut State, Out) -> StreamResult<Next> + Send + Sync + 'static,
1273 {
1274 let stage = Arc::new(f);
1275 self.via(Flow::from_transform(move |input| {
1276 let stage = Arc::clone(&stage);
1277 let decider = Arc::clone(&decider);
1278 let seed = seed.clone();
1279 let mut state = seed.clone();
1280 Box::new(input.filter_map(move |item| match item {
1281 Ok(item) => match call_supervised("stateful_map_result callback", || {
1282 stage(&mut state, item)
1283 }) {
1284 Ok(next) => Some(Ok(next)),
1285 Err(error) => match decide_supervision(&decider, &error) {
1286 SupervisionDirective::Stop => Some(Err(error)),
1287 SupervisionDirective::Resume => None,
1288 SupervisionDirective::Restart => {
1289 state = seed.clone();
1290 None
1291 }
1292 },
1293 },
1294 Err(error) => Some(Err(error)),
1295 }))
1296 }))
1297 }
1298
1299 #[must_use]
1300 pub fn stateful_map_concat<State, Next, F, I>(self, seed: State, f: F) -> Flow<In, Next, Mat>
1301 where
1302 State: Clone + Send + Sync + 'static,
1303 Next: Send + 'static,
1304 F: Fn(&mut State, Out) -> I + Send + Sync + 'static,
1305 I: IntoIterator<Item = Next>,
1306 I::IntoIter: Send + 'static,
1307 {
1308 let stage = Arc::new(f);
1309 self.via(Flow::from_transform(move |mut input| {
1310 let stage = Arc::clone(&stage);
1311 let mut state = seed.clone();
1312 let mut current = None::<I::IntoIter>;
1313 let mut done = false;
1314 Box::new(std::iter::from_fn(move || {
1315 loop {
1316 if let Some(iter) = &mut current {
1317 if let Some(item) = iter.next() {
1318 return Some(Ok(item));
1319 }
1320 current = None;
1321 }
1322
1323 if done {
1324 return None;
1325 }
1326
1327 match input.next()? {
1328 Ok(item) => current = Some(stage(&mut state, item).into_iter()),
1329 Err(error) => {
1330 done = true;
1331 return Some(Err(error));
1332 }
1333 }
1334 }
1335 }))
1336 }))
1337 }
1338
1339 #[must_use]
1340 pub fn try_stateful_map_concat<State, Next, F, I>(
1341 self,
1342 seed: State,
1343 f: F,
1344 ) -> Flow<In, Next, Mat>
1345 where
1346 State: Clone + Send + Sync + 'static,
1347 Next: Send + 'static,
1348 F: Fn(&mut State, Out) -> StreamResult<I> + Send + Sync + 'static,
1349 I: IntoIterator<Item = Next>,
1350 I::IntoIter: Send + 'static,
1351 {
1352 let stage = Arc::new(f);
1353 self.via(Flow::from_transform(move |mut input| {
1354 let stage = Arc::clone(&stage);
1355 let mut state = seed.clone();
1356 let mut current = None::<I::IntoIter>;
1357 let mut done = false;
1358 Box::new(std::iter::from_fn(move || {
1359 loop {
1360 if let Some(iter) = &mut current {
1361 if let Some(item) = iter.next() {
1362 return Some(Ok(item));
1363 }
1364 current = None;
1365 }
1366
1367 if done {
1368 return None;
1369 }
1370
1371 match input.next()? {
1372 Ok(item) => match stage(&mut state, item) {
1373 Ok(items) => current = Some(items.into_iter()),
1374 Err(error) => {
1375 done = true;
1376 return Some(Err(error));
1377 }
1378 },
1379 Err(error) => {
1380 done = true;
1381 return Some(Err(error));
1382 }
1383 }
1384 }
1385 }))
1386 }))
1387 }
1388
1389 #[must_use]
1391 #[deprecated(
1392 since = "0.9.0",
1393 note = "renamed to `try_stateful_map_concat` (idiomatic); the `_result` name still works"
1394 )]
1395 pub fn stateful_map_concat_result<State, Next, F, I>(
1396 self,
1397 seed: State,
1398 f: F,
1399 ) -> Flow<In, Next, Mat>
1400 where
1401 State: Clone + Send + Sync + 'static,
1402 Next: Send + 'static,
1403 F: Fn(&mut State, Out) -> StreamResult<I> + Send + Sync + 'static,
1404 I: IntoIterator<Item = Next>,
1405 I::IntoIter: Send + 'static,
1406 {
1407 self.try_stateful_map_concat(seed, f)
1408 }
1409
1410 #[must_use]
1411 pub fn stateful_map_concat_result_with_supervision<State, Next, F, I>(
1412 self,
1413 seed: State,
1414 f: F,
1415 decider: SupervisionDecider,
1416 ) -> Flow<In, Next, Mat>
1417 where
1418 State: Clone + Send + Sync + 'static,
1419 Next: Send + 'static,
1420 F: Fn(&mut State, Out) -> StreamResult<I> + Send + Sync + 'static,
1421 I: IntoIterator<Item = Next>,
1422 I::IntoIter: Send + 'static,
1423 {
1424 let stage = Arc::new(f);
1425 self.via(Flow::from_transform(move |mut input| {
1426 let stage = Arc::clone(&stage);
1427 let decider = Arc::clone(&decider);
1428 let seed = seed.clone();
1429 let mut state = seed.clone();
1430 let mut current = None::<I::IntoIter>;
1431 let mut done = false;
1432 Box::new(std::iter::from_fn(move || {
1433 loop {
1434 if let Some(iter) = &mut current {
1435 if let Some(item) = iter.next() {
1436 return Some(Ok(item));
1437 }
1438 current = None;
1439 }
1440
1441 if done {
1442 return None;
1443 }
1444
1445 match input.next()? {
1446 Ok(item) => {
1447 match call_supervised("stateful_map_concat_result callback", || {
1448 stage(&mut state, item)
1449 }) {
1450 Ok(items) => current = Some(items.into_iter()),
1451 Err(error) => match decide_supervision(&decider, &error) {
1452 SupervisionDirective::Stop => {
1453 done = true;
1454 return Some(Err(error));
1455 }
1456 SupervisionDirective::Resume => {}
1457 SupervisionDirective::Restart => state = seed.clone(),
1458 },
1459 }
1460 }
1461 Err(error) => {
1462 done = true;
1463 return Some(Err(error));
1464 }
1465 }
1466 }
1467 }))
1468 }))
1469 }
1470
1471 #[must_use]
1472 pub fn map_async<Next, F, Fut>(self, parallelism: usize, f: F) -> Flow<In, Next, Mat>
1476 where
1477 Next: Send + 'static,
1478 F: Fn(Out) -> Fut + Send + Sync + 'static,
1479 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1480 {
1481 assert!(
1482 parallelism > 0,
1483 "map_async parallelism must be greater than zero"
1484 );
1485 let stage = Arc::new(f);
1486 self.via(Flow::from_runtime_transform_with_hints(
1487 move |input, _materializer| {
1488 let stage = Arc::clone(&stage);
1489 Ok(map_async_ordered(input, parallelism, stage))
1490 },
1491 FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1492 ))
1493 }
1494
1495 #[must_use]
1496 pub fn map_async_with_supervision<Next, F, Fut>(
1501 self,
1502 parallelism: usize,
1503 f: F,
1504 decider: SupervisionDecider,
1505 ) -> Flow<In, Next, Mat>
1506 where
1507 Next: Send + 'static,
1508 F: Fn(Out) -> Fut + Send + Sync + 'static,
1509 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1510 {
1511 assert!(
1512 parallelism > 0,
1513 "map_async parallelism must be greater than zero"
1514 );
1515 let stage = Arc::new(f);
1516 self.via(Flow::from_runtime_transform(move |input, _materializer| {
1517 let stage = Arc::clone(&stage);
1518 let decider = Arc::clone(&decider);
1519 Ok(map_async_ordered_supervised(
1520 input,
1521 parallelism,
1522 stage,
1523 decider,
1524 ))
1525 }))
1526 }
1527
1528 #[must_use]
1529 pub fn map_async_unordered<Next, F, Fut>(self, parallelism: usize, f: F) -> Flow<In, Next, Mat>
1533 where
1534 Next: Send + 'static,
1535 F: Fn(Out) -> Fut + Send + Sync + 'static,
1536 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1537 {
1538 assert!(
1539 parallelism > 0,
1540 "map_async_unordered parallelism must be greater than zero"
1541 );
1542 let stage = Arc::new(f);
1543 self.via(Flow::from_runtime_transform_with_hints(
1544 move |input, _materializer| {
1545 let stage = Arc::clone(&stage);
1546 Ok(map_async_unordered(input, parallelism, stage))
1547 },
1548 FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1549 ))
1550 }
1551
1552 #[must_use]
1553 pub fn map_async_unordered_with_supervision<Next, F, Fut>(
1554 self,
1555 parallelism: usize,
1556 f: F,
1557 decider: SupervisionDecider,
1558 ) -> Flow<In, Next, Mat>
1559 where
1560 Next: Send + 'static,
1561 F: Fn(Out) -> Fut + Send + Sync + 'static,
1562 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1563 {
1564 assert!(
1565 parallelism > 0,
1566 "map_async_unordered parallelism must be greater than zero"
1567 );
1568 let stage = Arc::new(f);
1569 self.via(Flow::from_runtime_transform(move |input, _materializer| {
1570 let stage = Arc::clone(&stage);
1571 let decider = Arc::clone(&decider);
1572 Ok(map_async_unordered_supervised(
1573 input,
1574 parallelism,
1575 stage,
1576 decider,
1577 ))
1578 }))
1579 }
1580
1581 #[must_use]
1582 pub fn map_async_partitioned<Key, Next, Partition, F, Fut>(
1586 self,
1587 parallelism: usize,
1588 per_partition: usize,
1589 partition: Partition,
1590 f: F,
1591 ) -> Flow<In, Next, Mat>
1592 where
1593 Key: Clone + Eq + Hash + Send + 'static,
1594 Next: Send + 'static,
1595 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
1596 F: Fn(Out) -> Fut + Send + Sync + 'static,
1597 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
1598 {
1599 assert!(
1600 parallelism > 0,
1601 "map_async_partitioned parallelism must be greater than zero"
1602 );
1603 assert!(
1604 per_partition > 0,
1605 "map_async_partitioned per_partition must be greater than zero"
1606 );
1607 let partition = Arc::new(partition);
1608 let stage = Arc::new(f);
1609 self.via(Flow::from_runtime_transform_with_hints(
1610 move |input, _materializer| {
1611 let partition = Arc::clone(&partition);
1612 let stage = Arc::clone(&stage);
1613 Ok(map_async_partitioned(
1614 input,
1615 parallelism,
1616 per_partition,
1617 partition,
1618 stage,
1619 ))
1620 },
1621 FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
1622 ))
1623 }
1624
1625 #[must_use]
1626 pub fn take(self, n: usize) -> Flow<In, Out, Mat> {
1627 self.via(Flow::from_transform(move |input| Box::new(input.take(n))))
1628 }
1629
1630 #[must_use]
1631 pub fn drop(self, n: usize) -> Flow<In, Out, Mat> {
1632 self.via(Flow::from_transform(move |input| {
1633 let mut remaining = n;
1634 Box::new(input.filter_map(move |item| match item {
1635 Ok(_) if remaining > 0 => {
1636 remaining -= 1;
1637 None
1638 }
1639 other => Some(other),
1640 }))
1641 }))
1642 }
1643
1644 #[must_use]
1645 pub fn take_while<F>(self, predicate: F) -> Flow<In, Out, Mat>
1646 where
1647 F: Fn(&Out) -> bool + Send + Sync + 'static,
1648 {
1649 let predicate = Arc::new(predicate);
1650 self.via(Flow::from_transform(move |mut input| {
1651 let predicate = Arc::clone(&predicate);
1652 let mut open = true;
1653 Box::new(std::iter::from_fn(move || {
1654 if !open {
1655 return None;
1656 }
1657
1658 match input.next() {
1659 Some(Ok(item)) if predicate(&item) => Some(Ok(item)),
1660 Some(Ok(_)) | None => {
1661 open = false;
1662 None
1663 }
1664 Some(Err(error)) => Some(Err(error)),
1665 }
1666 }))
1667 }))
1668 }
1669
1670 #[must_use]
1671 pub fn drop_while<F>(self, predicate: F) -> Flow<In, Out, Mat>
1672 where
1673 F: Fn(&Out) -> bool + Send + Sync + 'static,
1674 {
1675 let predicate = Arc::new(predicate);
1676 self.via(Flow::from_transform(move |mut input| {
1677 let predicate = Arc::clone(&predicate);
1678 let mut dropping = true;
1679 Box::new(std::iter::from_fn(move || {
1680 loop {
1681 let next = input.next()?;
1682 match next {
1683 Ok(item) if dropping && predicate(&item) => continue,
1684 Ok(item) => {
1685 dropping = false;
1686 return Some(Ok(item));
1687 }
1688 Err(error) => return Some(Err(error)),
1689 }
1690 }
1691 }))
1692 }))
1693 }
1694
1695 #[must_use]
1696 pub fn limit(self, max: u64) -> Flow<In, Out, Mat> {
1697 self.via(Flow::from_transform(move |input| {
1698 let mut seen = 0_u64;
1699 Box::new(input.map(move |item| match item {
1700 Ok(item) if seen < max => {
1701 seen += 1;
1702 Ok(item)
1703 }
1704 Ok(_) => Err(StreamError::LimitExceeded { max }),
1705 Err(error) => Err(error),
1706 }))
1707 }))
1708 }
1709
1710 #[must_use]
1711 pub fn grouped(self, size: usize) -> Flow<In, Vec<Out>, Mat> {
1712 assert!(size > 0, "grouped size must be greater than zero");
1713 self.via(Flow::from_transform(move |mut input| {
1714 Box::new(std::iter::from_fn(move || {
1715 let mut group = Vec::with_capacity(size);
1716 while group.len() < size {
1717 match input.next() {
1718 Some(Ok(item)) => group.push(item),
1719 Some(Err(error)) => return Some(Err(error)),
1720 None => break,
1721 }
1722 }
1723
1724 if group.is_empty() {
1725 None
1726 } else {
1727 Some(Ok(group))
1728 }
1729 }))
1730 }))
1731 }
1732
1733 #[must_use]
1734 pub fn scan<State, F>(self, seed: State, f: F) -> Flow<In, State, Mat>
1735 where
1736 State: Clone + Send + Sync + 'static,
1737 F: Fn(State, Out) -> State + Send + Sync + 'static,
1738 {
1739 let stage = Arc::new(f);
1740 self.via(Flow::from_transform(move |mut input| {
1741 let stage = Arc::clone(&stage);
1742 let mut state = Some(seed.clone());
1743 let mut emit_seed = true;
1744 Box::new(std::iter::from_fn(move || {
1745 if emit_seed {
1746 emit_seed = false;
1747 return Some(Ok(state.as_ref().expect("scan state present").clone()));
1748 }
1749
1750 match input.next()? {
1751 Ok(item) => {
1752 let prev = state.take().expect("scan state present");
1755 let next = stage(prev, item);
1756 state = Some(next.clone());
1757 Some(Ok(next))
1758 }
1759 Err(error) => Some(Err(error)),
1760 }
1761 }))
1762 }))
1763 }
1764
1765 #[must_use]
1766 pub fn scan_async<State, F, Fut>(self, seed: State, f: F) -> Flow<In, State, Mat>
1767 where
1768 State: Clone + Send + Sync + 'static,
1769 F: Fn(State, Out) -> Fut + Send + Sync + 'static,
1770 Fut: Future<Output = StreamResult<State>> + Send + 'static,
1771 {
1772 let stage = Arc::new(f);
1773 self.via(Flow::from_transform(move |mut input| {
1774 let stage = Arc::clone(&stage);
1775 let mut state = Some(seed.clone());
1776 let mut emit_seed = true;
1777 let mut terminated = false;
1778 Box::new(std::iter::from_fn(move || {
1779 if terminated {
1780 return None;
1781 }
1782 if emit_seed {
1783 emit_seed = false;
1784 return Some(Ok(state
1785 .as_ref()
1786 .expect("scan_async state present")
1787 .clone()));
1788 }
1789
1790 match input.next()? {
1791 Ok(item) => {
1792 let prev = state.take().expect("scan_async state present");
1793 match catch_unwind_failed("scan_async factory", || stage(prev, item))
1794 .and_then(run_future_inline_or_spawn)
1795 {
1796 Ok(next) => {
1797 state = Some(next.clone());
1798 Some(Ok(next))
1799 }
1800 Err(error) => {
1801 terminated = true;
1802 Some(Err(error))
1803 }
1804 }
1805 }
1806 Err(error) => {
1807 terminated = true;
1808 Some(Err(error))
1809 }
1810 }
1811 }))
1812 }))
1813 }
1814
1815 #[must_use]
1816 pub fn try_scan<State, F>(self, seed: State, f: F) -> Flow<In, State, Mat>
1817 where
1818 State: Clone + Send + Sync + 'static,
1819 F: Fn(State, Out) -> StreamResult<State> + Send + Sync + 'static,
1820 {
1821 let stage = Arc::new(f);
1822 self.via(Flow::from_transform(move |mut input| {
1823 let stage = Arc::clone(&stage);
1824 let mut state = Some(seed.clone());
1825 let mut emit_seed = true;
1826 Box::new(std::iter::from_fn(move || {
1827 if emit_seed {
1828 emit_seed = false;
1829 return Some(Ok(state.as_ref().expect("scan state present").clone()));
1830 }
1831
1832 match input.next()? {
1833 Ok(item) => {
1834 let prev = state.take().expect("scan state present");
1835 match stage(prev, item) {
1836 Ok(next) => {
1837 state = Some(next.clone());
1838 Some(Ok(next))
1839 }
1840 Err(error) => Some(Err(error)),
1841 }
1842 }
1843 Err(error) => Some(Err(error)),
1844 }
1845 }))
1846 }))
1847 }
1848
1849 #[must_use]
1851 #[deprecated(
1852 since = "0.9.0",
1853 note = "renamed to `try_scan` (idiomatic); the `_result` name still works"
1854 )]
1855 pub fn scan_result<State, F>(self, seed: State, f: F) -> Flow<In, State, Mat>
1856 where
1857 State: Clone + Send + Sync + 'static,
1858 F: Fn(State, Out) -> StreamResult<State> + Send + Sync + 'static,
1859 {
1860 self.try_scan(seed, f)
1861 }
1862
1863 #[must_use]
1864 pub fn scan_result_with_supervision<State, F>(
1865 self,
1866 seed: State,
1867 f: F,
1868 decider: SupervisionDecider,
1869 ) -> Flow<In, State, Mat>
1870 where
1871 State: Clone + Send + Sync + 'static,
1872 F: Fn(State, Out) -> StreamResult<State> + Send + Sync + 'static,
1873 {
1874 let stage = Arc::new(f);
1875 self.via(Flow::from_transform(move |mut input| {
1876 let stage = Arc::clone(&stage);
1877 let decider = Arc::clone(&decider);
1878 let seed = seed.clone();
1879 let mut state = Some(seed.clone());
1880 let mut emit_seed = true;
1881 Box::new(std::iter::from_fn(move || {
1882 loop {
1883 if emit_seed {
1884 emit_seed = false;
1885 return Some(Ok(state.as_ref().expect("scan state present").clone()));
1886 }
1887
1888 match input.next()? {
1889 Ok(item) => {
1890 let prev = state.take().expect("scan state present");
1891 match call_supervised("scan_result callback", || {
1892 stage(prev.clone(), item)
1893 }) {
1894 Ok(next) => {
1895 state = Some(next.clone());
1896 return Some(Ok(next));
1897 }
1898 Err(error) => match decide_supervision(&decider, &error) {
1899 SupervisionDirective::Stop => return Some(Err(error)),
1900 SupervisionDirective::Resume => {
1901 state = Some(prev);
1902 }
1903 SupervisionDirective::Restart => {
1904 state = Some(seed.clone());
1905 emit_seed = true;
1906 }
1907 },
1908 }
1909 }
1910 Err(error) => return Some(Err(error)),
1911 }
1912 }
1913 }))
1914 }))
1915 }
1916
1917 #[must_use]
1918 pub fn sliding(self, size: usize, step: usize) -> Flow<In, Vec<Out>, Mat>
1919 where
1920 Out: Clone,
1921 {
1922 assert!(size > 0, "sliding size must be greater than zero");
1923 assert!(step > 0, "sliding step must be greater than zero");
1924 self.via(Flow::from_transform(move |mut input| {
1925 let mut buffer = VecDeque::with_capacity(size.max(step));
1931 let mut terminated = false;
1932
1933 Box::new(std::iter::from_fn(move || {
1934 if terminated {
1935 return None;
1936 }
1937
1938 loop {
1939 match input.next() {
1940 Some(Ok(item)) => {
1941 buffer.push_back(item);
1942 if buffer.len() < size {
1943 continue;
1944 } else if buffer.len() == size {
1945 return Some(Ok(buffer.iter().cloned().collect()));
1946 } else if step <= size {
1947 for _ in 0..step {
1948 buffer.pop_front();
1949 }
1950 if buffer.len() == size {
1951 return Some(Ok(buffer.iter().cloned().collect()));
1952 }
1953 } else if buffer.len() == step {
1954 buffer.clear();
1956 }
1957 }
1958 Some(Err(error)) => {
1959 terminated = true;
1960 return Some(Err(error));
1961 }
1962 None => {
1963 terminated = true;
1964 if !buffer.is_empty() && buffer.len() < size {
1965 return Some(Ok(buffer.iter().cloned().collect()));
1966 }
1967 return None;
1968 }
1969 }
1970 }
1971 }))
1972 }))
1973 }
1974
1975 #[must_use]
1976 pub fn fold<Acc, F>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
1977 where
1978 Acc: Clone + Send + Sync + 'static,
1979 F: Fn(Acc, Out) -> Acc + Send + Sync + 'static,
1980 {
1981 let stage = Arc::new(f);
1982 self.via(Flow::from_transform(move |input| {
1983 let stage = Arc::clone(&stage);
1984 let mut acc = zero.clone();
1985 for item in input {
1986 match item {
1987 Ok(item) => acc = stage(acc, item),
1988 Err(error) => return Box::new(std::iter::once(Err(error))),
1989 }
1990 }
1991 Box::new(std::iter::once(Ok(acc)))
1992 }))
1993 }
1994
1995 #[must_use]
1996 pub fn fold_async<Acc, F, Fut>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
1997 where
1998 Acc: Clone + Send + Sync + 'static,
1999 F: Fn(Acc, Out) -> Fut + Send + Sync + 'static,
2000 Fut: Future<Output = StreamResult<Acc>> + Send + 'static,
2001 {
2002 let stage = Arc::new(f);
2003 self.via(Flow::from_transform(move |mut input| {
2004 let stage = Arc::clone(&stage);
2005 let mut acc = Some(zero.clone());
2006 let mut done = false;
2007 Box::new(std::iter::from_fn(move || {
2008 if done {
2009 return None;
2010 }
2011 done = true;
2012
2013 for item in input.by_ref() {
2014 let item = match item {
2015 Ok(item) => item,
2016 Err(error) => return Some(Err(error)),
2017 };
2018 let current = acc.take().expect("fold_async accumulator present");
2019 match catch_unwind_failed("fold_async factory", || stage(current, item))
2020 .and_then(run_future_inline_or_spawn)
2021 {
2022 Ok(next) => acc = Some(next),
2023 Err(error) => return Some(Err(error)),
2024 }
2025 }
2026
2027 Some(Ok(acc.take().expect("fold_async accumulator present")))
2028 }))
2029 }))
2030 }
2031
2032 #[must_use]
2033 pub fn map_with_resource<Resource, Next, Create, F, Close>(
2034 self,
2035 create: Create,
2036 f: F,
2037 close: Close,
2038 ) -> Flow<In, Next, Mat>
2039 where
2040 Resource: Send + 'static,
2041 Next: Send + 'static,
2042 Create: Fn() -> StreamResult<Resource> + Send + Sync + 'static,
2043 F: Fn(&mut Resource, Out) -> StreamResult<Next> + Send + Sync + 'static,
2044 Close: Fn(Resource) -> StreamResult<Option<Next>> + Send + Sync + 'static,
2045 {
2046 let create = Arc::new(create);
2047 let stage = Arc::new(f);
2048 let close = Arc::new(close);
2049 self.via(Flow::from_transform(move |input| {
2050 Box::new(MapWithResourceStream {
2051 input,
2052 create: Arc::clone(&create),
2053 stage: Arc::clone(&stage),
2054 close: Arc::clone(&close),
2055 resource: None,
2056 created: false,
2057 pending_terminal: None,
2058 terminated: false,
2059 _marker: PhantomData,
2060 }) as BoxStream<Next>
2061 }))
2062 }
2063
2064 #[must_use]
2065 pub fn try_fold<Acc, F>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
2066 where
2067 Acc: Clone + Send + Sync + 'static,
2068 F: Fn(Acc, Out) -> StreamResult<Acc> + Send + Sync + 'static,
2069 {
2070 let stage = Arc::new(f);
2071 self.via(Flow::from_transform(move |input| {
2072 let stage = Arc::clone(&stage);
2073 let mut acc = zero.clone();
2074 for item in input {
2075 match item {
2076 Ok(item) => match stage(acc, item) {
2077 Ok(next) => acc = next,
2078 Err(error) => return Box::new(std::iter::once(Err(error))),
2079 },
2080 Err(error) => return Box::new(std::iter::once(Err(error))),
2081 }
2082 }
2083 Box::new(std::iter::once(Ok(acc)))
2084 }))
2085 }
2086
2087 #[must_use]
2089 #[deprecated(
2090 since = "0.9.0",
2091 note = "renamed to `try_fold` (idiomatic); the `_result` name still works"
2092 )]
2093 pub fn fold_result<Acc, F>(self, zero: Acc, f: F) -> Flow<In, Acc, Mat>
2094 where
2095 Acc: Clone + Send + Sync + 'static,
2096 F: Fn(Acc, Out) -> StreamResult<Acc> + Send + Sync + 'static,
2097 {
2098 self.try_fold(zero, f)
2099 }
2100
2101 #[must_use]
2102 pub fn fold_result_with_supervision<Acc, F>(
2103 self,
2104 zero: Acc,
2105 f: F,
2106 decider: SupervisionDecider,
2107 ) -> Flow<In, Acc, Mat>
2108 where
2109 Acc: Clone + Send + Sync + 'static,
2110 F: Fn(Acc, Out) -> StreamResult<Acc> + Send + Sync + 'static,
2111 {
2112 let stage = Arc::new(f);
2113 self.via(Flow::from_transform(move |input| {
2114 let stage = Arc::clone(&stage);
2115 let decider = Arc::clone(&decider);
2116 let mut acc = zero.clone();
2117 for item in input {
2118 match item {
2119 Ok(item) => {
2120 let previous = acc;
2121 match call_supervised("fold_result callback", || {
2122 stage(previous.clone(), item)
2123 }) {
2124 Ok(next) => acc = next,
2125 Err(error) => match decide_supervision(&decider, &error) {
2126 SupervisionDirective::Stop => {
2127 return Box::new(std::iter::once(Err(error)));
2128 }
2129 SupervisionDirective::Resume => acc = previous,
2130 SupervisionDirective::Restart => acc = zero.clone(),
2131 },
2132 }
2133 }
2134 Err(error) => return Box::new(std::iter::once(Err(error))),
2135 }
2136 }
2137 Box::new(std::iter::once(Ok(acc)))
2138 }))
2139 }
2140
2141 #[must_use]
2142 pub fn reduce<F>(self, f: F) -> Flow<In, Out, Mat>
2143 where
2144 F: Fn(Out, Out) -> Out + Send + Sync + 'static,
2145 {
2146 let stage = Arc::new(f);
2147 self.via(Flow::from_transform(move |mut input| {
2148 let Some(first) = input.next() else {
2149 return Box::new(std::iter::once(Err(StreamError::EmptyStream)));
2150 };
2151 let mut acc = match first {
2152 Ok(item) => item,
2153 Err(error) => return Box::new(std::iter::once(Err(error))),
2154 };
2155 for item in input {
2156 match item {
2157 Ok(item) => acc = stage(acc, item),
2158 Err(error) => return Box::new(std::iter::once(Err(error))),
2159 }
2160 }
2161 Box::new(std::iter::once(Ok(acc)))
2162 }))
2163 }
2164
2165 #[must_use]
2166 pub fn try_reduce<F>(self, f: F) -> Flow<In, Out, Mat>
2167 where
2168 Out: Clone,
2169 F: Fn(Out, Out) -> StreamResult<Out> + Send + Sync + 'static,
2170 {
2171 let stage = Arc::new(f);
2172 self.via(Flow::from_transform(move |mut input| {
2173 let Some(first) = input.next() else {
2174 return Box::new(std::iter::once(Err(StreamError::EmptyStream)));
2175 };
2176 let mut acc = match first {
2177 Ok(item) => item,
2178 Err(error) => return Box::new(std::iter::once(Err(error))),
2179 };
2180 for item in input {
2181 match item {
2182 Ok(item) => match stage(acc, item) {
2183 Ok(next) => acc = next,
2184 Err(error) => return Box::new(std::iter::once(Err(error))),
2185 },
2186 Err(error) => return Box::new(std::iter::once(Err(error))),
2187 }
2188 }
2189 Box::new(std::iter::once(Ok(acc)))
2190 }))
2191 }
2192
2193 #[must_use]
2195 #[deprecated(
2196 since = "0.9.0",
2197 note = "renamed to `try_reduce` (idiomatic); the `_result` name still works"
2198 )]
2199 pub fn reduce_result<F>(self, f: F) -> Flow<In, Out, Mat>
2200 where
2201 Out: Clone,
2202 F: Fn(Out, Out) -> StreamResult<Out> + Send + Sync + 'static,
2203 {
2204 self.try_reduce(f)
2205 }
2206
2207 #[must_use]
2208 pub fn reduce_result_with_supervision<F>(
2209 self,
2210 f: F,
2211 decider: SupervisionDecider,
2212 ) -> Flow<In, Out, Mat>
2213 where
2214 Out: Clone,
2215 F: Fn(Out, Out) -> StreamResult<Out> + Send + Sync + 'static,
2216 {
2217 let stage = Arc::new(f);
2218 self.via(Flow::from_transform(move |input| {
2219 let stage = Arc::clone(&stage);
2220 let decider = Arc::clone(&decider);
2221 let mut acc = None::<Out>;
2222 for item in input {
2223 match item {
2224 Ok(item) => {
2225 let Some(previous) = acc.take() else {
2226 acc = Some(item);
2227 continue;
2228 };
2229 match call_supervised("reduce_result callback", || {
2230 stage(previous.clone(), item)
2231 }) {
2232 Ok(next) => acc = Some(next),
2233 Err(error) => match decide_supervision(&decider, &error) {
2234 SupervisionDirective::Stop => {
2235 return Box::new(std::iter::once(Err(error)));
2236 }
2237 SupervisionDirective::Resume => acc = Some(previous),
2238 SupervisionDirective::Restart => acc = None,
2239 },
2240 }
2241 }
2242 Err(error) => return Box::new(std::iter::once(Err(error))),
2243 }
2244 }
2245 match acc {
2246 Some(acc) => Box::new(std::iter::once(Ok(acc))),
2247 None => Box::new(std::iter::once(Err(StreamError::EmptyStream))),
2248 }
2249 }))
2250 }
2251
2252 #[must_use]
2253 pub fn map_error<F>(self, f: F) -> Flow<In, Out, Mat>
2254 where
2255 F: Fn(StreamError) -> StreamError + Send + Sync + 'static,
2256 {
2257 let stage = Arc::new(f);
2258 self.via(Flow::from_transform(move |input| {
2259 let stage = Arc::clone(&stage);
2260 Box::new(input.map(move |item| item.map_err(|error| stage(error))))
2261 }))
2262 }
2263
2264 #[must_use]
2265 pub fn recover<F>(self, f: F) -> Flow<In, Out, Mat>
2266 where
2267 F: Fn(StreamError) -> Option<Out> + Send + Sync + 'static,
2268 {
2269 let stage = Arc::new(f);
2270 self.via(Flow::from_transform(move |mut input| {
2271 let stage = Arc::clone(&stage);
2272 let mut done = false;
2273 Box::new(std::iter::from_fn(move || {
2274 if done {
2275 return None;
2276 }
2277 match input.next()? {
2278 Ok(item) => Some(Ok(item)),
2279 Err(error) => {
2280 done = true;
2281 match stage(error.clone()) {
2282 Some(item) => Some(Ok(item)),
2283 None => Some(Err(error)),
2284 }
2285 }
2286 }
2287 }))
2288 }))
2289 }
2290
2291 #[must_use]
2292 pub fn recover_with<F>(self, f: F) -> Flow<In, Out, Mat>
2293 where
2294 F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2295 {
2296 self.recover_with_attempts(None, f)
2298 }
2299
2300 #[must_use]
2301 pub fn recover_with_retries<F>(self, retries: usize, f: F) -> Flow<In, Out, Mat>
2302 where
2303 F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2304 {
2305 self.recover_with_attempts(Some(retries), f)
2306 }
2307
2308 fn recover_with_attempts<F>(self, attempts: Option<usize>, f: F) -> Flow<In, Out, Mat>
2309 where
2310 F: Fn(StreamError) -> Option<Source<Out>> + Send + Sync + 'static,
2311 {
2312 let stage = Arc::new(f);
2313 self.via(Flow::from_runtime_transform(move |input, materializer| {
2314 let replacement_materializer =
2315 materializer.with_name_prefix(materializer.name_prefix().to_owned());
2316 let stage = Arc::clone(&stage);
2317 let mut remaining_retries = attempts;
2319 let mut current = input;
2320 let mut terminated = false;
2321
2322 Ok(Box::new(std::iter::from_fn(move || {
2323 if terminated {
2324 return None;
2325 }
2326
2327 loop {
2328 match current.next() {
2329 Some(Ok(item)) => return Some(Ok(item)),
2330 Some(Err(error)) if remaining_retries != Some(0) => {
2331 if let Some(remaining) = remaining_retries.as_mut() {
2332 *remaining -= 1;
2333 }
2334 match stage(error.clone()) {
2335 Some(source) => match Arc::clone(&source.factory)
2336 .create(&replacement_materializer)
2337 {
2338 Ok((stream, _)) => current = stream,
2339 Err(error) => {
2340 terminated = true;
2341 return Some(Err(error));
2342 }
2343 },
2344 None => {
2345 terminated = true;
2346 return Some(Err(error));
2347 }
2348 }
2349 }
2350 Some(Err(error)) => {
2351 terminated = true;
2352 return Some(Err(error));
2353 }
2354 None => {
2355 terminated = true;
2356 return None;
2357 }
2358 }
2359 }
2360 })) as BoxStream<Out>)
2361 }))
2362 }
2363
2364 #[must_use]
2365 pub fn on_error_complete(self) -> Flow<In, Out, Mat> {
2366 self.via(Flow::from_transform(move |mut input| {
2367 let mut done = false;
2368 Box::new(std::iter::from_fn(move || {
2369 if done {
2370 return None;
2371 }
2372 match input.next()? {
2373 Ok(item) => Some(Ok(item)),
2374 Err(_) => {
2375 done = true;
2376 None
2377 }
2378 }
2379 }))
2380 }))
2381 }
2382
2383 #[must_use]
2384 pub fn prefix_and_tail(self, n: usize) -> Flow<In, (Vec<Out>, Source<Out>), Mat> {
2385 self.via(Flow::from_runtime_transform(move |input, _materializer| {
2386 Ok(prefix_and_tail_stream(input, n))
2387 }))
2388 }
2389
2390 #[must_use]
2391 pub fn flat_map_prefix<Next, NextMat, F>(self, n: usize, f: F) -> Flow<In, Next, Mat>
2392 where
2393 Next: Send + 'static,
2394 NextMat: Send + 'static,
2395 F: Fn(Vec<Out>) -> Flow<Out, Next, NextMat> + Send + Sync + 'static,
2396 Out: Clone,
2397 {
2398 let stage = Arc::new(f);
2399 self.via(Flow::from_runtime_transform(
2400 move |mut input, materializer| {
2401 let mut prefix = Vec::with_capacity(n);
2402 while prefix.len() < n {
2403 match input.next() {
2404 Some(Ok(item)) => prefix.push(item),
2405 Some(Err(error)) => {
2406 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>);
2407 }
2408 None => break,
2409 }
2410 }
2411
2412 let flow = stage(prefix);
2413 let transform = flow.transform;
2414 let _ = (flow.materialize)()?;
2415 match transform {
2416 FlowTransform::Pure(transform) => Ok(transform(input)),
2417 FlowTransform::Runtime(transform) => transform(input, materializer),
2418 }
2419 },
2420 ))
2421 }
2422
2423 #[must_use]
2424 pub fn group_by<Key, F>(
2425 self,
2426 max_substreams: usize,
2427 f: F,
2428 allow_closed_substream_recreation: bool,
2429 ) -> Flow<In, Source<Out>, Mat>
2430 where
2431 Key: Clone + Eq + Hash + Send + 'static,
2432 F: Fn(&Out) -> Key + Send + Sync + 'static,
2433 Out: Clone,
2434 {
2435 self.group_by_with_batching(
2436 max_substreams,
2437 f,
2438 allow_closed_substream_recreation,
2439 GroupByBatchMode::Immediate,
2440 )
2441 }
2442
2443 pub(super) fn group_by_with_batching<Key, F>(
2444 self,
2445 max_substreams: usize,
2446 f: F,
2447 allow_closed_substream_recreation: bool,
2448 batch_mode: GroupByBatchMode,
2449 ) -> Flow<In, Source<Out>, Mat>
2450 where
2451 Key: Clone + Eq + Hash + Send + 'static,
2452 F: Fn(&Out) -> Key + Send + Sync + 'static,
2453 Out: Clone,
2454 {
2455 assert!(
2456 max_substreams > 0,
2457 "group_by max_substreams must be greater than zero"
2458 );
2459 let key_fn = Arc::new(f);
2460 self.via(Flow::from_runtime_transform(move |input, materializer| {
2461 Ok(group_by_stream(
2462 input,
2463 max_substreams,
2464 allow_closed_substream_recreation,
2465 Arc::clone(&key_fn),
2466 batch_mode,
2467 materializer,
2468 ))
2469 }))
2470 }
2471
2472 #[must_use]
2473 pub fn split_when<F>(self, predicate: F) -> Flow<In, Source<Out>, Mat>
2474 where
2475 F: Fn(&Out) -> bool + Send + Sync + 'static,
2476 Out: Clone,
2477 {
2478 let predicate = Arc::new(predicate);
2479 self.via(Flow::from_runtime_transform(move |input, materializer| {
2480 Ok(split_streams(
2481 input,
2482 SplitMode::When,
2483 Arc::clone(&predicate),
2484 materializer,
2485 ))
2486 }))
2487 }
2488
2489 #[must_use]
2490 pub fn split_after<F>(self, predicate: F) -> Flow<In, Source<Out>, Mat>
2491 where
2492 F: Fn(&Out) -> bool + Send + Sync + 'static,
2493 Out: Clone,
2494 {
2495 let predicate = Arc::new(predicate);
2496 self.via(Flow::from_runtime_transform(move |input, materializer| {
2497 Ok(split_streams(
2498 input,
2499 SplitMode::After,
2500 Arc::clone(&predicate),
2501 materializer,
2502 ))
2503 }))
2504 }
2505
2506 #[must_use]
2507 pub fn flat_map_concat<Next, NextMat, F>(self, f: F) -> Flow<In, Next, Mat>
2508 where
2509 Next: Send + 'static,
2510 NextMat: Send + 'static,
2511 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
2512 {
2513 let stage = Arc::new(f);
2514 self.via(Flow::from_runtime_transform(move |input, materializer| {
2515 Ok(flat_map_concat_stream(
2516 input,
2517 Arc::clone(&stage),
2518 materializer,
2519 ))
2520 }))
2521 }
2522
2523 #[must_use]
2524 pub fn flat_map_merge<Next, NextMat, F>(self, breadth: usize, f: F) -> Flow<In, Next, Mat>
2525 where
2526 Next: Send + 'static,
2527 NextMat: Send + 'static,
2528 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
2529 {
2530 assert!(
2531 breadth > 0,
2532 "flat_map_merge breadth must be greater than zero"
2533 );
2534 let stage = Arc::new(f);
2535 self.via(Flow::from_runtime_transform_with_hints(
2536 move |input, materializer| {
2537 Ok(flat_map_merge_stream(
2538 input,
2539 breadth,
2540 Arc::clone(&stage),
2541 materializer,
2542 ))
2543 },
2544 FlowHints::PRESERVES_TERMINAL_CONSUMER_BATCH,
2545 ))
2546 }
2547
2548 #[must_use]
2549 pub fn concat<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2550 where
2551 Mat2: Send + 'static,
2552 {
2553 self.concat_sources([that])
2554 }
2555
2556 #[must_use]
2557 pub fn concat_lazy<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2558 where
2559 Mat2: Send + 'static,
2560 {
2561 let that_factory = that.factory;
2562 self.via(Flow::from_runtime_transform(move |input, materializer| {
2563 let primary = input;
2564 Ok(concat_streams_lazy(
2565 primary,
2566 vec![Arc::clone(&that_factory)],
2567 materializer,
2568 ))
2569 }))
2570 }
2571
2572 #[must_use]
2573 pub fn concat_all_lazy<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
2574 where
2575 Mat2: Send + 'static,
2576 I: IntoIterator<Item = Source<Out, Mat2>>,
2577 {
2578 let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2579 self.via(Flow::from_runtime_transform(move |input, materializer| {
2580 Ok(concat_streams_lazy(
2581 input,
2582 source_factories.clone(),
2583 materializer,
2584 ))
2585 }))
2586 }
2587
2588 #[must_use]
2589 pub fn prepend<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2590 where
2591 Mat2: Send + 'static,
2592 {
2593 self.prepend_sources([that])
2594 }
2595
2596 #[must_use]
2597 pub fn prepend_lazy<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2598 where
2599 Mat2: Send + 'static,
2600 {
2601 self.prepend(that)
2602 }
2603
2604 #[must_use]
2605 pub fn or_else<Mat2>(self, secondary: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2606 where
2607 Mat2: Send + 'static,
2608 {
2609 let secondary_factory = secondary.factory;
2610 self.via(Flow::from_runtime_transform(move |input, materializer| {
2611 let secondary = match Arc::clone(&secondary_factory).create(materializer) {
2612 Ok((stream, _)) => stream,
2613 Err(error) => Box::new(std::iter::once(Err(error))) as BoxStream<Out>,
2614 };
2615 Ok(or_else_stream(input, secondary))
2616 }))
2617 }
2618
2619 #[must_use]
2620 pub fn interleave<Mat2>(
2621 self,
2622 that: Source<Out, Mat2>,
2623 segment_size: usize,
2624 ) -> Flow<In, Out, Mat>
2625 where
2626 Mat2: Send + 'static,
2627 {
2628 self.interleave_all([that], segment_size, false)
2629 }
2630
2631 #[must_use]
2632 pub fn interleave_all<Mat2, I>(
2633 self,
2634 those: I,
2635 segment_size: usize,
2636 eager_close: bool,
2637 ) -> Flow<In, Out, Mat>
2638 where
2639 Mat2: Send + 'static,
2640 I: IntoIterator<Item = Source<Out, Mat2>>,
2641 {
2642 let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2643 self.via(Flow::from_runtime_transform(move |input, materializer| {
2644 let mut streams = Vec::with_capacity(source_factories.len() + 1);
2645 streams.push(input);
2646 for factory in &source_factories {
2647 let stream = match Arc::clone(factory).create(materializer) {
2648 Ok((stream, _)) => stream,
2649 Err(error) => {
2650 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2651 }
2652 };
2653 streams.push(stream);
2654 }
2655 Ok(interleave_streams(streams, segment_size, eager_close))
2656 }))
2657 }
2658
2659 #[must_use]
2660 pub fn merge_sorted<Mat2>(self, that: Source<Out, Mat2>) -> Flow<In, Out, Mat>
2661 where
2662 Out: Ord,
2663 Mat2: Send + 'static,
2664 {
2665 let source_factory = that.factory;
2666 self.via(Flow::from_runtime_transform(move |input, materializer| {
2667 let other = match Arc::clone(&source_factory).create(materializer) {
2668 Ok((stream, _)) => stream,
2669 Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>),
2670 };
2671 Ok(merge_sorted_stream(input, other))
2672 }))
2673 }
2674
2675 #[must_use]
2676 pub fn merge_latest<Mat2>(
2677 self,
2678 that: Source<Out, Mat2>,
2679 eager_complete: bool,
2680 ) -> Flow<In, Vec<Out>, Mat>
2681 where
2682 Out: Clone,
2683 Mat2: Send + 'static,
2684 {
2685 let source_factory = that.factory;
2686 self.via(Flow::from_runtime_transform(move |input, materializer| {
2687 let other = match Arc::clone(&source_factory).create(materializer) {
2688 Ok((stream, _)) => stream,
2689 Err(error) => {
2690 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Vec<Out>>);
2691 }
2692 };
2693 Ok(merge_latest_streams(vec![input, other], eager_complete))
2694 }))
2695 }
2696
2697 #[must_use]
2698 pub fn merge_all<Mat2, I>(self, those: I, eager_complete: bool) -> Flow<In, Out, Mat>
2699 where
2700 Mat2: Send + 'static,
2701 I: IntoIterator<Item = Source<Out, Mat2>>,
2702 {
2703 let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
2704 self.via(Flow::from_runtime_transform(move |input, materializer| {
2705 let mut streams = Vec::with_capacity(source_factories.len() + 1);
2706 streams.push(input);
2707 for factory in &source_factories {
2708 let stream = match Arc::clone(factory).create(materializer) {
2709 Ok((stream, _)) => stream,
2710 Err(error) => {
2711 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
2712 }
2713 };
2714 streams.push(stream);
2715 }
2716 Ok(merge_streams(streams, eager_complete))
2717 }))
2718 }
2719
2720 #[must_use]
2721 pub fn zip_with<Mat2, Out2, Next, F>(
2722 self,
2723 that: Source<Out2, Mat2>,
2724 combine: F,
2725 ) -> Flow<In, Next, Mat>
2726 where
2727 Out2: Send + 'static,
2728 Next: Send + 'static,
2729 Mat2: Send + 'static,
2730 F: Fn(Out, Out2) -> Next + Send + Sync + 'static,
2731 {
2732 let source_factory = that.factory;
2733 let combine = Arc::new(combine);
2734 self.via(Flow::from_runtime_transform(move |input, materializer| {
2735 let other = match Arc::clone(&source_factory).create(materializer) {
2736 Ok((stream, _)) => stream,
2737 Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>),
2738 };
2739 let combine = Arc::clone(&combine);
2740 Ok(Box::new(
2741 zip_streams(input, other)
2742 .map(move |item| item.map(|(left, right)| combine(left, right))),
2743 ) as BoxStream<Next>)
2744 }))
2745 }
2746
2747 #[must_use]
2748 pub fn zip_latest<Mat2, Out2>(self, that: Source<Out2, Mat2>) -> Flow<In, (Out, Out2), Mat>
2749 where
2750 Out: Clone,
2751 Out2: Clone + Send + 'static,
2752 Mat2: Send + 'static,
2753 {
2754 self.zip_latest_with(that, true, |left, right| (left, right))
2755 }
2756
2757 #[must_use]
2758 pub fn zip_latest_with<Mat2, Out2, Next, F>(
2759 self,
2760 that: Source<Out2, Mat2>,
2761 eager_complete: bool,
2762 combine: F,
2763 ) -> Flow<In, Next, Mat>
2764 where
2765 Out: Clone,
2766 Out2: Clone + Send + 'static,
2767 Next: Send + 'static,
2768 Mat2: Send + 'static,
2769 F: Fn(Out, Out2) -> Next + Send + Sync + 'static,
2770 {
2771 let source_factory = that.factory;
2772 let combine = Arc::new(combine);
2773 self.via(Flow::from_runtime_transform(move |input, materializer| {
2774 let other = match Arc::clone(&source_factory).create(materializer) {
2775 Ok((stream, _)) => stream,
2776 Err(error) => return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Next>),
2777 };
2778 Ok(zip_latest_with_stream(
2779 input,
2780 other,
2781 eager_complete,
2782 Arc::clone(&combine),
2783 ))
2784 }))
2785 }
2786
2787 #[must_use]
2788 pub fn zip_with_index(self) -> Flow<In, (Out, u64), Mat> {
2789 self.via(Flow::from_runtime_transform(
2790 move |mut input, _materializer| {
2791 let mut index = 0_u64;
2792 Ok(Box::new(std::iter::from_fn(move || {
2793 input.next().map(|item| {
2794 item.map(|value| {
2795 let pair = (value, index);
2796 index = index.wrapping_add(1);
2797 pair
2798 })
2799 })
2800 })) as BoxStream<(Out, u64)>)
2801 },
2802 ))
2803 }
2804
2805 #[must_use]
2806 pub fn zip_all<Mat2, Out2>(
2807 self,
2808 that: Source<Out2, Mat2>,
2809 this_elem: Out,
2810 that_elem: Out2,
2811 ) -> Flow<In, (Out, Out2), Mat>
2812 where
2813 Out: Clone + Sync,
2814 Out2: Clone + Send + Sync + 'static,
2815 Mat2: Send + 'static,
2816 {
2817 let source_factory = that.factory;
2818 self.via(Flow::from_runtime_transform(move |input, materializer| {
2819 let other = match Arc::clone(&source_factory).create(materializer) {
2820 Ok((stream, _)) => stream,
2821 Err(error) => {
2822 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<(Out, Out2)>);
2823 }
2824 };
2825 Ok(zip_all_stream(
2826 input,
2827 other,
2828 this_elem.clone(),
2829 that_elem.clone(),
2830 ))
2831 }))
2832 }
2833
2834 #[must_use]
2835 pub fn also_to<SideMat>(self, sink: Sink<Out, SideMat>) -> Flow<In, Out, Mat>
2836 where
2837 Out: Clone,
2838 SideMat: Send + 'static,
2839 {
2840 self.via(Flow::from_runtime_transform(
2841 move |mut input: BoxStream<Out>, materializer| {
2842 let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 0)?;
2843 let mut sender = Some(side_sender);
2844 let side_mat = side_mat;
2845 Ok(Box::new(std::iter::from_fn(move || match input.next() {
2846 Some(Ok(item)) => {
2847 let _ = &side_mat;
2848 if sender
2849 .as_ref()
2850 .is_some_and(|sender| sender.send(Ok(item.clone())).is_err())
2851 {
2852 sender = None;
2853 return None;
2854 }
2855 Some(Ok(item))
2856 }
2857 Some(Err(error)) => {
2858 let _ = &side_mat;
2859 let _ = sender
2860 .as_ref()
2861 .and_then(|sender| sender.send(Err(error.clone())).ok());
2862 sender = None;
2863 Some(Err(error))
2864 }
2865 None => {
2866 let _ = &side_mat;
2867 sender = None;
2868 None
2869 }
2870 })) as BoxStream<Out>)
2871 },
2872 ))
2873 }
2874
2875 #[must_use]
2876 pub fn also_to_all<SideMat, I>(self, sinks: I) -> Flow<In, Out, Mat>
2877 where
2878 Out: Clone,
2879 SideMat: Send + 'static,
2880 I: IntoIterator<Item = Sink<Out, SideMat>>,
2881 {
2882 let sinks: Vec<_> = sinks.into_iter().collect();
2883 if sinks.is_empty() {
2884 return self;
2885 }
2886
2887 self.via(Flow::from_runtime_transform(
2888 move |mut input: BoxStream<Out>, materializer| {
2889 let mut sides = sinks
2890 .iter()
2891 .map(|sink| materialize_side_sink(sink, materializer, 0))
2892 .collect::<StreamResult<Vec<_>>>()?;
2893 Ok(Box::new(std::iter::from_fn(move || match input.next() {
2894 Some(Ok(item)) => {
2895 for (sender, _) in &sides {
2896 if sender.send(Ok(item.clone())).is_err() {
2897 sides.clear();
2898 return None;
2899 }
2900 }
2901 Some(Ok(item))
2902 }
2903 Some(Err(error)) => {
2904 for (sender, _) in &sides {
2905 let _ = sender.send(Err(error.clone())).ok();
2906 }
2907 sides.clear();
2908 Some(Err(error))
2909 }
2910 None => {
2911 sides.clear();
2912 None
2913 }
2914 })) as BoxStream<Out>)
2915 },
2916 ))
2917 }
2918
2919 #[must_use]
2920 pub fn divert_to<SideMat, F>(self, sink: Sink<Out, SideMat>, predicate: F) -> Flow<In, Out, Mat>
2921 where
2922 SideMat: Send + 'static,
2923 F: Fn(&Out) -> bool + Send + Sync + 'static,
2924 {
2925 let predicate = Arc::new(predicate);
2926 self.via(Flow::from_runtime_transform(
2927 move |mut input: BoxStream<Out>, materializer| {
2928 let predicate = Arc::clone(&predicate);
2929 let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 0)?;
2930 let mut sender = Some(side_sender);
2931 let side_mat = side_mat;
2932 Ok(Box::new(std::iter::from_fn(move || {
2933 loop {
2934 let _ = &side_mat;
2935 match input.next() {
2936 Some(Ok(item)) if predicate(&item) => {
2937 if sender
2938 .as_ref()
2939 .is_some_and(|sender| sender.send(Ok(item)).is_err())
2940 {
2941 sender = None;
2942 return None;
2943 }
2944 }
2945 Some(Ok(item)) => return Some(Ok(item)),
2946 Some(Err(error)) => {
2947 let _ = sender
2948 .as_ref()
2949 .and_then(|sender| sender.send(Err(error.clone())).ok());
2950 sender = None;
2951 return Some(Err(error));
2952 }
2953 None => {
2954 sender = None;
2955 return None;
2956 }
2957 }
2958 }
2959 })) as BoxStream<Out>)
2960 },
2961 ))
2962 }
2963
2964 #[must_use]
2965 pub fn wire_tap<SideMat>(self, sink: Sink<Out, SideMat>) -> Flow<In, Out, Mat>
2966 where
2967 Out: Clone,
2968 SideMat: Send + 'static,
2969 {
2970 self.via(Flow::from_runtime_transform(
2971 move |mut input: BoxStream<Out>, materializer| {
2972 let (side_sender, side_mat) = materialize_side_sink(&sink, materializer, 1)?;
2973 let mut sender = Some(side_sender);
2974 let side_mat = side_mat;
2975 Ok(Box::new(std::iter::from_fn(move || match input.next() {
2976 Some(Ok(item)) => {
2977 let _ = &side_mat;
2978 if let Some(side) = sender.as_ref() {
2979 match side.try_send(Ok(item.clone())) {
2980 Ok(()) | Err(std::sync::mpsc::TrySendError::Full(_)) => {}
2981 Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
2982 sender = None
2983 }
2984 }
2985 }
2986 Some(Ok(item))
2987 }
2988 Some(Err(error)) => {
2989 let _ = &side_mat;
2990 if let Some(side) = sender.as_ref() {
2991 match side.try_send(Err(error.clone())) {
2992 Ok(())
2993 | Err(std::sync::mpsc::TrySendError::Full(_))
2994 | Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {}
2995 }
2996 }
2997 sender = None;
2998 Some(Err(error))
2999 }
3000 None => {
3001 let _ = &side_mat;
3002 sender = None;
3003 None
3004 }
3005 })) as BoxStream<Out>)
3006 },
3007 ))
3008 }
3009
3010 fn concat_sources<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
3011 where
3012 Mat2: Send + 'static,
3013 I: IntoIterator<Item = Source<Out, Mat2>>,
3014 {
3015 let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
3016 self.via(Flow::from_runtime_transform(move |input, materializer| {
3017 let mut streams = Vec::with_capacity(source_factories.len() + 1);
3018 streams.push(input);
3019 for factory in &source_factories {
3020 let stream = match Arc::clone(factory).create(materializer) {
3021 Ok((stream, _)) => stream,
3022 Err(error) => {
3023 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
3024 }
3025 };
3026 streams.push(stream);
3027 }
3028 Ok(concat_streams(streams))
3029 }))
3030 }
3031
3032 fn prepend_sources<Mat2, I>(self, those: I) -> Flow<In, Out, Mat>
3033 where
3034 Mat2: Send + 'static,
3035 I: IntoIterator<Item = Source<Out, Mat2>>,
3036 {
3037 let source_factories: Vec<_> = those.into_iter().map(|source| source.factory).collect();
3038 self.via(Flow::from_runtime_transform(move |input, materializer| {
3039 let mut streams = Vec::with_capacity(source_factories.len() + 1);
3040 for factory in &source_factories {
3041 let stream = match Arc::clone(factory).create(materializer) {
3042 Ok((stream, _)) => stream,
3043 Err(error) => {
3044 return Ok(Box::new(std::iter::once(Err(error))) as BoxStream<Out>);
3045 }
3046 };
3047 streams.push(stream);
3048 }
3049 streams.push(input);
3050 Ok(concat_streams(streams))
3051 }))
3052 }
3053
3054 #[must_use]
3057 pub fn intersperse(self, inject: Out) -> Flow<In, Out, Mat>
3058 where
3059 Out: Clone + Sync,
3060 {
3061 let inject = Arc::new(inject);
3062 self.via(Flow::from_transform(move |mut input| {
3063 let inject = Arc::clone(&inject);
3064 let mut first = true;
3065 Box::new(std::iter::from_fn(move || {
3066 if first {
3067 first = false;
3068 match input.next() {
3069 None => None,
3070 Some(item) => {
3071 if item.is_err() {
3072 first = true;
3073 }
3074 Some(item)
3075 }
3076 }
3077 } else {
3078 match input.next() {
3079 None => None,
3080 Some(item) => {
3081 if item.is_err() {
3082 first = true;
3083 Some(item)
3084 } else {
3085 Some(Ok((*inject).clone()))
3086 }
3087 }
3088 }
3089 }
3090 }))
3091 }))
3092 }
3093
3094 #[must_use]
3095 pub fn flatten_optional<Inner>(self) -> Flow<In, Inner, Mat>
3096 where
3097 Out: Into<Option<Inner>>,
3098 Inner: Send + 'static,
3099 {
3100 self.filter_map(|item| item.into())
3101 }
3102
3103 #[must_use]
3104 pub fn grouped_weighted<F>(self, max_weight: usize, cost_fn: F) -> Flow<In, Vec<Out>, Mat>
3105 where
3106 F: Fn(&Out) -> usize + Send + Sync + 'static,
3107 {
3108 let cost_fn = Arc::new(cost_fn);
3109 self.via(Flow::from_transform(move |mut input| {
3110 let cost_fn = Arc::clone(&cost_fn);
3111 Box::new(std::iter::from_fn(move || {
3112 let mut group = Vec::new();
3113 let mut weight = 0usize;
3114 while weight < max_weight {
3115 match input.next() {
3116 Some(Ok(item)) => {
3117 let item_weight = cost_fn(&item);
3118 if weight > 0 && weight + item_weight > max_weight {
3119 group.push(item);
3120 break;
3121 }
3122 weight += item_weight;
3123 group.push(item);
3124 }
3125 Some(Err(error)) => return Some(Err(error)),
3126 None => break,
3127 }
3128 }
3129 if group.is_empty() {
3130 None
3131 } else {
3132 Some(Ok(group))
3133 }
3134 }))
3135 }))
3136 }
3137
3138 #[must_use]
3139 pub fn limit_weighted<F>(self, max_weight: usize, cost_fn: F) -> Flow<In, Out, Mat>
3140 where
3141 F: Fn(&Out) -> usize + Send + Sync + 'static,
3142 {
3143 let cost_fn = Arc::new(cost_fn);
3144 self.via(Flow::from_transform(move |mut input| {
3145 let cost_fn = Arc::clone(&cost_fn);
3146 let mut weight = 0usize;
3147 Box::new(std::iter::from_fn(move || match input.next()? {
3148 Ok(item) => {
3149 let item_weight = cost_fn(&item);
3150 if weight + item_weight > max_weight {
3151 Some(Err(StreamError::LimitExceeded {
3152 max: max_weight as u64,
3153 }))
3154 } else {
3155 weight += item_weight;
3156 Some(Ok(item))
3157 }
3158 }
3159 Err(error) => Some(Err(error)),
3160 }))
3161 }))
3162 }
3163
3164 #[must_use]
3165 pub fn contramap<NewIn, F>(self, f: F) -> Flow<NewIn, Out, Mat>
3166 where
3167 NewIn: Send + 'static,
3168 F: Fn(NewIn) -> In + Send + Sync + 'static,
3169 {
3170 let stage = Arc::new(f);
3171 Flow::from_transform(move |input| {
3172 let stage = Arc::clone(&stage);
3173 Box::new(input.map(move |item| item.map(|item| stage(item))))
3174 })
3175 .via_mat(self, Keep::right)
3176 }
3177
3178 #[must_use]
3179 pub fn monitor<F>(self, f: F) -> Flow<In, Out, Mat>
3180 where
3181 Out: Clone,
3182 F: Fn(&Out) + Send + Sync + 'static,
3183 {
3184 let stage = Arc::new(f);
3185 self.via(Flow::from_transform(move |input| {
3186 let stage = Arc::clone(&stage);
3187 Box::new(input.map(move |item| match item {
3188 Ok(item) => {
3189 stage(&item);
3190 Ok(item)
3191 }
3192 Err(error) => Err(error),
3193 }))
3194 }))
3195 }
3196
3197 #[must_use]
3198 pub fn watch_termination<CallbackMat, F>(self, materialize_callback: F) -> Flow<In, Out, Mat>
3199 where
3200 Mat: Clone,
3201 CallbackMat: Send + 'static,
3202 F: Fn(Mat) -> CallbackMat + Send + Sync + 'static,
3203 {
3204 let cb = Arc::new(materialize_callback);
3205 self.map_materialized_value(move |mat| {
3206 let _ = cb(mat.clone());
3207 mat
3208 })
3209 }
3210
3211 #[must_use]
3212 pub fn to<SinkMat>(self, sink: Sink<Out, SinkMat>) -> Sink<In, Mat>
3213 where
3214 SinkMat: Send + 'static,
3215 {
3216 self.to_mat(sink, Keep::left)
3217 }
3218
3219 #[must_use]
3220 pub fn to_mat<SinkMat, Combined, F>(
3221 self,
3222 sink: Sink<Out, SinkMat>,
3223 combine: F,
3224 ) -> Sink<In, Combined>
3225 where
3226 SinkMat: Send + 'static,
3227 Combined: Send + 'static,
3228 F: Fn(Mat, SinkMat) -> Combined + Send + Sync + 'static,
3229 {
3230 let transform = self.transform;
3231 let materialize = self.materialize;
3232 let combine = Arc::new(combine);
3233 Sink::from_runner(move |input, materializer| {
3234 let flow_mat = materialize()?;
3235 let input = match &transform {
3236 FlowTransform::Pure(transform) => transform(input),
3237 FlowTransform::Runtime(transform) => transform(input, materializer)?,
3238 };
3239 let sink_mat = sink.run(input, materializer)?;
3240 Ok(combine(flow_mat, sink_mat))
3241 })
3242 }
3243
3244 #[must_use]
3245 pub fn map_materialized_value<NextMat, F>(self, f: F) -> Flow<In, Out, NextMat>
3246 where
3247 NextMat: Send + 'static,
3248 F: Fn(Mat) -> NextMat + Send + Sync + 'static,
3249 {
3250 let transform = self.transform;
3251 let materialize = self.materialize;
3252 let hints = self.hints;
3253 let attributes = self.attributes;
3254 let f = Arc::new(f);
3255 match transform {
3256 FlowTransform::Pure(transform) => Flow::from_parts_with_hints(
3257 move |input| transform(input),
3258 move || {
3259 let mat = materialize()?;
3260 Ok(f(mat))
3261 },
3262 hints,
3263 )
3264 .with_attributes(attributes),
3265 FlowTransform::Runtime(transform) => Flow {
3266 transform: FlowTransform::Runtime(transform),
3267 materialize: Arc::new(move || {
3268 let mat = materialize()?;
3269 Ok(f(mat))
3270 }),
3271 hints,
3272 attributes,
3273 },
3274 }
3275 }
3276}
3277
3278struct MapWithResourceStream<In, Out, Resource, Create, F, Close>
3279where
3280 Create: Fn() -> StreamResult<Resource>,
3281 F: Fn(&mut Resource, In) -> StreamResult<Out>,
3282 Close: Fn(Resource) -> StreamResult<Option<Out>>,
3283{
3284 input: BoxStream<In>,
3285 create: Arc<Create>,
3286 stage: Arc<F>,
3287 close: Arc<Close>,
3288 resource: Option<Resource>,
3289 created: bool,
3290 pending_terminal: Option<StreamError>,
3291 terminated: bool,
3292 _marker: PhantomData<fn() -> Out>,
3293}
3294
3295impl<In, Out, Resource, Create, F, Close> MapWithResourceStream<In, Out, Resource, Create, F, Close>
3296where
3297 Create: Fn() -> StreamResult<Resource>,
3298 F: Fn(&mut Resource, In) -> StreamResult<Out>,
3299 Close: Fn(Resource) -> StreamResult<Option<Out>>,
3300{
3301 fn ensure_created(&mut self) -> StreamResult<()> {
3302 if self.created {
3303 return Ok(());
3304 }
3305 self.created = true;
3306 let resource = catch_unwind_failed("map_with_resource create", || (self.create)())
3307 .and_then(|result| result)?;
3308 self.resource = Some(resource);
3309 Ok(())
3310 }
3311
3312 fn close_resource(&mut self) -> StreamResult<Option<Out>> {
3313 match self.resource.take() {
3314 Some(resource) => {
3315 catch_unwind_failed("map_with_resource close", || (self.close)(resource))
3316 .and_then(|result| result)
3317 }
3318 None => Ok(None),
3319 }
3320 }
3321
3322 fn close_with_terminal(&mut self, terminal: Option<StreamError>) -> Option<StreamResult<Out>> {
3323 self.terminated = terminal.is_none();
3324 match self.close_resource() {
3325 Ok(Some(item)) => {
3326 self.pending_terminal = terminal;
3327 Some(Ok(item))
3328 }
3329 Ok(None) => match terminal {
3330 Some(error) => {
3331 self.terminated = true;
3332 Some(Err(error))
3333 }
3334 None => {
3335 self.terminated = true;
3336 None
3337 }
3338 },
3339 Err(error) => {
3340 self.terminated = true;
3341 Some(Err(terminal.unwrap_or(error)))
3342 }
3343 }
3344 }
3345}
3346
3347impl<In, Out, Resource, Create, F, Close> Iterator
3348 for MapWithResourceStream<In, Out, Resource, Create, F, Close>
3349where
3350 Create: Fn() -> StreamResult<Resource>,
3351 F: Fn(&mut Resource, In) -> StreamResult<Out>,
3352 Close: Fn(Resource) -> StreamResult<Option<Out>>,
3353{
3354 type Item = StreamResult<Out>;
3355
3356 fn next(&mut self) -> Option<Self::Item> {
3357 if let Some(error) = self.pending_terminal.take() {
3358 self.terminated = true;
3359 return Some(Err(error));
3360 }
3361 if self.terminated {
3362 return None;
3363 }
3364 if let Err(error) = self.ensure_created() {
3365 self.terminated = true;
3366 return Some(Err(error));
3367 }
3368
3369 match self.input.next() {
3370 Some(Ok(item)) => {
3371 let result = {
3372 let resource = self
3373 .resource
3374 .as_mut()
3375 .expect("map_with_resource resource is open");
3376 catch_unwind_failed("map_with_resource function", || {
3377 (self.stage)(resource, item)
3378 })
3379 .and_then(|result| result)
3380 };
3381 match result {
3382 Ok(item) => Some(Ok(item)),
3383 Err(error) => self.close_with_terminal(Some(error)),
3384 }
3385 }
3386 Some(Err(error)) => self.close_with_terminal(Some(error)),
3387 None => self.close_with_terminal(None),
3388 }
3389 }
3390}
3391
3392impl<I1: Send + 'static, O1: Send + 'static, I2: Send + 'static, O2: Send + 'static>
3393 BidiFlow<I1, O1, I2, O2>
3394{
3395 #[must_use]
3396 pub fn from_flows<Mat1, Mat2>(top: Flow<I1, O1, Mat1>, bottom: Flow<I2, O2, Mat2>) -> Self
3397 where
3398 Mat1: Send + 'static,
3399 Mat2: Send + 'static,
3400 {
3401 Self {
3402 top: top.map_materialized_value(|_| NotUsed),
3403 bottom: bottom.map_materialized_value(|_| NotUsed),
3404 attributes: Attributes::default(),
3405 }
3406 }
3407}
3408
3409impl<I1: Send + 'static, O1: Send + 'static, I2: Send + 'static, O2: Send + 'static>
3410 BidiFlow<I1, O1, I2, O2>
3411{
3412 #[must_use]
3413 pub fn attributes(&self) -> &Attributes {
3414 &self.attributes
3415 }
3416
3417 #[must_use]
3418 pub fn with_attributes(mut self, attributes: Attributes) -> Self {
3419 self.attributes = attributes;
3420 self
3421 }
3422
3423 #[must_use]
3424 pub fn add_attributes(mut self, attributes: Attributes) -> Self {
3425 self.attributes = self.attributes.and(attributes);
3426 self
3427 }
3428
3429 #[must_use]
3430 pub fn named(self, name: impl Into<String>) -> Self {
3431 self.add_attributes(Attributes::named(name))
3432 }
3433
3434 #[must_use]
3435 pub fn join<Mat2>(self, flow: Flow<O1, I2, Mat2>) -> Flow<I1, O2, NotUsed>
3436 where
3437 Mat2: Send + 'static,
3438 {
3439 self.top
3440 .via(flow)
3441 .via(self.bottom)
3442 .map_materialized_value(|_| NotUsed)
3443 .with_attributes(self.attributes)
3444 }
3445
3446 #[must_use]
3447 pub fn atop<OO1: Send + 'static, II2: Send + 'static>(
3448 self,
3449 bidi: BidiFlow<O1, OO1, II2, I2>,
3450 ) -> BidiFlow<I1, OO1, II2, O2> {
3451 BidiFlow {
3452 top: self.top.via(bidi.top).map_materialized_value(|_| NotUsed),
3453 bottom: bidi
3454 .bottom
3455 .via(self.bottom)
3456 .map_materialized_value(|_| NotUsed),
3457 attributes: self.attributes.and(bidi.attributes),
3458 }
3459 }
3460
3461 #[must_use]
3462 pub fn reversed(self) -> BidiFlow<I2, O2, I1, O1> {
3463 BidiFlow {
3464 top: self.bottom,
3465 bottom: self.top.map_materialized_value(|_| NotUsed),
3466 attributes: self.attributes,
3467 }
3468 }
3469}
3470
3471impl<In, Out, Resource, Create, F, Close> Drop
3472 for MapWithResourceStream<In, Out, Resource, Create, F, Close>
3473where
3474 Create: Fn() -> StreamResult<Resource>,
3475 F: Fn(&mut Resource, In) -> StreamResult<Out>,
3476 Close: Fn(Resource) -> StreamResult<Option<Out>>,
3477{
3478 fn drop(&mut self) {
3479 let _ = self.close_resource();
3480 }
3481}
3482
3483fn materialize_inner_flow<In, Out, InnerMat>(
3484 flow: Flow<In, Out, InnerMat>,
3485 input: BoxStream<In>,
3486 materializer: &Materializer,
3487) -> StreamResult<(BoxStream<Out>, InnerMat)>
3488where
3489 In: Send + 'static,
3490 Out: Send + 'static,
3491 InnerMat: Send + 'static,
3492{
3493 let mat = (flow.materialize)()?;
3494 let stream = match flow.transform {
3495 FlowTransform::Pure(transform) => transform(input),
3496 FlowTransform::Runtime(transform) => transform(input, materializer)?,
3497 };
3498 Ok((stream, mat))
3499}
3500
3501struct FutureFlowStream<In, Out, InnerMat, F, Fut> {
3502 future: Arc<F>,
3503 materializer: Materializer,
3504 input: Option<BoxStream<In>>,
3505 current: Option<BoxStream<Out>>,
3506 mat_sender: Option<oneshot::Sender<StreamResult<InnerMat>>>,
3507 initialized: bool,
3508 terminated: bool,
3509 _marker: PhantomData<fn() -> Fut>,
3510}
3511
3512impl<In, Out, InnerMat, F, Fut> FutureFlowStream<In, Out, InnerMat, F, Fut>
3513where
3514 In: Send + 'static,
3515 Out: Send + 'static,
3516 InnerMat: Send + 'static,
3517 F: Fn() -> Fut,
3518 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3519{
3520 fn complete_mat(&mut self, result: StreamResult<InnerMat>) {
3521 if let Some(sender) = self.mat_sender.take() {
3522 let _ = sender.send(result);
3523 }
3524 }
3525
3526 fn initialize(&mut self) -> StreamResult<()> {
3527 if self.initialized {
3528 return Ok(());
3529 }
3530 self.initialized = true;
3531 let flow = match catch_unwind_failed("future_flow factory", || (self.future)())
3532 .and_then(run_future_inline_or_spawn)
3533 {
3534 Ok(flow) => flow,
3535 Err(error) => {
3536 self.complete_mat(Err(error.clone()));
3537 return Err(error);
3538 }
3539 };
3540 let input = self.input.take().expect("future_flow input available");
3541 match materialize_inner_flow(flow, input, &self.materializer) {
3542 Ok((stream, mat)) => {
3543 self.current = Some(stream);
3544 self.complete_mat(Ok(mat));
3545 Ok(())
3546 }
3547 Err(error) => {
3548 self.complete_mat(Err(error.clone()));
3549 Err(error)
3550 }
3551 }
3552 }
3553}
3554
3555impl<In, Out, InnerMat, F, Fut> Iterator for FutureFlowStream<In, Out, InnerMat, F, Fut>
3556where
3557 In: Send + 'static,
3558 Out: Send + 'static,
3559 InnerMat: Send + 'static,
3560 F: Fn() -> Fut,
3561 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3562{
3563 type Item = StreamResult<Out>;
3564
3565 fn next(&mut self) -> Option<Self::Item> {
3566 if self.terminated {
3567 return None;
3568 }
3569 if let Err(error) = self.initialize() {
3570 self.terminated = true;
3571 return Some(Err(error));
3572 }
3573 match self
3574 .current
3575 .as_mut()
3576 .expect("future_flow current stream initialized")
3577 .next()
3578 {
3579 Some(Ok(item)) => Some(Ok(item)),
3580 Some(Err(error)) => {
3581 self.terminated = true;
3582 Some(Err(error))
3583 }
3584 None => {
3585 self.terminated = true;
3586 None
3587 }
3588 }
3589 }
3590}
3591
3592impl<In, Out, InnerMat, F, Fut> Drop for FutureFlowStream<In, Out, InnerMat, F, Fut> {
3593 fn drop(&mut self) {
3594 if !self.initialized
3595 && let Some(sender) = self.mat_sender.take()
3596 {
3597 let _ = sender.send(Err(StreamError::Failed(
3598 "future flow was never materialized".into(),
3599 )));
3600 }
3601 }
3602}
3603
3604struct LazyFutureFlowStream<In, Out, InnerMat, F, Fut> {
3605 create: Arc<F>,
3606 materializer: Materializer,
3607 input: Option<BoxStream<In>>,
3608 current: Option<BoxStream<Out>>,
3609 mat_sender: Option<oneshot::Sender<StreamResult<InnerMat>>>,
3610 initialized: bool,
3611 terminated: bool,
3612 _marker: PhantomData<fn() -> Fut>,
3613}
3614
3615impl<In, Out, InnerMat, F, Fut> LazyFutureFlowStream<In, Out, InnerMat, F, Fut>
3616where
3617 In: Send + 'static,
3618 Out: Send + 'static,
3619 InnerMat: Send + 'static,
3620 F: Fn() -> Fut,
3621 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3622{
3623 fn complete_mat(&mut self, result: StreamResult<InnerMat>) {
3624 if let Some(sender) = self.mat_sender.take() {
3625 let _ = sender.send(result);
3626 }
3627 }
3628
3629 fn initialize(&mut self) -> Result<bool, StreamError> {
3630 if self.initialized {
3631 return Ok(true);
3632 }
3633 self.initialized = true;
3634 let first = match self
3635 .input
3636 .as_mut()
3637 .expect("lazy_future_flow input available")
3638 .next()
3639 {
3640 Some(Ok(item)) => item,
3641 Some(Err(error)) => {
3642 self.complete_mat(Err(error.clone()));
3643 return Err(error);
3644 }
3645 None => {
3646 self.complete_mat(Err(StreamError::Failed(
3647 "lazy flow was never materialized".into(),
3648 )));
3649 self.terminated = true;
3650 return Ok(false);
3651 }
3652 };
3653
3654 let flow = match catch_unwind_failed("lazy_future_flow factory", || (self.create)())
3655 .and_then(run_future_inline_or_spawn)
3656 {
3657 Ok(flow) => flow,
3658 Err(error) => {
3659 self.complete_mat(Err(error.clone()));
3660 return Err(error);
3661 }
3662 };
3663 let input = prepend_first_stream(
3664 first,
3665 self.input
3666 .take()
3667 .expect("lazy_future_flow input available after first element"),
3668 );
3669 match materialize_inner_flow(flow, input, &self.materializer) {
3670 Ok((stream, mat)) => {
3671 self.current = Some(stream);
3672 self.complete_mat(Ok(mat));
3673 Ok(true)
3674 }
3675 Err(error) => {
3676 self.complete_mat(Err(error.clone()));
3677 Err(error)
3678 }
3679 }
3680 }
3681}
3682
3683impl<In, Out, InnerMat, F, Fut> Iterator for LazyFutureFlowStream<In, Out, InnerMat, F, Fut>
3684where
3685 In: Send + 'static,
3686 Out: Send + 'static,
3687 InnerMat: Send + 'static,
3688 F: Fn() -> Fut,
3689 Fut: Future<Output = StreamResult<Flow<In, Out, InnerMat>>> + Send + 'static,
3690{
3691 type Item = StreamResult<Out>;
3692
3693 fn next(&mut self) -> Option<Self::Item> {
3694 if self.terminated {
3695 return None;
3696 }
3697 match self.initialize() {
3698 Ok(true) => {}
3699 Ok(false) => return None,
3700 Err(error) => {
3701 self.terminated = true;
3702 return Some(Err(error));
3703 }
3704 }
3705 match self
3706 .current
3707 .as_mut()
3708 .expect("lazy_future_flow current stream initialized")
3709 .next()
3710 {
3711 Some(Ok(item)) => Some(Ok(item)),
3712 Some(Err(error)) => {
3713 self.terminated = true;
3714 Some(Err(error))
3715 }
3716 None => {
3717 self.terminated = true;
3718 None
3719 }
3720 }
3721 }
3722}
3723
3724impl<In, Out, InnerMat, F, Fut> Drop for LazyFutureFlowStream<In, Out, InnerMat, F, Fut> {
3725 fn drop(&mut self) {
3726 if !self.initialized
3727 && let Some(sender) = self.mat_sender.take()
3728 {
3729 let _ = sender.send(Err(StreamError::Failed(
3730 "lazy flow was never materialized".into(),
3731 )));
3732 }
3733 }
3734}
3735
3736fn prepend_first_stream<In>(first: In, mut rest: BoxStream<In>) -> BoxStream<In>
3737where
3738 In: Send + 'static,
3739{
3740 let mut first = Some(first);
3741 Box::new(std::iter::from_fn(move || {
3742 if let Some(item) = first.take() {
3743 Some(Ok(item))
3744 } else {
3745 rest.next()
3746 }
3747 }))
3748}
3749
3750pub(super) fn poll_once_or_pending<Fut, T>(future: Fut) -> Result<StreamResult<T>, Pin<Box<Fut>>>
3751where
3752 Fut: Future<Output = StreamResult<T>>,
3753{
3754 let mut future = Box::pin(future);
3755 let _guard = stream_tokio_runtime().enter();
3756 let waker = noop_waker();
3757 let mut cx = Context::from_waker(&waker);
3758 match catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(&mut cx))) {
3759 Ok(Poll::Ready(output)) => Ok(output),
3760 Ok(Poll::Pending) => Err(future),
3761 Err(_) => Ok(Err(StreamError::Failed("future task panicked".into()))),
3762 }
3763}
3764
3765pub(super) fn spawn_completion_task<Fut, T, Msg, Map>(
3766 task_id: usize,
3767 future: Pin<Box<Fut>>,
3768 sender: std::sync::mpsc::Sender<(usize, Msg)>,
3769 map: Map,
3770) -> AbortOnDropHandle<()>
3771where
3772 Fut: Future<Output = StreamResult<T>> + Send + 'static,
3773 T: Send + 'static,
3774 Msg: Send + 'static,
3775 Map: FnOnce(StreamResult<T>) -> Msg + Send + 'static,
3776{
3777 spawn_tokio_task(async move {
3778 let result = AssertUnwindSafe(future).catch_unwind().await;
3779 let message = match result {
3780 Ok(output) => map(output),
3781 Err(_) => map(Err(StreamError::Failed("future task panicked".into()))),
3782 };
3783 let _ = sender.send((task_id, message));
3784 })
3785}
3786
3787pub(super) fn recv_completion<Msg>(
3788 receiver: &std::sync::mpsc::Receiver<(usize, Msg)>,
3789) -> Option<(usize, Msg)> {
3790 let mut idle_spins = 0;
3791 loop {
3792 match receiver.try_recv() {
3793 Ok(message) => return Some(message),
3794 Err(std::sync::mpsc::TryRecvError::Disconnected) => return None,
3795 Err(std::sync::mpsc::TryRecvError::Empty) if idle_spins < STREAM_READY_SPINS => {
3796 idle_spins += STREAM_SPIN_BACKOFF;
3797 for _ in 0..STREAM_SPIN_BACKOFF {
3798 std::hint::spin_loop();
3799 }
3800 }
3801 Err(std::sync::mpsc::TryRecvError::Empty) => {
3802 idle_spins = 0;
3803 match receiver.recv_timeout(STREAM_MAX_PARK) {
3804 Ok(message) => return Some(message),
3805 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
3806 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return None,
3807 }
3808 }
3809 }
3810 }
3811}
3812
3813pub(super) fn run_future_inline_or_spawn<Fut, T>(future: Fut) -> StreamResult<T>
3814where
3815 Fut: Future<Output = StreamResult<T>> + Send + 'static,
3816 T: Send + 'static,
3817{
3818 match poll_once_or_pending(future) {
3819 Ok(result) => result,
3820 Err(future) => {
3821 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<T>)>();
3822 let _task = spawn_completion_task(0, future, sender, |result| result);
3823 recv_completion(&receiver)
3824 .map(|(_, result)| result)
3825 .unwrap_or_else(|| Err(StreamError::Failed("future task dropped".into())))
3826 }
3827 }
3828}
3829
3830fn map_async_ordered<Out, Next, F, Fut>(
3831 mut input: BoxStream<Out>,
3832 parallelism: usize,
3833 stage: Arc<F>,
3834) -> BoxStream<Next>
3835where
3836 Out: Send + 'static,
3837 Next: Send + 'static,
3838 F: Fn(Out) -> Fut + Send + Sync + 'static,
3839 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
3840{
3841 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
3842 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
3843 let mut next_index = 0_usize;
3844 let mut next_to_emit = 0_usize;
3845 let mut completed = BTreeMap::new();
3846 let mut input_done = false;
3847
3848 Box::new(std::iter::from_fn(move || {
3849 loop {
3850 if let Some(result) = completed.remove(&next_to_emit) {
3851 next_to_emit += 1;
3852 return Some(result);
3853 }
3854
3855 while tasks.len() + completed.len() < parallelism && !input_done {
3856 match input.next() {
3857 Some(Ok(item)) => {
3858 let index = next_index;
3859 next_index += 1;
3860 match poll_once_or_pending(stage(item)) {
3861 Ok(result) => {
3862 if index == next_to_emit {
3863 next_to_emit += 1;
3864 return Some(result);
3865 }
3866 completed.insert(index, result);
3867 }
3868 Err(future) => {
3869 tasks.insert(
3870 index,
3871 spawn_completion_task(
3872 index,
3873 future,
3874 sender.clone(),
3875 |result| result,
3876 ),
3877 );
3878 }
3879 }
3880 }
3881 Some(Err(error)) => {
3882 completed.insert(next_index, Err(error));
3883 next_index += 1;
3884 input_done = true;
3885 }
3886 None => input_done = true,
3887 }
3888 }
3889
3890 if let Some(result) = completed.remove(&next_to_emit) {
3891 next_to_emit += 1;
3892 return Some(result);
3893 }
3894
3895 if tasks.is_empty() {
3896 return None;
3897 }
3898
3899 if let Some((index, result)) = recv_completion(&receiver) {
3900 tasks.remove(&index);
3901 if index == next_to_emit {
3902 next_to_emit += 1;
3903 return Some(result);
3904 }
3905 completed.insert(index, result);
3906 }
3907 }
3908 }))
3909}
3910
3911fn supervise_async_result<Next>(
3912 result: StreamResult<Next>,
3913 decider: &SupervisionDecider,
3914) -> Option<StreamResult<Next>> {
3915 match result {
3916 Ok(item) => Some(Ok(item)),
3917 Err(error) => match decide_supervision(decider, &error) {
3918 SupervisionDirective::Stop => Some(Err(error)),
3919 SupervisionDirective::Resume | SupervisionDirective::Restart => None,
3920 },
3921 }
3922}
3923
3924fn map_async_ordered_supervised<Out, Next, F, Fut>(
3925 mut input: BoxStream<Out>,
3926 parallelism: usize,
3927 stage: Arc<F>,
3928 decider: SupervisionDecider,
3929) -> BoxStream<Next>
3930where
3931 Out: Send + 'static,
3932 Next: Send + 'static,
3933 F: Fn(Out) -> Fut + Send + Sync + 'static,
3934 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
3935{
3936 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
3937 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
3938 let mut next_index = 0_usize;
3939 let mut next_to_emit = 0_usize;
3940 let mut completed = BTreeMap::<usize, Option<StreamResult<Next>>>::new();
3941 let mut input_done = false;
3942
3943 Box::new(std::iter::from_fn(move || {
3944 loop {
3945 while let Some(result) = completed.remove(&next_to_emit) {
3946 next_to_emit += 1;
3947 if let Some(result) = result {
3948 return Some(result);
3949 }
3950 }
3951
3952 while tasks.len() + completed.len() < parallelism && !input_done {
3953 match input.next() {
3954 Some(Ok(item)) => {
3955 let index = next_index;
3956 next_index += 1;
3957 match catch_unwind(AssertUnwindSafe(|| poll_once_or_pending(stage(item)))) {
3958 Ok(Ok(result)) => {
3959 let result = supervise_async_result(result, &decider);
3960 if index == next_to_emit {
3961 next_to_emit += 1;
3962 if let Some(result) = result {
3963 return Some(result);
3964 }
3965 } else {
3966 completed.insert(index, result);
3967 }
3968 }
3969 Ok(Err(future)) => {
3970 tasks.insert(
3971 index,
3972 spawn_completion_task(
3973 index,
3974 future,
3975 sender.clone(),
3976 |result| result,
3977 ),
3978 );
3979 }
3980 Err(_) => {
3981 let error = panic_stream_error("map_async callback");
3982 let result = supervise_async_result(Err(error), &decider);
3983 if index == next_to_emit {
3984 next_to_emit += 1;
3985 if let Some(result) = result {
3986 return Some(result);
3987 }
3988 } else {
3989 completed.insert(index, result);
3990 }
3991 }
3992 }
3993 }
3994 Some(Err(error)) => {
3995 completed.insert(next_index, Some(Err(error)));
3996 next_index += 1;
3997 input_done = true;
3998 }
3999 None => input_done = true,
4000 }
4001 }
4002
4003 while let Some(result) = completed.remove(&next_to_emit) {
4004 next_to_emit += 1;
4005 if let Some(result) = result {
4006 return Some(result);
4007 }
4008 }
4009
4010 if tasks.is_empty() {
4011 return None;
4012 }
4013
4014 if let Some((index, result)) = recv_completion(&receiver) {
4015 tasks.remove(&index);
4016 let result = supervise_async_result(result, &decider);
4017 if index == next_to_emit {
4018 next_to_emit += 1;
4019 if let Some(result) = result {
4020 return Some(result);
4021 }
4022 } else {
4023 completed.insert(index, result);
4024 }
4025 }
4026 }
4027 }))
4028}
4029
4030fn concat_streams<Out>(streams: Vec<BoxStream<Out>>) -> BoxStream<Out>
4031where
4032 Out: Send + 'static,
4033{
4034 let mut streams: VecDeque<_> = streams.into();
4035 let mut current = streams.pop_front();
4036 Box::new(std::iter::from_fn(move || {
4037 loop {
4038 match current.as_mut() {
4039 Some(stream) => match stream.next() {
4040 Some(item) => return Some(item),
4041 None => current = streams.pop_front(),
4042 },
4043 None => return None,
4044 }
4045 }
4046 }))
4047}
4048
4049fn concat_streams_lazy<Out, Mat>(
4050 initial: BoxStream<Out>,
4051 factories: Vec<Arc<dyn SourceFactory<Out, Mat>>>,
4052 materializer: &Materializer,
4053) -> BoxStream<Out>
4054where
4055 Out: Send + 'static,
4056 Mat: Send + 'static,
4057{
4058 let mut current = Some(initial);
4059 let mut remaining: VecDeque<_> = factories.into();
4060 let materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
4061 Box::new(std::iter::from_fn(move || {
4062 loop {
4063 match current.as_mut() {
4064 Some(stream) => match stream.next() {
4065 Some(item) => return Some(item),
4066 None => {
4067 current = remaining.pop_front().map(|factory| {
4068 match factory.create(&materializer) {
4069 Ok((stream, _)) => stream,
4070 Err(error) => {
4071 Box::new(std::iter::once(Err(error))) as BoxStream<Out>
4072 }
4073 }
4074 });
4075 }
4076 },
4077 None => return None,
4078 }
4079 }
4080 }))
4081}
4082
4083fn or_else_stream<Out>(mut primary: BoxStream<Out>, mut secondary: BoxStream<Out>) -> BoxStream<Out>
4084where
4085 Out: Send + 'static,
4086{
4087 let mut primary_emitted = false;
4088 let mut using_secondary = false;
4089 Box::new(std::iter::from_fn(move || {
4090 loop {
4091 if using_secondary {
4092 return secondary.next();
4093 }
4094
4095 match primary.next() {
4096 Some(Ok(item)) => {
4097 primary_emitted = true;
4098 return Some(Ok(item));
4099 }
4100 Some(Err(error)) => return Some(Err(error)),
4101 None if primary_emitted => return None,
4102 None => using_secondary = true,
4103 }
4104 }
4105 }))
4106}
4107
4108fn interleave_streams<Out>(
4109 streams: Vec<BoxStream<Out>>,
4110 segment_size: usize,
4111 eager_close: bool,
4112) -> BoxStream<Out>
4113where
4114 Out: Send + 'static,
4115{
4116 if segment_size == 0 {
4117 return Box::new(std::iter::once(Err(StreamError::GraphValidation(
4118 "interleave segment size must be greater than zero".into(),
4119 ))));
4120 }
4121
4122 let mut streams: Vec<Option<BoxStream<Out>>> = streams.into_iter().map(Some).collect();
4123 let mut pending: Vec<Option<StreamResult<Out>>> = (0..streams.len()).map(|_| None).collect();
4124 let mut current = 0usize;
4125 let mut emitted = 0usize;
4126 Box::new(std::iter::from_fn(move || {
4127 loop {
4128 if streams.iter().all(Option::is_none) {
4129 return None;
4130 }
4131 if streams[current].is_none() {
4132 match next_active_stream(&streams, current) {
4133 Some(next) => {
4134 current = next;
4135 emitted = 0;
4136 }
4137 None => return None,
4138 }
4139 }
4140
4141 let Some(stream) = streams[current].as_mut() else {
4142 continue;
4143 };
4144 let next_item = pending[current].take().or_else(|| stream.next());
4145 match next_item {
4146 Some(Ok(item)) => {
4147 emitted += 1;
4148 if emitted == segment_size {
4149 emitted = 0;
4150 if let Some(next) = next_active_stream(&streams, current) {
4151 current = next;
4152 }
4153 }
4154 return Some(Ok(item));
4155 }
4156 Some(Err(error)) => return Some(Err(error)),
4157 None => {
4158 streams[current] = None;
4159 emitted = 0;
4160 if eager_close {
4161 return None;
4162 }
4163 match next_active_stream(&streams, current) {
4164 Some(next) => current = next,
4165 None => return None,
4166 }
4167 }
4168 }
4169 }
4170 }))
4171}
4172
4173fn next_active_stream<Out>(streams: &[Option<BoxStream<Out>>], current: usize) -> Option<usize>
4174where
4175 Out: Send + 'static,
4176{
4177 if streams.is_empty() {
4178 return None;
4179 }
4180 for offset in 1..=streams.len() {
4181 let index = (current + offset) % streams.len();
4182 if streams[index].is_some() {
4183 return Some(index);
4184 }
4185 }
4186 None
4187}
4188
4189fn materialize_side_sink<Out, Mat>(
4190 sink: &Sink<Out, Mat>,
4191 materializer: &Materializer,
4192 buffer: usize,
4193) -> StreamResult<(std::sync::mpsc::SyncSender<StreamResult<Out>>, Mat)>
4194where
4195 Out: Send + 'static,
4196 Mat: Send + 'static,
4197{
4198 let (sender, receiver) = std::sync::mpsc::sync_channel(buffer);
4199 let mat = sink.run(side_receiver_stream(receiver), materializer)?;
4200 Ok((sender, mat))
4201}
4202
4203fn side_receiver_stream<Out>(
4204 receiver: std::sync::mpsc::Receiver<StreamResult<Out>>,
4205) -> BoxStream<Out>
4206where
4207 Out: Send + 'static,
4208{
4209 Box::new(std::iter::from_fn(move || receiver.recv().ok()))
4210}
4211
4212#[derive(Clone)]
4213enum LiveSubstreamTerminal {
4214 Complete,
4215 Error(StreamError),
4216}
4217
4218const LIVE_SUBSTREAM_CAPACITY: usize = 256;
4219const LIVE_SUBSTREAM_BATCH: usize = 64;
4220const FLAT_MAP_MERGE_SUBSTREAM_WINDOW: usize = 64;
4221
4222struct LiveSubstreamShared<T> {
4223 state: Mutex<LiveSubstreamState<T>>,
4224 available: Condvar,
4225 cancelled: Arc<AtomicBool>,
4226 capacity: usize,
4227 batch_size: usize,
4228}
4229
4230struct LiveSubstreamState<T> {
4231 buffered: usize,
4232 batches: VecDeque<VecDeque<T>>,
4233 terminal: Option<LiveSubstreamTerminal>,
4234}
4235
4236impl<T> LiveSubstreamShared<T> {
4237 fn new() -> Arc<Self> {
4238 Self::with_capacity(LIVE_SUBSTREAM_CAPACITY)
4239 }
4240
4241 fn with_capacity(capacity: usize) -> Arc<Self> {
4242 Self::with_batching(capacity, LIVE_SUBSTREAM_BATCH)
4243 }
4244
4245 fn with_batching(capacity: usize, batch_size: usize) -> Arc<Self> {
4246 Arc::new(Self {
4247 state: Mutex::new(LiveSubstreamState {
4248 buffered: 0,
4249 batches: VecDeque::new(),
4250 terminal: None,
4251 }),
4252 available: Condvar::new(),
4253 cancelled: Arc::new(AtomicBool::new(false)),
4254 capacity,
4255 batch_size: batch_size.max(1),
4256 })
4257 }
4258}
4259
4260struct LiveSubstreamStream<T> {
4261 shared: Arc<LiveSubstreamShared<T>>,
4262 completion: Option<StreamCompletion<NotUsed>>,
4263 local_batch: VecDeque<T>,
4264}
4265
4266impl<T> Iterator for LiveSubstreamStream<T> {
4267 type Item = StreamResult<T>;
4268
4269 fn next(&mut self) -> Option<Self::Item> {
4270 if let Some(item) = self.local_batch.pop_front() {
4271 return Some(Ok(item));
4272 }
4273
4274 let mut state = self
4275 .shared
4276 .state
4277 .lock()
4278 .unwrap_or_else(|poison| poison.into_inner());
4279 loop {
4280 if let Some(mut batch) = state.batches.pop_front() {
4281 state.buffered -= batch.len();
4282 drop(state);
4283 self.shared.available.notify_all();
4284 let item = batch.pop_front().expect("live substream batch has an item");
4285 self.local_batch = batch;
4286 return Some(Ok(item));
4287 }
4288 if let Some(terminal) = state.terminal.clone() {
4289 return match terminal {
4290 LiveSubstreamTerminal::Complete => None,
4291 LiveSubstreamTerminal::Error(error) => Some(Err(error)),
4292 };
4293 }
4294 state = self
4295 .shared
4296 .available
4297 .wait(state)
4298 .unwrap_or_else(|poison| poison.into_inner());
4299 }
4300 }
4301}
4302
4303impl<T> Drop for LiveSubstreamStream<T> {
4304 fn drop(&mut self) {
4305 self.shared.cancelled.store(true, Ordering::SeqCst);
4306 self.shared.available.notify_all();
4307 let _ = self.completion.take();
4308 }
4309}
4310
4311fn push_live_substream_batch<T>(
4315 shared: &Arc<LiveSubstreamShared<T>>,
4316 batch: &mut VecDeque<T>,
4317) -> Result<(), ()> {
4318 while !batch.is_empty() {
4319 let mut state = shared
4320 .state
4321 .lock()
4322 .unwrap_or_else(|poison| poison.into_inner());
4323 while state.buffered >= shared.capacity && state.terminal.is_none() {
4324 if shared.cancelled.load(Ordering::SeqCst) {
4325 batch.clear();
4326 return Err(());
4327 }
4328 state = shared
4329 .available
4330 .wait(state)
4331 .unwrap_or_else(|poison| poison.into_inner());
4332 }
4333 if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
4334 batch.clear();
4335 return Err(());
4336 }
4337 let was_empty = state.buffered == 0;
4338 while state.buffered < shared.capacity && !batch.is_empty() {
4340 let item = batch.pop_front().expect("batch non-empty");
4341 if let Some(back) = state.batches.back_mut()
4342 && back.len() < shared.batch_size
4343 {
4344 back.push_back(item);
4345 } else {
4346 let mut new_batch =
4347 VecDeque::with_capacity(shared.batch_size.min(shared.capacity.max(1)));
4348 new_batch.push_back(item);
4349 state.batches.push_back(new_batch);
4350 }
4351 state.buffered += 1;
4352 }
4353 drop(state);
4354 if was_empty {
4355 shared.available.notify_all();
4356 }
4357 }
4358 Ok(())
4359}
4360
4361fn push_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>, item: T) -> Result<(), T> {
4362 let mut state = shared
4363 .state
4364 .lock()
4365 .unwrap_or_else(|poison| poison.into_inner());
4366 while state.buffered >= shared.capacity && state.terminal.is_none() {
4367 if shared.cancelled.load(Ordering::SeqCst) {
4368 return Err(item);
4369 }
4370 state = shared
4371 .available
4372 .wait(state)
4373 .unwrap_or_else(|poison| poison.into_inner());
4374 }
4375 if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
4376 return Err(item);
4377 }
4378 let was_empty = state.buffered == 0;
4379 if let Some(batch) = state.batches.back_mut()
4380 && batch.len() < shared.batch_size
4381 {
4382 batch.push_back(item);
4383 } else {
4384 let mut batch = VecDeque::with_capacity(shared.batch_size.min(shared.capacity.max(1)));
4385 batch.push_back(item);
4386 state.batches.push_back(batch);
4387 }
4388 state.buffered += 1;
4389 drop(state);
4390 if was_empty {
4391 shared.available.notify_all();
4392 }
4393 Ok(())
4394}
4395
4396fn complete_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>) {
4397 let mut state = shared
4398 .state
4399 .lock()
4400 .unwrap_or_else(|poison| poison.into_inner());
4401 if state.terminal.is_none() {
4402 state.terminal = Some(LiveSubstreamTerminal::Complete);
4403 }
4404 drop(state);
4405 shared.available.notify_all();
4406}
4407
4408fn fail_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>, error: StreamError) {
4409 let mut state = shared
4410 .state
4411 .lock()
4412 .unwrap_or_else(|poison| poison.into_inner());
4413 if state.terminal.is_none() {
4414 state.terminal = Some(LiveSubstreamTerminal::Error(error));
4415 }
4416 drop(state);
4417 shared.available.notify_all();
4418}
4419
4420fn cancel_live_substream<T>(shared: &Arc<LiveSubstreamShared<T>>) {
4421 fail_live_substream(shared, StreamError::Cancelled);
4422}
4423
4424fn source_from_live_substream<T>(shared: Arc<LiveSubstreamShared<T>>) -> Source<T>
4425where
4426 T: Send + 'static,
4427{
4428 let claimed = Arc::new(AtomicBool::new(false));
4429 Source::from_materialized_factory(move |_materializer| {
4430 if claimed.swap(true, Ordering::SeqCst) {
4431 return Err(StreamError::Failed(
4432 "substream source cannot be materialized more than once".into(),
4433 ));
4434 }
4435 Ok((
4436 Box::new(LiveSubstreamStream {
4437 shared: Arc::clone(&shared),
4438 completion: None,
4439 local_batch: VecDeque::new(),
4440 }) as BoxStream<T>,
4441 NotUsed,
4442 ))
4443 })
4444}
4445
4446fn source_from_once_stream<T>(stream: BoxStream<T>) -> Source<T>
4447where
4448 T: Send + 'static,
4449{
4450 let stream = Arc::new(Mutex::new(Some(stream)));
4451 Source::from_materialized_factory(move |_materializer| {
4452 let mut slot = stream.lock().unwrap_or_else(|poison| poison.into_inner());
4453 let stream = slot.take().ok_or_else(|| {
4454 StreamError::Failed("substream source cannot be materialized more than once".into())
4455 })?;
4456 Ok((stream, NotUsed))
4457 })
4458}
4459
4460fn prefix_and_tail_stream<Out>(
4461 input: BoxStream<Out>,
4462 n: usize,
4463) -> BoxStream<(Vec<Out>, Source<Out>)>
4464where
4465 Out: Send + 'static,
4466{
4467 let mut input = Some(input);
4468 let mut emitted = false;
4469 Box::new(std::iter::from_fn(move || {
4470 if emitted {
4471 return None;
4472 }
4473 emitted = true;
4474
4475 let mut prefix = Vec::with_capacity(n);
4476 while prefix.len() < n {
4477 match input
4478 .as_mut()
4479 .expect("prefix_and_tail input available")
4480 .next()
4481 {
4482 Some(Ok(item)) => prefix.push(item),
4483 Some(Err(error)) => return Some(Err(error)),
4484 None => return Some(Ok((prefix, Source::empty()))),
4485 }
4486 }
4487
4488 Some(Ok((
4489 prefix,
4490 source_from_once_stream(input.take().expect("tail input available")),
4491 )))
4492 }))
4493}
4494
4495struct GroupByWorkerGuard<Key, Out> {
4496 outer: Arc<LiveSubstreamShared<Source<Out>>>,
4497 active: HashMap<Key, Arc<LiveSubstreamShared<Out>>>,
4498 closed: HashSet<Key>,
4499 armed: bool,
4500}
4501
4502impl<Key, Out> GroupByWorkerGuard<Key, Out>
4503where
4504 Key: Eq + Hash,
4505{
4506 fn new(outer: Arc<LiveSubstreamShared<Source<Out>>>) -> Self {
4507 Self {
4508 outer,
4509 active: HashMap::new(),
4510 closed: HashSet::new(),
4511 armed: true,
4512 }
4513 }
4514
4515 fn disarm(&mut self) {
4516 self.armed = false;
4517 }
4518
4519 fn fail_all(&self, error: StreamError)
4520 where
4521 Out: Send + 'static,
4522 {
4523 fail_live_substream(&self.outer, error.clone());
4524 for substream in self.active.values() {
4525 fail_live_substream(substream, error.clone());
4526 }
4527 }
4528
4529 fn complete_all(&self)
4530 where
4531 Out: Send + 'static,
4532 {
4533 complete_live_substream(&self.outer);
4534 for substream in self.active.values() {
4535 complete_live_substream(substream);
4536 }
4537 }
4538
4539 fn cancel_all(&self)
4540 where
4541 Out: Send + 'static,
4542 {
4543 for substream in self.active.values() {
4544 cancel_live_substream(substream);
4545 }
4546 }
4547}
4548
4549impl<Key, Out> Drop for GroupByWorkerGuard<Key, Out> {
4550 fn drop(&mut self) {
4551 if self.armed {
4552 fail_live_substream(&self.outer, StreamError::AbruptTermination);
4553 for substream in self.active.values() {
4554 fail_live_substream(substream, StreamError::AbruptTermination);
4555 }
4556 }
4557 }
4558}
4559
4560fn group_by_flush_write_batch<Key, Out>(
4561 guard: &mut GroupByWorkerGuard<Key, Out>,
4562 wb_key: &mut Option<Key>,
4563 wb_sub: &mut Option<Arc<LiveSubstreamShared<Out>>>,
4564 wb_items: &mut VecDeque<Out>,
4565 allow_closed_substream_recreation: bool,
4566) where
4567 Key: Clone + Eq + Hash,
4568 Out: Send + 'static,
4569{
4570 if wb_items.is_empty() {
4571 return;
4572 }
4573 let key = wb_key.take().expect("wb_key set when wb_items non-empty");
4574 if let Some(ref sub) = *wb_sub {
4575 if push_live_substream_batch(sub, wb_items).is_err() {
4576 guard.active.remove(&key);
4577 if !allow_closed_substream_recreation {
4578 guard.closed.insert(key);
4579 }
4580 }
4581 } else {
4582 wb_items.clear();
4583 }
4584 *wb_sub = None;
4585}
4586
4587fn group_by_stream<Out, Key, F>(
4588 mut input: BoxStream<Out>,
4589 max_substreams: usize,
4590 allow_closed_substream_recreation: bool,
4591 key_fn: Arc<F>,
4592 batch_mode: GroupByBatchMode,
4593 materializer: &Materializer,
4594) -> BoxStream<Source<Out>>
4595where
4596 Out: Clone + Send + 'static,
4597 Key: Clone + Eq + Hash + Send + 'static,
4598 F: Fn(&Out) -> Key + Send + Sync + 'static,
4599{
4600 let outer = LiveSubstreamShared::new();
4601 let worker_outer = Arc::clone(&outer);
4602 let batch_repeated_keys = batch_mode == GroupByBatchMode::FiniteEagerNoRecreate;
4603 let completion = materializer.spawn_stream(move |cancelled| {
4604 let mut guard = GroupByWorkerGuard::new(worker_outer);
4605
4606 let mut wb_key: Option<Key> = None;
4611 let mut wb_sub: Option<Arc<LiveSubstreamShared<Out>>> = None;
4612 let mut wb_items: VecDeque<Out> = VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH);
4613
4614 while !cancelled.load(Ordering::SeqCst) {
4615 if guard.outer.cancelled.load(Ordering::SeqCst) {
4616 guard.disarm();
4617 return Ok(NotUsed);
4618 }
4619
4620 match input.next() {
4621 Some(Ok(item)) => {
4622 let key = match catch_unwind(AssertUnwindSafe(|| key_fn(&item))) {
4623 Ok(key) => key,
4624 Err(_panic) => {
4625 wb_items.clear();
4626 guard.fail_all(StreamError::AbruptTermination);
4627 guard.disarm();
4628 return Ok(NotUsed);
4629 }
4630 };
4631
4632 if let Some(current) = guard.active.get(&key)
4636 && current.cancelled.load(Ordering::SeqCst)
4637 {
4638 if wb_key.as_ref() == Some(&key) {
4641 wb_items.clear();
4642 wb_key = None;
4643 wb_sub = None;
4644 }
4645 guard.active.remove(&key);
4646 if !allow_closed_substream_recreation {
4647 guard.closed.insert(key.clone());
4648 }
4649 }
4650
4651 let mut item = item;
4652 if let Some(current) = guard.active.get(&key).cloned() {
4653 if !batch_repeated_keys {
4654 item = match push_live_substream(¤t, item) {
4655 Ok(()) => {
4656 continue;
4657 }
4658 Err(item) => item,
4659 };
4660 guard.active.remove(&key);
4661 if !allow_closed_substream_recreation {
4662 guard.closed.insert(key.clone());
4663 continue;
4664 }
4665 } else {
4666 if wb_key.as_ref() != Some(&key) {
4670 group_by_flush_write_batch(
4671 &mut guard,
4672 &mut wb_key,
4673 &mut wb_sub,
4674 &mut wb_items,
4675 allow_closed_substream_recreation,
4676 );
4677 wb_key = Some(key.clone());
4678 wb_sub = Some(current);
4679 }
4680 wb_items.push_back(item);
4681 if wb_items.len() >= LIVE_SUBSTREAM_BATCH {
4682 group_by_flush_write_batch(
4683 &mut guard,
4684 &mut wb_key,
4685 &mut wb_sub,
4686 &mut wb_items,
4687 allow_closed_substream_recreation,
4688 );
4689 }
4690 continue;
4691 }
4692 }
4693
4694 if !wb_items.is_empty() {
4697 group_by_flush_write_batch(
4698 &mut guard,
4699 &mut wb_key,
4700 &mut wb_sub,
4701 &mut wb_items,
4702 allow_closed_substream_recreation,
4703 );
4704 }
4705
4706 if guard.closed.contains(&key) {
4707 continue;
4708 }
4709
4710 if guard.active.len() + guard.closed.len() == max_substreams {
4711 let error = StreamError::Failed(format!(
4712 "group_by reached max_substreams ({max_substreams})"
4713 ));
4714 guard.fail_all(error.clone());
4715 guard.disarm();
4716 return Err(error);
4717 }
4718
4719 let substream = LiveSubstreamShared::with_capacity(LIVE_SUBSTREAM_CAPACITY);
4720 push_live_substream(&substream, item)
4721 .unwrap_or_else(|_| unreachable!("fresh group_by substream"));
4722 guard.active.insert(key.clone(), Arc::clone(&substream));
4723 if push_live_substream(
4724 &guard.outer,
4725 source_from_live_substream(Arc::clone(&substream)),
4726 )
4727 .is_err()
4728 {
4729 guard.cancel_all();
4730 cancel_live_substream(&substream);
4731 guard.disarm();
4732 return Ok(NotUsed);
4733 }
4734 }
4735 Some(Err(error)) => {
4736 wb_items.clear();
4737 guard.fail_all(error.clone());
4738 guard.disarm();
4739 return Err(error);
4740 }
4741 None => {
4742 group_by_flush_write_batch(
4743 &mut guard,
4744 &mut wb_key,
4745 &mut wb_sub,
4746 &mut wb_items,
4747 allow_closed_substream_recreation,
4748 );
4749 guard.complete_all();
4750 guard.disarm();
4751 return Ok(NotUsed);
4752 }
4753 }
4754 }
4755
4756 group_by_flush_write_batch(
4757 &mut guard,
4758 &mut wb_key,
4759 &mut wb_sub,
4760 &mut wb_items,
4761 allow_closed_substream_recreation,
4762 );
4763 guard.complete_all();
4764 guard.disarm();
4765 Ok(NotUsed)
4766 });
4767
4768 Box::new(LiveSubstreamStream {
4769 shared: outer,
4770 completion: Some(completion),
4771 local_batch: VecDeque::new(),
4772 })
4773}
4774
4775#[derive(Clone, Copy, Debug)]
4776enum SplitMode {
4777 When,
4778 After,
4779}
4780
4781#[cfg(test)]
4782struct SplitWorkerGuard<Out> {
4783 outer: Arc<LiveSubstreamShared<Source<Out>>>,
4784 current: Option<Arc<LiveSubstreamShared<Out>>>,
4785 armed: bool,
4786 pending: VecDeque<Out>,
4790}
4791
4792#[cfg(test)]
4793impl<Out> SplitWorkerGuard<Out> {
4794 fn new(outer: Arc<LiveSubstreamShared<Source<Out>>>) -> Self {
4795 Self {
4796 outer,
4797 current: None,
4798 armed: true,
4799 pending: VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH),
4800 }
4801 }
4802
4803 fn disarm(&mut self) {
4804 self.armed = false;
4805 }
4806
4807 fn open_segment(&mut self) -> Result<(), ()>
4810 where
4811 Out: Send + 'static,
4812 {
4813 let substream = LiveSubstreamShared::new();
4814 self.current = Some(Arc::clone(&substream));
4815 push_live_substream(&self.outer, source_from_live_substream(substream)).map_err(|_| ())
4816 }
4817
4818 fn flush_pending(&mut self) -> Result<(), ()>
4821 where
4822 Out: Send + 'static,
4823 {
4824 if self.pending.is_empty() {
4825 return Ok(());
4826 }
4827 match self.current {
4828 Some(ref current) => push_live_substream_batch(current, &mut self.pending),
4829 None => {
4830 self.pending.clear();
4831 Ok(())
4832 }
4833 }
4834 }
4835
4836 fn push_item(&mut self, item: Out) -> Result<(), ()>
4840 where
4841 Out: Send + 'static,
4842 {
4843 if self.current.is_none() {
4844 return Ok(());
4846 }
4847 self.pending.push_back(item);
4848 if self.pending.len() >= LIVE_SUBSTREAM_BATCH {
4849 self.flush_pending()
4850 } else {
4851 Ok(())
4852 }
4853 }
4854
4855 fn close_segment(&mut self)
4858 where
4859 Out: Send + 'static,
4860 {
4861 let _ = self.flush_pending();
4862 if let Some(current) = self.current.take() {
4863 complete_live_substream(¤t);
4864 }
4865 }
4866
4867 fn fail_current(&mut self, error: StreamError)
4868 where
4869 Out: Send + 'static,
4870 {
4871 self.pending.clear();
4872 if let Some(current) = self.current.take() {
4873 fail_live_substream(¤t, error);
4874 }
4875 }
4876
4877 fn fail_all(&mut self, error: StreamError)
4878 where
4879 Out: Send + 'static,
4880 {
4881 self.fail_current(error.clone());
4882 fail_live_substream(&self.outer, error);
4883 }
4884
4885 fn complete_all(&mut self)
4886 where
4887 Out: Send + 'static,
4888 {
4889 self.close_segment();
4890 complete_live_substream(&self.outer);
4891 }
4892}
4893
4894#[cfg(test)]
4895impl<Out> Drop for SplitWorkerGuard<Out> {
4896 fn drop(&mut self) {
4897 if self.armed {
4898 self.pending.clear();
4899 if let Some(current) = self.current.take() {
4900 fail_live_substream(¤t, StreamError::AbruptTermination);
4901 }
4902 fail_live_substream(&self.outer, StreamError::AbruptTermination);
4903 }
4904 }
4905}
4906
4907fn split_streams<Out, F>(
4908 input: BoxStream<Out>,
4909 mode: SplitMode,
4910 predicate: Arc<F>,
4911 materializer: &Materializer,
4912) -> BoxStream<Source<Out>>
4913where
4914 Out: Clone + Send + 'static,
4915 F: Fn(&Out) -> bool + Send + Sync + 'static,
4916{
4917 #[cfg(test)]
4918 if current_substream_mode() == SubstreamExecutorMode::LegacyOnly {
4919 return split_streams_legacy(input, mode, predicate, materializer);
4920 }
4921 let parent_cancelled = Arc::new(AtomicBool::new(false));
4922 split_streams_fast(input, mode, predicate, parent_cancelled, materializer)
4923}
4924
4925#[cfg(test)]
4926fn split_streams_legacy<Out, F>(
4927 mut input: BoxStream<Out>,
4928 mode: SplitMode,
4929 predicate: Arc<F>,
4930 materializer: &Materializer,
4931) -> BoxStream<Source<Out>>
4932where
4933 Out: Clone + Send + 'static,
4934 F: Fn(&Out) -> bool + Send + Sync + 'static,
4935{
4936 let outer = LiveSubstreamShared::new();
4937 let worker_outer = Arc::clone(&outer);
4938 let completion = materializer.spawn_stream(move |cancelled| {
4939 let mut guard = SplitWorkerGuard::new(Arc::clone(&worker_outer));
4940
4941 while !cancelled.load(Ordering::SeqCst) {
4942 if worker_outer.cancelled.load(Ordering::SeqCst) {
4943 guard.disarm();
4944 return Ok(NotUsed);
4945 }
4946
4947 match input.next() {
4948 Some(Ok(item)) => {
4949 let split = match catch_unwind(AssertUnwindSafe(|| predicate(&item))) {
4950 Ok(split) => split,
4951 Err(_panic) => {
4952 guard.fail_all(StreamError::AbruptTermination);
4953 guard.disarm();
4954 return Ok(NotUsed);
4955 }
4956 };
4957
4958 match mode {
4959 SplitMode::When => {
4960 if split && guard.current.is_some() {
4961 guard.close_segment();
4963 }
4964 if guard.current.is_none() && guard.open_segment().is_err() {
4966 guard.disarm();
4967 return Ok(NotUsed);
4968 }
4969 if guard.push_item(item).is_err() {
4970 guard.disarm();
4971 return Ok(NotUsed);
4972 }
4973 }
4974 SplitMode::After => {
4975 if guard.current.is_none() && guard.open_segment().is_err() {
4976 guard.disarm();
4977 return Ok(NotUsed);
4978 }
4979 if guard.push_item(item).is_err() {
4980 guard.disarm();
4981 return Ok(NotUsed);
4982 }
4983 if split {
4984 guard.close_segment();
4985 }
4986 }
4987 }
4988 }
4989 Some(Err(error)) => {
4990 guard.fail_all(error.clone());
4991 guard.disarm();
4992 return Err(error);
4993 }
4994 None => {
4995 guard.complete_all();
4996 guard.disarm();
4997 return Ok(NotUsed);
4998 }
4999 }
5000 }
5001
5002 guard.complete_all();
5003 guard.disarm();
5004 Ok(NotUsed)
5005 });
5006
5007 Box::new(LiveSubstreamStream {
5008 shared: outer,
5009 completion: Some(completion),
5010 local_batch: VecDeque::new(),
5011 })
5012}
5013
5014trait SplitConsumer<T: Send + 'static>: Send + 'static {
5019 fn push_item(&mut self, item: T) -> StreamResult<()>;
5020 fn complete(self: Box<Self>);
5021 fn fail(self: Box<Self>, error: StreamError);
5022}
5023
5024struct FoldConsumer<T, Acc> {
5025 acc: Option<Acc>,
5026 f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
5027 tx: futures::channel::oneshot::Sender<StreamResult<Acc>>,
5028}
5029
5030impl<T: Send + 'static, Acc: Send + 'static> SplitConsumer<T> for FoldConsumer<T, Acc> {
5031 fn push_item(&mut self, item: T) -> StreamResult<()> {
5032 let acc = self.acc.take().expect("FoldConsumer: push after done");
5033 self.acc = Some((self.f)(acc, item));
5034 Ok(())
5035 }
5036 fn complete(mut self: Box<Self>) {
5037 let acc = self
5038 .acc
5039 .take()
5040 .expect("FoldConsumer: complete called twice");
5041 let _ = self.tx.send(Ok(acc));
5042 }
5043 fn fail(mut self: Box<Self>, error: StreamError) {
5044 self.acc = None;
5045 let _ = self.tx.send(Err(error));
5046 }
5047}
5048
5049struct FoldResultConsumer<T, Acc> {
5050 acc: Option<Acc>,
5051 f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
5052 tx: futures::channel::oneshot::Sender<StreamResult<Acc>>,
5053}
5054
5055impl<T: Send + 'static, Acc: Send + 'static> SplitConsumer<T> for FoldResultConsumer<T, Acc> {
5056 fn push_item(&mut self, item: T) -> StreamResult<()> {
5057 let acc = self
5058 .acc
5059 .take()
5060 .expect("FoldResultConsumer: push after done");
5061 match (self.f)(acc, item) {
5062 Ok(new_acc) => {
5063 self.acc = Some(new_acc);
5064 Ok(())
5065 }
5066 Err(e) => Err(e),
5067 }
5068 }
5069 fn complete(mut self: Box<Self>) {
5070 let acc = self
5071 .acc
5072 .take()
5073 .expect("FoldResultConsumer: complete called twice");
5074 let _ = self.tx.send(Ok(acc));
5075 }
5076 fn fail(mut self: Box<Self>, error: StreamError) {
5077 self.acc = None;
5078 let _ = self.tx.send(Err(error));
5079 }
5080}
5081
5082struct CollectConsumer<T> {
5083 items: Vec<T>,
5084 tx: futures::channel::oneshot::Sender<StreamResult<Vec<T>>>,
5085}
5086
5087impl<T: Send + 'static> SplitConsumer<T> for CollectConsumer<T> {
5088 fn push_item(&mut self, item: T) -> StreamResult<()> {
5089 self.items.push(item);
5090 Ok(())
5091 }
5092 fn complete(self: Box<Self>) {
5093 let _ = self.tx.send(Ok(self.items));
5094 }
5095 fn fail(self: Box<Self>, error: StreamError) {
5096 let _ = self.tx.send(Err(error));
5097 }
5098}
5099
5100struct IgnoreConsumer<T> {
5101 tx: futures::channel::oneshot::Sender<StreamResult<NotUsed>>,
5102 _phantom: std::marker::PhantomData<fn(T)>,
5103}
5104
5105impl<T: Send + 'static> SplitConsumer<T> for IgnoreConsumer<T> {
5106 fn push_item(&mut self, _item: T) -> StreamResult<()> {
5107 Ok(())
5108 }
5109 fn complete(self: Box<Self>) {
5110 let _ = self.tx.send(Ok(NotUsed));
5111 }
5112 fn fail(self: Box<Self>, error: StreamError) {
5113 let _ = self.tx.send(Err(error));
5114 }
5115}
5116
5117struct TerminalDrainCancelGuard<T: 'static> {
5118 hook: Option<Arc<dyn TerminalSourceHookDyn<T>>>,
5119}
5120
5121impl<T: 'static> TerminalDrainCancelGuard<T> {
5122 fn new(hook: Arc<dyn TerminalSourceHookDyn<T>>) -> Self {
5123 Self { hook: Some(hook) }
5124 }
5125
5126 fn disarm(&mut self) {
5127 self.hook = None;
5128 }
5129}
5130
5131impl<T: 'static> Drop for TerminalDrainCancelGuard<T> {
5132 fn drop(&mut self) {
5133 if let Some(hook) = self.hook.take() {
5134 hook.cancel_terminal();
5135 }
5136 }
5137}
5138
5139fn terminal_drain_status<T: 'static>(
5140 hook: &Arc<dyn TerminalSourceHookDyn<T>>,
5141 materializer: &Materializer,
5142 cancelled: &Arc<AtomicBool>,
5143) -> StreamResult<()> {
5144 if materializer.is_shutdown() {
5145 hook.cancel_terminal();
5146 Err(StreamError::AbruptTermination)
5147 } else if cancelled.load(Ordering::SeqCst) {
5148 hook.cancel_terminal();
5149 Err(StreamError::Cancelled)
5150 } else {
5151 Ok(())
5152 }
5153}
5154
5155fn wait_direct_terminal_result<T: 'static, Acc: Send + 'static>(
5156 hook: &Arc<dyn TerminalSourceHookDyn<T>>,
5157 materializer: &Materializer,
5158 cancelled: &Arc<AtomicBool>,
5159 receiver: std::sync::mpsc::Receiver<StreamResult<Acc>>,
5160) -> StreamResult<Acc> {
5161 loop {
5162 if materializer.is_shutdown() {
5163 hook.cancel_terminal();
5164 return Err(StreamError::AbruptTermination);
5165 }
5166 if cancelled.load(Ordering::SeqCst) {
5167 hook.cancel_terminal();
5168 return Err(StreamError::Cancelled);
5169 }
5170 match receiver.recv_timeout(Duration::from_millis(1)) {
5171 Ok(result) => return result,
5172 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
5173 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
5174 return Err(StreamError::AbruptTermination);
5175 }
5176 }
5177 }
5178}
5179
5180fn register_direct_terminal<T: Send + 'static, Acc: Send + 'static>(
5181 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5182 materializer: &Materializer,
5183 consumer: Box<dyn TerminalSinkConsumerDyn<T>>,
5184 receiver: std::sync::mpsc::Receiver<StreamResult<Acc>>,
5185) -> Option<StreamResult<Box<dyn Any + Send>>> {
5186 if !hook.supports_direct_terminal() {
5187 return None;
5188 }
5189 let cancellation = StreamCancellation::for_external_completion();
5190 let cancelled = cancellation.cancelled();
5191 if let Err(error) = hook
5192 .try_register_direct_terminal(consumer, Arc::clone(&cancelled))
5193 .expect("direct terminal support advertised")
5194 {
5195 return Some(Err(error));
5196 }
5197 let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
5198 let (tx, rx) = futures::channel::oneshot::channel();
5199 thread::spawn(move || {
5200 let result = wait_direct_terminal_result(&hook, &worker_materializer, &cancelled, receiver);
5201 let _ = tx.send(result);
5202 });
5203 let completion = StreamCompletion::from_receiver(rx, Some(cancellation));
5204 Some(Ok(Box::new(completion)))
5205}
5206
5207struct DirectFoldConsumer<T, Acc> {
5208 acc: Option<Acc>,
5209 f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
5210 tx: std::sync::mpsc::Sender<StreamResult<Acc>>,
5211}
5212
5213impl<T: Send + 'static, Acc: Send + 'static> TerminalSinkConsumerDyn<T>
5214 for DirectFoldConsumer<T, Acc>
5215{
5216 fn on_item(&mut self, item: T) -> StreamResult<()> {
5217 let previous = self.acc.take().expect("fold accumulator present");
5218 match catch_unwind(AssertUnwindSafe(|| (self.f)(previous, item))) {
5219 Ok(next) => {
5220 self.acc = Some(next);
5221 Ok(())
5222 }
5223 Err(_) => Err(StreamError::AbruptTermination),
5224 }
5225 }
5226
5227 fn finish(mut self: Box<Self>, result: StreamResult<()>) {
5228 let result = result.map(|()| self.acc.take().expect("fold accumulator present"));
5229 let _ = self.tx.send(result);
5230 }
5231}
5232
5233struct DirectFoldResultConsumer<T, Acc> {
5234 acc: Option<Acc>,
5235 f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
5236 tx: std::sync::mpsc::Sender<StreamResult<Acc>>,
5237}
5238
5239impl<T: Send + 'static, Acc: Send + 'static> TerminalSinkConsumerDyn<T>
5240 for DirectFoldResultConsumer<T, Acc>
5241{
5242 fn on_item(&mut self, item: T) -> StreamResult<()> {
5243 let previous = self.acc.take().expect("fold accumulator present");
5244 let result = catch_unwind(AssertUnwindSafe(|| (self.f)(previous, item)))
5245 .unwrap_or(Err(StreamError::AbruptTermination));
5246 match result {
5247 Ok(next) => {
5248 self.acc = Some(next);
5249 Ok(())
5250 }
5251 Err(error) => Err(error),
5252 }
5253 }
5254
5255 fn finish(mut self: Box<Self>, result: StreamResult<()>) {
5256 let result = result.map(|()| self.acc.take().expect("fold accumulator present"));
5257 let _ = self.tx.send(result);
5258 }
5259}
5260
5261struct DirectCollectConsumer<T> {
5262 items: Vec<T>,
5263 tx: std::sync::mpsc::Sender<StreamResult<Vec<T>>>,
5264}
5265
5266impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectCollectConsumer<T> {
5267 fn on_item(&mut self, item: T) -> StreamResult<()> {
5268 self.items.push(item);
5269 Ok(())
5270 }
5271
5272 fn finish(self: Box<Self>, result: StreamResult<()>) {
5273 let Self { items, tx } = *self;
5274 let result = result.map(|()| items);
5275 let _ = tx.send(result);
5276 }
5277}
5278
5279struct DirectIgnoreConsumer {
5280 tx: std::sync::mpsc::Sender<StreamResult<NotUsed>>,
5281}
5282
5283impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectIgnoreConsumer {
5284 fn on_item(&mut self, _item: T) -> StreamResult<()> {
5285 Ok(())
5286 }
5287
5288 fn finish(self: Box<Self>, result: StreamResult<()>) {
5289 let result = result.map(|()| NotUsed);
5290 let _ = self.tx.send(result);
5291 }
5292}
5293
5294pub(super) struct ForeachDescriptor<T> {
5295 pub(super) f: Arc<dyn Fn(T) + Send + Sync + 'static>,
5296}
5297
5298struct DirectForeachConsumer<T> {
5299 f: Arc<dyn Fn(T) + Send + Sync + 'static>,
5300 tx: std::sync::mpsc::Sender<StreamResult<NotUsed>>,
5301}
5302
5303impl<T: Send + 'static> TerminalSinkConsumerDyn<T> for DirectForeachConsumer<T> {
5304 fn on_item(&mut self, item: T) -> StreamResult<()> {
5305 catch_unwind(AssertUnwindSafe(|| (self.f)(item)))
5306 .map_err(|_| StreamError::AbruptTermination)
5307 }
5308
5309 fn finish(self: Box<Self>, result: StreamResult<()>) {
5310 let result = result.map(|()| NotUsed);
5311 let _ = self.tx.send(result);
5312 }
5313}
5314
5315pub(super) struct FoldDescriptor<T, Acc> {
5318 pub(super) zero: Acc,
5319 pub(super) f: Arc<dyn Fn(Acc, T) -> Acc + Send + Sync + 'static>,
5320}
5321
5322impl<T: Send + 'static, Acc: Clone + Send + Sync + 'static> FoldFastPathDyn<T>
5323 for FoldDescriptor<T, Acc>
5324{
5325 fn try_register(
5326 &self,
5327 hook: Arc<dyn SplitSegmentHookDyn>,
5328 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5329 let hook_any = hook.as_any_arc();
5330 let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5331 if slot.claimed.swap(true, Ordering::SeqCst) {
5332 return Some(Err(StreamError::Failed(
5333 "substream source cannot be materialized more than once".into(),
5334 )));
5335 }
5336 let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Acc>>();
5337 {
5338 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5339
5340 if let Some(terminal) = state.terminal.take() {
5342 let buffer = std::mem::take(&mut state.buffer);
5343 state.consumer = SegmentConsumer::DirectTaken;
5344 drop(state);
5345 let mut acc = self.zero.clone();
5347 for item in buffer {
5348 acc = (self.f)(acc, item);
5349 }
5350 let result = match terminal {
5351 LiveSubstreamTerminal::Complete => Ok(acc),
5352 LiveSubstreamTerminal::Error(e) => Err(e),
5353 };
5354 let _ = tx.send(result);
5355 let completion = StreamCompletion::from_receiver(rx, None);
5356 return Some(Ok(Box::new(completion)));
5357 }
5358
5359 let consumer: Box<dyn SplitConsumer<T>> = Box::new(FoldConsumer {
5361 acc: Some(self.zero.clone()),
5362 f: Arc::clone(&self.f),
5363 tx,
5364 });
5365 state.consumer = SegmentConsumer::Direct(consumer);
5366 }
5367 slot.available.notify_all();
5368 let completion = StreamCompletion::from_receiver(rx, None);
5369 Some(Ok(Box::new(completion)))
5370 }
5371
5372 fn supports_terminal_drain(&self) -> bool {
5373 true
5374 }
5375
5376 fn try_register_direct_terminal(
5377 &self,
5378 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5379 materializer: &Materializer,
5380 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5381 let (tx, rx) = std::sync::mpsc::channel();
5382 let consumer = Box::new(DirectFoldConsumer {
5383 acc: Some(self.zero.clone()),
5384 f: Arc::clone(&self.f),
5385 tx,
5386 });
5387 register_direct_terminal(hook, materializer, consumer, rx)
5388 }
5389
5390 fn try_register_terminal_drain(
5391 &self,
5392 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5393 materializer: &Materializer,
5394 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5395 let zero = self.zero.clone();
5396 let f = Arc::clone(&self.f);
5397 let worker_materializer =
5398 materializer.with_name_prefix(materializer.name_prefix().to_owned());
5399 let completion = materializer.spawn_stream(move |cancelled| {
5400 let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5401 let mut acc = Some(zero);
5402 let mut batch = Vec::new();
5403 loop {
5404 let status =
5405 hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5406 for (index, item) in batch.drain(..).enumerate() {
5407 let previous = acc.take().expect("fold accumulator present");
5408 acc = Some(f(previous, item));
5409 if (index + 1) % 64 == 0 {
5410 terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5411 }
5412 }
5413 if matches!(status, TerminalSourceStatus::Completed) {
5414 guard.disarm();
5415 return Ok(acc.expect("fold accumulator present"));
5416 }
5417 }
5418 });
5419 Some(Ok(Box::new(completion)))
5420 }
5421}
5422
5423pub(super) struct FoldResultDescriptor<T, Acc> {
5424 pub(super) zero: Acc,
5425 pub(super) f: Arc<dyn Fn(Acc, T) -> StreamResult<Acc> + Send + Sync + 'static>,
5426}
5427
5428impl<T: Send + 'static, Acc: Clone + Send + Sync + 'static> FoldFastPathDyn<T>
5429 for FoldResultDescriptor<T, Acc>
5430{
5431 fn try_register(
5432 &self,
5433 hook: Arc<dyn SplitSegmentHookDyn>,
5434 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5435 let hook_any = hook.as_any_arc();
5436 let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5437 if slot.claimed.swap(true, Ordering::SeqCst) {
5438 return Some(Err(StreamError::Failed(
5439 "substream source cannot be materialized more than once".into(),
5440 )));
5441 }
5442 let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Acc>>();
5443 {
5444 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5445
5446 if let Some(terminal) = state.terminal.take() {
5448 let buffer = std::mem::take(&mut state.buffer);
5449 state.consumer = SegmentConsumer::DirectTaken;
5450 drop(state);
5451 let acc = self.zero.clone();
5452 let result = match terminal {
5453 LiveSubstreamTerminal::Complete => buffer
5454 .into_iter()
5455 .try_fold(acc, |a, item| (self.f)(a, item)),
5456 LiveSubstreamTerminal::Error(e) => Err(e),
5457 };
5458 let _ = tx.send(result);
5459 let completion = StreamCompletion::from_receiver(rx, None);
5460 return Some(Ok(Box::new(completion)));
5461 }
5462
5463 let consumer: Box<dyn SplitConsumer<T>> = Box::new(FoldResultConsumer {
5465 acc: Some(self.zero.clone()),
5466 f: Arc::clone(&self.f),
5467 tx,
5468 });
5469 state.consumer = SegmentConsumer::Direct(consumer);
5470 }
5471 slot.available.notify_all();
5472 let completion = StreamCompletion::from_receiver(rx, None);
5473 Some(Ok(Box::new(completion)))
5474 }
5475
5476 fn supports_terminal_drain(&self) -> bool {
5477 true
5478 }
5479
5480 fn try_register_direct_terminal(
5481 &self,
5482 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5483 materializer: &Materializer,
5484 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5485 let (tx, rx) = std::sync::mpsc::channel();
5486 let consumer = Box::new(DirectFoldResultConsumer {
5487 acc: Some(self.zero.clone()),
5488 f: Arc::clone(&self.f),
5489 tx,
5490 });
5491 register_direct_terminal(hook, materializer, consumer, rx)
5492 }
5493
5494 fn try_register_terminal_drain(
5495 &self,
5496 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5497 materializer: &Materializer,
5498 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5499 let zero = self.zero.clone();
5500 let f = Arc::clone(&self.f);
5501 let worker_materializer =
5502 materializer.with_name_prefix(materializer.name_prefix().to_owned());
5503 let completion = materializer.spawn_stream(move |cancelled| {
5504 let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5505 let mut acc = Some(zero);
5506 let mut batch = Vec::new();
5507 loop {
5508 let status =
5509 hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5510 for (index, item) in batch.drain(..).enumerate() {
5511 let previous = acc.take().expect("fold accumulator present");
5512 match f(previous, item) {
5513 Ok(next) => {
5514 acc = Some(next);
5515 }
5516 Err(error) => return Err(error),
5517 }
5518 if (index + 1) % 64 == 0 {
5519 terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5520 }
5521 }
5522 if matches!(status, TerminalSourceStatus::Completed) {
5523 guard.disarm();
5524 return Ok(acc.expect("fold accumulator present"));
5525 }
5526 }
5527 });
5528 Some(Ok(Box::new(completion)))
5529 }
5530}
5531
5532pub(super) struct CollectDescriptor<T> {
5533 pub(super) _phantom: std::marker::PhantomData<fn(T)>,
5534}
5535
5536impl<T: Send + 'static> FoldFastPathDyn<T> for CollectDescriptor<T> {
5537 fn try_register(
5538 &self,
5539 hook: Arc<dyn SplitSegmentHookDyn>,
5540 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5541 let hook_any = hook.as_any_arc();
5542 let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5543 if slot.claimed.swap(true, Ordering::SeqCst) {
5544 return Some(Err(StreamError::Failed(
5545 "substream source cannot be materialized more than once".into(),
5546 )));
5547 }
5548 let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<Vec<T>>>();
5549 {
5550 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5551
5552 if let Some(terminal) = state.terminal.take() {
5554 let buffer = std::mem::take(&mut state.buffer);
5555 state.consumer = SegmentConsumer::DirectTaken;
5556 drop(state);
5557 let items: Vec<T> = buffer.into_iter().collect();
5558 let result = match terminal {
5559 LiveSubstreamTerminal::Complete => Ok(items),
5560 LiveSubstreamTerminal::Error(e) => Err(e),
5561 };
5562 let _ = tx.send(result);
5563 let completion = StreamCompletion::from_receiver(rx, None);
5564 return Some(Ok(Box::new(completion)));
5565 }
5566
5567 let consumer: Box<dyn SplitConsumer<T>> = Box::new(CollectConsumer {
5569 items: Vec::new(),
5570 tx,
5571 });
5572 state.consumer = SegmentConsumer::Direct(consumer);
5573 }
5574 slot.available.notify_all();
5575 let completion = StreamCompletion::from_receiver(rx, None);
5576 Some(Ok(Box::new(completion)))
5577 }
5578
5579 fn supports_terminal_drain(&self) -> bool {
5580 true
5581 }
5582
5583 fn try_register_direct_terminal(
5584 &self,
5585 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5586 materializer: &Materializer,
5587 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5588 let (tx, rx) = std::sync::mpsc::channel();
5589 let consumer = Box::new(DirectCollectConsumer {
5590 items: Vec::new(),
5591 tx,
5592 });
5593 register_direct_terminal(hook, materializer, consumer, rx)
5594 }
5595
5596 fn try_register_terminal_drain(
5597 &self,
5598 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5599 materializer: &Materializer,
5600 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5601 let worker_materializer =
5602 materializer.with_name_prefix(materializer.name_prefix().to_owned());
5603 let completion = materializer.spawn_stream(move |cancelled| {
5604 let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5605 let mut items = Vec::new();
5606 let mut batch = Vec::new();
5607 loop {
5608 let status =
5609 hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5610 for (index, item) in batch.drain(..).enumerate() {
5611 items.push(item);
5612 if (index + 1) % 64 == 0 {
5613 terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5614 }
5615 }
5616 if matches!(status, TerminalSourceStatus::Completed) {
5617 guard.disarm();
5618 return Ok(items);
5619 }
5620 }
5621 });
5622 Some(Ok(Box::new(completion)))
5623 }
5624}
5625
5626pub(super) struct IgnoreDescriptor<T> {
5627 pub(super) _phantom: std::marker::PhantomData<fn(T)>,
5628}
5629
5630impl<T: Send + 'static> FoldFastPathDyn<T> for IgnoreDescriptor<T> {
5631 fn try_register(
5632 &self,
5633 hook: Arc<dyn SplitSegmentHookDyn>,
5634 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5635 let hook_any = hook.as_any_arc();
5636 let slot = hook_any.downcast::<SegmentConsumerSlot<T>>().ok()?;
5637 if slot.claimed.swap(true, Ordering::SeqCst) {
5638 return Some(Err(StreamError::Failed(
5639 "substream source cannot be materialized more than once".into(),
5640 )));
5641 }
5642 let (tx, rx) = futures::channel::oneshot::channel::<StreamResult<NotUsed>>();
5643 {
5644 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5645
5646 if let Some(terminal) = state.terminal.take() {
5648 state.consumer = SegmentConsumer::DirectTaken;
5649 drop(state);
5650 let result = match terminal {
5651 LiveSubstreamTerminal::Complete => Ok(NotUsed),
5652 LiveSubstreamTerminal::Error(e) => Err(e),
5653 };
5654 let _ = tx.send(result);
5655 let completion = StreamCompletion::from_receiver(rx, None);
5656 return Some(Ok(Box::new(completion)));
5657 }
5658
5659 let consumer: Box<dyn SplitConsumer<T>> = Box::new(IgnoreConsumer {
5661 tx,
5662 _phantom: std::marker::PhantomData,
5663 });
5664 state.consumer = SegmentConsumer::Direct(consumer);
5665 }
5666 slot.available.notify_all();
5667 let completion = StreamCompletion::from_receiver(rx, None);
5668 Some(Ok(Box::new(completion)))
5669 }
5670
5671 fn supports_terminal_drain(&self) -> bool {
5672 true
5673 }
5674
5675 fn try_register_direct_terminal(
5676 &self,
5677 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5678 materializer: &Materializer,
5679 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5680 let (tx, rx) = std::sync::mpsc::channel();
5681 let consumer = Box::new(DirectIgnoreConsumer { tx });
5682 register_direct_terminal(hook, materializer, consumer, rx)
5683 }
5684
5685 fn try_register_terminal_drain(
5686 &self,
5687 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5688 materializer: &Materializer,
5689 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5690 let worker_materializer =
5691 materializer.with_name_prefix(materializer.name_prefix().to_owned());
5692 let completion = materializer.spawn_stream(move |cancelled| {
5693 let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5694 let mut batch = Vec::new();
5695 loop {
5696 let status =
5697 hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5698 batch.clear();
5699 if matches!(status, TerminalSourceStatus::Completed) {
5700 guard.disarm();
5701 return Ok(NotUsed);
5702 }
5703 }
5704 });
5705 Some(Ok(Box::new(completion)))
5706 }
5707}
5708
5709impl<T: Send + 'static> FoldFastPathDyn<T> for ForeachDescriptor<T> {
5710 fn try_register(
5711 &self,
5712 _hook: Arc<dyn SplitSegmentHookDyn>,
5713 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5714 None
5715 }
5716
5717 fn supports_terminal_drain(&self) -> bool {
5718 true
5719 }
5720
5721 fn try_register_direct_terminal(
5722 &self,
5723 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5724 materializer: &Materializer,
5725 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5726 let (tx, rx) = std::sync::mpsc::channel();
5727 let consumer = Box::new(DirectForeachConsumer {
5728 f: Arc::clone(&self.f),
5729 tx,
5730 });
5731 register_direct_terminal(hook, materializer, consumer, rx)
5732 }
5733
5734 fn try_register_terminal_drain(
5735 &self,
5736 hook: Arc<dyn TerminalSourceHookDyn<T>>,
5737 materializer: &Materializer,
5738 ) -> Option<StreamResult<Box<dyn Any + Send>>> {
5739 let f = Arc::clone(&self.f);
5740 let worker_materializer =
5741 materializer.with_name_prefix(materializer.name_prefix().to_owned());
5742 let completion = materializer.spawn_stream(move |cancelled| {
5743 let mut guard = TerminalDrainCancelGuard::new(Arc::clone(&hook));
5744 let mut batch = Vec::new();
5745 loop {
5746 let status =
5747 hook.drain_terminal_batch(&worker_materializer, &cancelled, &mut batch)?;
5748 for (index, item) in batch.drain(..).enumerate() {
5749 f(item);
5750 if (index + 1) % 64 == 0 {
5751 terminal_drain_status(&hook, &worker_materializer, &cancelled)?;
5752 }
5753 }
5754 if matches!(status, TerminalSourceStatus::Completed) {
5755 guard.disarm();
5756 return Ok(NotUsed);
5757 }
5758 }
5759 });
5760 Some(Ok(Box::new(completion)))
5761 }
5762}
5763
5764enum SegmentConsumer<T: Send + 'static> {
5767 Pending,
5768 Direct(Box<dyn SplitConsumer<T>>),
5769 DirectTaken,
5770 Fallback,
5771}
5772
5773struct SegmentSlotState<T: Send + 'static> {
5774 buffer: VecDeque<T>,
5775 consumer: SegmentConsumer<T>,
5776 terminal: Option<LiveSubstreamTerminal>,
5777}
5778
5779pub(super) struct SegmentConsumerSlot<T: Send + 'static> {
5780 claimed: Arc<AtomicBool>,
5781 state: Mutex<SegmentSlotState<T>>,
5782 available: Condvar,
5783 parent_cancelled: Arc<AtomicBool>,
5784}
5785
5786impl<T: Clone + Send + 'static> SegmentConsumerSlot<T> {
5787 fn new(parent_cancelled: Arc<AtomicBool>) -> Arc<Self> {
5788 Arc::new(Self {
5789 claimed: Arc::new(AtomicBool::new(false)),
5790 state: Mutex::new(SegmentSlotState {
5791 buffer: VecDeque::new(),
5792 consumer: SegmentConsumer::Pending,
5793 terminal: None,
5794 }),
5795 available: Condvar::new(),
5796 parent_cancelled,
5797 })
5798 }
5799}
5800
5801impl<T: Clone + Send + 'static> SplitSegmentHookDyn for SegmentConsumerSlot<T> {
5802 fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
5803 self
5804 }
5805}
5806
5807impl<T: Clone + Send + 'static> SourceFactory<T, NotUsed> for SegmentConsumerSlot<T> {
5808 fn create(
5809 self: Arc<Self>,
5810 _materializer: &Materializer,
5811 ) -> StreamResult<(BoxStream<T>, NotUsed)> {
5812 if self.claimed.swap(true, Ordering::SeqCst) {
5813 return Err(StreamError::Failed(
5814 "substream source cannot be materialized more than once".into(),
5815 ));
5816 }
5817 {
5818 let mut state = self.state.lock().unwrap_or_else(|p| p.into_inner());
5819 state.consumer = SegmentConsumer::Fallback;
5820 }
5821 self.available.notify_all();
5822 let stream = FallbackSegmentStream {
5823 slot: Arc::clone(&self),
5824 };
5825 Ok((Box::new(stream), NotUsed))
5826 }
5827}
5828
5829struct FallbackSegmentStream<T: Send + 'static> {
5830 slot: Arc<SegmentConsumerSlot<T>>,
5831}
5832
5833impl<T: Clone + Send + 'static> Iterator for FallbackSegmentStream<T> {
5834 type Item = StreamResult<T>;
5835
5836 fn next(&mut self) -> Option<Self::Item> {
5837 let mut state = self.slot.state.lock().unwrap_or_else(|p| p.into_inner());
5838 loop {
5839 if let Some(item) = state.buffer.pop_front() {
5840 drop(state);
5841 self.slot.available.notify_all();
5842 return Some(Ok(item));
5843 }
5844 if let Some(terminal) = &state.terminal {
5845 return match terminal {
5846 LiveSubstreamTerminal::Complete => None,
5847 LiveSubstreamTerminal::Error(e) => Some(Err(e.clone())),
5848 };
5849 }
5850 if self.slot.parent_cancelled.load(Ordering::SeqCst) {
5851 return Some(Err(StreamError::AbruptTermination));
5852 }
5853 state = self
5854 .slot
5855 .available
5856 .wait(state)
5857 .unwrap_or_else(|p| p.into_inner());
5858 }
5859 }
5860}
5861
5862impl<T: Send + 'static> Drop for FallbackSegmentStream<T> {
5863 fn drop(&mut self) {
5864 let mut state = self.slot.state.lock().unwrap_or_else(|p| p.into_inner());
5865 if state.terminal.is_none() {
5866 state.terminal = Some(LiveSubstreamTerminal::Error(StreamError::Cancelled));
5867 }
5868 drop(state);
5869 self.slot.available.notify_all();
5870 }
5871}
5872
5873struct SplitFastWorkerGuard<Out: Send + 'static> {
5876 outer: Arc<LiveSubstreamShared<Source<Out>>>,
5877 current_slot: Option<Arc<SegmentConsumerSlot<Out>>>,
5878 current_consumer: Option<Box<dyn SplitConsumer<Out>>>,
5879 armed: bool,
5880 parent_cancelled: Arc<AtomicBool>,
5881 local_pending: VecDeque<Out>,
5882}
5883
5884impl<Out: Clone + Send + 'static> SplitFastWorkerGuard<Out> {
5885 fn new(
5886 outer: Arc<LiveSubstreamShared<Source<Out>>>,
5887 parent_cancelled: Arc<AtomicBool>,
5888 ) -> Self {
5889 Self {
5890 outer,
5891 current_slot: None,
5892 current_consumer: None,
5893 armed: true,
5894 parent_cancelled,
5895 local_pending: VecDeque::with_capacity(LIVE_SUBSTREAM_BATCH),
5896 }
5897 }
5898
5899 fn disarm(&mut self) {
5900 self.armed = false;
5901 }
5902
5903 fn open_segment(&mut self) -> Result<(), ()> {
5904 let slot = SegmentConsumerSlot::new(Arc::clone(&self.parent_cancelled));
5905 let factory: Arc<dyn SourceFactory<Out, NotUsed>> =
5906 Arc::clone(&slot) as Arc<dyn SourceFactory<Out, NotUsed>>;
5907 let hook: Arc<dyn SplitSegmentHookDyn> = Arc::clone(&slot) as Arc<dyn SplitSegmentHookDyn>;
5908 let source = Source {
5909 factory,
5910 terminal_factory: None,
5911 hints: SourceHints::default(),
5912 attributes: Attributes::default(),
5913 split_hook: Some(hook),
5914 };
5915 self.current_slot = Some(slot);
5916 push_live_substream(&self.outer, source).map_err(|_| ())
5917 }
5918
5919 fn push_item(&mut self, item: Out) -> Result<(), ()> {
5920 if let Some(ref mut consumer) = self.current_consumer {
5922 if let Err(e) = consumer.push_item(item) {
5923 let c = self.current_consumer.take().unwrap();
5924 c.fail(e);
5925 return Err(());
5926 }
5927 return Ok(());
5928 }
5929
5930 self.local_pending.push_back(item);
5932 if self.local_pending.len() >= LIVE_SUBSTREAM_BATCH {
5933 self.flush_pending()
5934 } else {
5935 Ok(())
5936 }
5937 }
5938
5939 fn flush_pending(&mut self) -> Result<(), ()> {
5942 if self.local_pending.is_empty() {
5943 return Ok(());
5944 }
5945
5946 if let Some(ref mut consumer) = self.current_consumer {
5948 for item in self.local_pending.drain(..) {
5949 if let Err(e) = consumer.push_item(item) {
5950 let c = self.current_consumer.take().unwrap();
5951 c.fail(e);
5952 return Err(());
5953 }
5954 }
5955 return Ok(());
5956 }
5957
5958 let slot = match &self.current_slot {
5959 Some(s) => Arc::clone(s),
5960 None => {
5961 self.local_pending.clear();
5962 return Ok(());
5963 }
5964 };
5965
5966 loop {
5967 if self.parent_cancelled.load(Ordering::SeqCst) {
5968 self.local_pending.clear();
5969 return Err(());
5970 }
5971
5972 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
5973
5974 if matches!(state.consumer, SegmentConsumer::Direct(_)) {
5975 let consumer =
5976 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
5977 SegmentConsumer::Direct(c) => c,
5978 _ => unreachable!(),
5979 };
5980 let mut drain_buf = std::mem::take(&mut state.buffer);
5981 drop(state);
5982 let mut consumer_box = consumer;
5983 for item in drain_buf.drain(..) {
5984 if let Err(e) = consumer_box.push_item(item) {
5985 consumer_box.fail(e);
5986 self.local_pending.clear();
5987 return Err(());
5988 }
5989 }
5990 for item in self.local_pending.drain(..) {
5991 if let Err(e) = consumer_box.push_item(item) {
5992 consumer_box.fail(e);
5993 return Err(());
5994 }
5995 }
5996 self.current_consumer = Some(consumer_box);
5997 return Ok(());
5998 }
5999
6000 if matches!(
6001 state.consumer,
6002 SegmentConsumer::Pending | SegmentConsumer::Fallback
6003 ) {
6004 let is_fallback = matches!(state.consumer, SegmentConsumer::Fallback);
6005 let cap = LIVE_SUBSTREAM_CAPACITY.saturating_sub(state.buffer.len());
6006 if cap > 0 {
6007 let to_flush = cap.min(self.local_pending.len());
6008 let items: Vec<Out> = self.local_pending.drain(..to_flush).collect();
6009 state.buffer.extend(items);
6010 let has_more = !self.local_pending.is_empty();
6011 drop(state);
6012 if is_fallback {
6013 slot.available.notify_all();
6014 }
6015 if !has_more {
6016 return Ok(());
6017 }
6018 continue;
6022 } else {
6023 state = slot
6025 .available
6026 .wait(state)
6027 .unwrap_or_else(|p| p.into_inner());
6028 continue;
6029 }
6030 }
6031
6032 return Ok(());
6034 }
6035 }
6036
6037 fn close_segment(&mut self) {
6038 let _ = self.flush_pending();
6040
6041 if let Some(consumer) = self.current_consumer.take() {
6043 consumer.complete();
6044 self.current_slot = None;
6045 return;
6046 }
6047
6048 let slot = match self.current_slot.take() {
6049 Some(s) => s,
6050 None => return,
6051 };
6052
6053 if self.parent_cancelled.load(Ordering::SeqCst) {
6054 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6055 if state.terminal.is_none() {
6056 state.terminal = Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
6057 }
6058 drop(state);
6059 slot.available.notify_all();
6060 return;
6061 }
6062
6063 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6064 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6065 SegmentConsumer::Direct(mut consumer_box) => {
6066 let mut drain_buf = std::mem::take(&mut state.buffer);
6067 drop(state);
6068 for item in drain_buf.drain(..) {
6069 if let Err(e) = consumer_box.push_item(item) {
6070 consumer_box.fail(e);
6071 return;
6072 }
6073 }
6074 consumer_box.complete();
6075 }
6076 SegmentConsumer::DirectTaken => {
6077 }
6079 SegmentConsumer::Fallback => {
6080 state.consumer = SegmentConsumer::DirectTaken;
6081 if state.terminal.is_none() {
6082 state.terminal = Some(LiveSubstreamTerminal::Complete);
6083 }
6084 drop(state);
6085 slot.available.notify_all();
6086 }
6087 SegmentConsumer::Pending => {
6088 state.consumer = SegmentConsumer::Pending;
6090 if state.terminal.is_none() {
6091 state.terminal = Some(LiveSubstreamTerminal::Complete);
6092 }
6093 drop(state);
6094 slot.available.notify_all();
6096 }
6097 }
6098 }
6099
6100 fn fail_segment(&mut self, error: StreamError) {
6101 self.local_pending.clear();
6103
6104 if let Some(consumer) = self.current_consumer.take() {
6105 consumer.fail(error);
6106 self.current_slot = None;
6107 return;
6108 }
6109 let slot = match self.current_slot.take() {
6110 Some(s) => s,
6111 None => return,
6112 };
6113 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6114 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6115 SegmentConsumer::Direct(c) => {
6116 drop(state);
6117 c.fail(error);
6118 }
6119 _ => {
6120 if state.terminal.is_none() {
6121 state.terminal = Some(LiveSubstreamTerminal::Error(error));
6122 }
6123 drop(state);
6124 slot.available.notify_all();
6125 }
6126 }
6127 }
6128
6129 fn fail_all(&mut self, error: StreamError) {
6130 self.fail_segment(error.clone());
6131 fail_live_substream(&self.outer, error);
6132 }
6133
6134 fn complete_all(&mut self) {
6135 self.close_segment();
6136 complete_live_substream(&self.outer);
6137 }
6138}
6139
6140impl<Out: Send + 'static> Drop for SplitFastWorkerGuard<Out> {
6141 fn drop(&mut self) {
6142 if self.armed {
6143 self.local_pending.clear(); if let Some(consumer) = self.current_consumer.take() {
6145 consumer.fail(StreamError::AbruptTermination);
6146 } else if let Some(slot) = self.current_slot.take() {
6147 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6148 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6149 SegmentConsumer::Direct(c) => {
6150 drop(state);
6151 c.fail(StreamError::AbruptTermination);
6152 }
6153 _ => {
6154 if state.terminal.is_none() {
6155 state.terminal =
6156 Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
6157 }
6158 drop(state);
6159 slot.available.notify_all();
6160 }
6161 }
6162 }
6163 fail_live_substream(&self.outer, StreamError::AbruptTermination);
6164 }
6165 }
6166}
6167
6168fn split_streams_fast<Out, F>(
6169 mut input: BoxStream<Out>,
6170 mode: SplitMode,
6171 predicate: Arc<F>,
6172 parent_cancelled: Arc<AtomicBool>,
6173 materializer: &Materializer,
6174) -> BoxStream<Source<Out>>
6175where
6176 Out: Clone + Send + 'static,
6177 F: Fn(&Out) -> bool + Send + Sync + 'static,
6178{
6179 let outer = LiveSubstreamShared::new();
6180 let worker_outer = Arc::clone(&outer);
6181 let completion = materializer.spawn_stream(move |cancelled| {
6182 let mut guard =
6183 SplitFastWorkerGuard::new(Arc::clone(&worker_outer), Arc::clone(&parent_cancelled));
6184
6185 while !cancelled.load(Ordering::SeqCst) {
6186 if worker_outer.cancelled.load(Ordering::SeqCst) {
6187 guard.disarm();
6188 return Ok(NotUsed);
6189 }
6190
6191 match input.next() {
6192 Some(Ok(item)) => {
6193 let split = match catch_unwind(AssertUnwindSafe(|| predicate(&item))) {
6194 Ok(s) => s,
6195 Err(_) => {
6196 guard.fail_all(StreamError::AbruptTermination);
6197 guard.disarm();
6198 return Ok(NotUsed);
6199 }
6200 };
6201
6202 match mode {
6203 SplitMode::When => {
6204 if split
6205 && (guard.current_slot.is_some()
6206 || guard.current_consumer.is_some())
6207 {
6208 guard.close_segment();
6209 }
6210 if guard.current_slot.is_none()
6211 && guard.current_consumer.is_none()
6212 && guard.open_segment().is_err()
6213 {
6214 guard.disarm();
6215 return Ok(NotUsed);
6216 }
6217 if guard.push_item(item).is_err() {
6218 guard.disarm();
6219 return Ok(NotUsed);
6220 }
6221 }
6222 SplitMode::After => {
6223 if guard.current_slot.is_none()
6224 && guard.current_consumer.is_none()
6225 && guard.open_segment().is_err()
6226 {
6227 guard.disarm();
6228 return Ok(NotUsed);
6229 }
6230 if guard.push_item(item).is_err() {
6231 guard.disarm();
6232 return Ok(NotUsed);
6233 }
6234 if split {
6235 guard.close_segment();
6236 }
6237 }
6238 }
6239 }
6240 Some(Err(error)) => {
6241 guard.fail_all(error.clone());
6242 guard.disarm();
6243 return Err(error);
6244 }
6245 None => {
6246 guard.complete_all();
6247 guard.disarm();
6248 return Ok(NotUsed);
6249 }
6250 }
6251 }
6252 guard.complete_all();
6253 guard.disarm();
6254 Ok(NotUsed)
6255 });
6256
6257 Box::new(LiveSubstreamStream {
6258 shared: outer,
6259 completion: Some(completion),
6260 local_batch: VecDeque::new(),
6261 })
6262}
6263
6264#[cfg(test)]
6266struct FlatMapMergeShared<T> {
6267 state: Mutex<FlatMapMergeState<T>>,
6268 available: Condvar,
6269 cancelled: Arc<AtomicBool>,
6270}
6271
6272#[cfg(test)]
6273struct FlatMapMergeState<T> {
6274 queued: VecDeque<(usize, T)>,
6275 window: HashMap<usize, usize>,
6276 active_streams: usize,
6277 input_done: bool,
6278 terminal: Option<StreamError>,
6279}
6280
6281#[cfg(test)]
6282impl<T> FlatMapMergeShared<T> {
6283 fn new() -> Arc<Self> {
6284 Arc::new(Self {
6285 state: Mutex::new(FlatMapMergeState {
6286 queued: VecDeque::new(),
6287 window: HashMap::new(),
6288 active_streams: 0,
6289 input_done: false,
6290 terminal: None,
6291 }),
6292 available: Condvar::new(),
6293 cancelled: Arc::new(AtomicBool::new(false)),
6294 })
6295 }
6296}
6297
6298#[cfg(test)]
6299struct FlatMapMergeStream<T> {
6300 shared: Arc<FlatMapMergeShared<T>>,
6301 completion: Option<StreamCompletion<NotUsed>>,
6302}
6303
6304#[cfg(test)]
6305impl<T> Iterator for FlatMapMergeStream<T> {
6306 type Item = StreamResult<T>;
6307
6308 fn next(&mut self) -> Option<Self::Item> {
6309 let mut state = self
6310 .shared
6311 .state
6312 .lock()
6313 .unwrap_or_else(|poison| poison.into_inner());
6314 loop {
6315 if let Some((stream_id, item)) = state.queued.pop_front() {
6316 if let Some(count) = state.window.get_mut(&stream_id) {
6319 *count -= 1;
6320 if *count == 0 {
6321 state.window.remove(&stream_id);
6322 }
6323 }
6324 drop(state);
6325 self.shared.available.notify_all();
6326 return Some(Ok(item));
6327 }
6328 if let Some(error) = state.terminal.clone() {
6329 return Some(Err(error));
6330 }
6331 if state.input_done && state.active_streams == 0 {
6332 return None;
6333 }
6334 state = self
6335 .shared
6336 .available
6337 .wait(state)
6338 .unwrap_or_else(|poison| poison.into_inner());
6339 }
6340 }
6341}
6342
6343#[cfg(test)]
6344impl<T> Drop for FlatMapMergeStream<T> {
6345 fn drop(&mut self) {
6346 self.shared.cancelled.store(true, Ordering::SeqCst);
6347 self.shared.available.notify_all();
6348 let _ = self.completion.take();
6349 }
6350}
6351
6352#[cfg(test)]
6353fn flat_map_merge_register_stream<T>(shared: &Arc<FlatMapMergeShared<T>>) {
6354 let mut state = shared
6355 .state
6356 .lock()
6357 .unwrap_or_else(|poison| poison.into_inner());
6358 state.active_streams += 1;
6359 drop(state);
6360 shared.available.notify_all();
6361}
6362
6363#[cfg(test)]
6364fn flat_map_merge_finish_stream<T>(
6365 shared: &Arc<FlatMapMergeShared<T>>,
6366 stream_id: usize,
6367 terminal: Result<(), StreamError>,
6368) {
6369 let mut state = shared
6370 .state
6371 .lock()
6372 .unwrap_or_else(|poison| poison.into_inner());
6373 state.active_streams = state.active_streams.saturating_sub(1);
6374 if let Err(error) = terminal
6375 && state.terminal.is_none()
6376 {
6377 state.terminal = Some(error);
6378 }
6379 state.window.entry(stream_id).or_default();
6381 drop(state);
6382 shared.available.notify_all();
6383}
6384
6385#[cfg(test)]
6386fn flat_map_merge_push<T>(
6387 shared: &Arc<FlatMapMergeShared<T>>,
6388 stream_id: usize,
6389 item: T,
6390) -> Result<(), T> {
6391 let mut state = shared
6394 .state
6395 .lock()
6396 .unwrap_or_else(|poison| poison.into_inner());
6397 while state.window.get(&stream_id).copied().unwrap_or(0) >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
6398 && state.terminal.is_none()
6399 {
6400 if shared.cancelled.load(Ordering::SeqCst) {
6401 return Err(item);
6402 }
6403 state = shared
6404 .available
6405 .wait(state)
6406 .unwrap_or_else(|poison| poison.into_inner());
6407 }
6408 if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
6409 return Err(item);
6410 }
6411 *state.window.entry(stream_id).or_insert(0) += 1;
6412 let was_empty = state.queued.is_empty();
6413 state.queued.push_back((stream_id, item));
6414 drop(state);
6415 if was_empty {
6416 shared.available.notify_all();
6417 }
6418 Ok(())
6419}
6420
6421#[cfg(test)]
6422struct FlatMapMergeCoordinatorGuard<T> {
6423 shared: Arc<FlatMapMergeShared<T>>,
6424 armed: bool,
6425}
6426
6427#[cfg(test)]
6428impl<T> FlatMapMergeCoordinatorGuard<T> {
6429 fn new(shared: Arc<FlatMapMergeShared<T>>) -> Self {
6430 Self {
6431 shared,
6432 armed: true,
6433 }
6434 }
6435
6436 fn finish(&mut self) {
6437 let mut state = self
6438 .shared
6439 .state
6440 .lock()
6441 .unwrap_or_else(|poison| poison.into_inner());
6442 state.input_done = true;
6443 drop(state);
6444 self.shared.available.notify_all();
6445 self.armed = false;
6446 }
6447}
6448
6449#[cfg(test)]
6450impl<T> Drop for FlatMapMergeCoordinatorGuard<T> {
6451 fn drop(&mut self) {
6452 if self.armed {
6453 let mut state = self
6454 .shared
6455 .state
6456 .lock()
6457 .unwrap_or_else(|poison| poison.into_inner());
6458 if state.terminal.is_none() {
6459 state.terminal = Some(StreamError::AbruptTermination);
6460 }
6461 state.input_done = true;
6462 drop(state);
6463 self.shared.cancelled.store(true, Ordering::SeqCst);
6464 self.shared.available.notify_all();
6465 }
6466 }
6467}
6468
6469#[cfg(test)]
6470struct FlatMapMergeWorkerGuard<T> {
6471 shared: Arc<FlatMapMergeShared<T>>,
6472 stream_id: usize,
6473 armed: bool,
6474}
6475
6476#[cfg(test)]
6477impl<T> FlatMapMergeWorkerGuard<T> {
6478 fn new(shared: Arc<FlatMapMergeShared<T>>, stream_id: usize) -> Self {
6479 Self {
6480 shared,
6481 stream_id,
6482 armed: true,
6483 }
6484 }
6485
6486 fn finish(&mut self, terminal: Result<(), StreamError>) {
6487 flat_map_merge_finish_stream(&self.shared, self.stream_id, terminal);
6488 self.armed = false;
6489 }
6490}
6491
6492#[cfg(test)]
6493impl<T> Drop for FlatMapMergeWorkerGuard<T> {
6494 fn drop(&mut self) {
6495 if self.armed {
6496 flat_map_merge_finish_stream(
6497 &self.shared,
6498 self.stream_id,
6499 Err(StreamError::AbruptTermination),
6500 );
6501 self.shared.cancelled.store(true, Ordering::SeqCst);
6502 }
6503 }
6504}
6505
6506#[cfg(test)]
6509#[derive(Clone, Copy, PartialEq, Eq, Debug)]
6510pub(crate) enum SubstreamExecutorMode {
6511 Auto,
6512 LegacyOnly,
6513 ReadyRingOnly,
6514 SplitSinkOnly,
6515}
6516
6517#[cfg(test)]
6518thread_local! {
6519 static SUBSTREAM_EXECUTOR_MODE: std::cell::Cell<SubstreamExecutorMode> =
6520 const { std::cell::Cell::new(SubstreamExecutorMode::Auto) };
6521}
6522
6523#[cfg(test)]
6524pub(crate) fn with_substream_mode<R>(mode: SubstreamExecutorMode, f: impl FnOnce() -> R) -> R {
6525 SUBSTREAM_EXECUTOR_MODE.with(|m| {
6526 let old = m.get();
6527 m.set(mode);
6528 let result = f();
6529 m.set(old);
6530 result
6531 })
6532}
6533
6534#[cfg(test)]
6535fn current_substream_mode() -> SubstreamExecutorMode {
6536 SUBSTREAM_EXECUTOR_MODE.with(|m| m.get())
6537}
6538
6539fn flat_map_merge_stream<Out, Next, NextMat, F>(
6542 input: BoxStream<Out>,
6543 breadth: usize,
6544 stage: Arc<F>,
6545 materializer: &Materializer,
6546) -> BoxStream<Next>
6547where
6548 Out: Send + 'static,
6549 Next: Send + 'static,
6550 NextMat: Send + 'static,
6551 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6552{
6553 #[cfg(test)]
6554 if current_substream_mode() == SubstreamExecutorMode::LegacyOnly {
6555 return flat_map_merge_stream_legacy(input, breadth, stage, materializer);
6556 }
6557 flat_map_merge_stream_ready(input, breadth, stage, materializer)
6558}
6559
6560#[cfg(test)]
6563fn flat_map_merge_stream_legacy<Out, Next, NextMat, F>(
6564 mut input: BoxStream<Out>,
6565 breadth: usize,
6566 stage: Arc<F>,
6567 materializer: &Materializer,
6568) -> BoxStream<Next>
6569where
6570 Out: Send + 'static,
6571 Next: Send + 'static,
6572 NextMat: Send + 'static,
6573 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6574{
6575 let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
6576 let shared = FlatMapMergeShared::new();
6577 let worker_shared = Arc::clone(&shared);
6578 let completion = materializer.spawn_stream(move |cancelled| {
6579 let mut next_id = 0usize;
6580 let mut guard = FlatMapMergeCoordinatorGuard::new(worker_shared);
6581 let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
6582
6583 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6584 workers.retain(|_, completion| completion.try_wait().is_none());
6585 {
6586 let mut state = guard
6587 .shared
6588 .state
6589 .lock()
6590 .unwrap_or_else(|poison| poison.into_inner());
6591 while state.active_streams >= breadth
6592 && state.terminal.is_none()
6593 && !cancelled.load(Ordering::SeqCst)
6594 && !guard.shared.cancelled.load(Ordering::SeqCst)
6595 {
6596 state = guard
6597 .shared
6598 .available
6599 .wait(state)
6600 .unwrap_or_else(|poison| poison.into_inner());
6601 }
6602 if state.terminal.is_some() {
6603 let error = state.terminal.clone().expect("terminal checked above");
6604 drop(state);
6605 guard.finish();
6606 return Err(error);
6607 }
6608 }
6609
6610 match input.next() {
6611 Some(Ok(item)) => {
6612 let source = stage(item);
6613 let (mut stream, _) =
6614 match Arc::clone(&source.factory).create(&worker_materializer) {
6615 Ok(parts) => parts,
6616 Err(error) => {
6617 let mut state = guard
6618 .shared
6619 .state
6620 .lock()
6621 .unwrap_or_else(|poison| poison.into_inner());
6622 if state.terminal.is_none() {
6623 state.terminal = Some(error.clone());
6624 }
6625 drop(state);
6626 guard.shared.available.notify_all();
6627 guard.finish();
6628 return Err(error);
6629 }
6630 };
6631 let stream_id = next_id;
6632 next_id += 1;
6633 flat_map_merge_register_stream(&guard.shared);
6634 let worker_shared = Arc::clone(&guard.shared);
6635 workers.insert(
6636 stream_id,
6637 worker_materializer.spawn_stream(move |inner_cancelled| {
6638 let mut worker_guard =
6639 FlatMapMergeWorkerGuard::new(Arc::clone(&worker_shared), stream_id);
6640 while !inner_cancelled.load(Ordering::SeqCst)
6641 && !worker_shared.cancelled.load(Ordering::SeqCst)
6642 {
6643 match stream.next() {
6644 Some(Ok(item)) => {
6645 if flat_map_merge_push(&worker_shared, stream_id, item)
6646 .is_err()
6647 {
6648 worker_guard.finish(Ok(()));
6649 return Ok(NotUsed);
6650 }
6651 }
6652 Some(Err(error)) => {
6653 worker_guard.finish(Err(error.clone()));
6654 return Err(error);
6655 }
6656 None => {
6657 worker_guard.finish(Ok(()));
6658 return Ok(NotUsed);
6659 }
6660 }
6661 }
6662 worker_guard.finish(Ok(()));
6663 Ok(NotUsed)
6664 }),
6665 );
6666 }
6667 Some(Err(error)) => {
6668 let mut state = guard
6669 .shared
6670 .state
6671 .lock()
6672 .unwrap_or_else(|poison| poison.into_inner());
6673 if state.terminal.is_none() {
6674 state.terminal = Some(error.clone());
6675 }
6676 drop(state);
6677 guard.shared.available.notify_all();
6678 guard.finish();
6679 return Err(error);
6680 }
6681 None => {
6682 guard.finish();
6683 break;
6684 }
6685 }
6686 }
6687
6688 if guard.armed {
6689 guard.finish();
6690 }
6691 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6699 workers.retain(|_, completion| completion.try_wait().is_none());
6700 let state = guard
6706 .shared
6707 .state
6708 .lock()
6709 .unwrap_or_else(|poison| poison.into_inner());
6710 if state.active_streams == 0 {
6711 return Ok(NotUsed);
6712 }
6713 drop(
6714 guard
6715 .shared
6716 .available
6717 .wait(state)
6718 .unwrap_or_else(|poison| poison.into_inner()),
6719 );
6720 }
6721 Ok(NotUsed)
6722 });
6723
6724 Box::new(FlatMapMergeStream {
6725 shared,
6726 completion: Some(completion),
6727 })
6728}
6729
6730const FLAT_MAP_MERGE_READY_BATCH: usize = 16;
6733
6734const FLAT_MAP_MERGE_INLINE_MICRO_MAX: usize = FLAT_MAP_MERGE_READY_BATCH;
6738
6739struct FlatMapMergeReadyShared<T> {
6740 coordinator: Mutex<FlatMapMergeReadyState<T>>,
6741 available: Condvar,
6742 cancelled: Arc<AtomicBool>,
6743}
6744
6745struct FlatMapMergeReadyState<T> {
6746 lanes: HashMap<usize, Arc<FlatMapMergeLane<T>>>,
6747 ready: std::collections::VecDeque<usize>,
6748 queued_items: usize,
6749 active_streams: usize,
6750 input_done: bool,
6751 terminal: Option<StreamError>,
6752 generation: u64,
6753}
6754
6755struct FlatMapMergeLane<T> {
6756 state: Mutex<FlatMapMergeLaneState<T>>,
6757 space_available: Condvar,
6758}
6759
6760struct FlatMapMergeLaneState<T> {
6761 buffer: std::collections::VecDeque<T>,
6762 in_ready_ring: bool,
6763 publishing: bool,
6764 closed: bool,
6765}
6766
6767struct FlatMapMergeReadyStream<T> {
6768 shared: Arc<FlatMapMergeReadyShared<T>>,
6769 completion: Option<StreamCompletion<NotUsed>>,
6770 local_batch: std::collections::VecDeque<T>,
6771}
6772
6773impl<T: Send + 'static> Iterator for FlatMapMergeReadyStream<T> {
6774 type Item = StreamResult<T>;
6775
6776 fn next(&mut self) -> Option<Self::Item> {
6777 if let Some(item) = self.local_batch.pop_front() {
6779 return Some(Ok(item));
6780 }
6781
6782 let mut coord = self
6784 .shared
6785 .coordinator
6786 .lock()
6787 .unwrap_or_else(|p| p.into_inner());
6788 loop {
6789 if let Some(error) = coord.terminal.clone() {
6791 return Some(Err(error));
6792 }
6793
6794 if let Some(lane_id) = coord.ready.pop_front() {
6796 let lane = coord.lanes.get(&lane_id).cloned();
6797 drop(coord);
6798
6799 if let Some(lane) = lane {
6800 let mut lane_state = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6802 while lane_state.publishing {
6803 lane_state = lane
6804 .space_available
6805 .wait(lane_state)
6806 .unwrap_or_else(|p| p.into_inner());
6807 }
6808 let drain_n = lane_state.buffer.len().min(FLAT_MAP_MERGE_READY_BATCH);
6809 let mut batch = std::collections::VecDeque::with_capacity(drain_n);
6810 for _ in 0..drain_n {
6811 if let Some(item) = lane_state.buffer.pop_front() {
6812 batch.push_back(item);
6813 }
6814 }
6815 let still_has_items = !lane_state.buffer.is_empty();
6816 let is_closed = lane_state.closed;
6817 if !still_has_items {
6818 lane_state.in_ready_ring = false;
6819 }
6820 drop(lane_state);
6822 if !batch.is_empty() {
6824 lane.space_available.notify_all();
6825 }
6826
6827 let freed = batch.len();
6829 let mut coord = self
6830 .shared
6831 .coordinator
6832 .lock()
6833 .unwrap_or_else(|p| p.into_inner());
6834 coord.queued_items = coord.queued_items.saturating_sub(freed);
6835 if still_has_items {
6836 coord.ready.push_back(lane_id);
6837 } else if is_closed {
6838 coord.lanes.remove(&lane_id);
6839 }
6840 coord.generation += 1;
6841 drop(coord);
6842 self.shared.available.notify_all();
6843
6844 let mut iter = batch.into_iter();
6846 if let Some(first) = iter.next() {
6847 self.local_batch.extend(iter);
6848 return Some(Ok(first));
6849 }
6850 }
6851 coord = self
6853 .shared
6854 .coordinator
6855 .lock()
6856 .unwrap_or_else(|p| p.into_inner());
6857 continue;
6858 }
6859
6860 if coord.terminal.is_none()
6862 && coord.input_done
6863 && coord.active_streams == 0
6864 && coord.queued_items == 0
6865 {
6866 return None;
6867 }
6868
6869 let seen = coord.generation;
6871 coord = self
6872 .shared
6873 .available
6874 .wait_while(coord, |s| {
6875 s.generation == seen
6876 && s.terminal.is_none()
6877 && s.ready.is_empty()
6878 && !(s.input_done && s.active_streams == 0 && s.queued_items == 0)
6879 })
6880 .unwrap_or_else(|p| p.into_inner());
6881 }
6882 }
6883}
6884
6885impl<T> Drop for FlatMapMergeReadyStream<T> {
6886 fn drop(&mut self) {
6887 let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
6888 let mut coord = self
6889 .shared
6890 .coordinator
6891 .lock()
6892 .unwrap_or_else(|p| p.into_inner());
6893 coord.generation += 1;
6894 coord.lanes.values().cloned().collect()
6895 };
6896 self.shared.cancelled.store(true, Ordering::SeqCst);
6897 self.shared.available.notify_all();
6898 for lane in lanes {
6899 lane.space_available.notify_all();
6900 }
6901 let _ = self.completion.take();
6902 }
6903}
6904
6905fn finish_lane_ready<T>(
6908 shared: &Arc<FlatMapMergeReadyShared<T>>,
6909 stream_id: usize,
6910 terminal: Result<(), StreamError>,
6911) {
6912 let is_error = terminal.is_err();
6913
6914 let lane_is_empty = {
6916 let coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6917 let lane_opt = coord.lanes.get(&stream_id).cloned();
6918 drop(coord);
6919 if let Some(lane) = lane_opt {
6920 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6921 ls.closed = true;
6922 ls.buffer.is_empty()
6923 } else {
6924 true
6925 }
6926 };
6927
6928 let lanes_to_notify: Vec<Arc<FlatMapMergeLane<T>>> = {
6930 let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6931 coord.active_streams = coord.active_streams.saturating_sub(1);
6932 if let Err(ref error) = terminal
6933 && coord.terminal.is_none()
6934 {
6935 coord.terminal = Some(error.clone());
6936 }
6937 if lane_is_empty {
6938 coord.lanes.remove(&stream_id);
6939 }
6940 coord.generation += 1;
6941 if is_error {
6942 coord.lanes.values().cloned().collect()
6943 } else {
6944 vec![]
6945 }
6946 };
6947
6948 if is_error {
6949 shared.cancelled.store(true, Ordering::SeqCst);
6950 for lane in &lanes_to_notify {
6951 lane.space_available.notify_all();
6952 }
6953 }
6954 shared.available.notify_all();
6955}
6956
6957struct FlatMapMergeReadyCoordinatorGuard<T> {
6958 shared: Arc<FlatMapMergeReadyShared<T>>,
6959 armed: bool,
6960}
6961
6962impl<T> FlatMapMergeReadyCoordinatorGuard<T> {
6963 fn new(shared: Arc<FlatMapMergeReadyShared<T>>) -> Self {
6964 Self {
6965 shared,
6966 armed: true,
6967 }
6968 }
6969
6970 fn finish(&mut self) {
6971 let mut coord = self
6972 .shared
6973 .coordinator
6974 .lock()
6975 .unwrap_or_else(|p| p.into_inner());
6976 coord.input_done = true;
6977 coord.generation += 1;
6978 drop(coord);
6979 self.shared.available.notify_all();
6980 self.armed = false;
6981 }
6982}
6983
6984impl<T> Drop for FlatMapMergeReadyCoordinatorGuard<T> {
6985 fn drop(&mut self) {
6986 if self.armed {
6987 let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
6988 let mut coord = self
6989 .shared
6990 .coordinator
6991 .lock()
6992 .unwrap_or_else(|p| p.into_inner());
6993 if coord.terminal.is_none() {
6994 coord.terminal = Some(StreamError::AbruptTermination);
6995 }
6996 coord.input_done = true;
6997 coord.generation += 1;
6998 coord.lanes.values().cloned().collect()
6999 };
7000 self.shared.cancelled.store(true, Ordering::SeqCst);
7001 self.shared.available.notify_all();
7002 for lane in lanes {
7003 lane.space_available.notify_all();
7004 }
7005 }
7006 }
7007}
7008
7009struct FlatMapMergeReadyWorkerGuard<T> {
7010 shared: Arc<FlatMapMergeReadyShared<T>>,
7011 stream_id: usize,
7012 armed: bool,
7013}
7014
7015impl<T> FlatMapMergeReadyWorkerGuard<T> {
7016 fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
7017 Self {
7018 shared,
7019 stream_id,
7020 armed: true,
7021 }
7022 }
7023
7024 fn finish(&mut self, terminal: Result<(), StreamError>) {
7025 finish_lane_ready(&self.shared, self.stream_id, terminal);
7026 self.armed = false;
7027 }
7028}
7029
7030impl<T> Drop for FlatMapMergeReadyWorkerGuard<T> {
7031 fn drop(&mut self) {
7032 if self.armed {
7033 finish_lane_ready(
7034 &self.shared,
7035 self.stream_id,
7036 Err(StreamError::AbruptTermination),
7037 );
7038 self.shared.cancelled.store(true, Ordering::SeqCst);
7039 }
7040 }
7041}
7042
7043fn publish_ready_batch<T>(
7058 shared: &Arc<FlatMapMergeReadyShared<T>>,
7059 stream_id: usize,
7060 lane: &Arc<FlatMapMergeLane<T>>,
7061 batch: std::collections::VecDeque<T>,
7062) {
7063 let batch_len = batch.len();
7064 if batch_len == 0 {
7065 return;
7066 }
7067
7068 let was_not_in_ready = {
7070 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
7071 let was_not = !ls.in_ready_ring;
7072 ls.publishing = true;
7073 ls.buffer.extend(batch);
7074 if !ls.buffer.is_empty() {
7075 ls.in_ready_ring = true;
7076 }
7077 was_not
7078 };
7079
7080 {
7082 let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
7083 coord.queued_items += batch_len;
7084 if was_not_in_ready {
7085 coord.ready.push_back(stream_id);
7086 }
7087 coord.generation += 1;
7088 }
7089 {
7090 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
7091 ls.publishing = false;
7092 }
7093 lane.space_available.notify_all();
7094 shared.available.notify_all();
7095}
7096
7097struct FlatMapMergeReadyInlineGuard<T> {
7107 shared: Arc<FlatMapMergeReadyShared<T>>,
7108 stream_id: usize,
7109 armed: bool,
7110}
7111
7112impl<T> FlatMapMergeReadyInlineGuard<T> {
7113 fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
7114 Self {
7115 shared,
7116 stream_id,
7117 armed: true,
7118 }
7119 }
7120
7121 fn finish(&mut self, terminal: Result<(), StreamError>) {
7122 finish_lane_ready(&self.shared, self.stream_id, terminal);
7123 self.armed = false;
7124 }
7125
7126 fn hand_off(&mut self) {
7127 self.armed = false;
7128 }
7129}
7130
7131impl<T> Drop for FlatMapMergeReadyInlineGuard<T> {
7132 fn drop(&mut self) {
7133 if self.armed {
7134 finish_lane_ready(
7135 &self.shared,
7136 self.stream_id,
7137 Err(StreamError::AbruptTermination),
7138 );
7139 }
7140 }
7141}
7142
7143fn flat_map_merge_stream_ready<Out, Next, NextMat, F>(
7144 mut input: BoxStream<Out>,
7145 breadth: usize,
7146 stage: Arc<F>,
7147 materializer: &Materializer,
7148) -> BoxStream<Next>
7149where
7150 Out: Send + 'static,
7151 Next: Send + 'static,
7152 NextMat: Send + 'static,
7153 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7154{
7155 let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7156 let shared: Arc<FlatMapMergeReadyShared<Next>> = Arc::new(FlatMapMergeReadyShared {
7157 coordinator: Mutex::new(FlatMapMergeReadyState {
7158 lanes: HashMap::new(),
7159 ready: std::collections::VecDeque::new(),
7160 queued_items: 0,
7161 active_streams: 0,
7162 input_done: false,
7163 terminal: None,
7164 generation: 0,
7165 }),
7166 available: Condvar::new(),
7167 cancelled: Arc::new(AtomicBool::new(false)),
7168 });
7169
7170 let worker_shared = Arc::clone(&shared);
7171 let completion = materializer.spawn_stream(move |cancelled| {
7172 let mut next_id = 0usize;
7173 let mut guard = FlatMapMergeReadyCoordinatorGuard::new(worker_shared);
7174 let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
7175
7176 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7177 workers.retain(|_, c| c.try_wait().is_none());
7178
7179 {
7181 let mut coord = guard
7182 .shared
7183 .coordinator
7184 .lock()
7185 .unwrap_or_else(|p| p.into_inner());
7186 loop {
7187 if coord.terminal.is_some()
7188 || coord.active_streams < breadth
7189 || cancelled.load(Ordering::SeqCst)
7190 || guard.shared.cancelled.load(Ordering::SeqCst)
7191 {
7192 break;
7193 }
7194 let seen = coord.generation;
7195 coord = guard
7196 .shared
7197 .available
7198 .wait_while(coord, |s| {
7199 s.generation == seen
7200 && s.terminal.is_none()
7201 && s.active_streams >= breadth
7202 && !cancelled.load(Ordering::SeqCst)
7203 && !guard.shared.cancelled.load(Ordering::SeqCst)
7204 })
7205 .unwrap_or_else(|p| p.into_inner());
7206 }
7207 if coord.terminal.is_some() {
7208 let error = coord.terminal.clone().expect("terminal checked");
7209 drop(coord);
7210 guard.finish();
7211 return Err(error);
7212 }
7213 }
7214
7215 match input.next() {
7216 Some(Ok(item)) => {
7217 let source = stage(item);
7219 let inline_hint = source.hints.inline_micro;
7221 let (stream, _) = match Arc::clone(&source.factory).create(&worker_materializer)
7222 {
7223 Ok(parts) => parts,
7224 Err(error) => {
7225 let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7226 let mut coord = guard
7227 .shared
7228 .coordinator
7229 .lock()
7230 .unwrap_or_else(|p| p.into_inner());
7231 if coord.terminal.is_none() {
7232 coord.terminal = Some(error.clone());
7233 }
7234 coord.generation += 1;
7235 coord.lanes.values().cloned().collect()
7236 };
7237 guard.shared.cancelled.store(true, Ordering::SeqCst);
7238 guard.shared.available.notify_all();
7239 for lane in lanes {
7240 lane.space_available.notify_all();
7241 }
7242 guard.finish();
7243 return Err(error);
7244 }
7245 };
7246
7247 let stream_id = next_id;
7248 next_id += 1;
7249 let mut stream = stream;
7250
7251 let lane: Arc<FlatMapMergeLane<Next>> = Arc::new(FlatMapMergeLane {
7253 state: Mutex::new(FlatMapMergeLaneState {
7254 buffer: std::collections::VecDeque::new(),
7255 in_ready_ring: false,
7256 publishing: false,
7257 closed: false,
7258 }),
7259 space_available: Condvar::new(),
7260 });
7261 {
7262 let mut coord = guard
7263 .shared
7264 .coordinator
7265 .lock()
7266 .unwrap_or_else(|p| p.into_inner());
7267 coord.lanes.insert(stream_id, Arc::clone(&lane));
7268 coord.active_streams += 1;
7269 }
7271
7272 let mut inline_guard =
7275 FlatMapMergeReadyInlineGuard::new(Arc::clone(&guard.shared), stream_id);
7276
7277 let is_inline = inline_hint
7278 .is_some_and(|h| h.max_success_items <= FLAT_MAP_MERGE_INLINE_MICRO_MAX);
7279
7280 if is_inline {
7281 let max_items = inline_hint.expect("checked").max_success_items;
7285
7286 if guard.shared.cancelled.load(Ordering::SeqCst)
7288 || cancelled.load(Ordering::SeqCst)
7289 {
7290 inline_guard.finish(Ok(()));
7291 continue;
7292 }
7293
7294 let max_pulls = max_items.saturating_add(1);
7297 let mut local_batch = std::collections::VecDeque::with_capacity(max_items);
7298 let mut terminal_result: Option<Result<(), StreamError>> = None;
7299
7300 for _ in 0..max_pulls {
7301 if guard.shared.cancelled.load(Ordering::SeqCst)
7302 || cancelled.load(Ordering::SeqCst)
7303 {
7304 break;
7305 }
7306 match stream.next() {
7307 Some(Ok(item)) => local_batch.push_back(item),
7308 Some(Err(e)) => {
7309 terminal_result = Some(Err(e));
7310 break;
7311 }
7312 None => {
7313 terminal_result = Some(Ok(()));
7314 break;
7315 }
7316 }
7317 }
7318
7319 if guard.shared.cancelled.load(Ordering::SeqCst)
7321 || cancelled.load(Ordering::SeqCst)
7322 {
7323 inline_guard.finish(Ok(()));
7324 continue;
7325 }
7326
7327 publish_ready_batch(&guard.shared, stream_id, &lane, local_batch);
7330
7331 match terminal_result {
7332 Some(Err(e)) => {
7333 inline_guard.finish(Err(e));
7336 }
7337 Some(Ok(())) | None => {
7338 inline_guard.finish(Ok(()));
7339 }
7340 }
7341 } else {
7342 inline_guard.hand_off();
7344 let worker_shared = Arc::clone(&guard.shared);
7345 let worker_lane = Arc::clone(&lane);
7346 workers.insert(
7347 stream_id,
7348 worker_materializer.spawn_stream(move |inner_cancelled| {
7349 let mut worker_guard = FlatMapMergeReadyWorkerGuard::new(
7350 Arc::clone(&worker_shared),
7351 stream_id,
7352 );
7353
7354 loop {
7355 if inner_cancelled.load(Ordering::SeqCst)
7356 || worker_shared.cancelled.load(Ordering::SeqCst)
7357 {
7358 worker_guard.finish(Ok(()));
7359 return Ok(NotUsed);
7360 }
7361
7362 let capacity;
7364 {
7365 let mut ls = worker_lane
7366 .state
7367 .lock()
7368 .unwrap_or_else(|p| p.into_inner());
7369 while ls.buffer.len() >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
7370 && !worker_shared.cancelled.load(Ordering::SeqCst)
7371 && !ls.closed
7372 && !inner_cancelled.load(Ordering::SeqCst)
7373 {
7374 ls = worker_lane
7375 .space_available
7376 .wait(ls)
7377 .unwrap_or_else(|p| p.into_inner());
7378 }
7379 if worker_shared.cancelled.load(Ordering::SeqCst)
7380 || ls.closed
7381 || inner_cancelled.load(Ordering::SeqCst)
7382 {
7383 drop(ls);
7384 worker_guard.finish(Ok(()));
7385 return Ok(NotUsed);
7386 }
7387 capacity =
7388 FLAT_MAP_MERGE_SUBSTREAM_WINDOW - ls.buffer.len();
7389 }
7390 let batch_size = capacity.min(FLAT_MAP_MERGE_READY_BATCH);
7394 let mut local_batch =
7395 std::collections::VecDeque::with_capacity(batch_size);
7396 let mut terminal_result: Option<Result<(), StreamError>> = None;
7397 for _ in 0..batch_size {
7398 if inner_cancelled.load(Ordering::SeqCst)
7399 || worker_shared.cancelled.load(Ordering::SeqCst)
7400 {
7401 break;
7402 }
7403 match stream.next() {
7404 Some(Ok(item)) => local_batch.push_back(item),
7405 Some(Err(e)) => {
7406 terminal_result = Some(Err(e));
7407 break;
7408 }
7409 None => {
7410 terminal_result = Some(Ok(()));
7411 break;
7412 }
7413 }
7414 }
7415 let batch_len = local_batch.len();
7416
7417 publish_ready_batch(
7420 &worker_shared,
7421 stream_id,
7422 &worker_lane,
7423 local_batch,
7424 );
7425
7426 match terminal_result {
7427 Some(Ok(())) => {
7428 worker_guard.finish(Ok(()));
7429 return Ok(NotUsed);
7430 }
7431 Some(Err(e)) => {
7432 worker_guard.finish(Err(e.clone()));
7433 return Err(e);
7434 }
7435 None => {
7436 if batch_len == 0 {
7437 worker_guard.finish(Ok(()));
7438 return Ok(NotUsed);
7439 }
7440 }
7441 }
7442 }
7443 }),
7444 );
7445 }
7446 }
7447 Some(Err(error)) => {
7448 let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7449 let mut coord = guard
7450 .shared
7451 .coordinator
7452 .lock()
7453 .unwrap_or_else(|p| p.into_inner());
7454 if coord.terminal.is_none() {
7455 coord.terminal = Some(error.clone());
7456 }
7457 coord.generation += 1;
7458 coord.lanes.values().cloned().collect()
7459 };
7460 guard.shared.cancelled.store(true, Ordering::SeqCst);
7461 guard.shared.available.notify_all();
7462 for lane in lanes {
7463 lane.space_available.notify_all();
7464 }
7465 guard.finish();
7466 return Err(error);
7467 }
7468 None => {
7469 guard.finish();
7470 break;
7471 }
7472 }
7473 }
7474
7475 if guard.armed {
7476 guard.finish();
7477 }
7478
7479 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7483 workers.retain(|_, c| c.try_wait().is_none());
7484 let coord = guard
7485 .shared
7486 .coordinator
7487 .lock()
7488 .unwrap_or_else(|p| p.into_inner());
7489 if coord.active_streams == 0 {
7490 return Ok(NotUsed);
7491 }
7492 let seen = coord.generation;
7493 drop(
7494 guard
7495 .shared
7496 .available
7497 .wait_while(coord, |s| s.generation == seen && s.active_streams > 0)
7498 .unwrap_or_else(|p| p.into_inner()),
7499 );
7500 }
7501 Ok(NotUsed)
7502 });
7503
7504 Box::new(FlatMapMergeReadyStream {
7505 shared,
7506 completion: Some(completion),
7507 local_batch: std::collections::VecDeque::new(),
7508 })
7509}
7510
7511fn flat_map_concat_stream<Out, Next, NextMat, F>(
7512 mut input: BoxStream<Out>,
7513 stage: Arc<F>,
7514 materializer: &Materializer,
7515) -> BoxStream<Next>
7516where
7517 Out: Send + 'static,
7518 Next: Send + 'static,
7519 NextMat: Send + 'static,
7520 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7521{
7522 let materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7523 let mut current: Option<BoxStream<Next>> = None;
7524 Box::new(std::iter::from_fn(move || {
7525 loop {
7526 if let Some(stream) = current.as_mut() {
7527 match stream.next() {
7528 Some(item) => return Some(item),
7529 None => current = None,
7530 }
7531 }
7532
7533 match input.next() {
7534 Some(Ok(item)) => {
7535 let source = stage(item);
7536 current = Some(match Arc::clone(&source.factory).create(&materializer) {
7537 Ok((stream, _)) => stream,
7538 Err(error) => Box::new(std::iter::once(Err(error))) as BoxStream<Next>,
7539 });
7540 }
7541 Some(Err(error)) => return Some(Err(error)),
7542 None => return None,
7543 }
7544 }
7545 }))
7546}
7547
7548fn map_async_unordered<Out, Next, F, Fut>(
7549 mut input: BoxStream<Out>,
7550 parallelism: usize,
7551 stage: Arc<F>,
7552) -> BoxStream<Next>
7553where
7554 Out: Send + 'static,
7555 Next: Send + 'static,
7556 F: Fn(Out) -> Fut + Send + Sync + 'static,
7557 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7558{
7559 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7560 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7561 let mut next_task_id = 0_usize;
7562 let mut input_done = false;
7563
7564 Box::new(std::iter::from_fn(move || {
7565 loop {
7566 while tasks.len() < parallelism && !input_done {
7567 match input.next() {
7568 Some(Ok(item)) => match poll_once_or_pending(stage(item)) {
7569 Ok(result) => return Some(result),
7570 Err(future) => {
7571 let task_id = next_task_id;
7572 next_task_id += 1;
7573 tasks.insert(
7574 task_id,
7575 spawn_completion_task(task_id, future, sender.clone(), |result| {
7576 result
7577 }),
7578 );
7579 }
7580 },
7581 Some(Err(error)) => {
7582 input_done = true;
7583 return Some(Err(error));
7584 }
7585 None => input_done = true,
7586 }
7587 }
7588
7589 if tasks.is_empty() {
7590 return None;
7591 }
7592
7593 if let Some((task_id, result)) = recv_completion(&receiver) {
7594 tasks.remove(&task_id);
7595 return Some(result);
7596 }
7597 }
7598 }))
7599}
7600
7601fn map_async_unordered_supervised<Out, Next, F, Fut>(
7602 mut input: BoxStream<Out>,
7603 parallelism: usize,
7604 stage: Arc<F>,
7605 decider: SupervisionDecider,
7606) -> BoxStream<Next>
7607where
7608 Out: Send + 'static,
7609 Next: Send + 'static,
7610 F: Fn(Out) -> Fut + Send + Sync + 'static,
7611 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7612{
7613 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7614 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7615 let mut next_task_id = 0_usize;
7616 let mut input_done = false;
7617
7618 Box::new(std::iter::from_fn(move || {
7619 loop {
7620 while tasks.len() < parallelism && !input_done {
7621 match input.next() {
7622 Some(Ok(item)) => {
7623 match catch_unwind(AssertUnwindSafe(|| poll_once_or_pending(stage(item)))) {
7624 Ok(Ok(result)) => {
7625 if let Some(result) = supervise_async_result(result, &decider) {
7626 return Some(result);
7627 }
7628 }
7629 Ok(Err(future)) => {
7630 let task_id = next_task_id;
7631 next_task_id += 1;
7632 tasks.insert(
7633 task_id,
7634 spawn_completion_task(
7635 task_id,
7636 future,
7637 sender.clone(),
7638 |result| result,
7639 ),
7640 );
7641 }
7642 Err(_) => {
7643 let error = panic_stream_error("map_async_unordered callback");
7644 if let Some(result) = supervise_async_result(Err(error), &decider) {
7645 return Some(result);
7646 }
7647 }
7648 }
7649 }
7650 Some(Err(error)) => {
7651 input_done = true;
7652 return Some(Err(error));
7653 }
7654 None => input_done = true,
7655 }
7656 }
7657
7658 if tasks.is_empty() {
7659 return None;
7660 }
7661
7662 if let Some((task_id, result)) = recv_completion(&receiver) {
7663 tasks.remove(&task_id);
7664 if let Some(result) = supervise_async_result(result, &decider) {
7665 return Some(result);
7666 }
7667 }
7668 }
7669 }))
7670}
7671
7672#[inline(always)]
7673fn map_async_partitioned_serial<Out, Key, Next, Partition, F, Fut>(
7674 mut input: BoxStream<Out>,
7675 partition: Arc<Partition>,
7676 stage: Arc<F>,
7677) -> BoxStream<Next>
7678where
7679 Out: Send + 'static,
7680 Key: Send + 'static,
7681 Next: Send + 'static,
7682 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7683 F: Fn(Out) -> Fut + Send + Sync + 'static,
7684 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7685{
7686 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7687 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(1);
7688 let mut next_index = 0_usize;
7689 Box::new(std::iter::from_fn(move || {
7690 if !tasks.is_empty()
7694 && let Some((task_id, result)) = recv_completion(&receiver)
7695 {
7696 tasks.remove(&task_id);
7697 return Some(result);
7698 }
7699
7700 let item = input.next()?;
7701 match item {
7702 Ok(item) => {
7703 let _ = partition(&item);
7704 let index = next_index;
7705 next_index += 1;
7706 Some(match poll_once_or_pending(stage(item)) {
7707 Ok(result) => result,
7708 Err(future) => {
7709 tasks.insert(
7710 index,
7711 spawn_completion_task(index, future, sender.clone(), |result| result),
7712 );
7713 let (task_id, result) =
7714 recv_completion(&receiver).expect("pending map_async task completion");
7715 tasks.remove(&task_id);
7716 result
7717 }
7718 })
7719 }
7720 Err(error) => Some(Err(error)),
7721 }
7722 }))
7723}
7724
7725#[inline(always)]
7726fn map_async_partitioned_scanning<Out, Key, Next, Partition, F, Fut>(
7727 mut input: BoxStream<Out>,
7728 parallelism: usize,
7729 per_partition: usize,
7730 partition: Arc<Partition>,
7731 stage: Arc<F>,
7732) -> BoxStream<Next>
7733where
7734 Out: Send + 'static,
7735 Key: Clone + Eq + Hash + Send + 'static,
7736 Next: Send + 'static,
7737 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7738 F: Fn(Out) -> Fut + Send + Sync + 'static,
7739 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7740{
7741 let (sender, receiver) = std::sync::mpsc::channel::<(usize, (Key, StreamResult<Next>))>();
7742 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7743 let mut active_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7744 let mut pending = VecDeque::<(usize, Key, Out)>::with_capacity(parallelism);
7745 let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7746 let mut next_index = 0_usize;
7747 let mut next_to_emit = 0_usize;
7748 let mut input_done = false;
7749
7750 Box::new(std::iter::from_fn(move || {
7751 loop {
7752 if let Some(result) = completed.remove(&next_to_emit) {
7753 next_to_emit += 1;
7754 return Some(result);
7755 }
7756
7757 while tasks.len() + completed.len() < parallelism {
7758 let next = pending
7759 .iter()
7760 .position(|(_, key, _)| {
7761 active_by_key.get(key).copied().unwrap_or(0) < per_partition
7762 })
7763 .and_then(|index| pending.remove(index))
7764 .or_else(|| {
7765 if input_done {
7766 return None;
7767 }
7768 match input.next() {
7769 Some(Ok(item)) => {
7770 let key = partition(&item);
7771 let index = next_index;
7772 next_index += 1;
7773 Some((index, key, item))
7774 }
7775 Some(Err(error)) => {
7776 completed.insert(next_index, Err(error));
7777 next_index += 1;
7778 input_done = true;
7779 None
7780 }
7781 None => {
7782 input_done = true;
7783 None
7784 }
7785 }
7786 });
7787
7788 let Some((index, key, item)) = next else {
7789 break;
7790 };
7791 if active_by_key.get(&key).copied().unwrap_or(0) >= per_partition {
7792 pending.push_back((index, key, item));
7793 if input_done || pending.len() >= parallelism {
7794 break;
7795 }
7796 continue;
7797 }
7798 *active_by_key.entry(key.clone()).or_default() += 1;
7799 match poll_once_or_pending(stage(item)) {
7800 Ok(result) => {
7801 if let Some(count) = active_by_key.get_mut(&key) {
7802 *count -= 1;
7803 if *count == 0 {
7804 active_by_key.remove(&key);
7805 }
7806 }
7807 completed.insert(index, result);
7808 }
7809 Err(future) => {
7810 tasks.insert(
7811 index,
7812 spawn_completion_task(index, future, sender.clone(), move |result| {
7813 (key, result)
7814 }),
7815 );
7816 }
7817 }
7818 }
7819
7820 if let Some(result) = completed.remove(&next_to_emit) {
7821 next_to_emit += 1;
7822 return Some(result);
7823 }
7824
7825 if tasks.is_empty() {
7826 return None;
7827 }
7828
7829 if let Some((index, (key, result))) = recv_completion(&receiver) {
7830 tasks.remove(&index);
7831 if let Some(count) = active_by_key.get_mut(&key) {
7832 *count -= 1;
7833 if *count == 0 {
7834 active_by_key.remove(&key);
7835 }
7836 }
7837 if index == next_to_emit {
7838 next_to_emit += 1;
7839 return Some(result);
7840 }
7841 completed.insert(index, result);
7842 }
7843 }
7844 }))
7845}
7846
7847fn map_async_partitioned<Out, Key, Next, Partition, F, Fut>(
7848 mut input: BoxStream<Out>,
7849 parallelism: usize,
7850 per_partition: usize,
7851 partition: Arc<Partition>,
7852 stage: Arc<F>,
7853) -> BoxStream<Next>
7854where
7855 Out: Send + 'static,
7856 Key: Clone + Eq + Hash + Send + 'static,
7857 Next: Send + 'static,
7858 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7859 F: Fn(Out) -> Fut + Send + Sync + 'static,
7860 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7861{
7862 if parallelism == 1 {
7863 return map_async_partitioned_serial(input, partition, stage);
7864 }
7865 if parallelism <= 4 {
7867 return map_async_partitioned_scanning(input, parallelism, per_partition, partition, stage);
7868 }
7869
7870 let (sender, receiver) = std::sync::mpsc::channel::<(usize, (usize, StreamResult<Next>))>();
7871 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7872 let mut slots_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7873 let mut slots = Vec::<PartitionSlot<Key, Out>>::with_capacity(parallelism);
7874 let mut free_slots = Vec::<usize>::new();
7875 let mut ready_slots = VecDeque::<usize>::with_capacity(parallelism);
7876 let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7877 let mut next_index = 0_usize;
7878 let mut next_to_emit = 0_usize;
7879 let mut input_done = false;
7880 Box::new(std::iter::from_fn(move || {
7881 loop {
7882 if let Some(result) = completed.remove(&next_to_emit) {
7883 next_to_emit += 1;
7884 return Some(result);
7885 }
7886
7887 while tasks.len() + completed.len() < parallelism {
7888 if let Some((index, slot, item)) =
7889 pop_ready_partition_slot(&mut slots, &mut ready_slots, per_partition)
7890 {
7891 match poll_once_or_pending(stage(item)) {
7892 Ok(result) => {
7893 let mut remove_empty = false;
7894 if let Some(state) = slots.get_mut(slot) {
7895 state.active -= 1;
7896 remove_empty = state.active == 0
7897 && state.queued.is_empty()
7898 && !state.in_ready_queue;
7899 }
7900 if remove_empty {
7901 retire_partition_slot(
7902 slot,
7903 &mut slots_by_key,
7904 &mut slots,
7905 &mut free_slots,
7906 );
7907 } else {
7908 ready_partition_slot(
7909 &mut slots,
7910 &mut ready_slots,
7911 slot,
7912 per_partition,
7913 );
7914 }
7915 if index == next_to_emit {
7916 next_to_emit += 1;
7917 return Some(result);
7918 }
7919 completed.insert(index, result);
7920 }
7921 Err(future) => {
7922 tasks.insert(
7923 index,
7924 spawn_completion_task(
7925 index,
7926 future,
7927 sender.clone(),
7928 move |result| (slot, result),
7929 ),
7930 );
7931 }
7932 }
7933 continue;
7934 }
7935
7936 if input_done {
7937 break;
7938 }
7939
7940 match input.next() {
7941 Some(Ok(item)) => {
7942 let key = partition(&item);
7943 let index = next_index;
7944 next_index += 1;
7945 let slot =
7946 partition_slot_for(key, &mut slots_by_key, &mut slots, &mut free_slots);
7947 let state = &mut slots[slot];
7948 if state.active < per_partition {
7949 match poll_once_or_pending(stage(item)) {
7950 Ok(result) => {
7951 if index == next_to_emit {
7952 next_to_emit += 1;
7953 if state.queued.is_empty() && !state.in_ready_queue {
7954 retire_partition_slot(
7955 slot,
7956 &mut slots_by_key,
7957 &mut slots,
7958 &mut free_slots,
7959 );
7960 }
7961 return Some(result);
7962 }
7963 completed.insert(index, result);
7964 if state.queued.is_empty() && !state.in_ready_queue {
7965 retire_partition_slot(
7966 slot,
7967 &mut slots_by_key,
7968 &mut slots,
7969 &mut free_slots,
7970 );
7971 }
7972 }
7973 Err(future) => {
7974 state.active += 1;
7975 tasks.insert(
7976 index,
7977 spawn_completion_task(
7978 index,
7979 future,
7980 sender.clone(),
7981 move |result| (slot, result),
7982 ),
7983 );
7984 }
7985 }
7986 } else {
7987 state.queued.push_back((index, item));
7988 }
7989 }
7990 Some(Err(error)) => {
7991 completed.insert(next_index, Err(error));
7992 next_index += 1;
7993 input_done = true;
7994 break;
7995 }
7996 None => {
7997 input_done = true;
7998 break;
7999 }
8000 }
8001 }
8002
8003 if let Some(result) = completed.remove(&next_to_emit) {
8004 next_to_emit += 1;
8005 return Some(result);
8006 }
8007
8008 if tasks.is_empty() {
8009 return None;
8010 }
8011
8012 if let Some((index, (slot, result))) = recv_completion(&receiver) {
8013 tasks.remove(&index);
8014 let mut remove_empty = false;
8015 if let Some(state) = slots.get_mut(slot) {
8016 state.active -= 1;
8017 remove_empty =
8018 state.active == 0 && state.queued.is_empty() && !state.in_ready_queue;
8019 }
8020 if remove_empty {
8021 retire_partition_slot(slot, &mut slots_by_key, &mut slots, &mut free_slots);
8022 } else {
8023 ready_partition_slot(&mut slots, &mut ready_slots, slot, per_partition);
8024 }
8025 if index == next_to_emit {
8026 next_to_emit += 1;
8027 return Some(result);
8028 }
8029 completed.insert(index, result);
8030 }
8031 }
8032 }))
8033}
8034
8035#[cfg(test)]
8036mod flat_map_merge_ready_ring_tests {
8037 use super::*;
8038 use std::sync::mpsc;
8039 use std::time::Duration;
8040
8041 fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
8042 let mut v = source.run_collect().unwrap();
8043 v.sort_unstable();
8044 v
8045 }
8046
8047 #[test]
8048 fn ready_ring_empty_upstream() {
8049 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8050 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8051 });
8052 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8053 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8054 });
8055 assert_eq!(legacy, ring);
8056 assert!(ring.is_empty());
8057 }
8058
8059 #[test]
8060 fn ready_ring_single_lane() {
8061 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8062 run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
8063 });
8064 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8065 run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
8066 });
8067 assert_eq!(legacy, ring);
8068 assert_eq!(ring, vec![42]);
8069 }
8070
8071 #[test]
8072 fn ready_ring_breadth_one_exact_order() {
8073 let make = || {
8075 crate::Source::from_iter(0_i32..5)
8076 .flat_map_merge(1, |x| crate::Source::single(x * 10))
8077 .run_collect()
8078 .unwrap()
8079 };
8080 let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8081 let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8082 legacy.sort_unstable();
8083 ring.sort_unstable();
8084 assert_eq!(legacy, ring);
8085 }
8086
8087 #[test]
8088 fn ready_ring_breadth_gt_input() {
8089 let make = || {
8090 run_sorted(
8091 crate::Source::from_iter(0_i32..3)
8092 .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2])),
8093 )
8094 };
8095 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8096 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8097 assert_eq!(legacy, ring);
8098 assert_eq!(ring.len(), 9);
8099 }
8100
8101 #[test]
8102 fn ready_ring_mixed_short_long() {
8103 let make = || {
8104 run_sorted(crate::Source::from_iter(0_i32..8).flat_map_merge(4, |x| {
8105 if x % 3 == 0 {
8106 crate::Source::from_iter(0..20_i32).map(move |i| x * 100 + i)
8107 } else {
8108 crate::Source::single(x)
8109 }
8110 }))
8111 };
8112 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8113 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8114 assert_eq!(legacy, ring);
8115 }
8116
8117 #[test]
8118 fn ready_ring_respects_breadth_bound() {
8119 use std::sync::atomic::{AtomicUsize, Ordering as Ord};
8120 let active = Arc::new(AtomicUsize::new(0));
8121 let max_active = Arc::new(AtomicUsize::new(0));
8122 let a2 = Arc::clone(&active);
8123 let m2 = Arc::clone(&max_active);
8124 let mut values = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8125 crate::Source::from_iter(0_i32..6)
8126 .flat_map_merge(2, move |item| {
8127 let a = Arc::clone(&a2);
8128 let m = Arc::clone(&m2);
8129 crate::Source::future(move || {
8130 let a = Arc::clone(&a);
8131 let m = Arc::clone(&m);
8132 async move {
8133 let now = a.fetch_add(1, Ord::SeqCst) + 1;
8134 let mut seen = m.load(Ord::SeqCst);
8135 while now > seen {
8136 match m.compare_exchange(seen, now, Ord::SeqCst, Ord::SeqCst) {
8137 Ok(_) => break,
8138 Err(v) => seen = v,
8139 }
8140 }
8141 thread::sleep(Duration::from_millis(20));
8142 a.fetch_sub(1, Ord::SeqCst);
8143 Ok(item)
8144 }
8145 })
8146 })
8147 .run_collect()
8148 .unwrap()
8149 });
8150 values.sort_unstable();
8151 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
8152 assert!(max_active.load(std::sync::atomic::Ordering::SeqCst) <= 2);
8153 }
8154
8155 #[test]
8156 fn ready_ring_fairness_slow_lane_not_starved() {
8157 use std::sync::atomic::{AtomicBool, Ordering as Ord};
8158 let slow_emitted = Arc::new(AtomicBool::new(false));
8159 let slow_flag = Arc::clone(&slow_emitted);
8160 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8161 crate::Source::from_iter(0_i32..3)
8162 .flat_map_merge(4, move |lane_id| {
8163 let flag = Arc::clone(&slow_flag);
8164 match lane_id {
8165 0 => crate::Source::from_iter(0_i32..50),
8166 1 => crate::Source::from_iter(100_i32..150),
8167 _ => crate::Source::future(move || {
8168 let flag = Arc::clone(&flag);
8169 async move {
8170 thread::sleep(Duration::from_millis(10));
8171 flag.store(true, Ord::SeqCst);
8172 Ok(999_i32)
8173 }
8174 }),
8175 }
8176 })
8177 .run_collect()
8178 .unwrap()
8179 });
8180 assert!(slow_emitted.load(std::sync::atomic::Ordering::SeqCst));
8181 assert!(results.contains(&999));
8182 assert_eq!(results.len(), 101);
8183 }
8184
8185 #[test]
8186 fn ready_ring_inner_failure_no_hang() {
8187 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8188 crate::Source::from_iter(0_i32..8)
8189 .flat_map_merge(4, |x| {
8190 if x == 3 {
8191 crate::Source::failed(StreamError::Failed("lane-fail".into()))
8192 } else {
8193 crate::Source::from_iter(0_i32..10)
8194 }
8195 })
8196 .run_collect()
8197 });
8198 assert_eq!(result, Err(StreamError::Failed("lane-fail".into())));
8199 }
8200
8201 #[test]
8202 fn ready_ring_factory_failure_propagates() {
8203 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8204 crate::Source::from_iter(0_i32..4)
8205 .flat_map_merge(2, |x| {
8206 if x == 2 {
8207 crate::Source::failed(StreamError::Failed("factory-fail".into()))
8208 } else {
8209 crate::Source::single(x)
8210 }
8211 })
8212 .run_collect()
8213 });
8214 assert!(result.is_err());
8215 }
8216
8217 #[test]
8218 fn ready_ring_closure_not_under_coordinator_lock() {
8219 let guard_mutex = Arc::new(std::sync::Mutex::<()>::new(()));
8220 let gm = Arc::clone(&guard_mutex);
8221 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8222 crate::Source::from_iter(0_i32..10)
8223 .flat_map_merge(4, move |x| {
8224 let _lock = gm.lock().unwrap();
8225 crate::Source::single(x)
8226 })
8227 .run_collect()
8228 .unwrap()
8229 });
8230 assert_eq!(results.len(), 10);
8231 }
8232
8233 #[test]
8240 fn ready_ring_bounded_memory_producer_blocks_at_window() {
8241 const WINDOW: usize = FLAT_MAP_MERGE_SUBSTREAM_WINDOW;
8242 const EXTRA: usize = 4;
8243 let (gate_tx, gate_rx) = mpsc::channel::<()>();
8246 let gate_tx = Arc::new(std::sync::Mutex::new(gate_tx));
8247 let gate_tx2 = Arc::clone(&gate_tx);
8248 let produced = Arc::new(std::sync::atomic::AtomicUsize::new(0));
8249 let prod2 = Arc::clone(&produced);
8250
8251 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8252 crate::Source::single(0_i32)
8253 .flat_map_merge(2, move |_| {
8254 let tx = Arc::clone(&gate_tx2);
8255 let prod = Arc::clone(&prod2);
8256 crate::Source::from_factory(move || {
8257 let tx = Arc::clone(&tx);
8258 let prod = Arc::clone(&prod);
8259 let mut i = 0_i32;
8260 Box::new(std::iter::from_fn(move || {
8261 if i as usize >= WINDOW + EXTRA {
8262 return None;
8263 }
8264 if i as usize == WINDOW {
8265 let _ = tx.lock().unwrap().send(());
8267 }
8268 prod.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
8269 i += 1;
8270 Some(Ok(i))
8271 }))
8272 })
8273 })
8274 .run_with(crate::Sink::queue())
8275 .unwrap()
8276 });
8277
8278 let mut total = 0;
8280 while queue.pull().unwrap().is_some() {
8281 total += 1;
8282 }
8283 let signal = gate_rx.recv_timeout(Duration::from_secs(1));
8285 assert!(signal.is_ok(), "producer never reached the window boundary");
8286 assert_eq!(total, WINDOW + EXTRA);
8287 }
8288
8289 #[test]
8290 fn ready_ring_cancellation_wakes_blocked_lanes() {
8291 let rt = crate::stream::runtime::Runtime::new();
8292 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8293 crate::Source::from_iter(0_i32..4)
8294 .flat_map_merge(4, |_| crate::Source::repeat(1_i32))
8295 .run_with_materializer(crate::Sink::queue(), &rt)
8296 .unwrap()
8297 });
8298 for _ in 0..8 {
8299 let _ = queue.pull();
8300 }
8301 drop(queue);
8302 rt.shutdown();
8303 }
8304
8305 #[test]
8306 fn ready_ring_lost_wakeup_stress() {
8307 for _ in 0..20 {
8308 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8309 crate::Source::from_iter(0_i32..50)
8310 .flat_map_merge(8, |item| {
8311 crate::Source::from_iter([item, item + 1, item + 2])
8312 })
8313 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8314 .unwrap()
8315 .wait()
8316 });
8317 assert_eq!(result, Ok(3825), "lost-wakeup stress: wrong sum");
8318 }
8319 }
8320
8321 #[test]
8322 fn ready_ring_concurrent_streams_lost_wakeup_stress() {
8323 const STREAMS: usize = 32;
8324 const ROUNDS: usize = 8;
8325 const EXPECTED: i64 = 998_080;
8326
8327 for _ in 0..ROUNDS {
8328 let barrier = Arc::new(std::sync::Barrier::new(STREAMS));
8329 let mut handles = Vec::with_capacity(STREAMS);
8330 for _ in 0..STREAMS {
8331 let barrier = Arc::clone(&barrier);
8332 handles.push(thread::spawn(move || {
8333 barrier.wait();
8334 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8335 crate::Source::from_iter(0_i64..32)
8336 .flat_map_merge(8, |item| {
8337 crate::Source::from_iter(item * 100..item * 100 + 20)
8338 })
8339 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8340 .unwrap()
8341 .wait()
8342 });
8343 assert_eq!(result, Ok(EXPECTED), "concurrent ready-ring sum");
8344 }));
8345 }
8346
8347 for handle in handles {
8348 handle.join().expect("ready-ring stress worker panicked");
8349 }
8350 }
8351 }
8352
8353 #[test]
8354 fn ready_ring_tail_loop_stress() {
8355 for _ in 0..20 {
8356 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8357 crate::Source::from_iter(0_i64..100)
8358 .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8359 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8360 .unwrap()
8361 .wait()
8362 });
8363 assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8364 }
8365 }
8366
8367 #[test]
8368 fn ready_ring_auto_mode_matches_ring() {
8369 let make = || {
8370 run_sorted(
8371 crate::Source::from_iter(0_i32..10)
8372 .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 100])),
8373 )
8374 };
8375 let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8376 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8377 assert_eq!(auto, ring);
8378 }
8379}
8380
8381#[cfg(test)]
8383mod inline_micro_source_tests {
8384 use super::*;
8385 use crate::stream::source::test_source_with_inline_micro_hint;
8386 use std::sync::mpsc;
8387
8388 fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
8389 let mut v = source.run_collect().unwrap();
8390 v.sort_unstable();
8391 v
8392 }
8393
8394 #[test]
8397 fn inline_empty_upstream() {
8398 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8399 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8400 });
8401 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8402 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8403 });
8404 assert_eq!(legacy, ring);
8405 assert!(ring.is_empty());
8406 }
8407
8408 #[test]
8409 fn inline_single_inner_source() {
8410 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8411 run_sorted(
8412 crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8413 )
8414 });
8415 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8416 run_sorted(
8417 crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8418 )
8419 });
8420 assert_eq!(legacy, ring);
8421 assert_eq!(ring, vec![198]);
8422 }
8423
8424 #[test]
8425 fn inline_breadth_one_exact_order() {
8426 let make = || {
8428 crate::Source::from_iter(0_i32..6)
8429 .flat_map_merge(1, |x| {
8430 crate::Source::from_iter([x * 10, x * 10 + 1, x * 10 + 2, x * 10 + 3])
8431 })
8432 .run_collect()
8433 .unwrap()
8434 };
8435 let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8436 let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8437 legacy.sort_unstable();
8438 ring.sort_unstable();
8439 assert_eq!(legacy, ring);
8440 }
8441
8442 #[test]
8443 fn inline_breadth_gt_input() {
8444 let make = || {
8445 run_sorted(
8446 crate::Source::from_iter(0_i32..3)
8447 .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8448 )
8449 };
8450 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8451 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8452 assert_eq!(legacy, ring);
8453 assert_eq!(ring.len(), 12);
8454 }
8455
8456 #[test]
8457 fn inline_mixed_empty_single_four_item() {
8458 let make = || {
8459 run_sorted(
8460 crate::Source::from_iter(0_i32..12).flat_map_merge(4, |x| match x % 3 {
8461 0 => crate::Source::empty(),
8462 1 => crate::Source::single(x),
8463 _ => crate::Source::from_iter([x, x + 100, x + 200, x + 300]),
8464 }),
8465 )
8466 };
8467 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8468 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8469 assert_eq!(legacy, ring);
8470 }
8471
8472 #[test]
8473 fn inline_2k_x4_b8_benchmark_shape() {
8474 let make = || {
8475 run_sorted(
8476 crate::Source::from_iter(0_i32..2_000).flat_map_merge(8, |item| {
8477 crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8478 }),
8479 )
8480 };
8481 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8482 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8483 assert_eq!(legacy, ring);
8484 assert_eq!(ring.len(), 8_000);
8485 }
8486
8487 #[test]
8490 fn inline_lost_wakeup_stress() {
8491 for _ in 0..20 {
8492 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8493 crate::Source::from_iter(0_i32..50)
8494 .flat_map_merge(8, |item| {
8495 crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8496 })
8497 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8498 .unwrap()
8499 .wait()
8500 });
8501 assert_eq!(result, Ok(5200), "lost-wakeup stress: wrong sum");
8505 }
8506 }
8507
8508 #[test]
8509 fn inline_tail_loop_stress() {
8510 for _ in 0..20 {
8511 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8512 crate::Source::from_iter(0_i64..100)
8513 .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8514 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8515 .unwrap()
8516 .wait()
8517 });
8518 assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8520 }
8521 }
8522
8523 #[test]
8526 fn inline_large_source_uses_worker_fallback() {
8527 let make = || {
8530 run_sorted(crate::Source::from_iter(0_i32..4).flat_map_merge(2, |x| {
8531 crate::Source::from_iter(0_i32..20).map(move |i| x * 100 + i)
8532 }))
8533 };
8534 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8535 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8536 assert_eq!(legacy, ring);
8537 assert_eq!(ring.len(), 80);
8538 }
8539
8540 #[test]
8543 fn inline_inner_error_before_any_item() {
8544 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8545 crate::Source::from_iter(0_i32..4)
8546 .flat_map_merge(2, |x| {
8547 if x == 1 {
8548 crate::Source::failed(StreamError::Failed("inline-err".into()))
8549 } else {
8550 crate::Source::from_iter([x, x + 1])
8551 }
8552 })
8553 .run_collect()
8554 });
8555 assert_eq!(result, Err(StreamError::Failed("inline-err".into())));
8556 }
8557
8558 #[test]
8559 fn inline_inner_error_after_some_items() {
8560 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8564 crate::Source::single(0_i32)
8565 .flat_map_merge(2, |_x| {
8566 test_source_with_inline_micro_hint(
8567 || {
8568 let mut count = 0;
8569 Box::new(std::iter::from_fn(move || {
8570 count += 1;
8571 match count {
8572 1 => Some(Ok(42_i32)),
8573 2 => Some(Err(StreamError::Failed("after-items".into()))),
8574 _ => None,
8575 }
8576 }))
8577 },
8578 1, )
8580 })
8581 .run_collect()
8582 });
8583 assert!(result.is_err());
8587 }
8588
8589 #[test]
8590 fn inline_worker_lane_fails_during_inline_drain() {
8591 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8594 crate::Source::from_iter(0_i32..10)
8595 .flat_map_merge(2, |x| {
8596 if x == 5 {
8597 crate::Source::failed(StreamError::Failed("worker-fail".into()))
8598 } else if x % 2 == 0 {
8599 crate::Source::from_iter([x, x + 1, x + 2])
8601 } else {
8602 crate::Source::from_iter(0_i32..40).map(move |i| x * 100 + i)
8604 }
8605 })
8606 .run_collect()
8607 });
8608 assert!(result.is_err());
8609 }
8610
8611 #[test]
8614 fn inline_cancellation_drop_output() {
8615 let rt = crate::stream::runtime::Runtime::new();
8618 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8619 crate::Source::from_iter(0_i32..100)
8620 .flat_map_merge(8, |x| {
8621 if x % 3 == 0 {
8622 crate::Source::from_iter([x, x + 1, x + 2, x + 3]) } else {
8624 crate::Source::repeat(x) }
8626 })
8627 .run_with_materializer(crate::Sink::queue(), &rt)
8628 .unwrap()
8629 });
8630 for _ in 0..16 {
8632 let _ = queue.pull();
8633 }
8634 drop(queue);
8635 rt.shutdown();
8636 }
8637
8638 #[test]
8641 fn inline_next_not_under_coordinator_lock() {
8642 let (tx, rx) = mpsc::channel::<()>();
8650 let tx = Arc::new(std::sync::Mutex::new(tx));
8651 let tx2 = Arc::clone(&tx);
8652
8653 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8654 crate::Source::from_iter(0_i32..2)
8655 .flat_map_merge(2, move |x| {
8656 let sig = Arc::clone(&tx2);
8657 if x == 0 {
8658 crate::Source::from_iter(0_i32..40)
8660 } else {
8661 test_source_with_inline_micro_hint(
8664 move || {
8665 let sig = Arc::clone(&sig);
8666 let mut emitted = false;
8667 Box::new(std::iter::from_fn(move || {
8668 if !emitted {
8669 emitted = true;
8670 let _ = sig.lock().unwrap().send(());
8671 Some(Ok(999_i32))
8672 } else {
8673 None
8674 }
8675 }))
8676 },
8677 1,
8678 )
8679 }
8680 })
8681 .run_collect()
8682 .unwrap()
8683 });
8684
8685 let signal = rx.recv_timeout(std::time::Duration::from_secs(5));
8687 assert!(signal.is_ok(), "inline next() never ran");
8688 assert!(results.contains(&999));
8689 assert_eq!(results.len(), 41);
8690 }
8691
8692 #[test]
8695 fn inline_auto_matches_readyring() {
8696 let make = || {
8697 run_sorted(
8698 crate::Source::from_iter(0_i32..20)
8699 .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8700 )
8701 };
8702 let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8703 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8704 assert_eq!(auto, ring);
8705 }
8706}
8707
8708#[cfg(test)]
8710mod split_sink_fast_path_tests {
8711 use super::*;
8712
8713 fn run_split_fold(
8715 items: Vec<u64>,
8716 split_mode: SplitMode,
8717 executor_mode: SubstreamExecutorMode,
8718 ) -> Vec<u64> {
8719 with_substream_mode(executor_mode, || {
8720 let source = match split_mode {
8721 SplitMode::When => crate::Source::from_iter(items).split_when(|x| x % 10 == 0),
8722 SplitMode::After => crate::Source::from_iter(items).split_after(|x| x % 10 == 0),
8723 };
8724 source
8725 .run_with(crate::Sink::fold(
8726 Vec::new(),
8727 |mut acc, seg: crate::Source<u64>| {
8728 let sum = seg
8729 .run_with(crate::Sink::fold(0u64, |a, x| a + x))
8730 .unwrap()
8731 .wait()
8732 .unwrap();
8733 acc.push(sum);
8734 acc
8735 },
8736 ))
8737 .unwrap()
8738 .wait()
8739 .unwrap()
8740 })
8741 }
8742
8743 fn run_split_collect_segments(
8745 items: Vec<u64>,
8746 split_mode: SplitMode,
8747 executor_mode: SubstreamExecutorMode,
8748 ) -> Vec<Vec<u64>> {
8749 with_substream_mode(executor_mode, || {
8750 let source = match split_mode {
8751 SplitMode::When => crate::Source::from_iter(items).split_when(move |x| x % 10 == 0),
8752 SplitMode::After => {
8753 crate::Source::from_iter(items).split_after(move |x| x % 10 == 0)
8754 }
8755 };
8756 source
8757 .run_with(crate::Sink::fold(
8758 Vec::new(),
8759 |mut acc, seg: crate::Source<u64>| {
8760 let v = seg
8761 .run_with(crate::Sink::collect())
8762 .unwrap()
8763 .wait()
8764 .unwrap();
8765 acc.push(v);
8766 acc
8767 },
8768 ))
8769 .unwrap()
8770 .wait()
8771 .unwrap()
8772 })
8773 }
8774
8775 #[test]
8778 fn split_fast_equivalence_empty_input() {
8779 for sm in [SplitMode::When, SplitMode::After] {
8780 let legacy = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::LegacyOnly);
8781 let fast = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::SplitSinkOnly);
8782 assert_eq!(legacy, fast, "empty input, mode {sm:?}");
8783 }
8784 }
8785
8786 #[test]
8787 fn split_fast_equivalence_no_boundaries() {
8788 let items: Vec<u64> = (1..=9).collect();
8789 for sm in [SplitMode::When, SplitMode::After] {
8790 let legacy =
8791 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8792 let fast =
8793 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8794 assert_eq!(legacy, fast, "no boundaries, mode {sm:?}");
8795 }
8796 }
8797
8798 #[test]
8799 fn split_fast_equivalence_first_element_boundary_when() {
8800 let items: Vec<u64> = vec![10, 1, 2, 3];
8802 let legacy = run_split_collect_segments(
8803 items.clone(),
8804 SplitMode::When,
8805 SubstreamExecutorMode::LegacyOnly,
8806 );
8807 let fast = run_split_collect_segments(
8808 items,
8809 SplitMode::When,
8810 SubstreamExecutorMode::SplitSinkOnly,
8811 );
8812 assert_eq!(legacy, fast);
8813 }
8814
8815 #[test]
8816 fn split_fast_equivalence_first_element_boundary_after() {
8817 let items: Vec<u64> = vec![10, 1, 2, 3];
8819 let legacy = run_split_collect_segments(
8820 items.clone(),
8821 SplitMode::After,
8822 SubstreamExecutorMode::LegacyOnly,
8823 );
8824 let fast = run_split_collect_segments(
8825 items,
8826 SplitMode::After,
8827 SubstreamExecutorMode::SplitSinkOnly,
8828 );
8829 assert_eq!(legacy, fast);
8830 }
8831
8832 #[test]
8833 fn split_fast_equivalence_consecutive_matches() {
8834 let items: Vec<u64> = vec![10, 20, 30, 1, 2, 40];
8835 for sm in [SplitMode::When, SplitMode::After] {
8836 let legacy =
8837 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8838 let fast =
8839 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8840 assert_eq!(legacy, fast, "consecutive matches, mode {sm:?}");
8841 }
8842 }
8843
8844 #[test]
8845 fn split_fast_equivalence_last_element_boundary() {
8846 let items: Vec<u64> = vec![1, 2, 3, 10];
8847 for sm in [SplitMode::When, SplitMode::After] {
8848 let legacy =
8849 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8850 let fast =
8851 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8852 assert_eq!(legacy, fast, "last element boundary, mode {sm:?}");
8853 }
8854 }
8855
8856 #[test]
8857 fn split_fast_equivalence_mixed() {
8858 let items: Vec<u64> = (0..50u64).collect();
8859 for sm in [SplitMode::When, SplitMode::After] {
8860 let legacy =
8861 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8862 let fast =
8863 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8864 assert_eq!(legacy, fast, "mixed 0..50, mode {sm:?}");
8865 }
8866 }
8867
8868 #[test]
8869 fn split_fast_equivalence_fold_sums() {
8870 let items: Vec<u64> = (0..50u64).collect();
8871 for sm in [SplitMode::When, SplitMode::After] {
8872 let legacy = run_split_fold(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8873 let fast = run_split_fold(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8874 assert_eq!(legacy, fast, "fold sums, mode {sm:?}");
8875 }
8876 }
8877
8878 #[test]
8879 fn split_fast_equivalence_with_collect() {
8880 let items: Vec<u64> = (0..312u64).collect();
8882 for sm in [SplitMode::When, SplitMode::After] {
8883 let legacy =
8884 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8885 let fast =
8886 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8887 assert_eq!(legacy, fast, "collect 312 items, mode {sm:?}");
8888 }
8889 }
8890
8891 #[test]
8894 fn split_fast_fold_result_equivalence() {
8895 let items: Vec<u64> = (0..50u64).collect();
8896 let run = |executor_mode| {
8897 with_substream_mode(executor_mode, || {
8898 crate::Source::from_iter(items.clone())
8899 .split_when(move |x| x % 10 == 0)
8900 .run_with(crate::Sink::fold(
8901 Vec::new(),
8902 |mut acc, seg: crate::Source<u64>| {
8903 let sum = seg
8904 .run_with(crate::Sink::fold_result(0u64, |a, x| Ok(a + x)))
8905 .unwrap()
8906 .wait()
8907 .unwrap();
8908 acc.push(sum);
8909 acc
8910 },
8911 ))
8912 .unwrap()
8913 .wait()
8914 .unwrap()
8915 })
8916 };
8917 assert_eq!(
8918 run(SubstreamExecutorMode::LegacyOnly),
8919 run(SubstreamExecutorMode::SplitSinkOnly)
8920 );
8921 }
8922
8923 #[test]
8926 fn split_fast_ignore_equivalence() {
8927 let items: Vec<u64> = (0..50u64).collect();
8928 let run = |executor_mode| {
8929 with_substream_mode(executor_mode, || {
8930 crate::Source::from_iter(items.clone())
8931 .split_when(move |x| x % 10 == 0)
8932 .run_with(crate::Sink::fold(0u64, |count, seg: crate::Source<u64>| {
8933 seg.run_with(crate::Sink::ignore()).unwrap().wait().unwrap();
8934 count + 1
8935 }))
8936 .unwrap()
8937 .wait()
8938 .unwrap()
8939 })
8940 };
8941 let legacy = run(SubstreamExecutorMode::LegacyOnly);
8942 let fast = run(SubstreamExecutorMode::SplitSinkOnly);
8943 assert_eq!(legacy, fast, "ignore segment counts must match");
8944 }
8945
8946 #[test]
8949 fn split_fast_one_shot_cannot_materialize_twice() {
8950 with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
8951 let materializer = crate::Runtime::default();
8952 let result = crate::Source::from_iter(1u64..=5)
8953 .split_when(|x| x % 3 == 0)
8954 .run_with(crate::Sink::fold(0u64, |_, seg: crate::Source<u64>| {
8955 let c1 = seg.clone().run_with(crate::Sink::fold(0u64, |a, x| a + x));
8957 let c2 = seg.run_with(crate::Sink::fold(0u64, |a, x| a + x));
8959 assert!(c1.is_ok(), "first materialization should succeed");
8960 assert!(c2.is_err(), "second materialization should fail: {c2:?}");
8961 let _ = c1.unwrap().wait();
8962 0u64
8963 }));
8964 let _ = result;
8965 let _ = &materializer;
8966 });
8967 }
8968
8969 #[test]
8972 fn split_fast_predicate_panic_both_modes() {
8973 for sm in [SplitMode::When, SplitMode::After] {
8976 let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
8977 let source = match sm {
8978 SplitMode::When => crate::Source::from_iter(0u64..10).split_when(|x| {
8979 if *x == 5 {
8980 panic!("test panic");
8981 }
8982 x % 3 == 0
8983 }),
8984 SplitMode::After => crate::Source::from_iter(0u64..10).split_after(|x| {
8985 if *x == 5 {
8986 panic!("test panic");
8987 }
8988 x % 3 == 0
8989 }),
8990 };
8991 source
8992 .run_with(crate::Sink::fold(
8993 Vec::<u64>::new(),
8994 |mut acc, seg: crate::Source<u64>| {
8995 let completion = seg.run_with(crate::Sink::ignore());
8997 if let Ok(c) = completion {
8998 let _ = c.wait();
8999 }
9000 acc.push(0u64);
9001 acc
9002 },
9003 ))
9004 .map(|c| c.wait())
9005 });
9006 let _ = result;
9008 }
9009 }
9010
9011 #[test]
9014 fn split_fast_stress_20x() {
9015 for i in 0..20 {
9016 let items: Vec<u64> = (0..10_000u64).collect();
9017 for sm in [SplitMode::When, SplitMode::After] {
9018 let fast = run_split_collect_segments(
9019 items.clone(),
9020 sm,
9021 SubstreamExecutorMode::SplitSinkOnly,
9022 );
9023 let legacy = run_split_collect_segments(
9024 items.clone(),
9025 sm,
9026 SubstreamExecutorMode::LegacyOnly,
9027 );
9028 assert_eq!(
9029 fast.len(),
9030 legacy.len(),
9031 "stress run {i} segment count mismatch, mode {sm:?}"
9032 );
9033 assert_eq!(
9034 fast.iter().flatten().sum::<u64>(),
9035 legacy.iter().flatten().sum::<u64>(),
9036 "stress run {i} sum mismatch, mode {sm:?}"
9037 );
9038 }
9039 }
9040 }
9041
9042 #[test]
9045 fn split_fast_auto_mode_matches_fast() {
9046 let items: Vec<u64> = (0..50u64).collect();
9047 for sm in [SplitMode::When, SplitMode::After] {
9048 let auto = run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::Auto);
9049 let fast =
9050 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
9051 assert_eq!(auto, fast, "auto == fast, mode {sm:?}");
9052 }
9053 }
9054
9055 #[test]
9058 fn split_fast_fallback_path_via_foreach() {
9059 use std::sync::atomic::{AtomicU64, Ordering as Ord};
9063 let total = Arc::new(AtomicU64::new(0));
9064 let total2 = Arc::clone(&total);
9065 let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9066 crate::Source::from_iter(0u64..30)
9067 .split_when(|x| x % 10 == 0)
9068 .run_with(crate::Sink::fold(
9069 Vec::new(),
9070 move |mut acc, seg: crate::Source<u64>| {
9071 let t = Arc::clone(&total2);
9072 seg.run_with(crate::Sink::foreach(move |x| {
9076 t.fetch_add(x, Ord::SeqCst);
9077 }))
9078 .unwrap()
9079 .wait()
9080 .unwrap();
9081 acc.push(1u64);
9082 acc
9083 },
9084 ))
9085 .unwrap()
9086 .wait()
9087 .unwrap()
9088 });
9089 assert_eq!(result.len(), 3, "should have 3 segments");
9090 assert_eq!(
9092 total.load(std::sync::atomic::Ordering::SeqCst),
9093 (0..30u64).sum::<u64>()
9094 );
9095 }
9096
9097 #[test]
9100 fn split_fast_liveness_segment_count_when() {
9101 let items: Vec<u64> = (0..30u64).collect();
9102 let fast = run_split_collect_segments(
9103 items.clone(),
9104 SplitMode::When,
9105 SubstreamExecutorMode::SplitSinkOnly,
9106 );
9107 let legacy =
9108 run_split_collect_segments(items, SplitMode::When, SubstreamExecutorMode::LegacyOnly);
9109 assert_eq!(fast.len(), legacy.len());
9110 }
9111
9112 #[test]
9113 fn split_fast_liveness_segment_count_after() {
9114 let items: Vec<u64> = (0..30u64).collect();
9115 let fast = run_split_collect_segments(
9116 items.clone(),
9117 SplitMode::After,
9118 SubstreamExecutorMode::SplitSinkOnly,
9119 );
9120 let legacy =
9121 run_split_collect_segments(items, SplitMode::After, SubstreamExecutorMode::LegacyOnly);
9122 assert_eq!(fast.len(), legacy.len());
9123 }
9124
9125 #[test]
9132 fn split_fast_bounded_memory_rendezvous() {
9133 use std::sync::{
9134 atomic::{AtomicBool, AtomicUsize, Ordering},
9135 mpsc,
9136 };
9137 use std::time::{Duration, Instant};
9138
9139 const CAPACITY: usize = LIVE_SUBSTREAM_CAPACITY;
9140 const BATCH: usize = LIVE_SUBSTREAM_BATCH;
9141 const TOTAL: usize = CAPACITY * 2;
9142 const MAX_IN_FLIGHT: usize = CAPACITY + BATCH;
9145
9146 let produced = Arc::new(AtomicUsize::new(0));
9147 let consumed = Arc::new(AtomicUsize::new(0));
9148 let bound_violated = Arc::new(AtomicBool::new(false));
9149
9150 let (item_tx, item_rx) = mpsc::channel::<u64>();
9152
9153 let prod_for_factory = Arc::clone(&produced);
9154 let prod_for_fold = Arc::clone(&produced);
9155 let cons_for_fold = Arc::clone(&consumed);
9156 let bv_for_fold = Arc::clone(&bound_violated);
9157
9158 let join = std::thread::spawn(move || {
9159 with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9160 crate::Source::from_factory(move || {
9161 let prod = Arc::clone(&prod_for_factory);
9162 let mut i = 0u64;
9163 Box::new(std::iter::from_fn(move || {
9164 if i as usize >= TOTAL {
9165 return None;
9166 }
9167 prod.fetch_add(1, Ordering::SeqCst);
9168 let val = i;
9169 i += 1;
9170 Some(Ok(val))
9171 }))
9172 })
9173 .split_when(|_| false)
9175 .run_with(crate::Sink::fold(
9176 0usize,
9177 move |count, seg: crate::Source<u64>| {
9178 let cons = Arc::clone(&cons_for_fold);
9179 let bv = Arc::clone(&bv_for_fold);
9180 let prod = Arc::clone(&prod_for_fold);
9181 let itx = item_tx.clone();
9183 seg.run_with(crate::Sink::foreach(move |x: u64| {
9186 let c = cons.fetch_add(1, Ordering::SeqCst) + 1;
9187 let p = prod.load(Ordering::SeqCst);
9188 if p > c + MAX_IN_FLIGHT {
9189 bv.store(true, Ordering::SeqCst);
9190 }
9191 let _ = itx.send(x);
9192 }))
9193 .unwrap()
9194 .wait()
9195 .unwrap();
9196 count + 1
9197 },
9198 ))
9199 .unwrap()
9200 .wait()
9201 .unwrap()
9202 })
9203 });
9204
9205 let mut received = Vec::with_capacity(TOTAL);
9210 let timeout = Duration::from_secs(60);
9211 let deadline = Instant::now() + timeout;
9212 for i in 0..TOTAL {
9213 let remaining = deadline.saturating_duration_since(Instant::now());
9214 if remaining == Duration::ZERO {
9215 panic!(
9216 "deadlock: received {} of {TOTAL} items within {timeout:?}",
9217 received.len()
9218 );
9219 }
9220 match item_rx.recv_timeout(remaining) {
9221 Ok(item) => received.push(item),
9222 Err(mpsc::RecvTimeoutError::Timeout) => {
9223 panic!(
9224 "deadlock: no item {i} before {timeout:?} rendezvous deadline; received {} of {TOTAL}",
9225 received.len()
9226 )
9227 }
9228 Err(mpsc::RecvTimeoutError::Disconnected) => {
9229 panic!("stream ended early at item {i}")
9230 }
9231 }
9232 }
9233
9234 let seg_count = join.join().expect("stream thread panicked");
9235
9236 assert!(
9238 !bound_violated.load(Ordering::SeqCst),
9239 "bound violated: producer ran >MAX_IN_FLIGHT={MAX_IN_FLIGHT} ahead of consumer"
9240 );
9241
9242 assert_eq!(seg_count, 1, "expected exactly 1 segment");
9244 assert_eq!(received.len(), TOTAL, "not all items received");
9245 let expected: Vec<u64> = (0..TOTAL as u64).collect();
9246 assert_eq!(received, expected, "items not in correct order");
9247 }
9248}