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