1use super::{BlockReason, Collector, CollectorFor, PushOutcome, RunnableChain};
17use crate::checkpoint::AckRef;
18use crate::deser::{Deserializer, EmitRecord, RecFamily};
19use crate::error::{DeserError, ErrorClass, 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>;
40
41 fn relieve(&mut self) -> Flow;
44
45 fn flush_terminal(&mut self) -> Flow;
48}
49
50#[derive(Debug, Default)]
53pub(crate) struct OpMeter {
54 handle: Option<Arc<OperatorMetrics>>,
55 records_in: u64,
56 records_out: u64,
57 filtered: u64,
58 skipped: u64,
59 record_errors: u64,
60}
61
62impl OpMeter {
63 pub(crate) fn new(handle: Option<Arc<OperatorMetrics>>) -> Self {
64 OpMeter {
65 handle,
66 ..Default::default()
67 }
68 }
69
70 #[inline(always)]
71 pub(crate) fn seen(&mut self) {
72 self.records_in += 1;
73 }
74
75 #[inline(always)]
76 pub(crate) fn out(&mut self) {
77 self.records_out += 1;
78 }
79
80 #[inline(always)]
81 pub(crate) fn out_n(&mut self, n: u64) {
82 self.records_out += n;
83 }
84
85 #[inline(always)]
86 pub(crate) fn filtered(&mut self) {
87 self.filtered += 1;
88 }
89
90 #[inline(always)]
91 pub(crate) fn skipped(&mut self) {
92 self.skipped += 1;
93 }
94
95 #[inline(always)]
96 pub(crate) fn record_error(&mut self) {
97 self.record_errors += 1;
98 }
99
100 pub(crate) fn flush(&mut self, elapsed: Duration) {
101 if let Some(h) = &self.handle {
102 h.batch(self.records_in, self.records_out, elapsed);
103 if self.filtered > 0 {
104 h.filtered(self.filtered);
105 }
106 if self.skipped > 0 {
107 h.skipped(self.skipped);
108 }
109 if self.record_errors > 0 {
110 h.errors(ErrorClass::RecordLevel, self.record_errors);
111 }
112 }
113 self.records_in = 0;
114 self.records_out = 0;
115 self.filtered = 0;
116 self.skipped = 0;
117 self.record_errors = 0;
118 }
119}
120
121#[derive(Debug, Default)]
124pub(crate) struct OpMeterSlot(pub(crate) OpMeter);
125
126impl Clone for OpMeterSlot {
127 fn clone(&self) -> Self {
128 OpMeterSlot(OpMeter::new(self.0.handle.clone()))
129 }
130}
131
132#[derive(Debug, Default)]
134pub(crate) struct FatalSlot(pub(crate) Option<FatalError>);
135
136impl Clone for FatalSlot {
137 fn clone(&self) -> Self {
138 FatalSlot(None)
139 }
140}
141
142#[derive(Clone, Debug)]
144pub struct Map<G, N> {
145 pub(crate) f: G,
146 pub(crate) next: N,
147 pub(crate) meter: OpMeterSlot,
148}
149
150impl<In, Out, G, N> Collector<In> for Map<G, N>
151where
152 G: FnMut(In) -> Out,
153 N: Collector<Out>,
154{
155 #[inline(always)]
156 fn push(&mut self, rec: Record<In>) -> Flow {
157 self.meter.0.seen();
158 self.meter.0.out();
159 self.next.push(rec.map(&mut self.f))
160 }
161}
162
163#[derive(Clone, Debug)]
166pub struct Filter<P, N> {
167 pub(crate) p: P,
168 pub(crate) next: N,
169 pub(crate) meter: OpMeterSlot,
170}
171
172impl<T, P, N> Collector<T> for Filter<P, N>
173where
174 P: FnMut(&T) -> bool,
175 N: Collector<T>,
176{
177 #[inline(always)]
178 fn push(&mut self, rec: Record<T>) -> Flow {
179 self.meter.0.seen();
180 if (self.p)(&rec.payload) {
181 self.meter.0.out();
182 self.next.push(rec)
183 } else {
184 self.meter.0.filtered();
185 Flow::Continue
186 }
187 }
188}
189
190#[derive(Clone, Debug)]
192pub struct Inspect<G, N> {
193 pub(crate) f: G,
194 pub(crate) next: N,
195}
196
197impl<T, G, N> Collector<T> for Inspect<G, N>
198where
199 G: FnMut(&T),
200 N: Collector<T>,
201{
202 #[inline(always)]
203 fn push(&mut self, rec: Record<T>) -> Flow {
204 (self.f)(&rec.payload);
205 self.next.push(rec)
206 }
207}
208
209static TRY_MAP_SKIP_WARN: RateLimit = RateLimit::new(5, Duration::from_secs(10));
210
211#[derive(Clone, Debug)]
215pub struct TryMap<G, N> {
216 pub(crate) f: G,
217 pub(crate) next: N,
218 pub(crate) policy: ErrorPolicy,
219 pub(crate) component: Arc<str>,
220 pub(crate) meter: OpMeterSlot,
221 pub(crate) fatal: FatalSlot,
222}
223
224impl<In, Out, E, G, N> Collector<In> for TryMap<G, N>
225where
226 G: FnMut(In) -> Result<Out, E>,
227 E: std::fmt::Display,
228 N: Collector<Out>,
229{
230 #[inline(always)]
231 fn push(&mut self, rec: Record<In>) -> Flow {
232 self.meter.0.seen();
233 if self.fatal.0.is_some() {
234 return Flow::Continue;
237 }
238 let Record { payload, meta, ack } = rec;
239 match (self.f)(payload) {
240 Ok(out) => {
241 self.meter.0.out();
242 self.next.push(Record {
243 payload: out,
244 meta,
245 ack,
246 })
247 }
248 Err(e) => {
249 match self.policy {
250 ErrorPolicy::Skip => {
251 self.meter.0.skipped();
252 self.meter.0.record_error();
253 crate::rate_limited_warn!(
254 TRY_MAP_SKIP_WARN,
255 component = &*self.component,
256 error = %e,
257 "record skipped by error policy"
258 );
259 }
260 _ => {
261 self.fatal.0 = Some(FatalError {
262 component: self.component.to_string(),
263 reason: e.to_string(),
264 });
265 }
266 }
267 Flow::Continue
268 }
269 }
270 }
271}
272
273pub struct Emitter<'a, OutF: RecFamily> {
278 next: &'a mut dyn CollectorFor<OutF>,
279 meta: RecordMeta,
280 ack: &'a AckRef,
281 emitted: u64,
282 flow: Flow,
283}
284
285impl<OutF: RecFamily> std::fmt::Debug for Emitter<'_, OutF> {
286 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
287 f.debug_struct("Emitter")
288 .field("emitted", &self.emitted)
289 .field("flow", &self.flow)
290 .finish_non_exhaustive()
291 }
292}
293
294impl<OutF: RecFamily> Emitter<'_, OutF> {
295 #[inline(always)]
298 pub fn emit<'buf>(&mut self, payload: OutF::Rec<'buf>) -> Flow {
299 let flow = self.next.push_rec(Record {
300 payload,
301 meta: self.meta,
302 ack: self.ack.clone(),
303 });
304 self.emitted += 1;
305 if self.flow != Flow::Blocked {
306 self.flow = flow;
307 }
308 flow
309 }
310
311 #[must_use]
313 pub fn meta(&self) -> RecordMeta {
314 self.meta
315 }
316}
317
318#[derive(Clone, Debug)]
322pub struct FlatMap<OutF: RecFamily, G, N> {
323 pub(crate) g: G,
324 pub(crate) next: N,
325 pub(crate) meter: OpMeterSlot,
326 pub(crate) _out: PhantomData<fn() -> OutF>,
327}
328
329impl<In, OutF, G, N> Collector<In> for FlatMap<OutF, G, N>
330where
331 OutF: RecFamily,
332 G: FnMut(In, &mut Emitter<'_, OutF>),
333 N: for<'buf> Collector<<OutF as RecFamily>::Rec<'buf>>,
334{
335 #[inline(always)]
336 fn push(&mut self, rec: Record<In>) -> Flow {
337 self.meter.0.seen();
338 let Record { payload, meta, ack } = rec;
339 let mut em = Emitter {
340 next: &mut self.next,
341 meta,
342 ack: &ack,
343 emitted: 0,
344 flow: Flow::Continue,
345 };
346 (self.g)(payload, &mut em);
347 let (emitted, flow) = (em.emitted, em.flow);
348 self.meter.0.out_n(emitted);
349 flow
350 }
351}
352
353impl<G, N: StageLifecycle> StageLifecycle for Map<G, N> {
356 fn on_batch_end(&mut self, elapsed: Duration) {
357 self.meter.0.flush(elapsed);
358 self.next.on_batch_end(elapsed);
359 }
360 fn take_fatal(&mut self) -> Option<FatalError> {
361 self.next.take_fatal()
362 }
363 fn relieve(&mut self) -> Flow {
364 self.next.relieve()
365 }
366 fn flush_terminal(&mut self) -> Flow {
367 self.next.flush_terminal()
368 }
369}
370
371impl<P, N: StageLifecycle> StageLifecycle for Filter<P, N> {
372 fn on_batch_end(&mut self, elapsed: Duration) {
373 self.meter.0.flush(elapsed);
374 self.next.on_batch_end(elapsed);
375 }
376 fn take_fatal(&mut self) -> Option<FatalError> {
377 self.next.take_fatal()
378 }
379 fn relieve(&mut self) -> Flow {
380 self.next.relieve()
381 }
382 fn flush_terminal(&mut self) -> Flow {
383 self.next.flush_terminal()
384 }
385}
386
387impl<G, N: StageLifecycle> StageLifecycle for Inspect<G, N> {
388 fn on_batch_end(&mut self, elapsed: Duration) {
389 self.next.on_batch_end(elapsed);
390 }
391 fn take_fatal(&mut self) -> Option<FatalError> {
392 self.next.take_fatal()
393 }
394 fn relieve(&mut self) -> Flow {
395 self.next.relieve()
396 }
397 fn flush_terminal(&mut self) -> Flow {
398 self.next.flush_terminal()
399 }
400}
401
402impl<G, N: StageLifecycle> StageLifecycle for TryMap<G, N> {
403 fn on_batch_end(&mut self, elapsed: Duration) {
404 self.meter.0.flush(elapsed);
405 self.next.on_batch_end(elapsed);
406 }
407 fn take_fatal(&mut self) -> Option<FatalError> {
408 self.fatal.0.take().or_else(|| self.next.take_fatal())
409 }
410 fn relieve(&mut self) -> Flow {
411 self.next.relieve()
412 }
413 fn flush_terminal(&mut self) -> Flow {
414 self.next.flush_terminal()
415 }
416}
417
418impl<OutF: RecFamily, G, N: StageLifecycle> StageLifecycle for FlatMap<OutF, G, N> {
419 fn on_batch_end(&mut self, elapsed: Duration) {
420 self.meter.0.flush(elapsed);
421 self.next.on_batch_end(elapsed);
422 }
423 fn take_fatal(&mut self) -> Option<FatalError> {
424 self.next.take_fatal()
425 }
426 fn relieve(&mut self) -> Flow {
427 self.next.relieve()
428 }
429 fn flush_terminal(&mut self) -> Flow {
430 self.next.flush_terminal()
431 }
432}
433
434struct OpsEmit<'a, Ops> {
439 ops: &'a mut Ops,
440 emitted: &'a mut u64,
441 flow: &'a mut Flow,
442}
443
444impl<'buf, T, Ops> EmitRecord<'buf, T> for OpsEmit<'_, Ops>
445where
446 Ops: Collector<T>,
447{
448 #[inline(always)]
449 fn emit(&mut self, rec: Record<T>) -> Flow {
450 *self.emitted += 1;
451 let flow = self.ops.push(rec);
452 if *self.flow != Flow::Blocked {
453 *self.flow = flow;
454 }
455 flow
456 }
457}
458
459static DESER_SKIP_WARN: RateLimit = RateLimit::new(5, Duration::from_secs(10));
460
461#[derive(Debug)]
467struct PendingPayload {
468 bytes: Vec<u8>,
469 key: Option<Vec<u8>>,
470 partition: crate::record::PartitionId,
471 offset: i64,
472 timestamp_ms: i64,
473}
474
475impl PendingPayload {
476 fn from_raw(raw: &RawPayload<'_>) -> Self {
477 PendingPayload {
478 bytes: raw.bytes.to_vec(),
479 key: raw.key.map(<[u8]>::to_vec),
480 partition: raw.partition,
481 offset: raw.offset,
482 timestamp_ms: raw.timestamp_ms,
483 }
484 }
485
486 fn as_raw(&self) -> RawPayload<'_> {
487 RawPayload {
488 bytes: &self.bytes,
489 key: self.key.as_deref(),
490 partition: self.partition,
491 offset: self.offset,
492 timestamp_ms: self.timestamp_ms,
493 }
494 }
495}
496
497enum Step {
499 Continue,
501 Backpressure,
504 NotReady,
507 Fatal(FatalError),
509}
510
511pub struct TypedChain<F: RecFamily, D, Ops> {
516 deser: D,
517 ops: Ops,
518 deser_policy: ErrorPolicy,
519 deser_metrics: Option<Arc<DeserMetrics>>,
520 cursor: usize,
522 mid_batch: bool,
523 pending: Option<PendingPayload>,
526 _family: PhantomData<fn() -> F>,
527}
528
529impl<F: RecFamily, D, Ops> TypedChain<F, D, Ops> {
530 pub(crate) fn new(
531 deser: D,
532 ops: Ops,
533 deser_policy: ErrorPolicy,
534 deser_metrics: Option<Arc<DeserMetrics>>,
535 ) -> Self {
536 TypedChain {
537 deser,
538 ops,
539 deser_policy,
540 deser_metrics,
541 cursor: 0,
542 mid_batch: false,
543 pending: None,
544 _family: PhantomData,
545 }
546 }
547}
548
549impl<F, D, Ops> TypedChain<F, D, Ops>
550where
551 F: RecFamily,
552 D: Deserializer<F>,
553 Ops: for<'buf> Collector<<F as RecFamily>::Rec<'buf>> + StageLifecycle + Send,
554{
555 fn deser_step(
557 &mut self,
558 raw: &RawPayload<'_>,
559 ack: &AckRef,
560 ok: &mut u64,
561 errors: &mut u64,
562 ) -> Step {
563 let mut flow = Flow::Continue;
564 let mut emitted = 0u64;
565 let result = self.deser.deserialize(
566 raw,
567 ack,
568 &mut OpsEmit {
569 ops: &mut self.ops,
570 emitted: &mut emitted,
571 flow: &mut flow,
572 },
573 );
574 *ok += emitted;
575 match result {
576 Err(DeserError::NotReady { .. }) => {
577 debug_assert_eq!(
578 emitted, 0,
579 "NotReady after emitting records would duplicate them on replay"
580 );
581 return Step::NotReady;
582 }
583 Err(e) => {
584 *errors += 1;
585 match self.deser_policy {
586 ErrorPolicy::Skip => {
587 crate::rate_limited_warn!(
588 DESER_SKIP_WARN,
589 partition = raw.partition.0,
590 offset = raw.offset,
591 error = %e,
592 "payload skipped by deserializer error policy"
593 );
594 }
595 _ => {
596 return Step::Fatal(FatalError {
597 component: "deserializer".to_string(),
598 reason: e.to_string(),
599 });
600 }
601 }
602 }
603 Ok(()) => {}
604 }
605 if let Some(fatal) = self.ops.take_fatal() {
606 return Step::Fatal(fatal);
607 }
608 if flow == Flow::Blocked || self.ops.relieve() == Flow::Blocked {
611 return Step::Backpressure;
612 }
613 Step::Continue
614 }
615}
616
617impl<F: RecFamily, D, Ops> std::fmt::Debug for TypedChain<F, D, Ops> {
618 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
619 f.debug_struct("TypedChain")
620 .field("cursor", &self.cursor)
621 .field("mid_batch", &self.mid_batch)
622 .finish_non_exhaustive()
623 }
624}
625
626impl<F, D, Ops> RunnableChain for TypedChain<F, D, Ops>
627where
628 F: RecFamily,
629 D: Deserializer<F>,
630 Ops: for<'buf> Collector<<F as RecFamily>::Rec<'buf>> + StageLifecycle + Send,
631{
632 fn push_batch<'buf>(&mut self, batch: &mut dyn PayloadBatch<'buf>, from: usize) -> PushOutcome {
633 if self.mid_batch {
634 debug_assert_eq!(
635 from, self.cursor,
636 "resume must continue where Blocked stopped"
637 );
638 } else {
639 debug_assert_eq!(from, 0, "fresh batches start at payload 0");
640 debug_assert!(
641 self.pending.is_none(),
642 "a not-ready payload must not survive its batch"
643 );
644 self.cursor = 0;
645 self.pending = None;
646 self.mid_batch = true;
647 }
648
649 if self.ops.relieve() == Flow::Blocked {
652 return PushOutcome::Blocked {
653 resume_at: self.cursor,
654 reason: BlockReason::Capacity,
655 };
656 }
657
658 let started = Instant::now();
659 let mut ok: u64 = 0;
660 let mut errors: u64 = 0;
661 let mut not_ready: u64 = 0;
662 let mut outcome: Option<PushOutcome> = None;
663
664 if let Some(p) = self.pending.take() {
667 let raw = p.as_raw();
668 match self.deser_step(&raw, batch.ack(), &mut ok, &mut errors) {
669 Step::Continue => self.cursor += 1,
670 Step::Backpressure => {
671 self.cursor += 1;
672 outcome = Some(PushOutcome::Blocked {
673 resume_at: self.cursor,
674 reason: BlockReason::Capacity,
675 });
676 }
677 Step::NotReady => {
678 not_ready += 1;
679 self.pending = Some(p);
680 outcome = Some(PushOutcome::Blocked {
681 resume_at: self.cursor,
682 reason: BlockReason::NotReady,
683 });
684 }
685 Step::Fatal(f) => outcome = Some(PushOutcome::Fatal(f)),
686 }
687 }
688
689 while outcome.is_none() {
690 let Some(raw) = batch.next_payload() else {
691 break;
692 };
693 match self.deser_step(&raw, batch.ack(), &mut ok, &mut errors) {
694 Step::Continue => self.cursor += 1,
695 Step::Backpressure => {
696 self.cursor += 1;
697 outcome = Some(PushOutcome::Blocked {
698 resume_at: self.cursor,
699 reason: BlockReason::Capacity,
700 });
701 }
702 Step::NotReady => {
703 not_ready += 1;
704 self.pending = Some(PendingPayload::from_raw(&raw));
705 outcome = Some(PushOutcome::Blocked {
706 resume_at: self.cursor,
707 reason: BlockReason::NotReady,
708 });
709 }
710 Step::Fatal(f) => outcome = Some(PushOutcome::Fatal(f)),
711 }
712 }
713
714 let elapsed = started.elapsed();
715 if let Some(m) = &self.deser_metrics {
716 m.batch(ok, errors, elapsed);
717 if errors > 0 && matches!(self.deser_policy, ErrorPolicy::Skip) {
718 m.dropped(errors);
719 }
720 if not_ready > 0 {
721 m.not_ready(not_ready);
722 }
723 }
724 self.ops.on_batch_end(elapsed);
725
726 match outcome {
727 Some(PushOutcome::Fatal(f)) => {
728 batch.ack().fail();
729 self.mid_batch = false;
730 self.pending = None;
731 PushOutcome::Fatal(f)
732 }
733 Some(blocked) => blocked,
734 None => {
735 self.mid_batch = false;
736 PushOutcome::Done
737 }
738 }
739 }
740
741 fn flush(&mut self) -> PushOutcome {
742 if let Some(fatal) = self.ops.take_fatal() {
743 return PushOutcome::Fatal(fatal);
744 }
745 match self.ops.flush_terminal() {
746 Flow::Continue => PushOutcome::Done,
747 Flow::Blocked => PushOutcome::Blocked {
748 resume_at: self.cursor,
749 reason: BlockReason::Capacity,
750 },
751 }
752 }
753
754 fn abandon_batch(&mut self) {
755 self.mid_batch = false;
761 self.cursor = 0;
762 self.pending = None;
763 }
764}