1use crate::batch::LogBatch;
2use crate::checkpoint::{
3 is_reserved_checkpoint_name, sweep_orphaned_temp_files, CheckpointData, CHECKPOINT_TEMP_MAX_AGE,
4};
5use crate::decompress::{FileDecompressor, NextLine, DEFAULT_MAX_LINE_SIZE};
6use anyhow::Context;
7use camino::Utf8PathBuf;
8use filenamegen::Glob;
9use futures::Stream;
10use notify::event::{CreateKind, ModifyKind};
11use notify::{Event, EventKind, Watcher};
12use serde::{Deserialize, Serialize};
13use std::collections::HashSet;
14use std::sync::atomic::{AtomicBool, Ordering};
15use std::sync::{Arc, Mutex};
16use std::time::Duration;
17use tokio::sync::Notify;
18use tracing::warn;
19
20fn default_pattern() -> String {
25 "*".to_string()
26}
27
28fn default_max_batch_size() -> usize {
29 100
30}
31
32fn default_max_batch_latency() -> Duration {
33 Duration::from_secs(1)
34}
35
36fn default_max_line_size() -> usize {
37 DEFAULT_MAX_LINE_SIZE
38}
39
40pub struct ConsumerConfig {
46 pub name: String,
49 pub max_batch_size: usize,
51 pub max_batch_latency: Duration,
53 pub checkpoint_name: Option<String>,
56 pub filter: Option<Box<dyn Fn(&serde_json::Value) -> anyhow::Result<bool> + Send>>,
59}
60
61impl ConsumerConfig {
62 pub fn new(name: impl Into<String>) -> Self {
63 Self {
64 name: name.into(),
65 max_batch_size: default_max_batch_size(),
66 max_batch_latency: default_max_batch_latency(),
67 checkpoint_name: None,
68 filter: None,
69 }
70 }
71
72 pub fn max_batch_size(mut self, size: usize) -> Self {
73 self.max_batch_size = size;
74 self
75 }
76
77 pub fn max_batch_latency(mut self, latency: Duration) -> Self {
78 self.max_batch_latency = latency;
79 self
80 }
81
82 pub fn checkpoint_name(mut self, name: impl Into<String>) -> Self {
83 self.checkpoint_name = Some(name.into());
84 self
85 }
86
87 pub fn filter<F>(mut self, f: F) -> Self
88 where
89 F: Fn(&serde_json::Value) -> anyhow::Result<bool> + Send + 'static,
90 {
91 self.filter = Some(Box::new(f));
92 self
93 }
94}
95
96pub struct MultiConsumerTailerConfig {
102 pub directory: Utf8PathBuf,
104 pub pattern: String,
106 pub poll_watcher: Option<Duration>,
108 pub tail: bool,
110 pub max_line_size: usize,
114 pub consumers: Vec<ConsumerConfig>,
116}
117
118impl MultiConsumerTailerConfig {
119 pub fn new(directory: Utf8PathBuf, consumers: Vec<ConsumerConfig>) -> Self {
120 Self {
121 directory,
122 pattern: default_pattern(),
123 poll_watcher: None,
124 tail: false,
125 max_line_size: default_max_line_size(),
126 consumers,
127 }
128 }
129
130 pub fn pattern(mut self, pattern: impl Into<String>) -> Self {
131 self.pattern = pattern.into();
132 self
133 }
134
135 pub fn poll_watcher(mut self, interval: Duration) -> Self {
136 self.poll_watcher = Some(interval);
137 self
138 }
139
140 pub fn tail(mut self, enable: bool) -> Self {
141 self.tail = enable;
142 self
143 }
144
145 pub fn max_line_size(mut self, size: usize) -> Self {
146 self.max_line_size = size;
147 self
148 }
149
150 pub async fn build(self) -> anyhow::Result<MultiConsumerTailer> {
152 for c in &self.consumers {
153 if let Some(name) = &c.checkpoint_name {
154 if is_reserved_checkpoint_name(name) {
155 anyhow::bail!(
156 "checkpoint_name {name:?} is not allowed because \
157 it would collide with temporary file names \
158 created during checkpoint writes"
159 );
160 }
161 }
162 }
163
164 let sweep_dir = self.directory.clone();
167 tokio::task::spawn_blocking(move || {
168 sweep_orphaned_temp_files(&sweep_dir, CHECKPOINT_TEMP_MAX_AGE)
169 })
170 .await
171 .ok();
172
173 let cp_paths: Vec<Option<Utf8PathBuf>> = self
176 .consumers
177 .iter()
178 .map(|c| {
179 c.checkpoint_name
180 .as_ref()
181 .map(|name| self.directory.join(format!(".{name}")))
182 })
183 .collect();
184
185 let mut consumer_checkpoints: Vec<Option<CheckpointData>> =
187 Vec::with_capacity(cp_paths.len());
188 let mut earliest_checkpoint: Option<CheckpointData> = None;
189
190 for cp_path in &cp_paths {
191 let cp = if self.tail {
192 resolve_tail_checkpoint(&self.directory, &self.pattern)?
193 } else if let Some(cp_path) = cp_path {
194 CheckpointData::load(cp_path).await?
195 } else {
196 None
197 };
198
199 match (&earliest_checkpoint, &cp) {
200 (None, Some(cp)) => {
201 earliest_checkpoint = Some(cp.clone());
202 }
203 (Some(existing), Some(cp)) => {
204 if cp.file < existing.file
205 || (cp.file == existing.file && cp.line < existing.line)
206 {
207 earliest_checkpoint = Some(cp.clone());
208 }
209 }
210 _ => {}
211 }
212
213 consumer_checkpoints.push(cp);
214 }
215
216 let closed = Arc::new(AtomicBool::new(false));
217 let close_notify = Arc::new(Notify::new());
218
219 let fs_notify = Arc::new(Notify::new());
220 let fs_notify_tx = fs_notify.clone();
221 let event_handler = move |res: Result<Event, _>| match res {
222 Ok(event) => match event.kind {
223 EventKind::Create(CreateKind::File) | EventKind::Modify(ModifyKind::Data(_)) => {
224 fs_notify_tx.notify_one();
225 }
226 _ => {}
227 },
228 Err(_) => {}
229 };
230 let mut watcher: Box<dyn Watcher + Send> = if let Some(interval) = self.poll_watcher {
231 Box::new(notify::PollWatcher::new(
232 event_handler,
233 notify::Config::default().with_poll_interval(interval),
234 )?)
235 } else {
236 Box::new(notify::recommended_watcher(event_handler)?)
237 };
238 watcher.watch(
239 &self.directory.clone().into_std_path_buf(),
240 notify::RecursiveMode::NonRecursive,
241 )?;
242
243 let shared = Arc::new(TailerShared {
244 closed,
245 close_notify,
246 });
247
248 let stream = make_multi_stream(
249 self.directory,
250 self.pattern,
251 self.max_line_size,
252 self.consumers,
253 earliest_checkpoint,
254 consumer_checkpoints,
255 cp_paths,
256 fs_notify,
257 shared.clone(),
258 );
259
260 Ok(MultiConsumerTailer {
261 close_handle: CloseHandle { shared },
262 _watcher: watcher,
263 stream: Box::pin(stream),
264 })
265 }
266}
267
268struct TailerShared {
273 closed: Arc<AtomicBool>,
274 close_notify: Arc<Notify>,
275}
276
277#[derive(Clone)]
279pub struct CloseHandle {
280 shared: Arc<TailerShared>,
281}
282
283impl CloseHandle {
284 pub fn close(&self) {
286 self.shared.closed.store(true, Ordering::SeqCst);
287 self.shared.close_notify.notify_waiters();
288 }
289}
290
291pub struct MultiConsumerTailer {
298 close_handle: CloseHandle,
299 _watcher: Box<dyn Watcher + Send>,
300 stream: std::pin::Pin<Box<dyn Stream<Item = anyhow::Result<Vec<LogBatch>>> + Send>>,
301}
302
303impl MultiConsumerTailer {
304 pub fn close_handle(&self) -> CloseHandle {
305 self.close_handle.clone()
306 }
307
308 pub fn close(&self) {
309 self.close_handle.close();
310 }
311}
312
313impl Stream for MultiConsumerTailer {
314 type Item = anyhow::Result<Vec<LogBatch>>;
315
316 fn poll_next(
317 mut self: std::pin::Pin<&mut Self>,
318 cx: &mut std::task::Context<'_>,
319 ) -> std::task::Poll<Option<Self::Item>> {
320 if self.close_handle.shared.closed.load(Ordering::SeqCst) {
321 return std::task::Poll::Ready(None);
322 }
323 self.stream.as_mut().poll_next(cx)
324 }
325}
326
327#[derive(Deserialize, Serialize)]
333pub struct LogTailerConfig {
334 pub directory: Utf8PathBuf,
335 #[serde(default = "default_pattern")]
336 pub pattern: String,
337 #[serde(default = "default_max_batch_size")]
338 pub max_batch_size: usize,
339 #[serde(default = "default_max_batch_latency", with = "duration_serde")]
340 pub max_batch_latency: Duration,
341 #[serde(default, skip_serializing_if = "Option::is_none")]
342 pub checkpoint_name: Option<String>,
343 #[serde(
344 default,
345 with = "duration_serde",
346 skip_serializing_if = "Option::is_none"
347 )]
348 pub poll_watcher: Option<Duration>,
349 #[serde(default)]
350 pub tail: bool,
351 #[serde(default = "default_max_line_size")]
352 pub max_line_size: usize,
353}
354
355impl LogTailerConfig {
356 pub fn new(directory: Utf8PathBuf) -> Self {
357 Self {
358 directory,
359 pattern: default_pattern(),
360 max_batch_size: default_max_batch_size(),
361 max_batch_latency: default_max_batch_latency(),
362 checkpoint_name: None,
363 poll_watcher: None,
364 tail: false,
365 max_line_size: default_max_line_size(),
366 }
367 }
368
369 pub fn pattern(mut self, pattern: impl Into<String>) -> Self {
370 self.pattern = pattern.into();
371 self
372 }
373
374 pub fn max_batch_size(mut self, size: usize) -> Self {
375 self.max_batch_size = size;
376 self
377 }
378
379 pub fn max_line_size(mut self, size: usize) -> Self {
380 self.max_line_size = size;
381 self
382 }
383
384 pub fn max_batch_latency(mut self, latency: Duration) -> Self {
385 self.max_batch_latency = latency;
386 self
387 }
388
389 pub fn checkpoint_name(mut self, name: impl Into<String>) -> Self {
390 self.checkpoint_name = Some(name.into());
391 self
392 }
393
394 pub fn poll_watcher(mut self, interval: Duration) -> Self {
395 self.poll_watcher = Some(interval);
396 self
397 }
398
399 pub fn tail(mut self, enable: bool) -> Self {
400 self.tail = enable;
401 self
402 }
403
404 pub async fn build(self) -> anyhow::Result<LogTailer> {
406 self.build_with_filter(None::<fn(&serde_json::Value) -> anyhow::Result<bool>>)
407 .await
408 }
409
410 pub async fn build_with_filter<F>(self, filter: Option<F>) -> anyhow::Result<LogTailer>
412 where
413 F: Fn(&serde_json::Value) -> anyhow::Result<bool> + Send + 'static,
414 {
415 let mut consumer = ConsumerConfig::new("default")
416 .max_batch_size(self.max_batch_size)
417 .max_batch_latency(self.max_batch_latency);
418 if let Some(name) = self.checkpoint_name.clone() {
419 consumer = consumer.checkpoint_name(name);
420 }
421 if let Some(f) = filter {
422 consumer = consumer.filter(f);
423 }
424
425 let multi_config = MultiConsumerTailerConfig {
426 directory: self.directory,
427 pattern: self.pattern,
428 poll_watcher: self.poll_watcher,
429 tail: self.tail,
430 max_line_size: self.max_line_size,
431 consumers: vec![consumer],
432 };
433
434 let multi = multi_config.build().await?;
435
436 Ok(LogTailer { inner: multi })
437 }
438}
439
440pub struct LogTailer {
445 inner: MultiConsumerTailer,
446}
447
448impl LogTailer {
449 pub fn close_handle(&self) -> CloseHandle {
450 self.inner.close_handle()
451 }
452
453 pub fn close(&self) {
454 self.inner.close();
455 }
456}
457
458impl Stream for LogTailer {
459 type Item = anyhow::Result<LogBatch>;
460
461 fn poll_next(
462 mut self: std::pin::Pin<&mut Self>,
463 cx: &mut std::task::Context<'_>,
464 ) -> std::task::Poll<Option<Self::Item>> {
465 use std::task::Poll;
466 match std::pin::Pin::new(&mut self.inner).poll_next(cx) {
469 Poll::Ready(Some(Ok(mut batches))) => Poll::Ready(Some(Ok(batches
470 .pop()
471 .expect("single consumer yields one batch")))),
472 Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
473 Poll::Ready(None) => Poll::Ready(None),
474 Poll::Pending => Poll::Pending,
475 }
476 }
477}
478
479fn resolve_tail_checkpoint(
484 directory: &Utf8PathBuf,
485 pattern: &str,
486) -> anyhow::Result<Option<CheckpointData>> {
487 let glob = Glob::new(pattern)?;
488 let mut files = vec![];
489 for path in glob.walk(directory) {
490 let path = directory.join(Utf8PathBuf::try_from(path).map_err(|e| anyhow::anyhow!("{e}"))?);
491 if path.is_file() {
492 files.push(path);
493 }
494 }
495 files.sort();
496 Ok(files.last().map(|f| CheckpointData {
497 file: f.to_string(),
498 line: 0,
499 }))
500}
501
502fn build_plan(
504 directory: &Utf8PathBuf,
505 pattern: &str,
506 checkpoint_paths: &[Option<Utf8PathBuf>],
507 last_processed: &Option<Utf8PathBuf>,
508 checkpoint: &Option<CheckpointData>,
509 bad_files: &HashSet<Utf8PathBuf>,
510) -> anyhow::Result<Vec<Utf8PathBuf>> {
511 let glob = Glob::new(pattern)?;
512 let mut result = vec![];
513 for path in glob.walk(directory) {
514 let path = directory.join(Utf8PathBuf::try_from(path).map_err(|e| anyhow::anyhow!("{e}"))?);
515 if checkpoint_paths.iter().any(|cp| cp.as_ref() == Some(&path)) {
517 continue;
518 }
519 if bad_files.contains(&path) {
522 continue;
523 }
524 if path.is_file() {
525 result.push(path);
526 }
527 }
528 result.sort();
529
530 if let Some(last) = last_processed {
531 result.retain(|item| item > last);
532 } else if let Some(cp) = checkpoint {
533 let cp_file = &cp.file;
534 result.retain(|item| item.as_str() >= cp_file.as_str());
535 }
536
537 Ok(result)
538}
539
540fn is_file_done(path: &Utf8PathBuf) -> bool {
541 path.metadata()
542 .map(|m| m.permissions().readonly())
543 .unwrap_or(false)
544}
545
546const FILTERED_PROGRESS_FLUSH_INTERVAL: usize = 10_000;
550
551#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
553struct ScanPosition {
554 file: Utf8PathBuf,
555 line: usize,
556}
557
558async fn flush_filtered_progress(
561 cp_paths: &[Option<Utf8PathBuf>],
562 last_handed_out: &[Option<ScanPosition>],
563 committed_through: &[Mutex<Option<ScanPosition>>],
564 file: &Utf8PathBuf,
565 line: usize,
566 last_persisted: &mut [Option<ScanPosition>],
567) -> anyhow::Result<()> {
568 let mut pending: Vec<(usize, Utf8PathBuf)> = vec![];
569 for (i, (cp_path, last)) in cp_paths.iter().zip(last_persisted.iter()).enumerate() {
570 let Some(cp_path) = cp_path else {
571 continue;
572 };
573 let blocked = {
576 let committed = committed_through[i]
577 .lock()
578 .expect("committed_through mutex poisoned");
579 last_handed_out[i] > *committed
580 };
581 if blocked {
582 continue;
583 }
584 let pos = ScanPosition {
585 file: file.clone(),
586 line,
587 };
588 if last.as_ref().is_some_and(|prev| pos <= *prev) {
591 continue;
592 }
593 pending.push((i, cp_path.clone()));
594 }
595 if pending.is_empty() {
596 return Ok(());
597 }
598
599 let segment = file.clone();
600 let written = tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<usize>> {
601 let mut written = Vec::with_capacity(pending.len());
602 for (i, cp_path) in pending {
603 CheckpointData::save_atomic(&cp_path, &segment, line)
609 .with_context(|| format!("persisting filtered-progress checkpoint {cp_path}"))?;
610 written.push(i);
611 }
612 Ok(written)
613 })
614 .await
615 .context("filtered-progress checkpoint task")??;
616
617 for i in written {
618 last_persisted[i] = Some(ScanPosition {
619 file: file.clone(),
620 line,
621 });
622 }
623 Ok(())
624}
625
626fn make_multi_stream(
635 directory: Utf8PathBuf,
636 pattern: String,
637 max_line_size: usize,
638 consumers: Vec<ConsumerConfig>,
639 earliest_checkpoint: Option<CheckpointData>,
640 mut consumer_checkpoints: Vec<Option<CheckpointData>>,
641 cp_paths: Vec<Option<Utf8PathBuf>>,
642 fs_notify: Arc<Notify>,
643 shared: Arc<TailerShared>,
644) -> impl Stream<Item = anyhow::Result<Vec<LogBatch>>> + Send {
645 let num_consumers = consumers.len();
646
647 let consumer_names: Vec<String> = consumers.iter().map(|c| c.name.clone()).collect();
649 let max_batch_sizes: Vec<usize> = consumers.iter().map(|c| c.max_batch_size).collect();
650 let max_batch_latencies: Vec<Duration> =
651 consumers.iter().map(|c| c.max_batch_latency).collect();
652 let filters: Vec<Option<Box<dyn Fn(&serde_json::Value) -> anyhow::Result<bool> + Send>>> =
653 consumers.into_iter().map(|c| c.filter).collect();
654
655 async_stream::try_stream! {
656 let mut last_processed: Option<Utf8PathBuf> = None;
657 let mut global_checkpoint = earliest_checkpoint;
658 let mut skip_lines: usize;
659 let retry_delay = Duration::from_millis(200);
660 let mut bad_files: HashSet<Utf8PathBuf> = HashSet::new();
667
668 let mut consumer_skip: Vec<usize> = vec![0; num_consumers];
674
675 let mut last_handed_out: Vec<Option<ScanPosition>> = vec![None; num_consumers];
678 let committed_through: Arc<Vec<Mutex<Option<ScanPosition>>>> =
682 Arc::new((0..num_consumers).map(|_| Mutex::new(None)).collect());
683
684 let mut last_persisted: Vec<Option<ScanPosition>> = consumer_checkpoints
687 .iter()
688 .map(|cp| {
689 cp.as_ref().map(|c| ScanPosition {
690 file: Utf8PathBuf::from(&c.file),
691 line: c.line,
692 })
693 })
694 .collect();
695
696 let mut lines_since_flush: usize = 0;
697
698 'outer: loop {
699 if shared.closed.load(Ordering::SeqCst) {
700 break;
701 }
702
703 let plan = build_plan(
704 &directory,
705 &pattern,
706 &cp_paths,
707 &last_processed,
708 &global_checkpoint,
709 &bad_files,
710 )?;
711
712 if plan.is_empty() {
713 tokio::select! {
714 _ = shared.close_notify.notified() => break,
715 _ = fs_notify.notified() => continue,
716 }
717 }
718
719 if let Some(cp) = &global_checkpoint {
721 if plan.first().map(|p| p.as_str()) == Some(cp.file.as_str()) {
722 skip_lines = cp.line;
723 } else {
724 skip_lines = 0;
725 }
726 } else {
727 skip_lines = 0;
728 }
729
730 for i in 0..num_consumers {
732 if let Some(cp) = &consumer_checkpoints[i] {
733 if plan.first().map(|p| p.as_str()) == Some(cp.file.as_str()) {
734 consumer_skip[i] = cp.line;
735 } else {
736 consumer_skip[i] = 0;
737 }
738 } else {
739 consumer_skip[i] = 0;
740 }
741 }
742 global_checkpoint.take();
743 for cp in consumer_checkpoints.iter_mut() {
744 cp.take();
745 }
746
747 let mut plan_index = 0;
748 let mut decomp: Option<FileDecompressor> = None;
749 let mut current_path: Option<&Utf8PathBuf> = None;
750 let mut last_lines_consumed: usize = 0;
751 let mut global_line_in_file: usize = skip_lines;
754
755 if let Some(path) = plan.get(plan_index) {
756 let path_std = path.as_std_path().to_owned();
757 decomp = Some(FileDecompressor::open_with_max_line_size(
758 &path_std,
759 max_line_size,
760 )?);
761 current_path = Some(path);
762 }
763
764 let mut batches: Vec<LogBatch> = (0..num_consumers)
769 .map(|i| LogBatch::with_consumer_name(consumer_names[i].clone()))
770 .collect();
771 let mut deadlines: Vec<Option<tokio::time::Instant>> = vec![None; num_consumers];
772
773 while decomp.is_some() {
774 if shared.closed.load(Ordering::SeqCst) {
775 break 'outer;
776 }
777
778 'fill: loop {
780 if shared.closed.load(Ordering::SeqCst) {
781 break 'outer;
782 }
783
784 let now = tokio::time::Instant::now();
786 let any_ready = (0..num_consumers).any(|i| {
787 !batches[i].is_empty()
788 && (batches[i].len() >= max_batch_sizes[i]
789 || deadlines[i].map_or(false, |d| now >= d))
790 });
791 if any_ready {
792 break 'fill;
793 }
794
795 let d = decomp.as_mut().expect("checked above");
796 let path = current_path.expect("set with decomp");
797
798 let mut advance_file: Option<bool> = None;
806
807 match d.next_line(skip_lines) {
808 Ok(NextLine::Line(line)) => {
809 match serde_json::from_str::<serde_json::Value>(&line.text) {
810 Ok(value) => {
811 for i in 0..num_consumers {
812 if global_line_in_file < consumer_skip[i] {
813 continue;
814 }
815 if let Some(ref f) = filters[i] {
816 if !f(&value)? {
817 continue;
818 }
819 }
820 batches[i].push_value(
821 value.clone(),
822 path,
823 line.byte_offset,
824 );
825 if cp_paths[i].is_some() {
826 last_handed_out[i] = Some(ScanPosition {
827 file: path.clone(),
828 line: d.lines_consumed,
829 });
830 }
831 if deadlines[i].is_none() {
833 deadlines[i] = Some(
834 tokio::time::Instant::now() + max_batch_latencies[i],
835 );
836 }
837 }
838 global_line_in_file += 1;
839 lines_since_flush += 1;
840 if lines_since_flush >= FILTERED_PROGRESS_FLUSH_INTERVAL {
841 flush_filtered_progress(
842 &cp_paths,
843 &last_handed_out,
844 committed_through.as_slice(),
845 path,
846 d.lines_consumed,
847 &mut last_persisted,
848 )
849 .await?;
850 lines_since_flush = 0;
851 }
852 }
853 Err(err) => {
854 warn!(
855 "Failed to parse a line from {path} (byte offset {}) \
856 as json: {err}. Skipping remainder of this file; \
857 move it aside if it is not a kumo-jsonl segment.",
858 line.byte_offset
859 );
860 advance_file = Some(true);
861 }
862 }
863 }
864 Ok(NextLine::Skipped { byte_offset, bytes }) => {
865 warn!(
871 "Skipping a record from {path} at byte offset {byte_offset} \
872 ({bytes} bytes) that exceeds max_line_size; the rest of the \
873 segment is still processed."
874 );
875 global_line_in_file += 1;
876 }
877 Ok(NextLine::None) => {
878 if is_file_done(path) {
880 if d.is_discarding_oversized_record() {
881 warn!(
887 "segment {path} ended while discarding a trailing record \
888 that exceeds max_line_size; the record is dropped and \
889 the segment is treated as complete."
890 );
891 } else if d.has_partial_data() {
892 warn!(
900 "unexpected EOF for {} with partial line data remaining; \
901 the writer exited without flushing the zstd stream. \
902 Discarding partial trailing data and advancing to next \
903 segment.",
904 path
905 );
906 }
907 advance_file = Some(false);
908 } else {
909 flush_filtered_progress(
912 &cp_paths,
913 &last_handed_out,
914 committed_through.as_slice(),
915 path,
916 d.lines_consumed,
917 &mut last_persisted,
918 )
919 .await?;
920 lines_since_flush = 0;
921 let earliest_deadline = (0..num_consumers)
924 .filter(|&i| !batches[i].is_empty())
925 .filter_map(|i| deadlines[i])
926 .min();
927
928 if let Some(deadline) = earliest_deadline {
929 let remaining = deadline.saturating_duration_since(
930 tokio::time::Instant::now(),
931 );
932 if remaining.is_zero() {
933 break 'fill;
934 }
935 tokio::select! {
936 _ = shared.close_notify.notified() => break 'outer,
937 _ = tokio::time::sleep(remaining.min(retry_delay)) => {},
938 _ = fs_notify.notified() => {},
939 }
940 } else {
941 tokio::select! {
943 _ = shared.close_notify.notified() => break 'outer,
944 _ = tokio::time::sleep(retry_delay) => {},
945 _ = fs_notify.notified() => {},
946 }
947 }
948 d.reset_eof();
949 }
950 }
951 Err(e) => {
952 warn!(
960 "Error decompressing {path}: {e:#}. Skipping this file; \
961 move it aside if it is not a kumo-jsonl segment."
962 );
963 advance_file = Some(true);
964 }
965 }
966
967 if let Some(mark_bad) = advance_file {
968 if mark_bad {
969 bad_files.insert(path.clone());
970 } else {
971 last_lines_consumed = d.lines_consumed;
972 last_processed = Some(path.clone());
973 flush_filtered_progress(
974 &cp_paths,
975 &last_handed_out,
976 committed_through.as_slice(),
977 path,
978 d.lines_consumed,
979 &mut last_persisted,
980 )
981 .await?;
982 lines_since_flush = 0;
983 }
984 skip_lines = 0;
985 global_line_in_file = 0;
986 for cs in consumer_skip.iter_mut() {
987 *cs = 0;
988 }
989 plan_index += 1;
990 if let Some(next_path) = plan.get(plan_index) {
991 let path_std = next_path.as_std_path().to_owned();
992 decomp = Some(FileDecompressor::open_with_max_line_size(
993 &path_std,
994 max_line_size,
995 )?);
996 current_path = Some(next_path);
997 continue 'fill;
998 } else {
999 decomp = None;
1000 current_path = None;
1001 break 'fill;
1002 }
1003 }
1004 }
1005
1006 let now = tokio::time::Instant::now();
1008 let mut ready: Vec<LogBatch> = Vec::new();
1009 for i in 0..num_consumers {
1010 let is_ready = !batches[i].is_empty()
1011 && (batches[i].len() >= max_batch_sizes[i]
1012 || deadlines[i].map_or(false, |d| now >= d)
1013 || decomp.is_none()); if !is_ready {
1016 continue;
1017 }
1018
1019 let mut batch = std::mem::replace(
1021 &mut batches[i],
1022 LogBatch::with_consumer_name(consumer_names[i].clone()),
1023 );
1024 deadlines[i] = None;
1025
1026 if let Some(ref cp_path) = cp_paths[i] {
1028 let (cp_file, cp_line) = if let Some(d) = &decomp {
1029 let path = current_path.expect("set with decomp");
1030 (path.clone(), d.lines_consumed)
1031 } else if let Some(last) = &last_processed {
1032 (last.clone(), last_lines_consumed)
1033 } else {
1034 unreachable!("non-empty batch without a source");
1035 };
1036 let cp_path = cp_path.clone();
1037 let committed_i = committed_through.clone();
1038 batch.set_commit_fn(Box::new(move || {
1039 CheckpointData::save_atomic(&cp_path, &cp_file, cp_line)?;
1043 let pos = ScanPosition {
1044 file: cp_file.clone(),
1045 line: cp_line,
1046 };
1047 let mut committed = committed_i[i]
1048 .lock()
1049 .expect("committed_through mutex poisoned");
1050 if committed.as_ref().map_or(true, |c| *c < pos) {
1051 *committed = Some(pos);
1052 }
1053 Ok(())
1054 }));
1055 }
1056 ready.push(batch);
1057 }
1058
1059 if !ready.is_empty() {
1060 skip_lines = 0;
1061 yield ready;
1062 }
1063 }
1064 }
1065 }
1066}