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 drop(
6106 slot.available
6107 .wait(state)
6108 .unwrap_or_else(|p| p.into_inner()),
6109 );
6110 continue;
6111 }
6112 }
6113
6114 return Ok(());
6116 }
6117 }
6118
6119 fn close_segment(&mut self) {
6120 let _ = self.flush_pending();
6122
6123 if let Some(consumer) = self.current_consumer.take() {
6125 consumer.complete();
6126 self.current_slot = None;
6127 return;
6128 }
6129
6130 let slot = match self.current_slot.take() {
6131 Some(s) => s,
6132 None => return,
6133 };
6134
6135 if self.parent_cancelled.load(Ordering::SeqCst) {
6136 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6137 if state.terminal.is_none() {
6138 state.terminal = Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
6139 }
6140 drop(state);
6141 slot.available.notify_all();
6142 return;
6143 }
6144
6145 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6146 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6147 SegmentConsumer::Direct(mut consumer_box) => {
6148 let mut drain_buf = std::mem::take(&mut state.buffer);
6149 drop(state);
6150 for item in drain_buf.drain(..) {
6151 if let Err(e) = consumer_box.push_item(item) {
6152 consumer_box.fail(e);
6153 return;
6154 }
6155 }
6156 consumer_box.complete();
6157 }
6158 SegmentConsumer::DirectTaken => {
6159 }
6161 SegmentConsumer::Fallback => {
6162 state.consumer = SegmentConsumer::DirectTaken;
6163 if state.terminal.is_none() {
6164 state.terminal = Some(LiveSubstreamTerminal::Complete);
6165 }
6166 drop(state);
6167 slot.available.notify_all();
6168 }
6169 SegmentConsumer::Pending => {
6170 state.consumer = SegmentConsumer::Pending;
6172 if state.terminal.is_none() {
6173 state.terminal = Some(LiveSubstreamTerminal::Complete);
6174 }
6175 drop(state);
6176 slot.available.notify_all();
6178 }
6179 }
6180 }
6181
6182 fn fail_segment(&mut self, error: StreamError) {
6183 self.local_pending.clear();
6185
6186 if let Some(consumer) = self.current_consumer.take() {
6187 consumer.fail(error);
6188 self.current_slot = None;
6189 return;
6190 }
6191 let slot = match self.current_slot.take() {
6192 Some(s) => s,
6193 None => return,
6194 };
6195 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6196 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6197 SegmentConsumer::Direct(c) => {
6198 drop(state);
6199 c.fail(error);
6200 }
6201 _ => {
6202 if state.terminal.is_none() {
6203 state.terminal = Some(LiveSubstreamTerminal::Error(error));
6204 }
6205 drop(state);
6206 slot.available.notify_all();
6207 }
6208 }
6209 }
6210
6211 fn fail_all(&mut self, error: StreamError) {
6212 self.fail_segment(error.clone());
6213 fail_live_substream(&self.outer, error);
6214 }
6215
6216 fn complete_all(&mut self) {
6217 self.close_segment();
6218 complete_live_substream(&self.outer);
6219 }
6220}
6221
6222impl<Out: Send + 'static> Drop for SplitFastWorkerGuard<Out> {
6223 fn drop(&mut self) {
6224 if self.armed {
6225 self.local_pending.clear(); if let Some(consumer) = self.current_consumer.take() {
6227 consumer.fail(StreamError::AbruptTermination);
6228 } else if let Some(slot) = self.current_slot.take() {
6229 let mut state = slot.state.lock().unwrap_or_else(|p| p.into_inner());
6230 match std::mem::replace(&mut state.consumer, SegmentConsumer::DirectTaken) {
6231 SegmentConsumer::Direct(c) => {
6232 drop(state);
6233 c.fail(StreamError::AbruptTermination);
6234 }
6235 _ => {
6236 if state.terminal.is_none() {
6237 state.terminal =
6238 Some(LiveSubstreamTerminal::Error(StreamError::AbruptTermination));
6239 }
6240 drop(state);
6241 slot.available.notify_all();
6242 }
6243 }
6244 }
6245 fail_live_substream(&self.outer, StreamError::AbruptTermination);
6246 }
6247 }
6248}
6249
6250fn split_streams_fast<Out, F>(
6251 mut input: BoxStream<Out>,
6252 mode: SplitMode,
6253 predicate: Arc<F>,
6254 parent_cancelled: Arc<AtomicBool>,
6255 materializer: &Materializer,
6256) -> BoxStream<Source<Out>>
6257where
6258 Out: Clone + Send + 'static,
6259 F: Fn(&Out) -> bool + Send + Sync + 'static,
6260{
6261 let outer = LiveSubstreamShared::new();
6262 let worker_outer = Arc::clone(&outer);
6263 let completion = materializer.spawn_stream(move |cancelled| {
6264 let mut guard =
6265 SplitFastWorkerGuard::new(Arc::clone(&worker_outer), Arc::clone(&parent_cancelled));
6266
6267 while !cancelled.load(Ordering::SeqCst) {
6268 if worker_outer.cancelled.load(Ordering::SeqCst) {
6269 guard.disarm();
6270 return Ok(NotUsed);
6271 }
6272
6273 match input.next() {
6274 Some(Ok(item)) => {
6275 let split = match catch_unwind(AssertUnwindSafe(|| predicate(&item))) {
6276 Ok(s) => s,
6277 Err(_) => {
6278 guard.fail_all(StreamError::AbruptTermination);
6279 guard.disarm();
6280 return Ok(NotUsed);
6281 }
6282 };
6283
6284 match mode {
6285 SplitMode::When => {
6286 if split
6287 && (guard.current_slot.is_some()
6288 || guard.current_consumer.is_some())
6289 {
6290 guard.close_segment();
6291 }
6292 if guard.current_slot.is_none()
6293 && guard.current_consumer.is_none()
6294 && guard.open_segment().is_err()
6295 {
6296 guard.disarm();
6297 return Ok(NotUsed);
6298 }
6299 if guard.push_item(item).is_err() {
6300 guard.disarm();
6301 return Ok(NotUsed);
6302 }
6303 }
6304 SplitMode::After => {
6305 if guard.current_slot.is_none()
6306 && guard.current_consumer.is_none()
6307 && guard.open_segment().is_err()
6308 {
6309 guard.disarm();
6310 return Ok(NotUsed);
6311 }
6312 if guard.push_item(item).is_err() {
6313 guard.disarm();
6314 return Ok(NotUsed);
6315 }
6316 if split {
6317 guard.close_segment();
6318 }
6319 }
6320 }
6321 }
6322 Some(Err(error)) => {
6323 guard.fail_all(error.clone());
6324 guard.disarm();
6325 return Err(error);
6326 }
6327 None => {
6328 guard.complete_all();
6329 guard.disarm();
6330 return Ok(NotUsed);
6331 }
6332 }
6333 }
6334 guard.complete_all();
6335 guard.disarm();
6336 Ok(NotUsed)
6337 });
6338
6339 Box::new(LiveSubstreamStream {
6340 shared: outer,
6341 completion: Some(completion),
6342 local_batch: VecDeque::new(),
6343 })
6344}
6345
6346#[cfg(test)]
6348struct FlatMapMergeShared<T> {
6349 state: Mutex<FlatMapMergeState<T>>,
6350 available: Condvar,
6351 cancelled: Arc<AtomicBool>,
6352}
6353
6354#[cfg(test)]
6355struct FlatMapMergeState<T> {
6356 queued: VecDeque<(usize, T)>,
6357 window: HashMap<usize, usize>,
6358 active_streams: usize,
6359 input_done: bool,
6360 terminal: Option<StreamError>,
6361}
6362
6363#[cfg(test)]
6364impl<T> FlatMapMergeShared<T> {
6365 fn new() -> Arc<Self> {
6366 Arc::new(Self {
6367 state: Mutex::new(FlatMapMergeState {
6368 queued: VecDeque::new(),
6369 window: HashMap::new(),
6370 active_streams: 0,
6371 input_done: false,
6372 terminal: None,
6373 }),
6374 available: Condvar::new(),
6375 cancelled: Arc::new(AtomicBool::new(false)),
6376 })
6377 }
6378}
6379
6380#[cfg(test)]
6381struct FlatMapMergeStream<T> {
6382 shared: Arc<FlatMapMergeShared<T>>,
6383 completion: Option<StreamCompletion<NotUsed>>,
6384}
6385
6386#[cfg(test)]
6387impl<T> Iterator for FlatMapMergeStream<T> {
6388 type Item = StreamResult<T>;
6389
6390 fn next(&mut self) -> Option<Self::Item> {
6391 let mut state = self
6392 .shared
6393 .state
6394 .lock()
6395 .unwrap_or_else(|poison| poison.into_inner());
6396 loop {
6397 if let Some((stream_id, item)) = state.queued.pop_front() {
6398 if let Some(count) = state.window.get_mut(&stream_id) {
6401 *count -= 1;
6402 if *count == 0 {
6403 state.window.remove(&stream_id);
6404 }
6405 }
6406 drop(state);
6407 self.shared.available.notify_all();
6408 return Some(Ok(item));
6409 }
6410 if let Some(error) = state.terminal.clone() {
6411 return Some(Err(error));
6412 }
6413 if state.input_done && state.active_streams == 0 {
6414 return None;
6415 }
6416 state = self
6417 .shared
6418 .available
6419 .wait(state)
6420 .unwrap_or_else(|poison| poison.into_inner());
6421 }
6422 }
6423}
6424
6425#[cfg(test)]
6426impl<T> Drop for FlatMapMergeStream<T> {
6427 fn drop(&mut self) {
6428 self.shared.cancelled.store(true, Ordering::SeqCst);
6429 self.shared.available.notify_all();
6430 let _ = self.completion.take();
6431 }
6432}
6433
6434#[cfg(test)]
6435fn flat_map_merge_register_stream<T>(shared: &Arc<FlatMapMergeShared<T>>) {
6436 let mut state = shared
6437 .state
6438 .lock()
6439 .unwrap_or_else(|poison| poison.into_inner());
6440 state.active_streams += 1;
6441 drop(state);
6442 shared.available.notify_all();
6443}
6444
6445#[cfg(test)]
6446fn flat_map_merge_finish_stream<T>(
6447 shared: &Arc<FlatMapMergeShared<T>>,
6448 stream_id: usize,
6449 terminal: Result<(), StreamError>,
6450) {
6451 let mut state = shared
6452 .state
6453 .lock()
6454 .unwrap_or_else(|poison| poison.into_inner());
6455 state.active_streams = state.active_streams.saturating_sub(1);
6456 if let Err(error) = terminal
6457 && state.terminal.is_none()
6458 {
6459 state.terminal = Some(error);
6460 }
6461 state.window.entry(stream_id).or_default();
6463 drop(state);
6464 shared.available.notify_all();
6465}
6466
6467#[cfg(test)]
6468fn flat_map_merge_push<T>(
6469 shared: &Arc<FlatMapMergeShared<T>>,
6470 stream_id: usize,
6471 item: T,
6472) -> Result<(), T> {
6473 let mut state = shared
6476 .state
6477 .lock()
6478 .unwrap_or_else(|poison| poison.into_inner());
6479 while state.window.get(&stream_id).copied().unwrap_or(0) >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
6480 && state.terminal.is_none()
6481 {
6482 if shared.cancelled.load(Ordering::SeqCst) {
6483 return Err(item);
6484 }
6485 state = shared
6486 .available
6487 .wait(state)
6488 .unwrap_or_else(|poison| poison.into_inner());
6489 }
6490 if shared.cancelled.load(Ordering::SeqCst) || state.terminal.is_some() {
6491 return Err(item);
6492 }
6493 *state.window.entry(stream_id).or_insert(0) += 1;
6494 let was_empty = state.queued.is_empty();
6495 state.queued.push_back((stream_id, item));
6496 drop(state);
6497 if was_empty {
6498 shared.available.notify_all();
6499 }
6500 Ok(())
6501}
6502
6503#[cfg(test)]
6504struct FlatMapMergeCoordinatorGuard<T> {
6505 shared: Arc<FlatMapMergeShared<T>>,
6506 armed: bool,
6507}
6508
6509#[cfg(test)]
6510impl<T> FlatMapMergeCoordinatorGuard<T> {
6511 fn new(shared: Arc<FlatMapMergeShared<T>>) -> Self {
6512 Self {
6513 shared,
6514 armed: true,
6515 }
6516 }
6517
6518 fn finish(&mut self) {
6519 let mut state = self
6520 .shared
6521 .state
6522 .lock()
6523 .unwrap_or_else(|poison| poison.into_inner());
6524 state.input_done = true;
6525 drop(state);
6526 self.shared.available.notify_all();
6527 self.armed = false;
6528 }
6529}
6530
6531#[cfg(test)]
6532impl<T> Drop for FlatMapMergeCoordinatorGuard<T> {
6533 fn drop(&mut self) {
6534 if self.armed {
6535 let mut state = self
6536 .shared
6537 .state
6538 .lock()
6539 .unwrap_or_else(|poison| poison.into_inner());
6540 if state.terminal.is_none() {
6541 state.terminal = Some(StreamError::AbruptTermination);
6542 }
6543 state.input_done = true;
6544 drop(state);
6545 self.shared.cancelled.store(true, Ordering::SeqCst);
6546 self.shared.available.notify_all();
6547 }
6548 }
6549}
6550
6551#[cfg(test)]
6552struct FlatMapMergeWorkerGuard<T> {
6553 shared: Arc<FlatMapMergeShared<T>>,
6554 stream_id: usize,
6555 armed: bool,
6556}
6557
6558#[cfg(test)]
6559impl<T> FlatMapMergeWorkerGuard<T> {
6560 fn new(shared: Arc<FlatMapMergeShared<T>>, stream_id: usize) -> Self {
6561 Self {
6562 shared,
6563 stream_id,
6564 armed: true,
6565 }
6566 }
6567
6568 fn finish(&mut self, terminal: Result<(), StreamError>) {
6569 flat_map_merge_finish_stream(&self.shared, self.stream_id, terminal);
6570 self.armed = false;
6571 }
6572}
6573
6574#[cfg(test)]
6575impl<T> Drop for FlatMapMergeWorkerGuard<T> {
6576 fn drop(&mut self) {
6577 if self.armed {
6578 flat_map_merge_finish_stream(
6579 &self.shared,
6580 self.stream_id,
6581 Err(StreamError::AbruptTermination),
6582 );
6583 self.shared.cancelled.store(true, Ordering::SeqCst);
6584 }
6585 }
6586}
6587
6588#[cfg(test)]
6591#[derive(Clone, Copy, PartialEq, Eq, Debug)]
6592pub(crate) enum SubstreamExecutorMode {
6593 Auto,
6594 LegacyOnly,
6595 ReadyRingOnly,
6596 SplitSinkOnly,
6597}
6598
6599#[cfg(test)]
6600thread_local! {
6601 static SUBSTREAM_EXECUTOR_MODE: std::cell::Cell<SubstreamExecutorMode> =
6602 const { std::cell::Cell::new(SubstreamExecutorMode::Auto) };
6603}
6604
6605#[cfg(test)]
6606pub(crate) fn with_substream_mode<R>(mode: SubstreamExecutorMode, f: impl FnOnce() -> R) -> R {
6607 SUBSTREAM_EXECUTOR_MODE.with(|m| {
6608 let old = m.get();
6609 m.set(mode);
6610 let result = f();
6611 m.set(old);
6612 result
6613 })
6614}
6615
6616#[cfg(test)]
6617fn current_substream_mode() -> SubstreamExecutorMode {
6618 SUBSTREAM_EXECUTOR_MODE.with(|m| m.get())
6619}
6620
6621fn flat_map_merge_stream<Out, Next, NextMat, F>(
6624 input: BoxStream<Out>,
6625 breadth: usize,
6626 stage: Arc<F>,
6627 materializer: &Materializer,
6628) -> BoxStream<Next>
6629where
6630 Out: Send + 'static,
6631 Next: Send + 'static,
6632 NextMat: Send + 'static,
6633 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6634{
6635 #[cfg(test)]
6636 if current_substream_mode() == SubstreamExecutorMode::LegacyOnly {
6637 return flat_map_merge_stream_legacy(input, breadth, stage, materializer);
6638 }
6639 flat_map_merge_stream_ready(input, breadth, stage, materializer)
6640}
6641
6642#[cfg(test)]
6645fn flat_map_merge_stream_legacy<Out, Next, NextMat, F>(
6646 mut input: BoxStream<Out>,
6647 breadth: usize,
6648 stage: Arc<F>,
6649 materializer: &Materializer,
6650) -> BoxStream<Next>
6651where
6652 Out: Send + 'static,
6653 Next: Send + 'static,
6654 NextMat: Send + 'static,
6655 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
6656{
6657 let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
6658 let shared = FlatMapMergeShared::new();
6659 let worker_shared = Arc::clone(&shared);
6660 let completion = materializer.spawn_stream(move |cancelled| {
6661 let mut next_id = 0usize;
6662 let mut guard = FlatMapMergeCoordinatorGuard::new(worker_shared);
6663 let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
6664
6665 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6666 workers.retain(|_, completion| completion.try_wait().is_none());
6667 {
6668 let mut state = guard
6669 .shared
6670 .state
6671 .lock()
6672 .unwrap_or_else(|poison| poison.into_inner());
6673 while state.active_streams >= breadth
6674 && state.terminal.is_none()
6675 && !cancelled.load(Ordering::SeqCst)
6676 && !guard.shared.cancelled.load(Ordering::SeqCst)
6677 {
6678 state = guard
6679 .shared
6680 .available
6681 .wait(state)
6682 .unwrap_or_else(|poison| poison.into_inner());
6683 }
6684 if state.terminal.is_some() {
6685 let error = state.terminal.clone().expect("terminal checked above");
6686 drop(state);
6687 guard.finish();
6688 return Err(error);
6689 }
6690 }
6691
6692 match input.next() {
6693 Some(Ok(item)) => {
6694 let source = stage(item);
6695 let (mut stream, _) =
6696 match Arc::clone(&source.factory).create(&worker_materializer) {
6697 Ok(parts) => parts,
6698 Err(error) => {
6699 let mut state = guard
6700 .shared
6701 .state
6702 .lock()
6703 .unwrap_or_else(|poison| poison.into_inner());
6704 if state.terminal.is_none() {
6705 state.terminal = Some(error.clone());
6706 }
6707 drop(state);
6708 guard.shared.available.notify_all();
6709 guard.finish();
6710 return Err(error);
6711 }
6712 };
6713 let stream_id = next_id;
6714 next_id += 1;
6715 flat_map_merge_register_stream(&guard.shared);
6716 let worker_shared = Arc::clone(&guard.shared);
6717 workers.insert(
6718 stream_id,
6719 worker_materializer.spawn_stream(move |inner_cancelled| {
6720 let mut worker_guard =
6721 FlatMapMergeWorkerGuard::new(Arc::clone(&worker_shared), stream_id);
6722 while !inner_cancelled.load(Ordering::SeqCst)
6723 && !worker_shared.cancelled.load(Ordering::SeqCst)
6724 {
6725 match stream.next() {
6726 Some(Ok(item)) => {
6727 if flat_map_merge_push(&worker_shared, stream_id, item)
6728 .is_err()
6729 {
6730 worker_guard.finish(Ok(()));
6731 return Ok(NotUsed);
6732 }
6733 }
6734 Some(Err(error)) => {
6735 worker_guard.finish(Err(error.clone()));
6736 return Err(error);
6737 }
6738 None => {
6739 worker_guard.finish(Ok(()));
6740 return Ok(NotUsed);
6741 }
6742 }
6743 }
6744 worker_guard.finish(Ok(()));
6745 Ok(NotUsed)
6746 }),
6747 );
6748 }
6749 Some(Err(error)) => {
6750 let mut state = guard
6751 .shared
6752 .state
6753 .lock()
6754 .unwrap_or_else(|poison| poison.into_inner());
6755 if state.terminal.is_none() {
6756 state.terminal = Some(error.clone());
6757 }
6758 drop(state);
6759 guard.shared.available.notify_all();
6760 guard.finish();
6761 return Err(error);
6762 }
6763 None => {
6764 guard.finish();
6765 break;
6766 }
6767 }
6768 }
6769
6770 if guard.armed {
6771 guard.finish();
6772 }
6773 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
6781 workers.retain(|_, completion| completion.try_wait().is_none());
6782 let state = guard
6788 .shared
6789 .state
6790 .lock()
6791 .unwrap_or_else(|poison| poison.into_inner());
6792 if state.active_streams == 0 {
6793 return Ok(NotUsed);
6794 }
6795 drop(
6796 guard
6797 .shared
6798 .available
6799 .wait(state)
6800 .unwrap_or_else(|poison| poison.into_inner()),
6801 );
6802 }
6803 Ok(NotUsed)
6804 });
6805
6806 Box::new(FlatMapMergeStream {
6807 shared,
6808 completion: Some(completion),
6809 })
6810}
6811
6812const FLAT_MAP_MERGE_READY_BATCH: usize = 16;
6815
6816const FLAT_MAP_MERGE_INLINE_MICRO_MAX: usize = FLAT_MAP_MERGE_READY_BATCH;
6820
6821struct FlatMapMergeReadyShared<T> {
6822 coordinator: Mutex<FlatMapMergeReadyState<T>>,
6823 available: Condvar,
6824 cancelled: Arc<AtomicBool>,
6825}
6826
6827struct FlatMapMergeReadyState<T> {
6828 lanes: HashMap<usize, Arc<FlatMapMergeLane<T>>>,
6829 ready: std::collections::VecDeque<usize>,
6830 queued_items: usize,
6831 active_streams: usize,
6832 input_done: bool,
6833 terminal: Option<StreamError>,
6834 generation: u64,
6835}
6836
6837struct FlatMapMergeLane<T> {
6838 state: Mutex<FlatMapMergeLaneState<T>>,
6839 space_available: Condvar,
6840}
6841
6842struct FlatMapMergeLaneState<T> {
6843 buffer: std::collections::VecDeque<T>,
6844 in_ready_ring: bool,
6845 publishing: bool,
6846 closed: bool,
6847}
6848
6849struct FlatMapMergeReadyStream<T> {
6850 shared: Arc<FlatMapMergeReadyShared<T>>,
6851 completion: Option<StreamCompletion<NotUsed>>,
6852 local_batch: std::collections::VecDeque<T>,
6853}
6854
6855impl<T: Send + 'static> Iterator for FlatMapMergeReadyStream<T> {
6856 type Item = StreamResult<T>;
6857
6858 fn next(&mut self) -> Option<Self::Item> {
6859 if let Some(item) = self.local_batch.pop_front() {
6861 return Some(Ok(item));
6862 }
6863
6864 let mut coord = self
6866 .shared
6867 .coordinator
6868 .lock()
6869 .unwrap_or_else(|p| p.into_inner());
6870 loop {
6871 if let Some(error) = coord.terminal.clone() {
6873 return Some(Err(error));
6874 }
6875
6876 if let Some(lane_id) = coord.ready.pop_front() {
6878 let lane = coord.lanes.get(&lane_id).cloned();
6879 drop(coord);
6880
6881 if let Some(lane) = lane {
6882 let mut lane_state = lane.state.lock().unwrap_or_else(|p| p.into_inner());
6884 while lane_state.publishing {
6885 lane_state = lane
6886 .space_available
6887 .wait(lane_state)
6888 .unwrap_or_else(|p| p.into_inner());
6889 }
6890 let drain_n = lane_state.buffer.len().min(FLAT_MAP_MERGE_READY_BATCH);
6891 let mut batch = std::collections::VecDeque::with_capacity(drain_n);
6892 for _ in 0..drain_n {
6893 if let Some(item) = lane_state.buffer.pop_front() {
6894 batch.push_back(item);
6895 }
6896 }
6897 let still_has_items = !lane_state.buffer.is_empty();
6898 let is_closed = lane_state.closed;
6899 if !still_has_items {
6900 lane_state.in_ready_ring = false;
6901 }
6902 drop(lane_state);
6904 if !batch.is_empty() {
6906 lane.space_available.notify_all();
6907 }
6908
6909 let freed = batch.len();
6911 let mut coord = self
6912 .shared
6913 .coordinator
6914 .lock()
6915 .unwrap_or_else(|p| p.into_inner());
6916 coord.queued_items = coord.queued_items.saturating_sub(freed);
6917 if still_has_items {
6918 coord.ready.push_back(lane_id);
6919 } else if is_closed {
6920 coord.lanes.remove(&lane_id);
6921 }
6922 coord.generation += 1;
6923 drop(coord);
6924 self.shared.available.notify_all();
6925
6926 let mut iter = batch.into_iter();
6928 if let Some(first) = iter.next() {
6929 self.local_batch.extend(iter);
6930 return Some(Ok(first));
6931 }
6932 }
6933 coord = self
6935 .shared
6936 .coordinator
6937 .lock()
6938 .unwrap_or_else(|p| p.into_inner());
6939 continue;
6940 }
6941
6942 if coord.terminal.is_none()
6944 && coord.input_done
6945 && coord.active_streams == 0
6946 && coord.queued_items == 0
6947 {
6948 return None;
6949 }
6950
6951 let seen = coord.generation;
6953 coord = self
6954 .shared
6955 .available
6956 .wait_while(coord, |s| {
6957 s.generation == seen
6958 && s.terminal.is_none()
6959 && s.ready.is_empty()
6960 && !(s.input_done && s.active_streams == 0 && s.queued_items == 0)
6961 })
6962 .unwrap_or_else(|p| p.into_inner());
6963 }
6964 }
6965}
6966
6967impl<T> Drop for FlatMapMergeReadyStream<T> {
6968 fn drop(&mut self) {
6969 let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
6970 let mut coord = self
6971 .shared
6972 .coordinator
6973 .lock()
6974 .unwrap_or_else(|p| p.into_inner());
6975 coord.generation += 1;
6976 coord.lanes.values().cloned().collect()
6977 };
6978 self.shared.cancelled.store(true, Ordering::SeqCst);
6979 self.shared.available.notify_all();
6980 for lane in lanes {
6981 lane.space_available.notify_all();
6982 }
6983 let _ = self.completion.take();
6984 }
6985}
6986
6987fn finish_lane_ready<T>(
6990 shared: &Arc<FlatMapMergeReadyShared<T>>,
6991 stream_id: usize,
6992 terminal: Result<(), StreamError>,
6993) {
6994 let is_error = terminal.is_err();
6995
6996 let lane_is_empty = {
6998 let coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
6999 let lane_opt = coord.lanes.get(&stream_id).cloned();
7000 drop(coord);
7001 if let Some(lane) = lane_opt {
7002 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
7003 ls.closed = true;
7004 ls.buffer.is_empty()
7005 } else {
7006 true
7007 }
7008 };
7009
7010 let lanes_to_notify: Vec<Arc<FlatMapMergeLane<T>>> = {
7012 let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
7013 coord.active_streams = coord.active_streams.saturating_sub(1);
7014 if let Err(ref error) = terminal
7015 && coord.terminal.is_none()
7016 {
7017 coord.terminal = Some(error.clone());
7018 }
7019 if lane_is_empty {
7020 coord.lanes.remove(&stream_id);
7021 }
7022 coord.generation += 1;
7023 if is_error {
7024 coord.lanes.values().cloned().collect()
7025 } else {
7026 vec![]
7027 }
7028 };
7029
7030 if is_error {
7031 shared.cancelled.store(true, Ordering::SeqCst);
7032 for lane in &lanes_to_notify {
7033 lane.space_available.notify_all();
7034 }
7035 }
7036 shared.available.notify_all();
7037}
7038
7039struct FlatMapMergeReadyCoordinatorGuard<T> {
7040 shared: Arc<FlatMapMergeReadyShared<T>>,
7041 armed: bool,
7042}
7043
7044impl<T> FlatMapMergeReadyCoordinatorGuard<T> {
7045 fn new(shared: Arc<FlatMapMergeReadyShared<T>>) -> Self {
7046 Self {
7047 shared,
7048 armed: true,
7049 }
7050 }
7051
7052 fn finish(&mut self) {
7053 let mut coord = self
7054 .shared
7055 .coordinator
7056 .lock()
7057 .unwrap_or_else(|p| p.into_inner());
7058 coord.input_done = true;
7059 coord.generation += 1;
7060 drop(coord);
7061 self.shared.available.notify_all();
7062 self.armed = false;
7063 }
7064}
7065
7066impl<T> Drop for FlatMapMergeReadyCoordinatorGuard<T> {
7067 fn drop(&mut self) {
7068 if self.armed {
7069 let lanes: Vec<Arc<FlatMapMergeLane<T>>> = {
7070 let mut coord = self
7071 .shared
7072 .coordinator
7073 .lock()
7074 .unwrap_or_else(|p| p.into_inner());
7075 if coord.terminal.is_none() {
7076 coord.terminal = Some(StreamError::AbruptTermination);
7077 }
7078 coord.input_done = true;
7079 coord.generation += 1;
7080 coord.lanes.values().cloned().collect()
7081 };
7082 self.shared.cancelled.store(true, Ordering::SeqCst);
7083 self.shared.available.notify_all();
7084 for lane in lanes {
7085 lane.space_available.notify_all();
7086 }
7087 }
7088 }
7089}
7090
7091struct FlatMapMergeReadyWorkerGuard<T> {
7092 shared: Arc<FlatMapMergeReadyShared<T>>,
7093 stream_id: usize,
7094 armed: bool,
7095}
7096
7097impl<T> FlatMapMergeReadyWorkerGuard<T> {
7098 fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
7099 Self {
7100 shared,
7101 stream_id,
7102 armed: true,
7103 }
7104 }
7105
7106 fn finish(&mut self, terminal: Result<(), StreamError>) {
7107 finish_lane_ready(&self.shared, self.stream_id, terminal);
7108 self.armed = false;
7109 }
7110}
7111
7112impl<T> Drop for FlatMapMergeReadyWorkerGuard<T> {
7113 fn drop(&mut self) {
7114 if self.armed {
7115 finish_lane_ready(
7116 &self.shared,
7117 self.stream_id,
7118 Err(StreamError::AbruptTermination),
7119 );
7120 self.shared.cancelled.store(true, Ordering::SeqCst);
7121 }
7122 }
7123}
7124
7125fn publish_ready_batch<T>(
7140 shared: &Arc<FlatMapMergeReadyShared<T>>,
7141 stream_id: usize,
7142 lane: &Arc<FlatMapMergeLane<T>>,
7143 batch: std::collections::VecDeque<T>,
7144) {
7145 let batch_len = batch.len();
7146 if batch_len == 0 {
7147 return;
7148 }
7149
7150 let was_not_in_ready = {
7152 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
7153 let was_not = !ls.in_ready_ring;
7154 ls.publishing = true;
7155 ls.buffer.extend(batch);
7156 if !ls.buffer.is_empty() {
7157 ls.in_ready_ring = true;
7158 }
7159 was_not
7160 };
7161
7162 {
7164 let mut coord = shared.coordinator.lock().unwrap_or_else(|p| p.into_inner());
7165 coord.queued_items += batch_len;
7166 if was_not_in_ready {
7167 coord.ready.push_back(stream_id);
7168 }
7169 coord.generation += 1;
7170 }
7171 {
7172 let mut ls = lane.state.lock().unwrap_or_else(|p| p.into_inner());
7173 ls.publishing = false;
7174 }
7175 lane.space_available.notify_all();
7176 shared.available.notify_all();
7177}
7178
7179struct FlatMapMergeReadyInlineGuard<T> {
7189 shared: Arc<FlatMapMergeReadyShared<T>>,
7190 stream_id: usize,
7191 armed: bool,
7192}
7193
7194impl<T> FlatMapMergeReadyInlineGuard<T> {
7195 fn new(shared: Arc<FlatMapMergeReadyShared<T>>, stream_id: usize) -> Self {
7196 Self {
7197 shared,
7198 stream_id,
7199 armed: true,
7200 }
7201 }
7202
7203 fn finish(&mut self, terminal: Result<(), StreamError>) {
7204 finish_lane_ready(&self.shared, self.stream_id, terminal);
7205 self.armed = false;
7206 }
7207
7208 fn hand_off(&mut self) {
7209 self.armed = false;
7210 }
7211}
7212
7213impl<T> Drop for FlatMapMergeReadyInlineGuard<T> {
7214 fn drop(&mut self) {
7215 if self.armed {
7216 finish_lane_ready(
7217 &self.shared,
7218 self.stream_id,
7219 Err(StreamError::AbruptTermination),
7220 );
7221 }
7222 }
7223}
7224
7225fn flat_map_merge_stream_ready<Out, Next, NextMat, F>(
7226 mut input: BoxStream<Out>,
7227 breadth: usize,
7228 stage: Arc<F>,
7229 materializer: &Materializer,
7230) -> BoxStream<Next>
7231where
7232 Out: Send + 'static,
7233 Next: Send + 'static,
7234 NextMat: Send + 'static,
7235 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7236{
7237 let worker_materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7238 let shared: Arc<FlatMapMergeReadyShared<Next>> = Arc::new(FlatMapMergeReadyShared {
7239 coordinator: Mutex::new(FlatMapMergeReadyState {
7240 lanes: HashMap::new(),
7241 ready: std::collections::VecDeque::new(),
7242 queued_items: 0,
7243 active_streams: 0,
7244 input_done: false,
7245 terminal: None,
7246 generation: 0,
7247 }),
7248 available: Condvar::new(),
7249 cancelled: Arc::new(AtomicBool::new(false)),
7250 });
7251
7252 let worker_shared = Arc::clone(&shared);
7253 let completion = materializer.spawn_stream(move |cancelled| {
7254 let mut next_id = 0usize;
7255 let mut guard = FlatMapMergeReadyCoordinatorGuard::new(worker_shared);
7256 let mut workers = HashMap::<usize, StreamCompletion<NotUsed>>::with_capacity(breadth);
7257
7258 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7259 workers.retain(|_, c| c.try_wait().is_none());
7260
7261 {
7263 let mut coord = guard
7264 .shared
7265 .coordinator
7266 .lock()
7267 .unwrap_or_else(|p| p.into_inner());
7268 loop {
7269 if coord.terminal.is_some()
7270 || coord.active_streams < breadth
7271 || cancelled.load(Ordering::SeqCst)
7272 || guard.shared.cancelled.load(Ordering::SeqCst)
7273 {
7274 break;
7275 }
7276 let seen = coord.generation;
7277 coord = guard
7278 .shared
7279 .available
7280 .wait_while(coord, |s| {
7281 s.generation == seen
7282 && s.terminal.is_none()
7283 && s.active_streams >= breadth
7284 && !cancelled.load(Ordering::SeqCst)
7285 && !guard.shared.cancelled.load(Ordering::SeqCst)
7286 })
7287 .unwrap_or_else(|p| p.into_inner());
7288 }
7289 if coord.terminal.is_some() {
7290 let error = coord.terminal.clone().expect("terminal checked");
7291 drop(coord);
7292 guard.finish();
7293 return Err(error);
7294 }
7295 }
7296
7297 match input.next() {
7298 Some(Ok(item)) => {
7299 let source = stage(item);
7301 let inline_hint = source.hints.inline_micro;
7303 let (stream, _) = match Arc::clone(&source.factory).create(&worker_materializer)
7304 {
7305 Ok(parts) => parts,
7306 Err(error) => {
7307 let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7308 let mut coord = guard
7309 .shared
7310 .coordinator
7311 .lock()
7312 .unwrap_or_else(|p| p.into_inner());
7313 if coord.terminal.is_none() {
7314 coord.terminal = Some(error.clone());
7315 }
7316 coord.generation += 1;
7317 coord.lanes.values().cloned().collect()
7318 };
7319 guard.shared.cancelled.store(true, Ordering::SeqCst);
7320 guard.shared.available.notify_all();
7321 for lane in lanes {
7322 lane.space_available.notify_all();
7323 }
7324 guard.finish();
7325 return Err(error);
7326 }
7327 };
7328
7329 let stream_id = next_id;
7330 next_id += 1;
7331 let mut stream = stream;
7332
7333 let lane: Arc<FlatMapMergeLane<Next>> = Arc::new(FlatMapMergeLane {
7335 state: Mutex::new(FlatMapMergeLaneState {
7336 buffer: std::collections::VecDeque::new(),
7337 in_ready_ring: false,
7338 publishing: false,
7339 closed: false,
7340 }),
7341 space_available: Condvar::new(),
7342 });
7343 {
7344 let mut coord = guard
7345 .shared
7346 .coordinator
7347 .lock()
7348 .unwrap_or_else(|p| p.into_inner());
7349 coord.lanes.insert(stream_id, Arc::clone(&lane));
7350 coord.active_streams += 1;
7351 }
7353
7354 let mut inline_guard =
7357 FlatMapMergeReadyInlineGuard::new(Arc::clone(&guard.shared), stream_id);
7358
7359 let is_inline = inline_hint
7360 .is_some_and(|h| h.max_success_items <= FLAT_MAP_MERGE_INLINE_MICRO_MAX);
7361
7362 if is_inline {
7363 let max_items = inline_hint.expect("checked").max_success_items;
7367
7368 if guard.shared.cancelled.load(Ordering::SeqCst)
7370 || cancelled.load(Ordering::SeqCst)
7371 {
7372 inline_guard.finish(Ok(()));
7373 continue;
7374 }
7375
7376 let max_pulls = max_items.saturating_add(1);
7379 let mut local_batch = std::collections::VecDeque::with_capacity(max_items);
7380 let mut terminal_result: Option<Result<(), StreamError>> = None;
7381
7382 for _ in 0..max_pulls {
7383 if guard.shared.cancelled.load(Ordering::SeqCst)
7384 || cancelled.load(Ordering::SeqCst)
7385 {
7386 break;
7387 }
7388 match stream.next() {
7389 Some(Ok(item)) => local_batch.push_back(item),
7390 Some(Err(e)) => {
7391 terminal_result = Some(Err(e));
7392 break;
7393 }
7394 None => {
7395 terminal_result = Some(Ok(()));
7396 break;
7397 }
7398 }
7399 }
7400
7401 if guard.shared.cancelled.load(Ordering::SeqCst)
7403 || cancelled.load(Ordering::SeqCst)
7404 {
7405 inline_guard.finish(Ok(()));
7406 continue;
7407 }
7408
7409 publish_ready_batch(&guard.shared, stream_id, &lane, local_batch);
7412
7413 match terminal_result {
7414 Some(Err(e)) => {
7415 inline_guard.finish(Err(e));
7418 }
7419 Some(Ok(())) | None => {
7420 inline_guard.finish(Ok(()));
7421 }
7422 }
7423 } else {
7424 inline_guard.hand_off();
7426 let worker_shared = Arc::clone(&guard.shared);
7427 let worker_lane = Arc::clone(&lane);
7428 workers.insert(
7429 stream_id,
7430 worker_materializer.spawn_stream(move |inner_cancelled| {
7431 let mut worker_guard = FlatMapMergeReadyWorkerGuard::new(
7432 Arc::clone(&worker_shared),
7433 stream_id,
7434 );
7435
7436 loop {
7437 if inner_cancelled.load(Ordering::SeqCst)
7438 || worker_shared.cancelled.load(Ordering::SeqCst)
7439 {
7440 worker_guard.finish(Ok(()));
7441 return Ok(NotUsed);
7442 }
7443
7444 let capacity;
7446 {
7447 let mut ls = worker_lane
7448 .state
7449 .lock()
7450 .unwrap_or_else(|p| p.into_inner());
7451 while ls.buffer.len() >= FLAT_MAP_MERGE_SUBSTREAM_WINDOW
7452 && !worker_shared.cancelled.load(Ordering::SeqCst)
7453 && !ls.closed
7454 && !inner_cancelled.load(Ordering::SeqCst)
7455 {
7456 ls = worker_lane
7457 .space_available
7458 .wait(ls)
7459 .unwrap_or_else(|p| p.into_inner());
7460 }
7461 if worker_shared.cancelled.load(Ordering::SeqCst)
7462 || ls.closed
7463 || inner_cancelled.load(Ordering::SeqCst)
7464 {
7465 drop(ls);
7466 worker_guard.finish(Ok(()));
7467 return Ok(NotUsed);
7468 }
7469 capacity =
7470 FLAT_MAP_MERGE_SUBSTREAM_WINDOW - ls.buffer.len();
7471 }
7472 let batch_size = capacity.min(FLAT_MAP_MERGE_READY_BATCH);
7476 let mut local_batch =
7477 std::collections::VecDeque::with_capacity(batch_size);
7478 let mut terminal_result: Option<Result<(), StreamError>> = None;
7479 for _ in 0..batch_size {
7480 if inner_cancelled.load(Ordering::SeqCst)
7481 || worker_shared.cancelled.load(Ordering::SeqCst)
7482 {
7483 break;
7484 }
7485 match stream.next() {
7486 Some(Ok(item)) => local_batch.push_back(item),
7487 Some(Err(e)) => {
7488 terminal_result = Some(Err(e));
7489 break;
7490 }
7491 None => {
7492 terminal_result = Some(Ok(()));
7493 break;
7494 }
7495 }
7496 }
7497 let batch_len = local_batch.len();
7498
7499 publish_ready_batch(
7502 &worker_shared,
7503 stream_id,
7504 &worker_lane,
7505 local_batch,
7506 );
7507
7508 match terminal_result {
7509 Some(Ok(())) => {
7510 worker_guard.finish(Ok(()));
7511 return Ok(NotUsed);
7512 }
7513 Some(Err(e)) => {
7514 worker_guard.finish(Err(e.clone()));
7515 return Err(e);
7516 }
7517 None => {
7518 if batch_len == 0 {
7519 worker_guard.finish(Ok(()));
7520 return Ok(NotUsed);
7521 }
7522 }
7523 }
7524 }
7525 }),
7526 );
7527 }
7528 }
7529 Some(Err(error)) => {
7530 let lanes: Vec<Arc<FlatMapMergeLane<Next>>> = {
7531 let mut coord = guard
7532 .shared
7533 .coordinator
7534 .lock()
7535 .unwrap_or_else(|p| p.into_inner());
7536 if coord.terminal.is_none() {
7537 coord.terminal = Some(error.clone());
7538 }
7539 coord.generation += 1;
7540 coord.lanes.values().cloned().collect()
7541 };
7542 guard.shared.cancelled.store(true, Ordering::SeqCst);
7543 guard.shared.available.notify_all();
7544 for lane in lanes {
7545 lane.space_available.notify_all();
7546 }
7547 guard.finish();
7548 return Err(error);
7549 }
7550 None => {
7551 guard.finish();
7552 break;
7553 }
7554 }
7555 }
7556
7557 if guard.armed {
7558 guard.finish();
7559 }
7560
7561 while !cancelled.load(Ordering::SeqCst) && !guard.shared.cancelled.load(Ordering::SeqCst) {
7565 workers.retain(|_, c| c.try_wait().is_none());
7566 let coord = guard
7567 .shared
7568 .coordinator
7569 .lock()
7570 .unwrap_or_else(|p| p.into_inner());
7571 if coord.active_streams == 0 {
7572 return Ok(NotUsed);
7573 }
7574 let seen = coord.generation;
7575 drop(
7576 guard
7577 .shared
7578 .available
7579 .wait_while(coord, |s| s.generation == seen && s.active_streams > 0)
7580 .unwrap_or_else(|p| p.into_inner()),
7581 );
7582 }
7583 Ok(NotUsed)
7584 });
7585
7586 Box::new(FlatMapMergeReadyStream {
7587 shared,
7588 completion: Some(completion),
7589 local_batch: std::collections::VecDeque::new(),
7590 })
7591}
7592
7593fn flat_map_concat_stream<Out, Next, NextMat, F>(
7594 mut input: BoxStream<Out>,
7595 stage: Arc<F>,
7596 materializer: &Materializer,
7597) -> BoxStream<Next>
7598where
7599 Out: Send + 'static,
7600 Next: Send + 'static,
7601 NextMat: Send + 'static,
7602 F: Fn(Out) -> Source<Next, NextMat> + Send + Sync + 'static,
7603{
7604 let materializer = materializer.with_name_prefix(materializer.name_prefix().to_owned());
7605 let mut current: Option<BoxStream<Next>> = None;
7606 Box::new(std::iter::from_fn(move || {
7607 loop {
7608 if let Some(stream) = current.as_mut() {
7609 match stream.next() {
7610 Some(item) => return Some(item),
7611 None => current = None,
7612 }
7613 }
7614
7615 match input.next() {
7616 Some(Ok(item)) => {
7617 let source = stage(item);
7618 current = Some(match Arc::clone(&source.factory).create(&materializer) {
7619 Ok((stream, _)) => stream,
7620 Err(error) => Box::new(std::iter::once(Err(error))) as BoxStream<Next>,
7621 });
7622 }
7623 Some(Err(error)) => return Some(Err(error)),
7624 None => return None,
7625 }
7626 }
7627 }))
7628}
7629
7630fn map_async_unordered<Out, Next, F, Fut>(
7631 mut input: BoxStream<Out>,
7632 parallelism: usize,
7633 stage: Arc<F>,
7634) -> BoxStream<Next>
7635where
7636 Out: Send + 'static,
7637 Next: Send + 'static,
7638 F: Fn(Out) -> Fut + Send + Sync + 'static,
7639 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7640{
7641 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7642 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7643 let mut next_task_id = 0_usize;
7644 let mut input_done = false;
7645
7646 Box::new(std::iter::from_fn(move || {
7647 loop {
7648 while tasks.len() < parallelism && !input_done {
7649 match input.next() {
7650 Some(Ok(item)) => match poll_once_or_pending(stage(item)) {
7651 Ok(result) => return Some(result),
7652 Err(future) => {
7653 let task_id = next_task_id;
7654 next_task_id += 1;
7655 tasks.insert(
7656 task_id,
7657 spawn_completion_task(task_id, future, sender.clone(), |result| {
7658 result
7659 }),
7660 );
7661 }
7662 },
7663 Some(Err(error)) => {
7664 input_done = true;
7665 return Some(Err(error));
7666 }
7667 None => input_done = true,
7668 }
7669 }
7670
7671 if tasks.is_empty() {
7672 return None;
7673 }
7674
7675 if let Some((task_id, result)) = recv_completion(&receiver) {
7676 tasks.remove(&task_id);
7677 return Some(result);
7678 }
7679 }
7680 }))
7681}
7682
7683fn map_async_unordered_supervised<Out, Next, F, Fut>(
7684 mut input: BoxStream<Out>,
7685 parallelism: usize,
7686 stage: Arc<F>,
7687 decider: SupervisionDecider,
7688) -> BoxStream<Next>
7689where
7690 Out: Send + 'static,
7691 Next: Send + 'static,
7692 F: Fn(Out) -> Fut + Send + Sync + 'static,
7693 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7694{
7695 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7696 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7697 let mut next_task_id = 0_usize;
7698 let mut input_done = false;
7699
7700 Box::new(std::iter::from_fn(move || {
7701 loop {
7702 while tasks.len() < parallelism && !input_done {
7703 match input.next() {
7704 Some(Ok(item)) => {
7705 match catch_unwind(AssertUnwindSafe(|| poll_once_or_pending(stage(item)))) {
7706 Ok(Ok(result)) => {
7707 if let Some(result) = supervise_async_result(result, &decider) {
7708 return Some(result);
7709 }
7710 }
7711 Ok(Err(future)) => {
7712 let task_id = next_task_id;
7713 next_task_id += 1;
7714 tasks.insert(
7715 task_id,
7716 spawn_completion_task(
7717 task_id,
7718 future,
7719 sender.clone(),
7720 |result| result,
7721 ),
7722 );
7723 }
7724 Err(_) => {
7725 let error = panic_stream_error("map_async_unordered callback");
7726 if let Some(result) = supervise_async_result(Err(error), &decider) {
7727 return Some(result);
7728 }
7729 }
7730 }
7731 }
7732 Some(Err(error)) => {
7733 input_done = true;
7734 return Some(Err(error));
7735 }
7736 None => input_done = true,
7737 }
7738 }
7739
7740 if tasks.is_empty() {
7741 return None;
7742 }
7743
7744 if let Some((task_id, result)) = recv_completion(&receiver) {
7745 tasks.remove(&task_id);
7746 if let Some(result) = supervise_async_result(result, &decider) {
7747 return Some(result);
7748 }
7749 }
7750 }
7751 }))
7752}
7753
7754#[inline(always)]
7755fn map_async_partitioned_serial<Out, Key, Next, Partition, F, Fut>(
7756 mut input: BoxStream<Out>,
7757 partition: Arc<Partition>,
7758 stage: Arc<F>,
7759) -> BoxStream<Next>
7760where
7761 Out: Send + 'static,
7762 Key: Send + 'static,
7763 Next: Send + 'static,
7764 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7765 F: Fn(Out) -> Fut + Send + Sync + 'static,
7766 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7767{
7768 let (sender, receiver) = std::sync::mpsc::channel::<(usize, StreamResult<Next>)>();
7769 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(1);
7770 let mut next_index = 0_usize;
7771 Box::new(std::iter::from_fn(move || {
7772 if !tasks.is_empty()
7776 && let Some((task_id, result)) = recv_completion(&receiver)
7777 {
7778 tasks.remove(&task_id);
7779 return Some(result);
7780 }
7781
7782 let item = input.next()?;
7783 match item {
7784 Ok(item) => {
7785 let _ = partition(&item);
7786 let index = next_index;
7787 next_index += 1;
7788 Some(match poll_once_or_pending(stage(item)) {
7789 Ok(result) => result,
7790 Err(future) => {
7791 tasks.insert(
7792 index,
7793 spawn_completion_task(index, future, sender.clone(), |result| result),
7794 );
7795 let (task_id, result) =
7796 recv_completion(&receiver).expect("pending map_async task completion");
7797 tasks.remove(&task_id);
7798 result
7799 }
7800 })
7801 }
7802 Err(error) => Some(Err(error)),
7803 }
7804 }))
7805}
7806
7807#[inline(always)]
7808fn map_async_partitioned_scanning<Out, Key, Next, Partition, F, Fut>(
7809 mut input: BoxStream<Out>,
7810 parallelism: usize,
7811 per_partition: usize,
7812 partition: Arc<Partition>,
7813 stage: Arc<F>,
7814) -> BoxStream<Next>
7815where
7816 Out: Send + 'static,
7817 Key: Clone + Eq + Hash + Send + 'static,
7818 Next: Send + 'static,
7819 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7820 F: Fn(Out) -> Fut + Send + Sync + 'static,
7821 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7822{
7823 let (sender, receiver) = std::sync::mpsc::channel::<(usize, (Key, StreamResult<Next>))>();
7824 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7825 let mut active_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7826 let mut pending = VecDeque::<(usize, Key, Out)>::with_capacity(parallelism);
7827 let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7828 let mut next_index = 0_usize;
7829 let mut next_to_emit = 0_usize;
7830 let mut input_done = false;
7831
7832 Box::new(std::iter::from_fn(move || {
7833 loop {
7834 if let Some(result) = completed.remove(&next_to_emit) {
7835 next_to_emit += 1;
7836 return Some(result);
7837 }
7838
7839 while tasks.len() + completed.len() < parallelism {
7840 let next = pending
7841 .iter()
7842 .position(|(_, key, _)| {
7843 active_by_key.get(key).copied().unwrap_or(0) < per_partition
7844 })
7845 .and_then(|index| pending.remove(index))
7846 .or_else(|| {
7847 if input_done {
7848 return None;
7849 }
7850 match input.next() {
7851 Some(Ok(item)) => {
7852 let key = partition(&item);
7853 let index = next_index;
7854 next_index += 1;
7855 Some((index, key, item))
7856 }
7857 Some(Err(error)) => {
7858 completed.insert(next_index, Err(error));
7859 next_index += 1;
7860 input_done = true;
7861 None
7862 }
7863 None => {
7864 input_done = true;
7865 None
7866 }
7867 }
7868 });
7869
7870 let Some((index, key, item)) = next else {
7871 break;
7872 };
7873 if active_by_key.get(&key).copied().unwrap_or(0) >= per_partition {
7874 pending.push_back((index, key, item));
7875 if input_done || pending.len() >= parallelism {
7876 break;
7877 }
7878 continue;
7879 }
7880 *active_by_key.entry(key.clone()).or_default() += 1;
7881 match poll_once_or_pending(stage(item)) {
7882 Ok(result) => {
7883 if let Some(count) = active_by_key.get_mut(&key) {
7884 *count -= 1;
7885 if *count == 0 {
7886 active_by_key.remove(&key);
7887 }
7888 }
7889 completed.insert(index, result);
7890 }
7891 Err(future) => {
7892 tasks.insert(
7893 index,
7894 spawn_completion_task(index, future, sender.clone(), move |result| {
7895 (key, result)
7896 }),
7897 );
7898 }
7899 }
7900 }
7901
7902 if let Some(result) = completed.remove(&next_to_emit) {
7903 next_to_emit += 1;
7904 return Some(result);
7905 }
7906
7907 if tasks.is_empty() {
7908 return None;
7909 }
7910
7911 if let Some((index, (key, result))) = recv_completion(&receiver) {
7912 tasks.remove(&index);
7913 if let Some(count) = active_by_key.get_mut(&key) {
7914 *count -= 1;
7915 if *count == 0 {
7916 active_by_key.remove(&key);
7917 }
7918 }
7919 if index == next_to_emit {
7920 next_to_emit += 1;
7921 return Some(result);
7922 }
7923 completed.insert(index, result);
7924 }
7925 }
7926 }))
7927}
7928
7929fn map_async_partitioned<Out, Key, Next, Partition, F, Fut>(
7930 mut input: BoxStream<Out>,
7931 parallelism: usize,
7932 per_partition: usize,
7933 partition: Arc<Partition>,
7934 stage: Arc<F>,
7935) -> BoxStream<Next>
7936where
7937 Out: Send + 'static,
7938 Key: Clone + Eq + Hash + Send + 'static,
7939 Next: Send + 'static,
7940 Partition: Fn(&Out) -> Key + Send + Sync + 'static,
7941 F: Fn(Out) -> Fut + Send + Sync + 'static,
7942 Fut: Future<Output = StreamResult<Next>> + Send + 'static,
7943{
7944 if parallelism == 1 {
7945 return map_async_partitioned_serial(input, partition, stage);
7946 }
7947 if parallelism <= 4 {
7949 return map_async_partitioned_scanning(input, parallelism, per_partition, partition, stage);
7950 }
7951
7952 let (sender, receiver) = std::sync::mpsc::channel::<(usize, (usize, StreamResult<Next>))>();
7953 let mut tasks = HashMap::<usize, AbortOnDropHandle<()>>::with_capacity(parallelism);
7954 let mut slots_by_key = HashMap::<Key, usize>::with_capacity(parallelism);
7955 let mut slots = Vec::<PartitionSlot<Key, Out>>::with_capacity(parallelism);
7956 let mut free_slots = Vec::<usize>::new();
7957 let mut ready_slots = VecDeque::<usize>::with_capacity(parallelism);
7958 let mut completed = BTreeMap::<usize, StreamResult<Next>>::new();
7959 let mut next_index = 0_usize;
7960 let mut next_to_emit = 0_usize;
7961 let mut input_done = false;
7962 Box::new(std::iter::from_fn(move || {
7963 loop {
7964 if let Some(result) = completed.remove(&next_to_emit) {
7965 next_to_emit += 1;
7966 return Some(result);
7967 }
7968
7969 while tasks.len() + completed.len() < parallelism {
7970 if let Some((index, slot, item)) =
7971 pop_ready_partition_slot(&mut slots, &mut ready_slots, per_partition)
7972 {
7973 match poll_once_or_pending(stage(item)) {
7974 Ok(result) => {
7975 let mut remove_empty = false;
7976 if let Some(state) = slots.get_mut(slot) {
7977 state.active -= 1;
7978 remove_empty = state.active == 0
7979 && state.queued.is_empty()
7980 && !state.in_ready_queue;
7981 }
7982 if remove_empty {
7983 retire_partition_slot(
7984 slot,
7985 &mut slots_by_key,
7986 &mut slots,
7987 &mut free_slots,
7988 );
7989 } else {
7990 ready_partition_slot(
7991 &mut slots,
7992 &mut ready_slots,
7993 slot,
7994 per_partition,
7995 );
7996 }
7997 if index == next_to_emit {
7998 next_to_emit += 1;
7999 return Some(result);
8000 }
8001 completed.insert(index, result);
8002 }
8003 Err(future) => {
8004 tasks.insert(
8005 index,
8006 spawn_completion_task(
8007 index,
8008 future,
8009 sender.clone(),
8010 move |result| (slot, result),
8011 ),
8012 );
8013 }
8014 }
8015 continue;
8016 }
8017
8018 if input_done {
8019 break;
8020 }
8021
8022 match input.next() {
8023 Some(Ok(item)) => {
8024 let key = partition(&item);
8025 let index = next_index;
8026 next_index += 1;
8027 let slot =
8028 partition_slot_for(key, &mut slots_by_key, &mut slots, &mut free_slots);
8029 let state = &mut slots[slot];
8030 if state.active < per_partition {
8031 match poll_once_or_pending(stage(item)) {
8032 Ok(result) => {
8033 if index == next_to_emit {
8034 next_to_emit += 1;
8035 if state.queued.is_empty() && !state.in_ready_queue {
8036 retire_partition_slot(
8037 slot,
8038 &mut slots_by_key,
8039 &mut slots,
8040 &mut free_slots,
8041 );
8042 }
8043 return Some(result);
8044 }
8045 completed.insert(index, result);
8046 if state.queued.is_empty() && !state.in_ready_queue {
8047 retire_partition_slot(
8048 slot,
8049 &mut slots_by_key,
8050 &mut slots,
8051 &mut free_slots,
8052 );
8053 }
8054 }
8055 Err(future) => {
8056 state.active += 1;
8057 tasks.insert(
8058 index,
8059 spawn_completion_task(
8060 index,
8061 future,
8062 sender.clone(),
8063 move |result| (slot, result),
8064 ),
8065 );
8066 }
8067 }
8068 } else {
8069 state.queued.push_back((index, item));
8070 }
8071 }
8072 Some(Err(error)) => {
8073 completed.insert(next_index, Err(error));
8074 next_index += 1;
8075 input_done = true;
8076 break;
8077 }
8078 None => {
8079 input_done = true;
8080 break;
8081 }
8082 }
8083 }
8084
8085 if let Some(result) = completed.remove(&next_to_emit) {
8086 next_to_emit += 1;
8087 return Some(result);
8088 }
8089
8090 if tasks.is_empty() {
8091 return None;
8092 }
8093
8094 if let Some((index, (slot, result))) = recv_completion(&receiver) {
8095 tasks.remove(&index);
8096 let mut remove_empty = false;
8097 if let Some(state) = slots.get_mut(slot) {
8098 state.active -= 1;
8099 remove_empty =
8100 state.active == 0 && state.queued.is_empty() && !state.in_ready_queue;
8101 }
8102 if remove_empty {
8103 retire_partition_slot(slot, &mut slots_by_key, &mut slots, &mut free_slots);
8104 } else {
8105 ready_partition_slot(&mut slots, &mut ready_slots, slot, per_partition);
8106 }
8107 if index == next_to_emit {
8108 next_to_emit += 1;
8109 return Some(result);
8110 }
8111 completed.insert(index, result);
8112 }
8113 }
8114 }))
8115}
8116
8117#[cfg(test)]
8118mod materialized_factory_tests {
8119 use super::*;
8120 use crate::Source;
8121 use std::panic::{AssertUnwindSafe, catch_unwind};
8122 use std::sync::atomic::{AtomicUsize, Ordering};
8123
8124 #[test]
8125 fn materialized_factory_panic_does_not_poison_reused_blueprint() {
8126 let calls = Arc::new(AtomicUsize::new(0));
8127 let flow = Flow::<i32, i32, usize>::from_materialized_factory({
8128 let calls = Arc::clone(&calls);
8129 move || {
8130 if calls.fetch_add(1, Ordering::SeqCst) == 0 {
8131 panic!("factory panics once");
8132 }
8133 let transform: PureTransform<i32, i32> = Arc::new(|input| input);
8134 (transform, 1)
8135 }
8136 });
8137
8138 let first = catch_unwind(AssertUnwindSafe(|| {
8139 let _ = Source::single(1_i32)
8140 .via(flow.clone())
8141 .run_collect()
8142 .expect("first materialization starts");
8143 }));
8144 assert!(first.is_err());
8145
8146 let values = Source::single(2_i32)
8147 .via(flow)
8148 .run_collect()
8149 .expect("second materialization succeeds");
8150 assert_eq!(values, vec![2]);
8151 }
8152}
8153
8154#[cfg(test)]
8155mod flat_map_merge_ready_ring_tests {
8156 use super::*;
8157 use std::sync::mpsc;
8158 use std::time::Duration;
8159
8160 fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
8161 let mut v = source.run_collect().unwrap();
8162 v.sort_unstable();
8163 v
8164 }
8165
8166 #[test]
8167 fn ready_ring_empty_upstream() {
8168 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8169 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8170 });
8171 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8172 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8173 });
8174 assert_eq!(legacy, ring);
8175 assert!(ring.is_empty());
8176 }
8177
8178 #[test]
8179 fn ready_ring_single_lane() {
8180 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8181 run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
8182 });
8183 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8184 run_sorted(crate::Source::single(42_i32).flat_map_merge(4, crate::Source::single))
8185 });
8186 assert_eq!(legacy, ring);
8187 assert_eq!(ring, vec![42]);
8188 }
8189
8190 #[test]
8191 fn ready_ring_breadth_one_exact_order() {
8192 let make = || {
8194 crate::Source::from_iter(0_i32..5)
8195 .flat_map_merge(1, |x| crate::Source::single(x * 10))
8196 .run_collect()
8197 .unwrap()
8198 };
8199 let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8200 let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8201 legacy.sort_unstable();
8202 ring.sort_unstable();
8203 assert_eq!(legacy, ring);
8204 }
8205
8206 #[test]
8207 fn ready_ring_breadth_gt_input() {
8208 let make = || {
8209 run_sorted(
8210 crate::Source::from_iter(0_i32..3)
8211 .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2])),
8212 )
8213 };
8214 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8215 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8216 assert_eq!(legacy, ring);
8217 assert_eq!(ring.len(), 9);
8218 }
8219
8220 #[test]
8221 fn ready_ring_mixed_short_long() {
8222 let make = || {
8223 run_sorted(crate::Source::from_iter(0_i32..8).flat_map_merge(4, |x| {
8224 if x % 3 == 0 {
8225 crate::Source::from_iter(0..20_i32).map(move |i| x * 100 + i)
8226 } else {
8227 crate::Source::single(x)
8228 }
8229 }))
8230 };
8231 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8232 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8233 assert_eq!(legacy, ring);
8234 }
8235
8236 #[test]
8237 fn ready_ring_respects_breadth_bound() {
8238 use std::sync::atomic::{AtomicUsize, Ordering as Ord};
8239 let active = Arc::new(AtomicUsize::new(0));
8240 let max_active = Arc::new(AtomicUsize::new(0));
8241 let a2 = Arc::clone(&active);
8242 let m2 = Arc::clone(&max_active);
8243 let mut values = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8244 crate::Source::from_iter(0_i32..6)
8245 .flat_map_merge(2, move |item| {
8246 let a = Arc::clone(&a2);
8247 let m = Arc::clone(&m2);
8248 crate::Source::future(move || {
8249 let a = Arc::clone(&a);
8250 let m = Arc::clone(&m);
8251 async move {
8252 let now = a.fetch_add(1, Ord::SeqCst) + 1;
8253 let mut seen = m.load(Ord::SeqCst);
8254 while now > seen {
8255 match m.compare_exchange(seen, now, Ord::SeqCst, Ord::SeqCst) {
8256 Ok(_) => break,
8257 Err(v) => seen = v,
8258 }
8259 }
8260 thread::sleep(Duration::from_millis(20));
8261 a.fetch_sub(1, Ord::SeqCst);
8262 Ok(item)
8263 }
8264 })
8265 })
8266 .run_collect()
8267 .unwrap()
8268 });
8269 values.sort_unstable();
8270 assert_eq!(values, vec![0, 1, 2, 3, 4, 5]);
8271 assert!(max_active.load(std::sync::atomic::Ordering::SeqCst) <= 2);
8272 }
8273
8274 #[test]
8275 fn ready_ring_fairness_slow_lane_not_starved() {
8276 use std::sync::atomic::{AtomicBool, Ordering as Ord};
8277 let slow_emitted = Arc::new(AtomicBool::new(false));
8278 let slow_flag = Arc::clone(&slow_emitted);
8279 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8280 crate::Source::from_iter(0_i32..3)
8281 .flat_map_merge(4, move |lane_id| {
8282 let flag = Arc::clone(&slow_flag);
8283 match lane_id {
8284 0 => crate::Source::from_iter(0_i32..50),
8285 1 => crate::Source::from_iter(100_i32..150),
8286 _ => crate::Source::future(move || {
8287 let flag = Arc::clone(&flag);
8288 async move {
8289 thread::sleep(Duration::from_millis(10));
8290 flag.store(true, Ord::SeqCst);
8291 Ok(999_i32)
8292 }
8293 }),
8294 }
8295 })
8296 .run_collect()
8297 .unwrap()
8298 });
8299 assert!(slow_emitted.load(std::sync::atomic::Ordering::SeqCst));
8300 assert!(results.contains(&999));
8301 assert_eq!(results.len(), 101);
8302 }
8303
8304 #[test]
8305 fn ready_ring_inner_failure_no_hang() {
8306 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8307 crate::Source::from_iter(0_i32..8)
8308 .flat_map_merge(4, |x| {
8309 if x == 3 {
8310 crate::Source::failed(StreamError::Failed("lane-fail".into()))
8311 } else {
8312 crate::Source::from_iter(0_i32..10)
8313 }
8314 })
8315 .run_collect()
8316 });
8317 assert_eq!(result, Err(StreamError::Failed("lane-fail".into())));
8318 }
8319
8320 #[test]
8321 fn ready_ring_factory_failure_propagates() {
8322 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8323 crate::Source::from_iter(0_i32..4)
8324 .flat_map_merge(2, |x| {
8325 if x == 2 {
8326 crate::Source::failed(StreamError::Failed("factory-fail".into()))
8327 } else {
8328 crate::Source::single(x)
8329 }
8330 })
8331 .run_collect()
8332 });
8333 assert!(result.is_err());
8334 }
8335
8336 #[test]
8337 fn ready_ring_closure_not_under_coordinator_lock() {
8338 let guard_mutex = Arc::new(std::sync::Mutex::<()>::new(()));
8339 let gm = Arc::clone(&guard_mutex);
8340 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8341 crate::Source::from_iter(0_i32..10)
8342 .flat_map_merge(4, move |x| {
8343 let _lock = gm.lock().unwrap();
8344 crate::Source::single(x)
8345 })
8346 .run_collect()
8347 .unwrap()
8348 });
8349 assert_eq!(results.len(), 10);
8350 }
8351
8352 #[test]
8359 fn ready_ring_bounded_memory_producer_blocks_at_window() {
8360 const WINDOW: usize = FLAT_MAP_MERGE_SUBSTREAM_WINDOW;
8361 const EXTRA: usize = 4;
8362 let (gate_tx, gate_rx) = mpsc::channel::<()>();
8365 let gate_tx = Arc::new(std::sync::Mutex::new(gate_tx));
8366 let gate_tx2 = Arc::clone(&gate_tx);
8367 let produced = Arc::new(std::sync::atomic::AtomicUsize::new(0));
8368 let prod2 = Arc::clone(&produced);
8369
8370 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8371 crate::Source::single(0_i32)
8372 .flat_map_merge(2, move |_| {
8373 let tx = Arc::clone(&gate_tx2);
8374 let prod = Arc::clone(&prod2);
8375 crate::Source::from_factory(move || {
8376 let tx = Arc::clone(&tx);
8377 let prod = Arc::clone(&prod);
8378 let mut i = 0_i32;
8379 Box::new(std::iter::from_fn(move || {
8380 if i as usize >= WINDOW + EXTRA {
8381 return None;
8382 }
8383 if i as usize == WINDOW {
8384 let _ = tx.lock().unwrap().send(());
8386 }
8387 prod.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
8388 i += 1;
8389 Some(Ok(i))
8390 }))
8391 })
8392 })
8393 .run_with(crate::Sink::queue())
8394 .unwrap()
8395 });
8396
8397 let mut total = 0;
8399 while queue.pull().unwrap().is_some() {
8400 total += 1;
8401 }
8402 let signal = gate_rx.recv_timeout(Duration::from_secs(1));
8404 assert!(signal.is_ok(), "producer never reached the window boundary");
8405 assert_eq!(total, WINDOW + EXTRA);
8406 }
8407
8408 #[test]
8409 fn ready_ring_cancellation_wakes_blocked_lanes() {
8410 let rt = crate::stream::runtime::Runtime::new();
8411 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8412 crate::Source::from_iter(0_i32..4)
8413 .flat_map_merge(4, |_| crate::Source::repeat(1_i32))
8414 .run_with_materializer(crate::Sink::queue(), &rt)
8415 .unwrap()
8416 });
8417 for _ in 0..8 {
8418 let _ = queue.pull();
8419 }
8420 drop(queue);
8421 rt.shutdown();
8422 }
8423
8424 #[test]
8425 fn ready_ring_lost_wakeup_stress() {
8426 for _ in 0..20 {
8427 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8428 crate::Source::from_iter(0_i32..50)
8429 .flat_map_merge(8, |item| {
8430 crate::Source::from_iter([item, item + 1, item + 2])
8431 })
8432 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8433 .unwrap()
8434 .wait()
8435 });
8436 assert_eq!(result, Ok(3825), "lost-wakeup stress: wrong sum");
8437 }
8438 }
8439
8440 #[test]
8441 fn ready_ring_concurrent_streams_lost_wakeup_stress() {
8442 const STREAMS: usize = 32;
8443 const ROUNDS: usize = 8;
8444 const EXPECTED: i64 = 998_080;
8445
8446 for _ in 0..ROUNDS {
8447 let barrier = Arc::new(std::sync::Barrier::new(STREAMS));
8448 let mut handles = Vec::with_capacity(STREAMS);
8449 for _ in 0..STREAMS {
8450 let barrier = Arc::clone(&barrier);
8451 handles.push(thread::spawn(move || {
8452 barrier.wait();
8453 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8454 crate::Source::from_iter(0_i64..32)
8455 .flat_map_merge(8, |item| {
8456 crate::Source::from_iter(item * 100..item * 100 + 20)
8457 })
8458 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8459 .unwrap()
8460 .wait()
8461 });
8462 assert_eq!(result, Ok(EXPECTED), "concurrent ready-ring sum");
8463 }));
8464 }
8465
8466 for handle in handles {
8467 handle.join().expect("ready-ring stress worker panicked");
8468 }
8469 }
8470 }
8471
8472 #[test]
8473 fn ready_ring_tail_loop_stress() {
8474 for _ in 0..20 {
8475 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8476 crate::Source::from_iter(0_i64..100)
8477 .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8478 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8479 .unwrap()
8480 .wait()
8481 });
8482 assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8483 }
8484 }
8485
8486 #[test]
8487 fn ready_ring_auto_mode_matches_ring() {
8488 let make = || {
8489 run_sorted(
8490 crate::Source::from_iter(0_i32..10)
8491 .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 100])),
8492 )
8493 };
8494 let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8495 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8496 assert_eq!(auto, ring);
8497 }
8498}
8499
8500#[cfg(test)]
8502mod inline_micro_source_tests {
8503 use super::*;
8504 use crate::stream::source::test_source_with_inline_micro_hint;
8505 use std::sync::mpsc;
8506
8507 fn run_sorted<T: Ord + Send + 'static>(source: crate::Source<T>) -> Vec<T> {
8508 let mut v = source.run_collect().unwrap();
8509 v.sort_unstable();
8510 v
8511 }
8512
8513 #[test]
8516 fn inline_empty_upstream() {
8517 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8518 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8519 });
8520 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8521 run_sorted(crate::Source::<i32>::empty().flat_map_merge(4, crate::Source::single))
8522 });
8523 assert_eq!(legacy, ring);
8524 assert!(ring.is_empty());
8525 }
8526
8527 #[test]
8528 fn inline_single_inner_source() {
8529 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, || {
8530 run_sorted(
8531 crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8532 )
8533 });
8534 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8535 run_sorted(
8536 crate::Source::single(99_i32).flat_map_merge(4, |x| crate::Source::single(x * 2)),
8537 )
8538 });
8539 assert_eq!(legacy, ring);
8540 assert_eq!(ring, vec![198]);
8541 }
8542
8543 #[test]
8544 fn inline_breadth_one_exact_order() {
8545 let make = || {
8547 crate::Source::from_iter(0_i32..6)
8548 .flat_map_merge(1, |x| {
8549 crate::Source::from_iter([x * 10, x * 10 + 1, x * 10 + 2, x * 10 + 3])
8550 })
8551 .run_collect()
8552 .unwrap()
8553 };
8554 let mut legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8555 let mut ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8556 legacy.sort_unstable();
8557 ring.sort_unstable();
8558 assert_eq!(legacy, ring);
8559 }
8560
8561 #[test]
8562 fn inline_breadth_gt_input() {
8563 let make = || {
8564 run_sorted(
8565 crate::Source::from_iter(0_i32..3)
8566 .flat_map_merge(100, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8567 )
8568 };
8569 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8570 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8571 assert_eq!(legacy, ring);
8572 assert_eq!(ring.len(), 12);
8573 }
8574
8575 #[test]
8576 fn inline_mixed_empty_single_four_item() {
8577 let make = || {
8578 run_sorted(
8579 crate::Source::from_iter(0_i32..12).flat_map_merge(4, |x| match x % 3 {
8580 0 => crate::Source::empty(),
8581 1 => crate::Source::single(x),
8582 _ => crate::Source::from_iter([x, x + 100, x + 200, x + 300]),
8583 }),
8584 )
8585 };
8586 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8587 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8588 assert_eq!(legacy, ring);
8589 }
8590
8591 #[test]
8592 fn inline_2k_x4_b8_benchmark_shape() {
8593 let make = || {
8594 run_sorted(
8595 crate::Source::from_iter(0_i32..2_000).flat_map_merge(8, |item| {
8596 crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8597 }),
8598 )
8599 };
8600 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8601 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8602 assert_eq!(legacy, ring);
8603 assert_eq!(ring.len(), 8_000);
8604 }
8605
8606 #[test]
8609 fn inline_lost_wakeup_stress() {
8610 for _ in 0..20 {
8611 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8612 crate::Source::from_iter(0_i32..50)
8613 .flat_map_merge(8, |item| {
8614 crate::Source::from_iter([item, item + 1, item + 2, item + 3])
8615 })
8616 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v as i64))
8617 .unwrap()
8618 .wait()
8619 });
8620 assert_eq!(result, Ok(5200), "lost-wakeup stress: wrong sum");
8624 }
8625 }
8626
8627 #[test]
8628 fn inline_tail_loop_stress() {
8629 for _ in 0..20 {
8630 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8631 crate::Source::from_iter(0_i64..100)
8632 .flat_map_merge(16, |item| crate::Source::from_iter([item, item + 1000]))
8633 .run_with(crate::Sink::fold(0i64, |acc, v| acc + v))
8634 .unwrap()
8635 .wait()
8636 });
8637 assert_eq!(result, Ok(109_900), "tail-loop stress: wrong sum");
8639 }
8640 }
8641
8642 #[test]
8645 fn inline_large_source_uses_worker_fallback() {
8646 let make = || {
8649 run_sorted(crate::Source::from_iter(0_i32..4).flat_map_merge(2, |x| {
8650 crate::Source::from_iter(0_i32..20).map(move |i| x * 100 + i)
8651 }))
8652 };
8653 let legacy = with_substream_mode(SubstreamExecutorMode::LegacyOnly, make);
8654 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8655 assert_eq!(legacy, ring);
8656 assert_eq!(ring.len(), 80);
8657 }
8658
8659 #[test]
8662 fn inline_inner_error_before_any_item() {
8663 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8664 crate::Source::from_iter(0_i32..4)
8665 .flat_map_merge(2, |x| {
8666 if x == 1 {
8667 crate::Source::failed(StreamError::Failed("inline-err".into()))
8668 } else {
8669 crate::Source::from_iter([x, x + 1])
8670 }
8671 })
8672 .run_collect()
8673 });
8674 assert_eq!(result, Err(StreamError::Failed("inline-err".into())));
8675 }
8676
8677 #[test]
8678 fn inline_inner_error_after_some_items() {
8679 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8683 crate::Source::single(0_i32)
8684 .flat_map_merge(2, |_x| {
8685 test_source_with_inline_micro_hint(
8686 || {
8687 let mut count = 0;
8688 Box::new(std::iter::from_fn(move || {
8689 count += 1;
8690 match count {
8691 1 => Some(Ok(42_i32)),
8692 2 => Some(Err(StreamError::Failed("after-items".into()))),
8693 _ => None,
8694 }
8695 }))
8696 },
8697 1, )
8699 })
8700 .run_collect()
8701 });
8702 assert!(result.is_err());
8706 }
8707
8708 #[test]
8709 fn inline_worker_lane_fails_during_inline_drain() {
8710 let result = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8713 crate::Source::from_iter(0_i32..10)
8714 .flat_map_merge(2, |x| {
8715 if x == 5 {
8716 crate::Source::failed(StreamError::Failed("worker-fail".into()))
8717 } else if x % 2 == 0 {
8718 crate::Source::from_iter([x, x + 1, x + 2])
8720 } else {
8721 crate::Source::from_iter(0_i32..40).map(move |i| x * 100 + i)
8723 }
8724 })
8725 .run_collect()
8726 });
8727 assert!(result.is_err());
8728 }
8729
8730 #[test]
8733 fn inline_cancellation_drop_output() {
8734 let rt = crate::stream::runtime::Runtime::new();
8737 let queue = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8738 crate::Source::from_iter(0_i32..100)
8739 .flat_map_merge(8, |x| {
8740 if x % 3 == 0 {
8741 crate::Source::from_iter([x, x + 1, x + 2, x + 3]) } else {
8743 crate::Source::repeat(x) }
8745 })
8746 .run_with_materializer(crate::Sink::queue(), &rt)
8747 .unwrap()
8748 });
8749 for _ in 0..16 {
8751 let _ = queue.pull();
8752 }
8753 drop(queue);
8754 rt.shutdown();
8755 }
8756
8757 #[test]
8760 fn inline_next_not_under_coordinator_lock() {
8761 let (tx, rx) = mpsc::channel::<()>();
8769 let tx = Arc::new(std::sync::Mutex::new(tx));
8770 let tx2 = Arc::clone(&tx);
8771
8772 let results = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, || {
8773 crate::Source::from_iter(0_i32..2)
8774 .flat_map_merge(2, move |x| {
8775 let sig = Arc::clone(&tx2);
8776 if x == 0 {
8777 crate::Source::from_iter(0_i32..40)
8779 } else {
8780 test_source_with_inline_micro_hint(
8783 move || {
8784 let sig = Arc::clone(&sig);
8785 let mut emitted = false;
8786 Box::new(std::iter::from_fn(move || {
8787 if !emitted {
8788 emitted = true;
8789 let _ = sig.lock().unwrap().send(());
8790 Some(Ok(999_i32))
8791 } else {
8792 None
8793 }
8794 }))
8795 },
8796 1,
8797 )
8798 }
8799 })
8800 .run_collect()
8801 .unwrap()
8802 });
8803
8804 let signal = rx.recv_timeout(std::time::Duration::from_secs(5));
8806 assert!(signal.is_ok(), "inline next() never ran");
8807 assert!(results.contains(&999));
8808 assert_eq!(results.len(), 41);
8809 }
8810
8811 #[test]
8814 fn inline_auto_matches_readyring() {
8815 let make = || {
8816 run_sorted(
8817 crate::Source::from_iter(0_i32..20)
8818 .flat_map_merge(4, |x| crate::Source::from_iter([x, x + 1, x + 2, x + 3])),
8819 )
8820 };
8821 let auto = with_substream_mode(SubstreamExecutorMode::Auto, make);
8822 let ring = with_substream_mode(SubstreamExecutorMode::ReadyRingOnly, make);
8823 assert_eq!(auto, ring);
8824 }
8825}
8826
8827#[cfg(test)]
8829mod split_sink_fast_path_tests {
8830 use super::*;
8831
8832 fn run_split_fold(
8834 items: Vec<u64>,
8835 split_mode: SplitMode,
8836 executor_mode: SubstreamExecutorMode,
8837 ) -> Vec<u64> {
8838 with_substream_mode(executor_mode, || {
8839 let source = match split_mode {
8840 SplitMode::When => crate::Source::from_iter(items).split_when(|x| x % 10 == 0),
8841 SplitMode::After => crate::Source::from_iter(items).split_after(|x| x % 10 == 0),
8842 };
8843 source
8844 .run_with(crate::Sink::fold(
8845 Vec::new(),
8846 |mut acc, seg: crate::Source<u64>| {
8847 let sum = seg
8848 .run_with(crate::Sink::fold(0u64, |a, x| a + x))
8849 .unwrap()
8850 .wait()
8851 .unwrap();
8852 acc.push(sum);
8853 acc
8854 },
8855 ))
8856 .unwrap()
8857 .wait()
8858 .unwrap()
8859 })
8860 }
8861
8862 fn run_split_collect_segments(
8864 items: Vec<u64>,
8865 split_mode: SplitMode,
8866 executor_mode: SubstreamExecutorMode,
8867 ) -> Vec<Vec<u64>> {
8868 with_substream_mode(executor_mode, || {
8869 let source = match split_mode {
8870 SplitMode::When => crate::Source::from_iter(items).split_when(move |x| x % 10 == 0),
8871 SplitMode::After => {
8872 crate::Source::from_iter(items).split_after(move |x| x % 10 == 0)
8873 }
8874 };
8875 source
8876 .run_with(crate::Sink::fold(
8877 Vec::new(),
8878 |mut acc, seg: crate::Source<u64>| {
8879 let v = seg
8880 .run_with(crate::Sink::collect())
8881 .unwrap()
8882 .wait()
8883 .unwrap();
8884 acc.push(v);
8885 acc
8886 },
8887 ))
8888 .unwrap()
8889 .wait()
8890 .unwrap()
8891 })
8892 }
8893
8894 #[test]
8897 fn split_fast_equivalence_empty_input() {
8898 for sm in [SplitMode::When, SplitMode::After] {
8899 let legacy = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::LegacyOnly);
8900 let fast = run_split_collect_segments(vec![], sm, SubstreamExecutorMode::SplitSinkOnly);
8901 assert_eq!(legacy, fast, "empty input, mode {sm:?}");
8902 }
8903 }
8904
8905 #[test]
8906 fn split_fast_equivalence_no_boundaries() {
8907 let items: Vec<u64> = (1..=9).collect();
8908 for sm in [SplitMode::When, SplitMode::After] {
8909 let legacy =
8910 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8911 let fast =
8912 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8913 assert_eq!(legacy, fast, "no boundaries, mode {sm:?}");
8914 }
8915 }
8916
8917 #[test]
8918 fn split_fast_equivalence_first_element_boundary_when() {
8919 let items: Vec<u64> = vec![10, 1, 2, 3];
8921 let legacy = run_split_collect_segments(
8922 items.clone(),
8923 SplitMode::When,
8924 SubstreamExecutorMode::LegacyOnly,
8925 );
8926 let fast = run_split_collect_segments(
8927 items,
8928 SplitMode::When,
8929 SubstreamExecutorMode::SplitSinkOnly,
8930 );
8931 assert_eq!(legacy, fast);
8932 }
8933
8934 #[test]
8935 fn split_fast_equivalence_first_element_boundary_after() {
8936 let items: Vec<u64> = vec![10, 1, 2, 3];
8938 let legacy = run_split_collect_segments(
8939 items.clone(),
8940 SplitMode::After,
8941 SubstreamExecutorMode::LegacyOnly,
8942 );
8943 let fast = run_split_collect_segments(
8944 items,
8945 SplitMode::After,
8946 SubstreamExecutorMode::SplitSinkOnly,
8947 );
8948 assert_eq!(legacy, fast);
8949 }
8950
8951 #[test]
8952 fn split_fast_equivalence_consecutive_matches() {
8953 let items: Vec<u64> = vec![10, 20, 30, 1, 2, 40];
8954 for sm in [SplitMode::When, SplitMode::After] {
8955 let legacy =
8956 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8957 let fast =
8958 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8959 assert_eq!(legacy, fast, "consecutive matches, mode {sm:?}");
8960 }
8961 }
8962
8963 #[test]
8964 fn split_fast_equivalence_last_element_boundary() {
8965 let items: Vec<u64> = vec![1, 2, 3, 10];
8966 for sm in [SplitMode::When, SplitMode::After] {
8967 let legacy =
8968 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8969 let fast =
8970 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8971 assert_eq!(legacy, fast, "last element boundary, mode {sm:?}");
8972 }
8973 }
8974
8975 #[test]
8976 fn split_fast_equivalence_mixed() {
8977 let items: Vec<u64> = (0..50u64).collect();
8978 for sm in [SplitMode::When, SplitMode::After] {
8979 let legacy =
8980 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8981 let fast =
8982 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8983 assert_eq!(legacy, fast, "mixed 0..50, mode {sm:?}");
8984 }
8985 }
8986
8987 #[test]
8988 fn split_fast_equivalence_fold_sums() {
8989 let items: Vec<u64> = (0..50u64).collect();
8990 for sm in [SplitMode::When, SplitMode::After] {
8991 let legacy = run_split_fold(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
8992 let fast = run_split_fold(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
8993 assert_eq!(legacy, fast, "fold sums, mode {sm:?}");
8994 }
8995 }
8996
8997 #[test]
8998 fn split_fast_equivalence_with_collect() {
8999 let items: Vec<u64> = (0..312u64).collect();
9001 for sm in [SplitMode::When, SplitMode::After] {
9002 let legacy =
9003 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::LegacyOnly);
9004 let fast =
9005 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
9006 assert_eq!(legacy, fast, "collect 312 items, mode {sm:?}");
9007 }
9008 }
9009
9010 #[test]
9013 fn split_fast_fold_result_equivalence() {
9014 let items: Vec<u64> = (0..50u64).collect();
9015 let run = |executor_mode| {
9016 with_substream_mode(executor_mode, || {
9017 crate::Source::from_iter(items.clone())
9018 .split_when(move |x| x % 10 == 0)
9019 .run_with(crate::Sink::fold(
9020 Vec::new(),
9021 |mut acc, seg: crate::Source<u64>| {
9022 let sum = seg
9023 .run_with(crate::Sink::fold_result(0u64, |a, x| Ok(a + x)))
9024 .unwrap()
9025 .wait()
9026 .unwrap();
9027 acc.push(sum);
9028 acc
9029 },
9030 ))
9031 .unwrap()
9032 .wait()
9033 .unwrap()
9034 })
9035 };
9036 assert_eq!(
9037 run(SubstreamExecutorMode::LegacyOnly),
9038 run(SubstreamExecutorMode::SplitSinkOnly)
9039 );
9040 }
9041
9042 #[test]
9045 fn split_fast_ignore_equivalence() {
9046 let items: Vec<u64> = (0..50u64).collect();
9047 let run = |executor_mode| {
9048 with_substream_mode(executor_mode, || {
9049 crate::Source::from_iter(items.clone())
9050 .split_when(move |x| x % 10 == 0)
9051 .run_with(crate::Sink::fold(0u64, |count, seg: crate::Source<u64>| {
9052 seg.run_with(crate::Sink::ignore()).unwrap().wait().unwrap();
9053 count + 1
9054 }))
9055 .unwrap()
9056 .wait()
9057 .unwrap()
9058 })
9059 };
9060 let legacy = run(SubstreamExecutorMode::LegacyOnly);
9061 let fast = run(SubstreamExecutorMode::SplitSinkOnly);
9062 assert_eq!(legacy, fast, "ignore segment counts must match");
9063 }
9064
9065 #[test]
9068 fn split_fast_one_shot_cannot_materialize_twice() {
9069 with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9070 let materializer = crate::Runtime::default();
9071 let result = crate::Source::from_iter(1u64..=5)
9072 .split_when(|x| x % 3 == 0)
9073 .run_with(crate::Sink::fold(0u64, |_, seg: crate::Source<u64>| {
9074 let c1 = seg.clone().run_with(crate::Sink::fold(0u64, |a, x| a + x));
9076 let c2 = seg.run_with(crate::Sink::fold(0u64, |a, x| a + x));
9078 assert!(c1.is_ok(), "first materialization should succeed");
9079 assert!(c2.is_err(), "second materialization should fail: {c2:?}");
9080 let _ = c1.unwrap().wait();
9081 0u64
9082 }));
9083 let _ = result;
9084 let _ = &materializer;
9085 });
9086 }
9087
9088 #[test]
9091 fn split_fast_predicate_panic_both_modes() {
9092 for sm in [SplitMode::When, SplitMode::After] {
9095 let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9096 let source = match sm {
9097 SplitMode::When => crate::Source::from_iter(0u64..10).split_when(|x| {
9098 if *x == 5 {
9099 panic!("test panic");
9100 }
9101 x % 3 == 0
9102 }),
9103 SplitMode::After => crate::Source::from_iter(0u64..10).split_after(|x| {
9104 if *x == 5 {
9105 panic!("test panic");
9106 }
9107 x % 3 == 0
9108 }),
9109 };
9110 source
9111 .run_with(crate::Sink::fold(
9112 Vec::<u64>::new(),
9113 |mut acc, seg: crate::Source<u64>| {
9114 let completion = seg.run_with(crate::Sink::ignore());
9116 if let Ok(c) = completion {
9117 let _ = c.wait();
9118 }
9119 acc.push(0u64);
9120 acc
9121 },
9122 ))
9123 .map(|c| c.wait())
9124 });
9125 let _ = result;
9127 }
9128 }
9129
9130 #[test]
9133 fn split_fast_stress_20x() {
9134 for i in 0..20 {
9135 let items: Vec<u64> = (0..10_000u64).collect();
9136 for sm in [SplitMode::When, SplitMode::After] {
9137 let fast = run_split_collect_segments(
9138 items.clone(),
9139 sm,
9140 SubstreamExecutorMode::SplitSinkOnly,
9141 );
9142 let legacy = run_split_collect_segments(
9143 items.clone(),
9144 sm,
9145 SubstreamExecutorMode::LegacyOnly,
9146 );
9147 assert_eq!(
9148 fast.len(),
9149 legacy.len(),
9150 "stress run {i} segment count mismatch, mode {sm:?}"
9151 );
9152 assert_eq!(
9153 fast.iter().flatten().sum::<u64>(),
9154 legacy.iter().flatten().sum::<u64>(),
9155 "stress run {i} sum mismatch, mode {sm:?}"
9156 );
9157 }
9158 }
9159 }
9160
9161 #[test]
9164 fn split_fast_auto_mode_matches_fast() {
9165 let items: Vec<u64> = (0..50u64).collect();
9166 for sm in [SplitMode::When, SplitMode::After] {
9167 let auto = run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::Auto);
9168 let fast =
9169 run_split_collect_segments(items.clone(), sm, SubstreamExecutorMode::SplitSinkOnly);
9170 assert_eq!(auto, fast, "auto == fast, mode {sm:?}");
9171 }
9172 }
9173
9174 #[test]
9177 fn split_fast_fallback_path_via_foreach() {
9178 use std::sync::atomic::{AtomicU64, Ordering as Ord};
9182 let total = Arc::new(AtomicU64::new(0));
9183 let total2 = Arc::clone(&total);
9184 let result = with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9185 crate::Source::from_iter(0u64..30)
9186 .split_when(|x| x % 10 == 0)
9187 .run_with(crate::Sink::fold(
9188 Vec::new(),
9189 move |mut acc, seg: crate::Source<u64>| {
9190 let t = Arc::clone(&total2);
9191 seg.run_with(crate::Sink::foreach(move |x| {
9195 t.fetch_add(x, Ord::SeqCst);
9196 }))
9197 .unwrap()
9198 .wait()
9199 .unwrap();
9200 acc.push(1u64);
9201 acc
9202 },
9203 ))
9204 .unwrap()
9205 .wait()
9206 .unwrap()
9207 });
9208 assert_eq!(result.len(), 3, "should have 3 segments");
9209 assert_eq!(
9211 total.load(std::sync::atomic::Ordering::SeqCst),
9212 (0..30u64).sum::<u64>()
9213 );
9214 }
9215
9216 #[test]
9219 fn split_fast_liveness_segment_count_when() {
9220 let items: Vec<u64> = (0..30u64).collect();
9221 let fast = run_split_collect_segments(
9222 items.clone(),
9223 SplitMode::When,
9224 SubstreamExecutorMode::SplitSinkOnly,
9225 );
9226 let legacy =
9227 run_split_collect_segments(items, SplitMode::When, SubstreamExecutorMode::LegacyOnly);
9228 assert_eq!(fast.len(), legacy.len());
9229 }
9230
9231 #[test]
9232 fn split_fast_liveness_segment_count_after() {
9233 let items: Vec<u64> = (0..30u64).collect();
9234 let fast = run_split_collect_segments(
9235 items.clone(),
9236 SplitMode::After,
9237 SubstreamExecutorMode::SplitSinkOnly,
9238 );
9239 let legacy =
9240 run_split_collect_segments(items, SplitMode::After, SubstreamExecutorMode::LegacyOnly);
9241 assert_eq!(fast.len(), legacy.len());
9242 }
9243
9244 #[test]
9251 fn split_fast_bounded_memory_rendezvous() {
9252 use std::sync::{
9253 atomic::{AtomicBool, AtomicUsize, Ordering},
9254 mpsc,
9255 };
9256 use std::time::{Duration, Instant};
9257
9258 const CAPACITY: usize = LIVE_SUBSTREAM_CAPACITY;
9259 const BATCH: usize = LIVE_SUBSTREAM_BATCH;
9260 const TOTAL: usize = CAPACITY * 2;
9261 const MAX_IN_FLIGHT: usize = CAPACITY + BATCH;
9264
9265 let produced = Arc::new(AtomicUsize::new(0));
9266 let consumed = Arc::new(AtomicUsize::new(0));
9267 let bound_violated = Arc::new(AtomicBool::new(false));
9268
9269 let (item_tx, item_rx) = mpsc::channel::<u64>();
9271
9272 let prod_for_factory = Arc::clone(&produced);
9273 let prod_for_fold = Arc::clone(&produced);
9274 let cons_for_fold = Arc::clone(&consumed);
9275 let bv_for_fold = Arc::clone(&bound_violated);
9276
9277 let join = std::thread::spawn(move || {
9278 with_substream_mode(SubstreamExecutorMode::SplitSinkOnly, || {
9279 crate::Source::from_factory(move || {
9280 let prod = Arc::clone(&prod_for_factory);
9281 let mut i = 0u64;
9282 Box::new(std::iter::from_fn(move || {
9283 if i as usize >= TOTAL {
9284 return None;
9285 }
9286 prod.fetch_add(1, Ordering::SeqCst);
9287 let val = i;
9288 i += 1;
9289 Some(Ok(val))
9290 }))
9291 })
9292 .split_when(|_| false)
9294 .run_with(crate::Sink::fold(
9295 0usize,
9296 move |count, seg: crate::Source<u64>| {
9297 let cons = Arc::clone(&cons_for_fold);
9298 let bv = Arc::clone(&bv_for_fold);
9299 let prod = Arc::clone(&prod_for_fold);
9300 let itx = item_tx.clone();
9302 seg.run_with(crate::Sink::foreach(move |x: u64| {
9305 let c = cons.fetch_add(1, Ordering::SeqCst) + 1;
9306 let p = prod.load(Ordering::SeqCst);
9307 if p > c + MAX_IN_FLIGHT {
9308 bv.store(true, Ordering::SeqCst);
9309 }
9310 let _ = itx.send(x);
9311 }))
9312 .unwrap()
9313 .wait()
9314 .unwrap();
9315 count + 1
9316 },
9317 ))
9318 .unwrap()
9319 .wait()
9320 .unwrap()
9321 })
9322 });
9323
9324 let mut received = Vec::with_capacity(TOTAL);
9329 let timeout = Duration::from_secs(60);
9330 let deadline = Instant::now() + timeout;
9331 for i in 0..TOTAL {
9332 let remaining = deadline.saturating_duration_since(Instant::now());
9333 if remaining == Duration::ZERO {
9334 panic!(
9335 "deadlock: received {} of {TOTAL} items within {timeout:?}",
9336 received.len()
9337 );
9338 }
9339 match item_rx.recv_timeout(remaining) {
9340 Ok(item) => received.push(item),
9341 Err(mpsc::RecvTimeoutError::Timeout) => {
9342 panic!(
9343 "deadlock: no item {i} before {timeout:?} rendezvous deadline; received {} of {TOTAL}",
9344 received.len()
9345 )
9346 }
9347 Err(mpsc::RecvTimeoutError::Disconnected) => {
9348 panic!("stream ended early at item {i}")
9349 }
9350 }
9351 }
9352
9353 let seg_count = join.join().expect("stream thread panicked");
9354
9355 assert!(
9357 !bound_violated.load(Ordering::SeqCst),
9358 "bound violated: producer ran >MAX_IN_FLIGHT={MAX_IN_FLIGHT} ahead of consumer"
9359 );
9360
9361 assert_eq!(seg_count, 1, "expected exactly 1 segment");
9363 assert_eq!(received.len(), TOTAL, "not all items received");
9364 let expected: Vec<u64> = (0..TOTAL as u64).collect();
9365 assert_eq!(received, expected, "items not in correct order");
9366 }
9367}