1use crate::config::Compression;
25use crate::fetch::{ChunkMsg, SplitFailure};
26use crate::framer::{Codec, FramerFactory, ObjectFramer};
27use crate::metrics::S3Metrics;
28use crate::offset::{MAX_RECORD_INDEX, Position};
29use crate::split_ctx::{PoisonKind, PoisonReport, SplitTracker};
30use spate_core::checkpoint::{AckIssuer, AckRef};
31use spate_core::coordination::{ControlWaker, SplitId};
32use spate_core::error::{ErrorClass, SourceError};
33use spate_core::record::{PartitionId, RawPayload};
34use spate_core::source::{LaneId, PayloadBatch, SourceLane};
35use std::sync::Arc;
36use std::time::{Duration, Instant};
37use tokio::sync::mpsc;
38use tokio::sync::mpsc::error::TryRecvError;
39
40struct CurrentObject {
42 ordinal: u32,
43 key: String,
44 next_record: u64,
46 event_time_ms: i64,
48}
49
50struct HeldRecord {
52 offset: i64,
53 event_time_ms: i64,
54 bytes: Vec<u8>,
55}
56
57pub struct S3Lane {
60 id: LaneId,
61 partition: PartitionId,
62 rx: mpsc::Receiver<ChunkMsg>,
63 handle: tokio::runtime::Handle,
64 issuer: AckIssuer,
65 compression: Compression,
66 framer: ObjectFramer,
67 resume: Option<Position>,
69 pending_discard: u64,
72 current: Option<CurrentObject>,
73 pending_end: bool,
76 held: Vec<HeldRecord>,
77 split: SplitId,
79 tracker: Arc<SplitTracker>,
82 poison_tx: std::sync::mpsc::Sender<PoisonReport>,
84 waker: ControlWaker,
88 poisoned: bool,
91 watermark_candidate: i64,
96 failed: Option<(ErrorClass, String)>,
98 metrics: Option<S3Metrics>,
99 decoded_reported: u64,
101}
102
103impl std::fmt::Debug for S3Lane {
104 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105 f.debug_struct("S3Lane")
106 .field("id", &self.id)
107 .field("partition", &self.partition)
108 .finish()
109 }
110}
111
112impl S3Lane {
113 #[expect(
114 clippy::too_many_arguments,
115 reason = "assembled in one place by the source's lane builder"
116 )]
117 pub(crate) fn new(
118 id: LaneId,
119 partition: PartitionId,
120 rx: mpsc::Receiver<ChunkMsg>,
121 handle: tokio::runtime::Handle,
122 issuer: AckIssuer,
123 compression: Compression,
124 make_framer: FramerFactory,
125 resume: Option<Position>,
126 split: SplitId,
127 tracker: Arc<SplitTracker>,
128 poison_tx: std::sync::mpsc::Sender<PoisonReport>,
129 waker: ControlWaker,
130 metrics: Option<S3Metrics>,
131 ) -> S3Lane {
132 let watermark_candidate = resume.map_or(0, |p| {
133 p.encode()
134 .expect("a decoded resume position always re-encodes")
135 });
136 S3Lane {
137 id,
138 partition,
139 rx,
140 handle,
141 issuer,
142 compression,
143 framer: ObjectFramer::new(make_framer),
144 resume,
145 pending_discard: 0,
146 current: None,
147 pending_end: false,
148 held: Vec::new(),
149 split,
150 tracker,
151 poison_tx,
152 waker,
153 poisoned: false,
154 watermark_candidate,
155 failed: None,
156 metrics,
157 decoded_reported: 0,
158 }
159 }
160
161 fn fail(&mut self, class: ErrorClass, reason: String) -> SourceError {
163 self.failed = Some((class, reason.clone()));
164 SourceError::Client { class, reason }
165 }
166
167 fn poison(&mut self, kind: PoisonKind, reason: String) {
172 self.held.clear();
173 self.current = None;
174 self.pending_end = false;
175 self.pending_discard = 0;
176 self.poisoned = true;
177 let _ = self.poison_tx.send(PoisonReport {
179 split: self.split.clone(),
180 kind,
181 reason,
182 });
183 self.waker.wake();
186 }
187
188 fn drain_framer(&mut self, max_records: usize) -> Result<(), SourceError> {
193 while self.held.len() < max_records {
194 let Some(bytes) = self.framer.pop_record() else {
195 break;
196 };
197 if self.pending_discard > 0 {
198 self.pending_discard -= 1;
199 continue;
200 }
201 let cur = self
202 .current
203 .as_mut()
204 .expect("framed records only exist within an object");
205 if cur.next_record > MAX_RECORD_INDEX {
206 let key = cur.key.clone();
209 self.poison(
210 PoisonKind::Undecodable,
211 format!(
212 "object \"{key}\" holds more than {} records, the composite-offset \
213 limit per object",
214 MAX_RECORD_INDEX + 1
215 ),
216 );
217 return Ok(());
218 }
219 let pos = Position {
220 ordinal: cur.ordinal,
221 record: cur.next_record,
222 };
223 let offset = match pos.encode() {
224 Ok(o) => o,
225 Err(e) => return Err(self.fail(ErrorClass::Fatal, e.to_string())),
226 };
227 cur.next_record += 1;
228 debug_assert!(
229 offset >= self.watermark_candidate,
230 "offsets must be monotonic"
231 );
232 self.watermark_candidate = offset + 1;
233 self.held.push(HeldRecord {
234 offset,
235 event_time_ms: cur.event_time_ms,
236 bytes,
237 });
238 }
239 Ok(())
240 }
241
242 fn finalize_object(&mut self) {
246 debug_assert!(self.pending_end && self.framer.queued() == 0);
247 let cur = self.current.take().expect("finalize without an object");
248 self.pending_end = false;
249 if self.pending_discard > 0 {
250 self.poison(
253 PoisonKind::EtagDrift,
254 format!(
255 "object \"{}\" ended {} records short of its committed position — \
256 its content changed under the checkpoint (the key set and object \
257 contents must stay frozen for the backfill's lifetime)",
258 cur.key, self.pending_discard
259 ),
260 );
261 return;
262 }
263 self.tracker.object_done();
264 if let Some(m) = &self.metrics {
265 m.objects_completed.increment(1);
266 m.objects_remaining.decrement(1.0);
267 }
268 }
269
270 fn on_msg(&mut self, msg: ChunkMsg) -> Result<(), SourceError> {
272 match msg {
273 ChunkMsg::ObjectStart {
274 ordinal,
275 key,
276 last_modified_ms,
277 } => {
278 debug_assert!(
279 self.current.is_none() && !self.pending_end,
280 "ObjectStart while the previous object is open"
281 );
282 let codec = Codec::resolve(self.compression, &key);
283 if let Err(e) = self.framer.begin_object(codec) {
284 self.poison(
285 PoisonKind::Undecodable,
286 format!("starting decode of \"{key}\": {e}"),
287 );
288 return Ok(());
289 }
290 let discard = match self.resume.take() {
294 Some(pos) if pos.ordinal == ordinal => pos.record,
295 Some(pos) => {
296 return Err(self.fail(
297 ErrorClass::Fatal,
298 format!(
299 "fetcher started at ordinal {ordinal} but the committed \
300 resume position is ordinal {} — internal wiring bug",
301 pos.ordinal
302 ),
303 ));
304 }
305 None => 0,
306 };
307 self.pending_discard = discard;
308 self.current = Some(CurrentObject {
309 ordinal,
310 key,
311 next_record: discard,
312 event_time_ms: last_modified_ms,
313 });
314 Ok(())
315 }
316 ChunkMsg::Chunk(bytes) => {
317 if let Some(m) = &self.metrics {
318 m.bytes_read.increment(bytes.len() as u64);
319 }
320 if let Err(e) = self.framer.push_chunk(&bytes) {
321 let key = self
322 .current
323 .as_ref()
324 .map_or_else(String::new, |c| c.key.clone());
325 self.poison(
326 PoisonKind::Undecodable,
327 format!("decoding \"{key}\": {e} (corrupt or truncated object?)"),
328 );
329 }
330 Ok(())
331 }
332 ChunkMsg::ObjectEnd => {
333 if let Err(e) = self.framer.finish_object() {
334 let key = self
335 .current
336 .as_ref()
337 .map_or_else(String::new, |c| c.key.clone());
338 self.poison(
339 PoisonKind::Undecodable,
340 format!("finishing decode of \"{key}\": {e} (truncated object?)"),
341 );
342 return Ok(());
343 }
344 self.pending_end = true;
345 Ok(())
346 }
347 ChunkMsg::LaneFailed(SplitFailure::Poison(kind, reason)) => {
348 self.poison(kind, reason);
349 Ok(())
350 }
351 ChunkMsg::LaneFailed(SplitFailure::Fatal(e)) => {
352 let (class, reason) = match e {
353 SourceError::Client { class, reason } => (class, reason),
354 other => (ErrorClass::Fatal, other.to_string()),
355 };
356 Err(self.fail(class, reason))
357 }
358 }
359 }
360
361 fn report_decoded(&mut self) {
363 if let Some(m) = &self.metrics {
364 let total = self.framer.decoded_bytes();
365 m.bytes_decoded
366 .increment(total.saturating_sub(self.decoded_reported));
367 self.decoded_reported = total;
368 }
369 }
370}
371
372impl SourceLane for S3Lane {
373 type Batch<'a> = S3Batch<'a>;
374
375 fn id(&self) -> LaneId {
376 self.id
377 }
378
379 fn partition(&self) -> PartitionId {
380 self.partition
381 }
382
383 fn poll(
384 &mut self,
385 max_records: usize,
386 timeout: Duration,
387 ) -> Result<Option<S3Batch<'_>>, SourceError> {
388 if let Some((class, reason)) = &self.failed {
389 return Err(SourceError::Client {
390 class: *class,
391 reason: reason.clone(),
392 });
393 }
394 if self.poisoned {
395 std::thread::sleep(timeout);
399 return Ok(None);
400 }
401 self.held.clear();
402 let deadline = Instant::now() + timeout;
403 let mut recv_attempted = false;
410
411 loop {
412 self.drain_framer(max_records)?;
413 if self.pending_end && !self.poisoned && self.framer.queued() == 0 {
414 self.finalize_object();
415 }
416 if self.poisoned {
417 return Ok(None);
418 }
419 if self.held.len() >= max_records {
420 break;
421 }
422 debug_assert!(!self.pending_end, "pending object end past the drain");
426 if recv_attempted && Instant::now() >= deadline {
435 if self.held.is_empty() {
436 return Ok(None);
437 }
438 break;
439 }
440 recv_attempted = true;
441 let msg = match self.rx.try_recv() {
442 Ok(m) => Some(m),
443 Err(TryRecvError::Empty) => {
444 if !self.held.is_empty() {
445 break; }
447 let remaining = deadline.saturating_duration_since(Instant::now());
448 if remaining.is_zero() {
449 return Ok(None);
450 }
451 let rx = &mut self.rx;
452 match self
456 .handle
457 .block_on(async { tokio::time::timeout(remaining, rx.recv()).await })
458 {
459 Ok(Some(m)) => Some(m),
460 Ok(None) => None,
461 Err(_) => return Ok(None), }
463 }
464 Err(TryRecvError::Disconnected) => None,
465 };
466 match msg {
467 Some(m) => {
468 self.on_msg(m)?;
469 if self.poisoned {
470 return Ok(None);
471 }
472 }
473 None => {
474 if self.current.is_some() {
478 return Err(self.fail(
479 ErrorClass::Fatal,
480 "object stream ended mid-object: the fetcher terminated \
481 unexpectedly"
482 .into(),
483 ));
484 }
485 if self.framer.queued() > 0 {
486 continue; }
488 if !self.held.is_empty() {
489 break; }
491 self.tracker.set_terminal(self.watermark_candidate);
497 self.waker.wake();
502 return Ok(None);
503 }
504 }
505 }
506
507 self.report_decoded();
508 let last = self
509 .held
510 .last()
511 .expect("non-empty batch at emit time")
512 .offset;
513 let ack = self.issuer.issue(self.partition, last);
514 Ok(Some(S3Batch {
515 records: &self.held,
516 idx: 0,
517 partition: self.partition,
518 ack,
519 }))
520 }
521}
522
523pub struct S3Batch<'a> {
525 records: &'a [HeldRecord],
526 idx: usize,
527 partition: PartitionId,
528 ack: AckRef,
529}
530
531impl std::fmt::Debug for S3Batch<'_> {
532 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
533 f.debug_struct("S3Batch")
534 .field("records", &self.records.len())
535 .field("idx", &self.idx)
536 .finish()
537 }
538}
539
540impl<'a> PayloadBatch<'a> for S3Batch<'a> {
541 fn next_payload(&mut self) -> Option<RawPayload<'a>> {
542 let rec = self.records.get(self.idx)?;
543 self.idx += 1;
544 Some(RawPayload {
545 bytes: &rec.bytes,
546 key: None,
547 partition: self.partition,
548 offset: rec.offset,
549 timestamp_ms: rec.event_time_ms,
550 })
551 }
552
553 fn ack(&self) -> &AckRef {
554 &self.ack
555 }
556}
557
558#[cfg(test)]
559mod tests {
560 use super::*;
561 use crate::testutil::TestLineFramer;
562 use spate_core::checkpoint::Checkpointer;
563
564 fn runtime() -> tokio::runtime::Runtime {
565 tokio::runtime::Builder::new_multi_thread()
566 .worker_threads(1)
567 .enable_all()
568 .build()
569 .unwrap()
570 }
571
572 struct LaneRig {
573 lane: S3Lane,
574 tx: Option<mpsc::Sender<ChunkMsg>>,
575 tracker: Arc<SplitTracker>,
576 poison_rx: std::sync::mpsc::Receiver<PoisonReport>,
577 _rt: tokio::runtime::Runtime,
578 }
579
580 fn rig(resume: Option<Position>) -> LaneRig {
581 let rt = runtime();
582 let (tx, rx) = mpsc::channel(64);
583 let checkpointer = Checkpointer::new();
584 let tracker = Arc::new(SplitTracker::new());
585 let (poison_tx, poison_rx) = std::sync::mpsc::channel();
586 let lane = S3Lane::new(
587 LaneId(0),
588 PartitionId(0),
589 rx,
590 rt.handle().clone(),
591 checkpointer.handle(),
592 Compression::Auto,
593 Arc::new(|| Box::new(TestLineFramer::new(1 << 20))),
594 resume,
595 SplitId::new("s3-test").unwrap(),
596 Arc::clone(&tracker),
597 poison_tx,
598 ControlWaker::inert(),
599 None,
600 );
601 LaneRig {
602 lane,
603 tx: Some(tx),
604 tracker,
605 poison_rx,
606 _rt: rt,
607 }
608 }
609
610 fn send(rig: &LaneRig, msg: ChunkMsg) {
611 rig.tx.as_ref().unwrap().try_send(msg).unwrap();
612 }
613
614 fn start(rig: &LaneRig, ordinal: u32, key: &str) {
615 send(
616 rig,
617 ChunkMsg::ObjectStart {
618 ordinal,
619 key: key.into(),
620 last_modified_ms: 1_000,
621 },
622 );
623 }
624
625 fn poll_batch(lane: &mut S3Lane, max: usize) -> Option<Vec<(i64, Vec<u8>)>> {
627 let batch = lane
628 .poll(max, Duration::from_millis(50))
629 .expect("poll must succeed");
630 batch.map(|mut b| {
631 let mut out = Vec::new();
632 while let Some(p) = b.next_payload() {
633 out.push((p.offset, p.bytes.to_vec()));
634 }
635 out
636 })
637 }
638
639 #[test]
640 fn frames_records_with_composite_offsets_across_objects() {
641 let mut r = rig(None);
642 start(&r, 0, "p/a.ndjson");
643 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"a1\na2\n")));
644 send(&r, ChunkMsg::ObjectEnd);
645 start(&r, 1, "p/b.ndjson");
646 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"b1\n")));
647 send(&r, ChunkMsg::ObjectEnd);
648 let records = poll_batch(&mut r.lane, 512).unwrap();
649 let expect = |ord: u32, rec: u64| {
650 Position {
651 ordinal: ord,
652 record: rec,
653 }
654 .encode()
655 .unwrap()
656 };
657 assert_eq!(
658 records,
659 vec![
660 (expect(0, 0), b"a1".to_vec()),
661 (expect(0, 1), b"a2".to_vec()),
662 (expect(1, 0), b"b1".to_vec()),
663 ],
664 "a batch may span objects; offsets stay monotonic"
665 );
666 }
667
668 #[test]
669 fn terminal_watermark_is_recorded_on_the_poll_after_the_final_batch() {
670 let mut r = rig(None);
671 start(&r, 0, "p/a.ndjson");
672 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"only\n")));
673 send(&r, ChunkMsg::ObjectEnd);
674 r.tx.take(); let records = poll_batch(&mut r.lane, 512).unwrap();
676 assert_eq!(records.len(), 1);
677 assert!(
678 r.tracker.terminal().is_none(),
679 "no terminal while the final batch is being handed out"
680 );
681 assert!(poll_batch(&mut r.lane, 512).is_none());
682 let expected = Position {
683 ordinal: 0,
684 record: 1,
685 }
686 .encode()
687 .unwrap();
688 assert_eq!(
689 r.tracker.terminal(),
690 Some(expected),
691 "terminal = one past the last emitted record, decided by the None poll"
692 );
693 }
694
695 #[test]
696 fn empty_input_terminates_at_the_resume_watermark() {
697 let resume = Position {
700 ordinal: 2,
701 record: 3,
702 };
703 let mut r = rig(Some(resume));
704 r.tx.take(); assert!(poll_batch(&mut r.lane, 512).is_none());
706 assert_eq!(r.tracker.terminal(), Some(resume.encode().unwrap()));
707 }
708
709 #[test]
710 fn resume_discards_the_committed_record_count() {
711 let mut r = rig(Some(Position {
713 ordinal: 0,
714 record: 2,
715 }));
716 start(&r, 0, "p/a.ndjson");
717 send(
718 &r,
719 ChunkMsg::Chunk(bytes::Bytes::from_static(b"r0\nr1\nr2\nr3\n")),
720 );
721 send(&r, ChunkMsg::ObjectEnd);
722 let records = poll_batch(&mut r.lane, 512).unwrap();
723 let expect = |rec: u64| {
724 Position {
725 ordinal: 0,
726 record: rec,
727 }
728 .encode()
729 .unwrap()
730 };
731 assert_eq!(
732 records,
733 vec![(expect(2), b"r2".to_vec()), (expect(3), b"r3".to_vec())],
734 "replayed committed records are discarded; indexes continue"
735 );
736 }
737
738 #[test]
739 fn resume_exactly_at_object_end_advances_cleanly() {
740 let mut r = rig(Some(Position {
742 ordinal: 0,
743 record: 2,
744 }));
745 start(&r, 0, "p/a.ndjson");
746 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"r0\nr1\n")));
747 send(&r, ChunkMsg::ObjectEnd);
748 start(&r, 1, "p/b.ndjson");
749 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"next\n")));
750 send(&r, ChunkMsg::ObjectEnd);
751 let records = poll_batch(&mut r.lane, 512).unwrap();
752 assert_eq!(
753 records,
754 vec![(
755 Position {
756 ordinal: 1,
757 record: 0
758 }
759 .encode()
760 .unwrap(),
761 b"next".to_vec()
762 )]
763 );
764 }
765
766 #[test]
767 fn resume_object_shorter_than_committed_poisons_the_split() {
768 let mut r = rig(Some(Position {
771 ordinal: 0,
772 record: 3,
773 }));
774 start(&r, 0, "p/a.ndjson");
775 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"only\n")));
776 send(&r, ChunkMsg::ObjectEnd);
777 assert!(poll_batch(&mut r.lane, 512).is_none());
778 let report = r.poison_rx.try_recv().expect("a poison report");
779 assert!(report.reason.contains("short"), "{}", report.reason);
780 assert_eq!(report.split.as_str(), "s3-test");
781 assert!(poll_batch(&mut r.lane, 512).is_none());
784 assert!(r.poison_rx.try_recv().is_err(), "poison reports once");
785 assert!(r.tracker.terminal().is_none());
786 }
787
788 #[test]
789 fn fetcher_poison_goes_quiescent_and_discards_undelivered_records() {
790 let mut r = rig(None);
791 start(&r, 0, "p/a.ndjson");
792 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"a1\n")));
793 send(
794 &r,
795 ChunkMsg::LaneFailed(SplitFailure::Poison(
796 PoisonKind::NotFound,
797 "object vanished".into(),
798 )),
799 );
800 assert!(poll_batch(&mut r.lane, 512).is_none());
805 let report = r.poison_rx.try_recv().expect("a poison report");
806 assert!(report.reason.contains("vanished"), "{}", report.reason);
807 }
808
809 #[test]
810 fn fetcher_fatal_failure_is_terminal() {
811 let mut r = rig(None);
812 send(
813 &r,
814 ChunkMsg::LaneFailed(SplitFailure::Fatal(SourceError::Client {
815 class: ErrorClass::Fatal,
816 reason: "access denied".into(),
817 })),
818 );
819 let err = r.lane.poll(512, Duration::from_millis(50)).unwrap_err();
820 assert!(err.to_string().contains("access denied"), "{err}");
821 let again = r.lane.poll(512, Duration::from_millis(50)).unwrap_err();
822 assert!(again.to_string().contains("access denied"), "sticky");
823 assert!(r.poison_rx.try_recv().is_err(), "fatal is not poison");
824 }
825
826 #[test]
827 fn mid_object_disconnect_is_fatal() {
828 let mut r = rig(None);
829 start(&r, 0, "p/a.ndjson");
830 send(&r, ChunkMsg::Chunk(bytes::Bytes::from_static(b"partial")));
831 r.tx.take(); let err = r.lane.poll(512, Duration::from_millis(50)).unwrap_err();
835 assert!(err.to_string().contains("mid-object"), "{err}");
836 assert!(
837 r.tracker.terminal().is_none(),
838 "a dead lane never terminates"
839 );
840 }
841
842 #[test]
843 fn batch_respects_max_records_and_continues() {
844 let mut r = rig(None);
845 start(&r, 0, "p/a.ndjson");
846 send(
847 &r,
848 ChunkMsg::Chunk(bytes::Bytes::from_static(b"1\n2\n3\n4\n5\n")),
849 );
850 send(&r, ChunkMsg::ObjectEnd);
851 let first = poll_batch(&mut r.lane, 2).unwrap();
852 assert_eq!(first.len(), 2);
853 let second = poll_batch(&mut r.lane, 2).unwrap();
854 assert_eq!(second.len(), 2);
855 let third = poll_batch(&mut r.lane, 2).unwrap();
856 assert_eq!(third.len(), 1);
857 assert_eq!(third[0].1, b"5".to_vec());
858 }
859
860 #[test]
861 fn idle_poll_times_out_with_none() {
862 let mut r = rig(None);
863 let started = Instant::now();
864 let polled = r.lane.poll(512, Duration::from_millis(60)).unwrap();
865 assert!(polled.is_none());
866 assert!(
867 started.elapsed() >= Duration::from_millis(50),
868 "idle poll must block up to the timeout, not busy-spin"
869 );
870 assert!(r.tracker.terminal().is_none(), "idle is not end-of-input");
871 }
872}