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