1use super::chain::{
71 FatalSlot, Filter, FlatMap, Inspect, Map, OpMeter, OpMeterSlot, StageLifecycle, TryMap,
72 TypedChain,
73};
74use super::handoff::{ChunkConfig, SinkHandoff};
75use super::split::{ErasedBranch, Sink, SinkCtx, SplitEmitter, SplitTerminal, new_branch};
76use super::{Collector, Emitter, RunnableChain};
77use crate::backpressure::InflightBudget;
78use crate::deser::{Deserializer, Owned, RecFamily};
79use crate::error::ErrorPolicy;
80use crate::metrics::{ComponentLabels, DeserMetrics, OperatorMetrics};
81use crate::sink::{RecordRouter, RowEncoder, ShardQueues};
82use std::marker::PhantomData;
83use std::sync::Arc;
84
85pub trait MapFn<In, Out>: FnMut(In) -> Out {}
90impl<G, In, Out> MapFn<In, Out> for G where G: FnMut(In) -> Out {}
91
92pub trait TryMapFn<In, Out, Err>: FnMut(In) -> Result<Out, Err> {}
94impl<G, In, Out, Err> TryMapFn<In, Out, Err> for G where G: FnMut(In) -> Result<Out, Err> {}
95
96pub trait Assemble<Term> {
101 type Out;
103 fn assemble(&self, term: Term) -> Self::Out;
105}
106
107#[derive(Clone, Copy, Debug, Default)]
109pub struct Root;
110
111impl<T> Assemble<T> for Root {
112 type Out = T;
113 fn assemble(&self, term: T) -> T {
114 term
115 }
116}
117
118#[derive(Clone, Debug)]
120pub struct MapPart<Prev, G> {
121 prev: Prev,
122 f: G,
123 meter: OpMeterSlot,
124}
125
126impl<Prev, G: Clone, Term> Assemble<Term> for MapPart<Prev, G>
127where
128 Prev: Assemble<Map<G, Term>>,
129{
130 type Out = Prev::Out;
131 fn assemble(&self, term: Term) -> Self::Out {
132 self.prev.assemble(Map {
133 f: self.f.clone(),
134 next: term,
135 meter: self.meter.clone(),
136 })
137 }
138}
139
140#[derive(Clone, Debug)]
142pub struct FilterPart<Prev, P> {
143 prev: Prev,
144 p: P,
145 meter: OpMeterSlot,
146}
147
148impl<Prev, P: Clone, Term> Assemble<Term> for FilterPart<Prev, P>
149where
150 Prev: Assemble<Filter<P, Term>>,
151{
152 type Out = Prev::Out;
153 fn assemble(&self, term: Term) -> Self::Out {
154 self.prev.assemble(Filter {
155 p: self.p.clone(),
156 next: term,
157 meter: self.meter.clone(),
158 })
159 }
160}
161
162#[derive(Clone, Debug)]
164pub struct InspectPart<Prev, G> {
165 prev: Prev,
166 f: G,
167}
168
169impl<Prev, G: Clone, Term> Assemble<Term> for InspectPart<Prev, G>
170where
171 Prev: Assemble<Inspect<G, Term>>,
172{
173 type Out = Prev::Out;
174 fn assemble(&self, term: Term) -> Self::Out {
175 self.prev.assemble(Inspect {
176 f: self.f.clone(),
177 next: term,
178 })
179 }
180}
181
182#[derive(Clone, Debug)]
184pub struct TryMapPart<Prev, G> {
185 prev: Prev,
186 f: G,
187 policy: ErrorPolicy,
188 component: Arc<str>,
189 meter: OpMeterSlot,
190}
191
192impl<Prev, G: Clone, Term> Assemble<Term> for TryMapPart<Prev, G>
193where
194 Prev: Assemble<TryMap<G, Term>>,
195{
196 type Out = Prev::Out;
197 fn assemble(&self, term: Term) -> Self::Out {
198 self.prev.assemble(TryMap {
199 f: self.f.clone(),
200 next: term,
201 policy: self.policy,
202 component: Arc::clone(&self.component),
203 meter: self.meter.clone(),
204 fatal: FatalSlot(None),
205 })
206 }
207}
208
209#[derive(Clone, Debug)]
211pub struct FlatMapPart<OutF: RecFamily, Prev, G> {
212 prev: Prev,
213 g: G,
214 meter: OpMeterSlot,
215 _out: PhantomData<fn() -> OutF>,
216}
217
218impl<OutF: RecFamily, Prev, G: Clone, Term> Assemble<Term> for FlatMapPart<OutF, Prev, G>
219where
220 Prev: Assemble<FlatMap<OutF, G, Term>>,
221{
222 type Out = Prev::Out;
223 fn assemble(&self, term: Term) -> Self::Out {
224 self.prev.assemble(FlatMap {
225 g: self.g.clone(),
226 next: term,
227 meter: self.meter.clone(),
228 _out: PhantomData,
229 })
230 }
231}
232
233#[derive(Clone, Debug)]
234struct MetricsSpec {
235 pipeline: String,
236 component: String,
237 deser: Arc<DeserMetrics>,
238}
239
240impl MetricsSpec {
241 fn op_handle(&self, idx: usize, kind: &'static str) -> Arc<OperatorMetrics> {
242 let labels = ComponentLabels::new(
243 self.pipeline.clone(),
244 format!("{}.{idx}_{kind}", self.component),
245 kind,
246 );
247 Arc::new(OperatorMetrics::new(&labels))
248 }
249}
250
251fn meter_for(metrics: &Option<MetricsSpec>, idx: usize, kind: &'static str) -> OpMeterSlot {
252 OpMeterSlot(OpMeter::new(
253 metrics.as_ref().map(|m| m.op_handle(idx, kind)),
254 ))
255}
256
257#[derive(Clone, Debug)]
262pub struct ChainBuilder<DF: RecFamily, CurF: RecFamily, D, P> {
263 deser: D,
264 parts: P,
265 deser_policy: ErrorPolicy,
266 metrics: Option<MetricsSpec>,
267 stage_idx: usize,
268 _fam: PhantomData<fn() -> (DF, CurF)>,
269}
270
271pub fn chain<F: RecFamily, D: Deserializer<F>>(deser: D) -> ChainBuilder<F, F, D, Root> {
273 ChainBuilder {
274 deser,
275 parts: Root,
276 deser_policy: ErrorPolicy::Skip,
277 metrics: None,
278 stage_idx: 0,
279 _fam: PhantomData,
280 }
281}
282
283pub fn chain_owned<T, D>(deser: D) -> ChainBuilder<Owned<T>, Owned<T>, D, Root>
285where
286 T: Send + 'static,
287 D: Deserializer<Owned<T>>,
288{
289 chain(deser)
290}
291
292impl<DF: RecFamily, CurF: RecFamily, D, P> ChainBuilder<DF, CurF, D, P> {
293 #[must_use]
300 pub fn with_metrics(
301 mut self,
302 pipeline: impl Into<String>,
303 component: impl Into<String>,
304 ) -> Self {
305 assert_eq!(
306 self.stage_idx, 0,
307 "with_metrics must be called before stages are added"
308 );
309 let pipeline = pipeline.into();
310 let component = component.into();
311 let deser_labels = ComponentLabels::new(
312 pipeline.clone(),
313 format!("{component}.deserializer"),
314 "deserializer",
315 );
316 self.metrics = Some(MetricsSpec {
317 pipeline,
318 component,
319 deser: Arc::new(DeserMetrics::new(&deser_labels)),
320 });
321 self
322 }
323
324 #[must_use]
326 pub fn deser_error_policy(mut self, policy: ErrorPolicy) -> Self {
327 self.deser_policy = policy;
328 self
329 }
330
331 #[must_use]
335 pub fn map_rec<NF, G>(self, f: G) -> ChainBuilder<DF, NF, D, MapPart<P, G>>
336 where
337 NF: RecFamily,
338 G: for<'buf> MapFn<CurF::Rec<'buf>, NF::Rec<'buf>>,
339 {
340 let Self {
341 deser,
342 parts,
343 deser_policy,
344 metrics,
345 stage_idx,
346 _fam,
347 } = self;
348 let meter = meter_for(&metrics, stage_idx, "map");
349 ChainBuilder {
350 deser,
351 parts: MapPart {
352 prev: parts,
353 f,
354 meter,
355 },
356 deser_policy,
357 metrics,
358 stage_idx: stage_idx + 1,
359 _fam: PhantomData,
360 }
361 }
362
363 #[must_use]
366 pub fn try_map_rec<NF, G, E>(
367 self,
368 f: G,
369 policy: ErrorPolicy,
370 ) -> ChainBuilder<DF, NF, D, TryMapPart<P, G>>
371 where
372 NF: RecFamily,
373 G: for<'buf> TryMapFn<CurF::Rec<'buf>, NF::Rec<'buf>, E>,
374 E: std::fmt::Display,
375 {
376 let Self {
377 deser,
378 parts,
379 deser_policy,
380 metrics,
381 stage_idx,
382 _fam,
383 } = self;
384 let meter = meter_for(&metrics, stage_idx, "try_map");
385 ChainBuilder {
386 deser,
387 parts: TryMapPart {
388 prev: parts,
389 f,
390 policy,
391 component: Arc::from(format!("try_map_{stage_idx}")),
392 meter,
393 },
394 deser_policy,
395 metrics,
396 stage_idx: stage_idx + 1,
397 _fam: PhantomData,
398 }
399 }
400
401 #[must_use]
403 pub fn filter<Pr>(self, p: Pr) -> ChainBuilder<DF, CurF, D, FilterPart<P, Pr>>
404 where
405 Pr: for<'buf> FnMut(&CurF::Rec<'buf>) -> bool,
406 {
407 let Self {
408 deser,
409 parts,
410 deser_policy,
411 metrics,
412 stage_idx,
413 _fam,
414 } = self;
415 let meter = meter_for(&metrics, stage_idx, "filter");
416 ChainBuilder {
417 deser,
418 parts: FilterPart {
419 prev: parts,
420 p,
421 meter,
422 },
423 deser_policy,
424 metrics,
425 stage_idx: stage_idx + 1,
426 _fam: PhantomData,
427 }
428 }
429
430 #[must_use]
432 pub fn inspect<G>(self, f: G) -> ChainBuilder<DF, CurF, D, InspectPart<P, G>>
433 where
434 G: for<'buf> FnMut(&CurF::Rec<'buf>),
435 {
436 let Self {
437 deser,
438 parts,
439 deser_policy,
440 metrics,
441 stage_idx,
442 _fam,
443 } = self;
444 ChainBuilder {
445 deser,
446 parts: InspectPart { prev: parts, f },
447 deser_policy,
448 metrics,
449 stage_idx: stage_idx + 1,
450 _fam: PhantomData,
451 }
452 }
453
454 #[must_use]
457 pub fn flat_map<OutF, G>(self, g: G) -> ChainBuilder<DF, OutF, D, FlatMapPart<OutF, P, G>>
458 where
459 OutF: RecFamily,
460 G: for<'buf> FnMut(CurF::Rec<'buf>, &mut Emitter<'_, OutF>),
461 {
462 let Self {
463 deser,
464 parts,
465 deser_policy,
466 metrics,
467 stage_idx,
468 _fam,
469 } = self;
470 let meter = meter_for(&metrics, stage_idx, "flat_map");
471 ChainBuilder {
472 deser,
473 parts: FlatMapPart {
474 prev: parts,
475 g,
476 meter,
477 _out: PhantomData,
478 },
479 deser_policy,
480 metrics,
481 stage_idx: stage_idx + 1,
482 _fam: PhantomData,
483 }
484 }
485
486 #[must_use]
494 pub fn sink<E, R>(
495 self,
496 encoder: E,
497 router: R,
498 cfg: ChunkConfig,
499 queues: ShardQueues,
500 budget: Arc<InflightBudget>,
501 ) -> SinkedChain<DF, CurF, D, P, E, R> {
502 let handoff_meter = meter_for(&self.metrics, self.stage_idx, "sink_handoff");
503 SinkedChain {
504 builder: self,
505 encoder,
506 router,
507 cfg,
508 queues,
509 budget,
510 handoff_meter,
511 }
512 }
513
514 #[must_use]
525 pub fn split(self, unmatched: ErrorPolicy) -> SplitBuilder<DF, CurF, D, P> {
526 SplitBuilder {
527 builder: self,
528 unmatched,
529 branches: Vec::new(),
530 next_idx: 0,
531 }
532 }
533}
534
535impl<DF: RecFamily, T: Send + 'static, D, P> ChainBuilder<DF, Owned<T>, D, P> {
537 #[must_use]
540 pub fn map<U, G>(self, f: G) -> ChainBuilder<DF, Owned<U>, D, MapPart<P, G>>
541 where
542 U: Send + 'static,
543 G: FnMut(T) -> U,
544 {
545 let Self {
546 deser,
547 parts,
548 deser_policy,
549 metrics,
550 stage_idx,
551 _fam,
552 } = self;
553 let meter = meter_for(&metrics, stage_idx, "map");
554 ChainBuilder {
555 deser,
556 parts: MapPart {
557 prev: parts,
558 f,
559 meter,
560 },
561 deser_policy,
562 metrics,
563 stage_idx: stage_idx + 1,
564 _fam: PhantomData,
565 }
566 }
567
568 #[must_use]
571 pub fn try_map<U, G, E>(
572 self,
573 f: G,
574 policy: ErrorPolicy,
575 ) -> ChainBuilder<DF, Owned<U>, D, TryMapPart<P, G>>
576 where
577 U: Send + 'static,
578 G: FnMut(T) -> Result<U, E>,
579 E: std::fmt::Display,
580 {
581 let Self {
582 deser,
583 parts,
584 deser_policy,
585 metrics,
586 stage_idx,
587 _fam,
588 } = self;
589 let meter = meter_for(&metrics, stage_idx, "try_map");
590 ChainBuilder {
591 deser,
592 parts: TryMapPart {
593 prev: parts,
594 f,
595 policy,
596 component: Arc::from(format!("try_map_{stage_idx}")),
597 meter,
598 },
599 deser_policy,
600 metrics,
601 stage_idx: stage_idx + 1,
602 _fam: PhantomData,
603 }
604 }
605}
606
607#[derive(Clone, Debug)]
610pub struct SinkedChain<DF: RecFamily, CurF: RecFamily, D, P, E, R> {
611 builder: ChainBuilder<DF, CurF, D, P>,
612 encoder: E,
613 router: R,
614 cfg: ChunkConfig,
615 queues: ShardQueues,
616 budget: Arc<InflightBudget>,
617 handoff_meter: OpMeterSlot,
618}
619
620impl<DF, CurF, D, P, E, R> SinkedChain<DF, CurF, D, P, E, R>
621where
622 DF: RecFamily,
623 CurF: RecFamily,
624 D: Deserializer<DF> + 'static,
625 P: Assemble<SinkHandoff<CurF, E, R>>,
626 P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
627 E: RowEncoder<CurF> + Clone + 'static,
628 R: RecordRouter<CurF> + Send + 'static,
629{
630 #[must_use]
635 pub fn build(self) -> Box<dyn RunnableChain> {
636 let SinkedChain {
637 builder,
638 encoder,
639 router,
640 cfg,
641 queues,
642 budget,
643 handoff_meter,
644 } = self;
645 let term = SinkHandoff::new(
646 encoder,
647 router,
648 queues,
649 budget,
650 cfg,
651 handoff_meter,
652 Arc::from("sink_handoff"),
653 );
654 let ops = builder.parts.assemble(term);
655 Box::new(TypedChain::<DF, D, _>::new(
656 builder.deser,
657 ops,
658 builder.deser_policy,
659 builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
660 ))
661 }
662
663 #[must_use]
666 pub fn build_factory(self) -> ChainFactory<DF, CurF, D, P, E, R>
667 where
668 D: Clone,
669 E: Clone,
670 R: Clone,
671 {
672 ChainFactory { spec: self }
673 }
674}
675
676#[derive(Clone, Debug)]
679pub struct ChainFactory<DF: RecFamily, CurF: RecFamily, D, P, E, R> {
680 spec: SinkedChain<DF, CurF, D, P, E, R>,
681}
682
683impl<DF, CurF, D, P, E, R> ChainFactory<DF, CurF, D, P, E, R>
684where
685 DF: RecFamily,
686 CurF: RecFamily,
687 D: Deserializer<DF> + Clone + 'static,
688 P: Assemble<SinkHandoff<CurF, E, R>>,
689 P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
690 E: RowEncoder<CurF> + Clone + 'static,
691 R: RecordRouter<CurF> + Clone + Send + 'static,
692{
693 #[must_use]
695 pub fn make(&self) -> Box<dyn RunnableChain> {
696 let spec = &self.spec;
697 let term = SinkHandoff::new(
698 spec.encoder.clone(),
699 spec.router.clone(),
700 spec.queues.clone(),
701 Arc::clone(&spec.budget),
702 spec.cfg,
703 spec.handoff_meter.clone(),
704 Arc::from("sink_handoff"),
705 );
706 let ops = spec.builder.parts.assemble(term);
707 Box::new(TypedChain::<DF, D, _>::new(
708 spec.builder.deser.clone(),
709 ops,
710 spec.builder.deser_policy,
711 spec.builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
712 ))
713 }
714}
715
716pub struct SplitBuilder<DF: RecFamily, CurF: RecFamily, D, P> {
721 builder: ChainBuilder<DF, CurF, D, P>,
722 unmatched: ErrorPolicy,
723 branches: Vec<Box<dyn ErasedBranch>>,
724 next_idx: usize,
725}
726
727impl<DF: RecFamily, CurF: RecFamily, D, P> std::fmt::Debug for SplitBuilder<DF, CurF, D, P> {
728 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
729 f.debug_struct("SplitBuilder")
730 .field("branches", &self.branches.len())
731 .field("unmatched", &self.unmatched)
732 .finish_non_exhaustive()
733 }
734}
735
736impl<DF: RecFamily, CurF: RecFamily, D, P> SplitBuilder<DF, CurF, D, P> {
737 #[must_use = "a dropped Sink<F> handle leaves its branch permanently unreachable"]
744 pub fn add<F, E, R>(&mut self, encoder: E, router: R, sink: SinkCtx) -> Sink<F>
745 where
746 F: RecFamily + 'static,
747 E: RowEncoder<F> + Clone + Send + 'static,
748 R: RecordRouter<F> + 'static,
749 {
750 let idx = self.next_idx;
751 self.next_idx += 1;
752 let meter = OpMeterSlot(OpMeter::new(self.builder.metrics.as_ref().map(|m| {
755 let labels = ComponentLabels::new(
756 m.pipeline.clone(),
757 format!("{}.sink.{}", m.component, sink.name),
758 "sink_handoff",
759 );
760 Arc::new(OperatorMetrics::new(&labels))
761 })));
762 let component: Arc<str> = Arc::from(format!("sink.{}", sink.name));
763 let branch = new_branch::<F, E, R>(
764 encoder,
765 router,
766 sink.queues,
767 sink.budget,
768 sink.chunk,
769 meter,
770 component,
771 );
772 self.branches.push(branch);
773 Sink::new(idx)
774 }
775
776 #[must_use]
781 pub fn route<G>(self, route: G) -> RoutedSplit<DF, CurF, D, P, G>
782 where
783 G: for<'buf> FnMut(CurF::Rec<'buf>, &mut SplitEmitter<'_>) + Send + 'static,
784 {
785 let handoff_meter = meter_for(&self.builder.metrics, self.builder.stage_idx, "split");
786 RoutedSplit {
787 builder: self.builder,
788 unmatched: self.unmatched,
789 branches: self.branches,
790 route,
791 handoff_meter,
792 }
793 }
794}
795
796pub struct RoutedSplit<DF: RecFamily, CurF: RecFamily, D, P, G> {
800 builder: ChainBuilder<DF, CurF, D, P>,
801 unmatched: ErrorPolicy,
802 branches: Vec<Box<dyn ErasedBranch>>,
803 route: G,
804 handoff_meter: OpMeterSlot,
805}
806
807impl<DF: RecFamily, CurF: RecFamily, D, P, G> std::fmt::Debug for RoutedSplit<DF, CurF, D, P, G> {
808 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
809 f.debug_struct("RoutedSplit")
810 .field("branches", &self.branches.len())
811 .field("unmatched", &self.unmatched)
812 .finish_non_exhaustive()
813 }
814}
815
816impl<DF, CurF, D, P, G> RoutedSplit<DF, CurF, D, P, G>
817where
818 DF: RecFamily,
819 CurF: RecFamily,
820 D: Deserializer<DF> + 'static,
821 G: for<'buf> FnMut(CurF::Rec<'buf>, &mut SplitEmitter<'_>) + Send + 'static,
822 P: Assemble<SplitTerminal<CurF, G>>,
823 P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
824{
825 #[must_use]
827 pub fn build(self) -> Box<dyn RunnableChain> {
828 let RoutedSplit {
829 builder,
830 unmatched,
831 branches,
832 route,
833 handoff_meter,
834 } = self;
835 let term = SplitTerminal::new(
836 route,
837 branches,
838 unmatched,
839 handoff_meter,
840 Arc::from("split"),
841 );
842 let ops = builder.parts.assemble(term);
843 Box::new(TypedChain::<DF, D, _>::new(
844 builder.deser,
845 ops,
846 builder.deser_policy,
847 builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
848 ))
849 }
850}