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