1use super::{BlockReason, Collector, CollectorFor, PushOutcome, RunnableChain};
17use crate::checkpoint::AckRef;
18use crate::deser::{Deserializer, EmitRecord, RecFamily};
19use crate::error::{DeserError, ErrorPolicy, FatalError};
20use crate::metrics::{DeserMetrics, OperatorMetrics};
21use crate::record::{Flow, RawPayload, Record, RecordMeta};
22use crate::source::PayloadBatch;
23use crate::telemetry::RateLimit;
24use std::marker::PhantomData;
25use std::sync::Arc;
26use std::time::{Duration, Instant};
27
28pub trait StageLifecycle {
32 fn on_batch_end(&mut self, elapsed: Duration);
37
38 fn take_fatal(&mut self) -> Option<FatalError>;
43
44 fn relieve(&mut self) -> Flow;
47
48 fn flush_terminal(&mut self) -> Flow;
51}
52
53#[derive(Debug, Default)]
56pub(crate) struct OpMeter {
57 handle: Option<Arc<OperatorMetrics>>,
58 tripped: std::sync::atomic::AtomicBool,
62 records_in: u64,
63 records_out: u64,
64 filtered: u64,
65 skipped: u64,
66 unrouted: u64,
67 record_errors: u64,
68}
69
70impl OpMeter {
71 pub(crate) fn new(handle: Option<Arc<OperatorMetrics>>) -> Self {
72 OpMeter {
73 handle,
74 ..Default::default()
75 }
76 }
77
78 #[inline(always)]
79 pub(crate) fn seen(&mut self) {
80 self.records_in += 1;
81 }
82
83 #[inline(always)]
84 pub(crate) fn out(&mut self) {
85 self.records_out += 1;
86 }
87
88 #[inline(always)]
89 pub(crate) fn out_n(&mut self, n: u64) {
90 self.records_out += n;
91 }
92
93 #[inline(always)]
94 pub(crate) fn filtered(&mut self) {
95 self.filtered += 1;
96 }
97
98 #[inline(always)]
99 pub(crate) fn skipped(&mut self) {
100 self.skipped += 1;
101 }
102
103 #[inline(always)]
104 pub(crate) fn unrouted(&mut self) {
105 self.unrouted += 1;
106 }
107
108 #[inline(always)]
109 pub(crate) fn record_error(&mut self) {
110 self.record_errors += 1;
111 }
112
113 pub(crate) fn fatal_error(&self) {
117 if self
118 .tripped
119 .swap(true, std::sync::atomic::Ordering::Relaxed)
120 {
121 return;
122 }
123 if let Some(h) = &self.handle {
124 h.fatal_error();
125 }
126 }
127
128 pub(crate) fn flush(&mut self, elapsed: Duration) {
129 if let Some(h) = &self.handle {
130 h.batch(self.records_in, self.records_out, elapsed);
131 if self.filtered > 0 {
132 h.filtered(self.filtered);
133 }
134 if self.skipped > 0 {
135 h.skipped(self.skipped);
136 }
137 if self.unrouted > 0 {
138 h.unrouted(self.unrouted);
139 }
140 if self.record_errors > 0 {
141 h.record_errors(self.record_errors);
142 }
143 }
144 self.records_in = 0;
145 self.records_out = 0;
146 self.filtered = 0;
147 self.skipped = 0;
148 self.unrouted = 0;
149 self.record_errors = 0;
150 }
151}
152
153#[derive(Debug, Default)]
156pub(crate) struct OpMeterSlot(pub(crate) OpMeter);
157
158impl Clone for OpMeterSlot {
159 fn clone(&self) -> Self {
160 OpMeterSlot(OpMeter::new(self.0.handle.clone()))
161 }
162}
163
164#[derive(Debug, Default)]
166pub(crate) struct FatalSlot(pub(crate) Option<FatalError>);
167
168impl Clone for FatalSlot {
169 fn clone(&self) -> Self {
170 FatalSlot(None)
171 }
172}
173
174#[derive(Clone, Debug)]
176pub struct Map<G, N> {
177 pub(crate) f: G,
178 pub(crate) next: N,
179 pub(crate) meter: OpMeterSlot,
180}
181
182impl<In, Out, G, N> Collector<In> for Map<G, N>
183where
184 G: FnMut(In) -> Out,
185 N: Collector<Out>,
186{
187 #[inline(always)]
188 fn push(&mut self, rec: Record<In>) -> Flow {
189 self.meter.0.seen();
190 self.meter.0.out();
191 self.next.push(rec.map(&mut self.f))
192 }
193}
194
195#[derive(Clone, Debug)]
198pub struct Filter<P, N> {
199 pub(crate) p: P,
200 pub(crate) next: N,
201 pub(crate) meter: OpMeterSlot,
202}
203
204impl<T, P, N> Collector<T> for Filter<P, N>
205where
206 P: FnMut(&T) -> bool,
207 N: Collector<T>,
208{
209 #[inline(always)]
210 fn push(&mut self, rec: Record<T>) -> Flow {
211 self.meter.0.seen();
212 if (self.p)(&rec.payload) {
213 self.meter.0.out();
214 self.next.push(rec)
215 } else {
216 self.meter.0.filtered();
217 Flow::Continue
218 }
219 }
220}
221
222#[derive(Clone, Debug)]
224pub struct Inspect<G, N> {
225 pub(crate) f: G,
226 pub(crate) next: N,
227}
228
229impl<T, G, N> Collector<T> for Inspect<G, N>
230where
231 G: FnMut(&T),
232 N: Collector<T>,
233{
234 #[inline(always)]
235 fn push(&mut self, rec: Record<T>) -> Flow {
236 (self.f)(&rec.payload);
237 self.next.push(rec)
238 }
239}
240
241static TRY_MAP_SKIP_WARN: RateLimit = RateLimit::new(5, Duration::from_secs(10));
242
243#[derive(Clone, Debug)]
247pub struct TryMap<G, N> {
248 pub(crate) f: G,
249 pub(crate) next: N,
250 pub(crate) policy: ErrorPolicy,
251 pub(crate) component: Arc<str>,
252 pub(crate) meter: OpMeterSlot,
253 pub(crate) fatal: FatalSlot,
254}
255
256impl<In, Out, E, G, N> Collector<In> for TryMap<G, N>
257where
258 G: FnMut(In) -> Result<Out, E>,
259 E: std::fmt::Display,
260 N: Collector<Out>,
261{
262 #[inline(always)]
263 fn push(&mut self, rec: Record<In>) -> Flow {
264 self.meter.0.seen();
265 if self.fatal.0.is_some() {
266 return Flow::Continue;
269 }
270 let Record { payload, meta, ack } = rec;
271 match (self.f)(payload) {
272 Ok(out) => {
273 self.meter.0.out();
274 self.next.push(Record {
275 payload: out,
276 meta,
277 ack,
278 })
279 }
280 Err(e) => {
281 match self.policy {
282 ErrorPolicy::Skip => {
283 self.meter.0.skipped();
284 self.meter.0.record_error();
285 crate::rate_limited_warn!(
286 TRY_MAP_SKIP_WARN,
287 component = &*self.component,
288 error = %e,
289 "record skipped by error policy"
290 );
291 }
292 _ => {
293 self.meter.0.fatal_error();
294 self.fatal.0 = Some(FatalError {
295 component: self.component.to_string(),
296 reason: e.to_string(),
297 });
298 }
299 }
300 Flow::Continue
301 }
302 }
303 }
304}
305
306pub struct Emitter<'a, OutF: RecFamily> {
311 next: &'a mut dyn CollectorFor<OutF>,
312 meta: RecordMeta,
313 ack: &'a AckRef,
314 emitted: u64,
315 flow: Flow,
316}
317
318impl<OutF: RecFamily> std::fmt::Debug for Emitter<'_, OutF> {
319 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320 f.debug_struct("Emitter")
321 .field("emitted", &self.emitted)
322 .field("flow", &self.flow)
323 .finish_non_exhaustive()
324 }
325}
326
327impl<OutF: RecFamily> Emitter<'_, OutF> {
328 #[inline(always)]
331 pub fn emit<'buf>(&mut self, payload: OutF::Rec<'buf>) -> Flow {
332 let flow = self.next.push_rec(Record {
333 payload,
334 meta: self.meta,
335 ack: self.ack.clone(),
336 });
337 self.emitted += 1;
338 if self.flow != Flow::Blocked {
339 self.flow = flow;
340 }
341 flow
342 }
343
344 #[must_use]
346 pub fn meta(&self) -> RecordMeta {
347 self.meta
348 }
349}
350
351#[derive(Clone, Debug)]
355pub struct FlatMap<OutF: RecFamily, G, N> {
356 pub(crate) g: G,
357 pub(crate) next: N,
358 pub(crate) meter: OpMeterSlot,
359 pub(crate) _out: PhantomData<fn() -> OutF>,
360}
361
362impl<In, OutF, G, N> Collector<In> for FlatMap<OutF, G, N>
363where
364 OutF: RecFamily,
365 G: FnMut(In, &mut Emitter<'_, OutF>),
366 N: for<'buf> Collector<<OutF as RecFamily>::Rec<'buf>>,
367{
368 #[inline(always)]
369 fn push(&mut self, rec: Record<In>) -> Flow {
370 self.meter.0.seen();
371 let Record { payload, meta, ack } = rec;
372 let mut em = Emitter {
373 next: &mut self.next,
374 meta,
375 ack: &ack,
376 emitted: 0,
377 flow: Flow::Continue,
378 };
379 (self.g)(payload, &mut em);
380 let (emitted, flow) = (em.emitted, em.flow);
381 self.meter.0.out_n(emitted);
382 flow
383 }
384}
385
386impl<G, N: StageLifecycle> StageLifecycle for Map<G, N> {
389 fn on_batch_end(&mut self, elapsed: Duration) {
390 self.meter.0.flush(elapsed);
391 self.next.on_batch_end(elapsed);
392 }
393 fn take_fatal(&mut self) -> Option<FatalError> {
394 self.next.take_fatal()
395 }
396 fn relieve(&mut self) -> Flow {
397 self.next.relieve()
398 }
399 fn flush_terminal(&mut self) -> Flow {
400 self.next.flush_terminal()
401 }
402}
403
404impl<P, N: StageLifecycle> StageLifecycle for Filter<P, N> {
405 fn on_batch_end(&mut self, elapsed: Duration) {
406 self.meter.0.flush(elapsed);
407 self.next.on_batch_end(elapsed);
408 }
409 fn take_fatal(&mut self) -> Option<FatalError> {
410 self.next.take_fatal()
411 }
412 fn relieve(&mut self) -> Flow {
413 self.next.relieve()
414 }
415 fn flush_terminal(&mut self) -> Flow {
416 self.next.flush_terminal()
417 }
418}
419
420impl<G, N: StageLifecycle> StageLifecycle for Inspect<G, N> {
421 fn on_batch_end(&mut self, elapsed: Duration) {
422 self.next.on_batch_end(elapsed);
423 }
424 fn take_fatal(&mut self) -> Option<FatalError> {
425 self.next.take_fatal()
426 }
427 fn relieve(&mut self) -> Flow {
428 self.next.relieve()
429 }
430 fn flush_terminal(&mut self) -> Flow {
431 self.next.flush_terminal()
432 }
433}
434
435impl<G, N: StageLifecycle> StageLifecycle for TryMap<G, N> {
436 fn on_batch_end(&mut self, elapsed: Duration) {
437 self.meter.0.flush(elapsed);
438 self.next.on_batch_end(elapsed);
439 }
440 fn take_fatal(&mut self) -> Option<FatalError> {
441 let downstream = self.next.take_fatal();
446 let own = self.fatal.0.take();
447 downstream.or(own)
448 }
449 fn relieve(&mut self) -> Flow {
450 self.next.relieve()
451 }
452 fn flush_terminal(&mut self) -> Flow {
453 self.next.flush_terminal()
454 }
455}
456
457impl<OutF: RecFamily, G, N: StageLifecycle> StageLifecycle for FlatMap<OutF, G, N> {
458 fn on_batch_end(&mut self, elapsed: Duration) {
459 self.meter.0.flush(elapsed);
460 self.next.on_batch_end(elapsed);
461 }
462 fn take_fatal(&mut self) -> Option<FatalError> {
463 self.next.take_fatal()
464 }
465 fn relieve(&mut self) -> Flow {
466 self.next.relieve()
467 }
468 fn flush_terminal(&mut self) -> Flow {
469 self.next.flush_terminal()
470 }
471}
472
473struct OpsEmit<'a, Ops> {
478 ops: &'a mut Ops,
479 emitted: &'a mut u64,
480 flow: &'a mut Flow,
481}
482
483impl<'buf, T, Ops> EmitRecord<'buf, T> for OpsEmit<'_, Ops>
484where
485 Ops: Collector<T>,
486{
487 #[inline(always)]
488 fn emit(&mut self, rec: Record<T>) -> Flow {
489 *self.emitted += 1;
490 let flow = self.ops.push(rec);
491 if *self.flow != Flow::Blocked {
492 *self.flow = flow;
493 }
494 flow
495 }
496}
497
498static DESER_SKIP_WARN: RateLimit = RateLimit::new(5, Duration::from_secs(10));
499
500#[derive(Debug)]
506struct PendingPayload {
507 bytes: Vec<u8>,
508 key: Option<Vec<u8>>,
509 partition: crate::record::PartitionId,
510 offset: i64,
511 timestamp_ms: i64,
512}
513
514impl PendingPayload {
515 fn from_raw(raw: &RawPayload<'_>) -> Self {
516 PendingPayload {
517 bytes: raw.bytes.to_vec(),
518 key: raw.key.map(<[u8]>::to_vec),
519 partition: raw.partition,
520 offset: raw.offset,
521 timestamp_ms: raw.timestamp_ms,
522 }
523 }
524
525 fn as_raw(&self) -> RawPayload<'_> {
526 RawPayload {
527 bytes: &self.bytes,
528 key: self.key.as_deref(),
529 partition: self.partition,
530 offset: self.offset,
531 timestamp_ms: self.timestamp_ms,
532 }
533 }
534}
535
536enum Step {
538 Continue,
540 Backpressure,
543 NotReady,
546 Fatal(FatalError),
548}
549
550pub struct TypedChain<F: RecFamily, D, Ops> {
555 deser: D,
556 ops: Ops,
557 deser_policy: ErrorPolicy,
558 deser_metrics: Option<Arc<DeserMetrics>>,
559 cursor: usize,
561 mid_batch: bool,
562 pending: Option<PendingPayload>,
565 _family: PhantomData<fn() -> F>,
566}
567
568impl<F: RecFamily, D, Ops> TypedChain<F, D, Ops> {
569 pub(crate) fn new(
570 deser: D,
571 ops: Ops,
572 deser_policy: ErrorPolicy,
573 deser_metrics: Option<Arc<DeserMetrics>>,
574 ) -> Self {
575 TypedChain {
576 deser,
577 ops,
578 deser_policy,
579 deser_metrics,
580 cursor: 0,
581 mid_batch: false,
582 pending: None,
583 _family: PhantomData,
584 }
585 }
586}
587
588impl<F, D, Ops> TypedChain<F, D, Ops>
589where
590 F: RecFamily,
591 D: Deserializer<F>,
592 Ops: for<'buf> Collector<<F as RecFamily>::Rec<'buf>> + StageLifecycle + Send,
593{
594 fn deser_step(
596 &mut self,
597 raw: &RawPayload<'_>,
598 ack: &AckRef,
599 ok: &mut u64,
600 errors: &mut u64,
601 fatal: &mut u64,
602 ) -> Step {
603 let mut flow = Flow::Continue;
604 let mut emitted = 0u64;
605 let result = self.deser.deserialize(
606 raw,
607 ack,
608 &mut OpsEmit {
609 ops: &mut self.ops,
610 emitted: &mut emitted,
611 flow: &mut flow,
612 },
613 );
614 *ok += emitted;
615 match result {
616 Err(DeserError::NotReady { .. }) => {
617 debug_assert_eq!(
618 emitted, 0,
619 "NotReady after emitting records would duplicate them on replay"
620 );
621 return Step::NotReady;
622 }
623 Err(DeserError::Fatal { reason }) => {
624 *fatal += 1;
625 return Step::Fatal(self.ops.take_fatal().unwrap_or(FatalError {
626 component: "deserializer".to_string(),
627 reason,
628 }));
629 }
630 Err(e) => {
631 *errors += 1;
632 match self.deser_policy {
633 ErrorPolicy::Skip => {
634 crate::rate_limited_warn!(
635 DESER_SKIP_WARN,
636 partition = raw.partition.0,
637 offset = raw.offset,
638 error = %e,
639 "payload skipped by deserializer error policy"
640 );
641 }
642 _ => {
643 return Step::Fatal(self.ops.take_fatal().unwrap_or_else(|| FatalError {
647 component: "deserializer".to_string(),
648 reason: e.to_string(),
649 }));
650 }
651 }
652 }
653 Ok(()) => {}
654 }
655 if let Some(fatal) = self.ops.take_fatal() {
659 return Step::Fatal(fatal);
660 }
661 if flow == Flow::Blocked || self.ops.relieve() == Flow::Blocked {
664 return Step::Backpressure;
665 }
666 Step::Continue
667 }
668}
669
670impl<F: RecFamily, D, Ops> std::fmt::Debug for TypedChain<F, D, Ops> {
671 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
672 f.debug_struct("TypedChain")
673 .field("cursor", &self.cursor)
674 .field("mid_batch", &self.mid_batch)
675 .finish_non_exhaustive()
676 }
677}
678
679impl<F, D, Ops> RunnableChain for TypedChain<F, D, Ops>
680where
681 F: RecFamily,
682 D: Deserializer<F>,
683 Ops: for<'buf> Collector<<F as RecFamily>::Rec<'buf>> + StageLifecycle + Send,
684{
685 fn push_batch<'buf>(&mut self, batch: &mut dyn PayloadBatch<'buf>, from: usize) -> PushOutcome {
686 if self.mid_batch {
687 debug_assert_eq!(
688 from, self.cursor,
689 "resume must continue where Blocked stopped"
690 );
691 } else {
692 debug_assert_eq!(from, 0, "fresh batches start at payload 0");
693 debug_assert!(
694 self.pending.is_none(),
695 "a not-ready payload must not survive its batch"
696 );
697 self.cursor = 0;
698 self.pending = None;
699 self.mid_batch = true;
700 }
701
702 if self.ops.relieve() == Flow::Blocked {
705 return PushOutcome::Blocked {
706 resume_at: self.cursor,
707 reason: BlockReason::Capacity,
708 };
709 }
710
711 let started = Instant::now();
712 let mut ok: u64 = 0;
713 let mut errors: u64 = 0;
714 let mut not_ready: u64 = 0;
715 let mut fatal: u64 = 0;
717 let mut outcome: Option<PushOutcome> = None;
718
719 if let Some(p) = self.pending.take() {
722 let raw = p.as_raw();
723 match self.deser_step(&raw, batch.ack(), &mut ok, &mut errors, &mut fatal) {
724 Step::Continue => self.cursor += 1,
725 Step::Backpressure => {
726 self.cursor += 1;
727 outcome = Some(PushOutcome::Blocked {
728 resume_at: self.cursor,
729 reason: BlockReason::Capacity,
730 });
731 }
732 Step::NotReady => {
733 not_ready += 1;
734 self.pending = Some(p);
735 outcome = Some(PushOutcome::Blocked {
736 resume_at: self.cursor,
737 reason: BlockReason::NotReady,
738 });
739 }
740 Step::Fatal(f) => outcome = Some(PushOutcome::Fatal(f)),
741 }
742 }
743
744 while outcome.is_none() {
745 let Some(raw) = batch.next_payload() else {
746 break;
747 };
748 match self.deser_step(&raw, batch.ack(), &mut ok, &mut errors, &mut fatal) {
749 Step::Continue => self.cursor += 1,
750 Step::Backpressure => {
751 self.cursor += 1;
752 outcome = Some(PushOutcome::Blocked {
753 resume_at: self.cursor,
754 reason: BlockReason::Capacity,
755 });
756 }
757 Step::NotReady => {
758 not_ready += 1;
759 self.pending = Some(PendingPayload::from_raw(&raw));
760 outcome = Some(PushOutcome::Blocked {
761 resume_at: self.cursor,
762 reason: BlockReason::NotReady,
763 });
764 }
765 Step::Fatal(f) => outcome = Some(PushOutcome::Fatal(f)),
766 }
767 }
768
769 let elapsed = started.elapsed();
770 if let Some(m) = &self.deser_metrics {
771 m.batch(ok, errors + fatal, elapsed);
772 if errors > 0 && matches!(self.deser_policy, ErrorPolicy::Skip) {
773 m.dropped(errors);
774 }
775 if not_ready > 0 {
776 m.not_ready(not_ready);
777 }
778 }
779 self.ops.on_batch_end(elapsed);
780
781 match outcome {
782 Some(PushOutcome::Fatal(f)) => {
783 batch.ack().fail();
784 self.mid_batch = false;
785 self.pending = None;
786 PushOutcome::Fatal(f)
787 }
788 Some(blocked) => blocked,
789 None => {
790 self.mid_batch = false;
791 PushOutcome::Done
792 }
793 }
794 }
795
796 fn flush(&mut self) -> PushOutcome {
797 if let Some(fatal) = self.ops.take_fatal() {
798 return PushOutcome::Fatal(fatal);
799 }
800 let flow = self.ops.flush_terminal();
801 if let Some(fatal) = self.ops.take_fatal() {
804 return PushOutcome::Fatal(fatal);
805 }
806 match flow {
807 Flow::Continue => PushOutcome::Done,
808 Flow::Blocked => PushOutcome::Blocked {
809 resume_at: self.cursor,
810 reason: BlockReason::Capacity,
811 },
812 }
813 }
814
815 fn abandon_batch(&mut self) {
816 self.mid_batch = false;
822 self.cursor = 0;
823 self.pending = None;
824 }
825}