use super::chain::{
FatalSlot, Filter, FlatMap, Inspect, Map, OpMeter, OpMeterSlot, StageLifecycle, TryMap,
TypedChain,
};
use super::handoff::{ChunkConfig, SinkHandoff};
use super::split::{ErasedBranch, Sink, SinkCtx, SplitEmitter, SplitTerminal, new_branch};
use super::{Collector, Emitter, RunnableChain};
use crate::backpressure::InflightBudget;
use crate::deser::{Deserializer, Owned, RecFamily};
use crate::error::ErrorPolicy;
use crate::metrics::{ComponentLabels, DeserMetrics, OperatorMetrics};
use crate::sink::{RecordRouter, RowEncoder, ShardQueues};
use std::marker::PhantomData;
use std::sync::Arc;
pub trait MapFn<In, Out>: FnMut(In) -> Out {}
impl<G, In, Out> MapFn<In, Out> for G where G: FnMut(In) -> Out {}
pub trait TryMapFn<In, Out, Err>: FnMut(In) -> Result<Out, Err> {}
impl<G, In, Out, Err> TryMapFn<In, Out, Err> for G where G: FnMut(In) -> Result<Out, Err> {}
pub trait Assemble<Term> {
type Out;
fn assemble(&self, term: Term) -> Self::Out;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct Root;
impl<T> Assemble<T> for Root {
type Out = T;
fn assemble(&self, term: T) -> T {
term
}
}
#[derive(Clone, Debug)]
pub struct MapPart<Prev, G> {
prev: Prev,
f: G,
meter: OpMeterSlot,
}
impl<Prev, G: Clone, Term> Assemble<Term> for MapPart<Prev, G>
where
Prev: Assemble<Map<G, Term>>,
{
type Out = Prev::Out;
fn assemble(&self, term: Term) -> Self::Out {
self.prev.assemble(Map {
f: self.f.clone(),
next: term,
meter: self.meter.clone(),
})
}
}
#[derive(Clone, Debug)]
pub struct FilterPart<Prev, P> {
prev: Prev,
p: P,
meter: OpMeterSlot,
}
impl<Prev, P: Clone, Term> Assemble<Term> for FilterPart<Prev, P>
where
Prev: Assemble<Filter<P, Term>>,
{
type Out = Prev::Out;
fn assemble(&self, term: Term) -> Self::Out {
self.prev.assemble(Filter {
p: self.p.clone(),
next: term,
meter: self.meter.clone(),
})
}
}
#[derive(Clone, Debug)]
pub struct InspectPart<Prev, G> {
prev: Prev,
f: G,
}
impl<Prev, G: Clone, Term> Assemble<Term> for InspectPart<Prev, G>
where
Prev: Assemble<Inspect<G, Term>>,
{
type Out = Prev::Out;
fn assemble(&self, term: Term) -> Self::Out {
self.prev.assemble(Inspect {
f: self.f.clone(),
next: term,
})
}
}
#[derive(Clone, Debug)]
pub struct TryMapPart<Prev, G> {
prev: Prev,
f: G,
policy: ErrorPolicy,
component: Arc<str>,
meter: OpMeterSlot,
}
impl<Prev, G: Clone, Term> Assemble<Term> for TryMapPart<Prev, G>
where
Prev: Assemble<TryMap<G, Term>>,
{
type Out = Prev::Out;
fn assemble(&self, term: Term) -> Self::Out {
self.prev.assemble(TryMap {
f: self.f.clone(),
next: term,
policy: self.policy,
component: Arc::clone(&self.component),
meter: self.meter.clone(),
fatal: FatalSlot(None),
})
}
}
#[derive(Clone, Debug)]
pub struct FlatMapPart<OutF: RecFamily, Prev, G> {
prev: Prev,
g: G,
meter: OpMeterSlot,
_out: PhantomData<fn() -> OutF>,
}
impl<OutF: RecFamily, Prev, G: Clone, Term> Assemble<Term> for FlatMapPart<OutF, Prev, G>
where
Prev: Assemble<FlatMap<OutF, G, Term>>,
{
type Out = Prev::Out;
fn assemble(&self, term: Term) -> Self::Out {
self.prev.assemble(FlatMap {
g: self.g.clone(),
next: term,
meter: self.meter.clone(),
_out: PhantomData,
})
}
}
#[derive(Clone, Debug)]
struct MetricsSpec {
pipeline: String,
component: String,
deser: Arc<DeserMetrics>,
}
impl MetricsSpec {
fn op_handle(&self, idx: usize, kind: &'static str) -> Arc<OperatorMetrics> {
let labels = ComponentLabels::new(
self.pipeline.clone(),
format!("{}.{idx}_{kind}", self.component),
kind,
);
Arc::new(OperatorMetrics::new(&labels))
}
}
fn meter_for(metrics: &Option<MetricsSpec>, idx: usize, kind: &'static str) -> OpMeterSlot {
OpMeterSlot(OpMeter::new(
metrics.as_ref().map(|m| m.op_handle(idx, kind)),
))
}
#[derive(Clone, Debug)]
pub struct ChainBuilder<DF: RecFamily, CurF: RecFamily, D, P> {
deser: D,
parts: P,
deser_policy: ErrorPolicy,
metrics: Option<MetricsSpec>,
stage_idx: usize,
_fam: PhantomData<fn() -> (DF, CurF)>,
}
pub fn chain<F: RecFamily, D: Deserializer<F>>(deser: D) -> ChainBuilder<F, F, D, Root> {
ChainBuilder {
deser,
parts: Root,
deser_policy: ErrorPolicy::Skip,
metrics: None,
stage_idx: 0,
_fam: PhantomData,
}
}
pub fn chain_owned<T, D>(deser: D) -> ChainBuilder<Owned<T>, Owned<T>, D, Root>
where
T: Send + 'static,
D: Deserializer<Owned<T>>,
{
chain(deser)
}
impl<DF: RecFamily, CurF: RecFamily, D, P> ChainBuilder<DF, CurF, D, P> {
#[must_use]
pub fn with_metrics(
mut self,
pipeline: impl Into<String>,
component: impl Into<String>,
) -> Self {
assert_eq!(
self.stage_idx, 0,
"with_metrics must be called before stages are added"
);
let pipeline = pipeline.into();
let component = component.into();
let deser_labels = ComponentLabels::new(
pipeline.clone(),
format!("{component}.deserializer"),
"deserializer",
);
self.metrics = Some(MetricsSpec {
pipeline,
component,
deser: Arc::new(DeserMetrics::new(&deser_labels)),
});
self
}
#[must_use]
pub fn deser_error_policy(mut self, policy: ErrorPolicy) -> Self {
self.deser_policy = policy;
self
}
#[must_use]
pub fn map_rec<NF, G>(self, f: G) -> ChainBuilder<DF, NF, D, MapPart<P, G>>
where
NF: RecFamily,
G: for<'buf> MapFn<CurF::Rec<'buf>, NF::Rec<'buf>>,
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "map");
ChainBuilder {
deser,
parts: MapPart {
prev: parts,
f,
meter,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn try_map_rec<NF, G, E>(
self,
f: G,
policy: ErrorPolicy,
) -> ChainBuilder<DF, NF, D, TryMapPart<P, G>>
where
NF: RecFamily,
G: for<'buf> TryMapFn<CurF::Rec<'buf>, NF::Rec<'buf>, E>,
E: std::fmt::Display,
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "try_map");
ChainBuilder {
deser,
parts: TryMapPart {
prev: parts,
f,
policy,
component: Arc::from(format!("try_map_{stage_idx}")),
meter,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn filter<Pr>(self, p: Pr) -> ChainBuilder<DF, CurF, D, FilterPart<P, Pr>>
where
Pr: for<'buf> FnMut(&CurF::Rec<'buf>) -> bool,
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "filter");
ChainBuilder {
deser,
parts: FilterPart {
prev: parts,
p,
meter,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn inspect<G>(self, f: G) -> ChainBuilder<DF, CurF, D, InspectPart<P, G>>
where
G: for<'buf> FnMut(&CurF::Rec<'buf>),
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
ChainBuilder {
deser,
parts: InspectPart { prev: parts, f },
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn flat_map<OutF, G>(self, g: G) -> ChainBuilder<DF, OutF, D, FlatMapPart<OutF, P, G>>
where
OutF: RecFamily,
G: for<'buf> FnMut(CurF::Rec<'buf>, &mut Emitter<'_, OutF>),
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "flat_map");
ChainBuilder {
deser,
parts: FlatMapPart {
prev: parts,
g,
meter,
_out: PhantomData,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn sink<E, R>(
self,
encoder: E,
router: R,
cfg: ChunkConfig,
queues: ShardQueues,
budget: Arc<InflightBudget>,
) -> SinkedChain<DF, CurF, D, P, E, R> {
let handoff_meter = meter_for(&self.metrics, self.stage_idx, "sink_handoff");
SinkedChain {
builder: self,
encoder,
router,
cfg,
queues,
budget,
handoff_meter,
}
}
#[must_use]
pub fn split(self, unmatched: ErrorPolicy) -> SplitBuilder<DF, CurF, D, P> {
SplitBuilder {
builder: self,
unmatched,
branches: Vec::new(),
next_idx: 0,
}
}
}
impl<DF: RecFamily, T: Send + 'static, D, P> ChainBuilder<DF, Owned<T>, D, P> {
#[must_use]
pub fn map<U, G>(self, f: G) -> ChainBuilder<DF, Owned<U>, D, MapPart<P, G>>
where
U: Send + 'static,
G: FnMut(T) -> U,
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "map");
ChainBuilder {
deser,
parts: MapPart {
prev: parts,
f,
meter,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
#[must_use]
pub fn try_map<U, G, E>(
self,
f: G,
policy: ErrorPolicy,
) -> ChainBuilder<DF, Owned<U>, D, TryMapPart<P, G>>
where
U: Send + 'static,
G: FnMut(T) -> Result<U, E>,
E: std::fmt::Display,
{
let Self {
deser,
parts,
deser_policy,
metrics,
stage_idx,
_fam,
} = self;
let meter = meter_for(&metrics, stage_idx, "try_map");
ChainBuilder {
deser,
parts: TryMapPart {
prev: parts,
f,
policy,
component: Arc::from(format!("try_map_{stage_idx}")),
meter,
},
deser_policy,
metrics,
stage_idx: stage_idx + 1,
_fam: PhantomData,
}
}
}
#[derive(Clone, Debug)]
pub struct SinkedChain<DF: RecFamily, CurF: RecFamily, D, P, E, R> {
builder: ChainBuilder<DF, CurF, D, P>,
encoder: E,
router: R,
cfg: ChunkConfig,
queues: ShardQueues,
budget: Arc<InflightBudget>,
handoff_meter: OpMeterSlot,
}
impl<DF, CurF, D, P, E, R> SinkedChain<DF, CurF, D, P, E, R>
where
DF: RecFamily,
CurF: RecFamily,
D: Deserializer<DF> + 'static,
P: Assemble<SinkHandoff<CurF, E, R>>,
P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
E: RowEncoder<CurF> + Clone + 'static,
R: RecordRouter<CurF> + Send + 'static,
{
#[must_use]
pub fn build(self) -> Box<dyn RunnableChain> {
let SinkedChain {
builder,
encoder,
router,
cfg,
queues,
budget,
handoff_meter,
} = self;
let term = SinkHandoff::new(
encoder,
router,
queues,
budget,
cfg,
handoff_meter,
Arc::from("sink_handoff"),
);
let ops = builder.parts.assemble(term);
Box::new(TypedChain::<DF, D, _>::new(
builder.deser,
ops,
builder.deser_policy,
builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
))
}
#[must_use]
pub fn build_factory(self) -> ChainFactory<DF, CurF, D, P, E, R>
where
D: Clone,
E: Clone,
R: Clone,
{
ChainFactory { spec: self }
}
}
#[derive(Clone, Debug)]
pub struct ChainFactory<DF: RecFamily, CurF: RecFamily, D, P, E, R> {
spec: SinkedChain<DF, CurF, D, P, E, R>,
}
impl<DF, CurF, D, P, E, R> ChainFactory<DF, CurF, D, P, E, R>
where
DF: RecFamily,
CurF: RecFamily,
D: Deserializer<DF> + Clone + 'static,
P: Assemble<SinkHandoff<CurF, E, R>>,
P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
E: RowEncoder<CurF> + Clone + 'static,
R: RecordRouter<CurF> + Clone + Send + 'static,
{
#[must_use]
pub fn make(&self) -> Box<dyn RunnableChain> {
let spec = &self.spec;
let term = SinkHandoff::new(
spec.encoder.clone(),
spec.router.clone(),
spec.queues.clone(),
Arc::clone(&spec.budget),
spec.cfg,
spec.handoff_meter.clone(),
Arc::from("sink_handoff"),
);
let ops = spec.builder.parts.assemble(term);
Box::new(TypedChain::<DF, D, _>::new(
spec.builder.deser.clone(),
ops,
spec.builder.deser_policy,
spec.builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
))
}
}
pub struct SplitBuilder<DF: RecFamily, CurF: RecFamily, D, P> {
builder: ChainBuilder<DF, CurF, D, P>,
unmatched: ErrorPolicy,
branches: Vec<Box<dyn ErasedBranch>>,
next_idx: usize,
}
impl<DF: RecFamily, CurF: RecFamily, D, P> std::fmt::Debug for SplitBuilder<DF, CurF, D, P> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SplitBuilder")
.field("branches", &self.branches.len())
.field("unmatched", &self.unmatched)
.finish_non_exhaustive()
}
}
impl<DF: RecFamily, CurF: RecFamily, D, P> SplitBuilder<DF, CurF, D, P> {
#[must_use = "a dropped Sink<F> handle leaves its branch permanently unreachable"]
pub fn add<F, E, R>(&mut self, encoder: E, router: R, sink: SinkCtx) -> Sink<F>
where
F: RecFamily + 'static,
E: RowEncoder<F> + Clone + Send + 'static,
R: RecordRouter<F> + 'static,
{
let idx = self.next_idx;
self.next_idx += 1;
let meter = OpMeterSlot(OpMeter::new(self.builder.metrics.as_ref().map(|m| {
let labels = ComponentLabels::new(
m.pipeline.clone(),
format!("{}.sink.{}", m.component, sink.name),
"sink_handoff",
);
Arc::new(OperatorMetrics::new(&labels))
})));
let component: Arc<str> = Arc::from(format!("sink.{}", sink.name));
let branch = new_branch::<F, E, R>(
encoder,
router,
sink.queues,
sink.budget,
sink.chunk,
meter,
component,
);
self.branches.push(branch);
Sink::new(idx)
}
#[must_use]
pub fn route<G>(self, route: G) -> RoutedSplit<DF, CurF, D, P, G>
where
G: for<'buf> FnMut(CurF::Rec<'buf>, &mut SplitEmitter<'_>) + Send + 'static,
{
let handoff_meter = meter_for(&self.builder.metrics, self.builder.stage_idx, "split");
RoutedSplit {
builder: self.builder,
unmatched: self.unmatched,
branches: self.branches,
route,
handoff_meter,
}
}
}
pub struct RoutedSplit<DF: RecFamily, CurF: RecFamily, D, P, G> {
builder: ChainBuilder<DF, CurF, D, P>,
unmatched: ErrorPolicy,
branches: Vec<Box<dyn ErasedBranch>>,
route: G,
handoff_meter: OpMeterSlot,
}
impl<DF: RecFamily, CurF: RecFamily, D, P, G> std::fmt::Debug for RoutedSplit<DF, CurF, D, P, G> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RoutedSplit")
.field("branches", &self.branches.len())
.field("unmatched", &self.unmatched)
.finish_non_exhaustive()
}
}
impl<DF, CurF, D, P, G> RoutedSplit<DF, CurF, D, P, G>
where
DF: RecFamily,
CurF: RecFamily,
D: Deserializer<DF> + 'static,
G: for<'buf> FnMut(CurF::Rec<'buf>, &mut SplitEmitter<'_>) + Send + 'static,
P: Assemble<SplitTerminal<CurF, G>>,
P::Out: for<'buf> Collector<<DF as RecFamily>::Rec<'buf>> + StageLifecycle + Send + 'static,
{
#[must_use]
pub fn build(self) -> Box<dyn RunnableChain> {
let RoutedSplit {
builder,
unmatched,
branches,
route,
handoff_meter,
} = self;
let term = SplitTerminal::new(
route,
branches,
unmatched,
handoff_meter,
Arc::from("split"),
);
let ops = builder.parts.assemble(term);
Box::new(TypedChain::<DF, D, _>::new(
builder.deser,
ops,
builder.deser_policy,
builder.metrics.as_ref().map(|m| Arc::clone(&m.deser)),
))
}
}