diff --git a/crates/kumo-log-tailer/src/batch.rs b/crates/kumo-log-tailer/src/batch.rs new file mode 100644 index 00000000..a4e72244 --- /dev/null +++ b/crates/kumo-log-tailer/src/batch.rs @@ -0,0 +1,99 @@ +use camino::Utf8PathBuf; + +/// Metadata about a single line within a [`LogBatch`]. +#[derive(Debug, Clone)] +pub struct LineInfo { + /// The byte offset of this line within the decompressed stream + /// of the segment file. + pub byte_offset: u64, +} + +/// A batch of log records yielded by the tailer stream. +/// +/// Contains the raw line strings along with metadata about which +/// segment file they came from and the byte offset of each line. +#[derive(Debug, Clone)] +pub struct LogBatch { + /// The path of the segment file these records were read from. + segment: Utf8PathBuf, + /// The raw log lines. + lines: Vec, + /// Per-line metadata (byte offset, etc.), parallel to `lines`. + line_info: Vec, +} + +impl LogBatch { + /// Create a new empty batch for the given segment file. + pub fn new(segment: Utf8PathBuf) -> Self { + Self { + segment, + lines: Vec::new(), + line_info: Vec::new(), + } + } + + /// Add a line to the batch along with its byte offset in the + /// decompressed stream. + pub fn push(&mut self, line: String, byte_offset: u64) { + self.lines.push(line); + self.line_info.push(LineInfo { byte_offset }); + } + + /// The number of records in this batch. + pub fn len(&self) -> usize { + self.lines.len() + } + + /// Whether the batch is empty. + pub fn is_empty(&self) -> bool { + self.lines.is_empty() + } + + /// The path of the segment file these records were read from. + pub fn segment(&self) -> &Utf8PathBuf { + &self.segment + } + + /// Per-line metadata, parallel to the lines returned by + /// [`AsRef<[String]>`]. + pub fn line_info(&self) -> &[LineInfo] { + &self.line_info + } +} + +impl AsRef<[String]> for LogBatch { + fn as_ref(&self) -> &[String] { + &self.lines + } +} + +impl TryFrom<&LogBatch> for Vec { + type Error = anyhow::Error; + + fn try_from(batch: &LogBatch) -> anyhow::Result { + let segment = batch.segment(); + batch + .lines + .iter() + .zip(batch.line_info.iter()) + .map(|(line, info)| { + serde_json::from_str(line).map_err(|err| { + anyhow::anyhow!( + "Failed to parse a line from {segment} (byte offset {}) \ + as json: {err}. Is the file corrupt? You may need to move \ + the file aside to make progress", + info.byte_offset + ) + }) + }) + .collect() + } +} + +impl TryFrom for Vec { + type Error = anyhow::Error; + + fn try_from(batch: LogBatch) -> anyhow::Result { + Vec::try_from(&batch) + } +} diff --git a/crates/kumo-log-tailer/src/decompress.rs b/crates/kumo-log-tailer/src/decompress.rs index 6615b244..410c7ba4 100644 --- a/crates/kumo-log-tailer/src/decompress.rs +++ b/crates/kumo-log-tailer/src/decompress.rs @@ -7,6 +7,13 @@ use zstd_safe::{DCtx, InBuffer, OutBuffer}; #[error("{}", zstd_safe::get_error_name(*.0))] pub struct ZStdError(pub usize); +/// A line extracted from the decompressed stream, along with its +/// byte offset in the decompressed data. +pub struct DecompressedLine { + pub text: String, + pub byte_offset: u64, +} + /// State for incremental zstd decompression and line extraction from a single file. pub struct FileDecompressor { file: BufReader, @@ -23,9 +30,13 @@ pub struct FileDecompressor { /// Equals skip_before + number of lines actually returned to caller. pub lines_consumed: usize, /// Buffered lines that have been extracted but not yet consumed. - pending_lines: VecDeque, + pending_lines: VecDeque, /// Whether we've seen EOF on the compressed input. saw_eof: bool, + /// Cumulative byte offset in the decompressed stream. + /// Tracks the position of `line_start` relative to the start + /// of the decompressed output. + decompressed_offset: u64, } impl FileDecompressor { @@ -55,6 +66,7 @@ impl FileDecompressor { lines_consumed: 0, pending_lines: VecDeque::new(), saw_eof: false, + decompressed_offset: 0, }) } @@ -66,7 +78,7 @@ impl FileDecompressor { /// - `Ok(Some(line))` — a complete line was extracted. /// - `Ok(None)` — no more data available right now. The caller should check /// if the file is done or retry later. - pub fn next_line(&mut self, skip_before: usize) -> anyhow::Result> { + pub fn next_line(&mut self, skip_before: usize) -> anyhow::Result> { // Return a buffered line if available if let Some(line) = self.pending_lines.pop_front() { self.lines_consumed += 1; @@ -116,12 +128,19 @@ impl FileDecompressor { while let Some(idx) = memchr::memchr(b'\n', &self.out_buffer[self.line_start..self.out_pos]) { + let line_byte_offset = self.decompressed_offset; if self.lines_decompressed >= skip_before { let this_line = &self.out_buffer[self.line_start..self.line_start + idx]; let line = String::from_utf8_lossy(this_line).into_owned(); - self.pending_lines.push_back(line); + self.pending_lines.push_back(DecompressedLine { + text: line, + byte_offset: line_byte_offset, + }); } - self.line_start += idx + 1; + // Advance past the line content + newline + let consumed = idx + 1; + self.decompressed_offset += consumed as u64; + self.line_start += consumed; self.lines_decompressed += 1; } diff --git a/crates/kumo-log-tailer/src/lib.rs b/crates/kumo-log-tailer/src/lib.rs index f1dde16b..21b2c73a 100644 --- a/crates/kumo-log-tailer/src/lib.rs +++ b/crates/kumo-log-tailer/src/lib.rs @@ -1,8 +1,10 @@ +pub mod batch; pub mod checkpoint; pub mod decompress; #[cfg(feature = "lua")] pub mod lua; pub mod tailer; +pub use batch::LogBatch; pub use checkpoint::CheckpointData; pub use tailer::{CloseHandle, LogTailer, LogTailerConfig}; diff --git a/crates/kumo-log-tailer/src/lua.rs b/crates/kumo-log-tailer/src/lua.rs index 7bba7b6a..2c7b75fb 100644 --- a/crates/kumo-log-tailer/src/lua.rs +++ b/crates/kumo-log-tailer/src/lua.rs @@ -27,31 +27,18 @@ impl LuaLogTailer { /// polls the underlying stream for the next batch. async fn batches(lua: Lua, this: UserDataRef, _: ()) -> mlua::Result { let stream = this.stream.clone(); - let close_handle = this.close_handle.clone(); lua.create_async_function(move |lua, ()| { let stream = stream.clone(); - let close_handle = close_handle.clone(); async move { let mut guard = stream.lock().await; match guard.next().await { Some(Ok(batch)) => { + let values: Vec = + (&batch).try_into().map_err(any_err)?; let table = lua.create_table()?; let options = config::serialize_options(); - let file_name = close_handle - .current_file() - .await - .map(|p| p.to_string()) - .unwrap_or_else(|| "".to_string()); - for (i, record) in batch.into_iter().enumerate() { - let json_value: serde_json::Value = serde_json::from_str(&record) - .map_err(|err| { - any_err(format!( - "Failed to parse a line from {file_name} as json: \ - {err}. Is the file corrupt? You may need to move \ - the file aside to make progress" - )) - })?; - let lua_value = lua.to_value_with(&json_value, options)?; + for (i, value) in values.into_iter().enumerate() { + let lua_value = lua.to_value_with(&value, options)?; table.raw_set(i + 1, lua_value)?; } Ok(mlua::Value::Table(table)) diff --git a/crates/kumo-log-tailer/src/tailer.rs b/crates/kumo-log-tailer/src/tailer.rs index 112d3d62..d8b58f36 100644 --- a/crates/kumo-log-tailer/src/tailer.rs +++ b/crates/kumo-log-tailer/src/tailer.rs @@ -1,3 +1,4 @@ +use crate::batch::LogBatch; use crate::checkpoint::CheckpointData; use crate::decompress::FileDecompressor; use camino::Utf8PathBuf; @@ -253,7 +254,7 @@ impl CloseHandle { pub struct LogTailer { close_handle: CloseHandle, _watcher: Box, - stream: std::pin::Pin>> + Send>>, + stream: std::pin::Pin> + Send>>, } impl LogTailer { @@ -272,7 +273,7 @@ impl LogTailer { } impl Stream for LogTailer { - type Item = anyhow::Result>; + type Item = anyhow::Result; fn poll_next( mut self: std::pin::Pin<&mut Self>, @@ -348,7 +349,7 @@ fn make_stream( shared: Arc, pos_file: Arc>>, pos_line: Arc, -) -> impl Stream>> + Send { +) -> impl Stream> + Send { async_stream::try_stream! { let mut last_processed: Option = None; let mut checkpoint = initial_checkpoint; @@ -419,7 +420,7 @@ fn make_stream( } } - let mut batch = Vec::new(); + let mut batch = LogBatch::new(path.clone()); let mut hit_eof = false; let batch_deadline = tokio::time::Instant::now() + max_batch_latency; @@ -431,7 +432,7 @@ fn make_stream( match decomp.next_line(skip_lines) { Ok(Some(line)) => { - batch.push(line); + batch.push(line.text, line.byte_offset); if batch.len() >= max_batch_size { break; } @@ -467,7 +468,7 @@ fn make_stream( loop { match decomp.next_line(skip_lines) { Ok(Some(line)) => { - batch.push(line); + batch.push(line.text, line.byte_offset); if batch.len() >= max_batch_size { break; } diff --git a/crates/kumo-log-tailer/tests/integration.rs b/crates/kumo-log-tailer/tests/integration.rs index 9e583f4d..8c2ca140 100644 --- a/crates/kumo-log-tailer/tests/integration.rs +++ b/crates/kumo-log-tailer/tests/integration.rs @@ -62,7 +62,7 @@ async fn test_checkpoint_resume_one_at_a_time() { .unwrap_or_else(|| panic!("expected a batch on iteration {i}")) .unwrap_or_else(|e| panic!("expected Ok batch on iteration {i}: {e}")); k9::assert_equal!(batch.len(), 1); - all_records.push(batch[0].clone()); + all_records.push(batch.as_ref()[0].clone()); tailer.as_mut().close().await.unwrap(); } @@ -125,7 +125,7 @@ async fn test_multiple_files_in_order() { batch = tailer.next() => { match batch { Some(Ok(records)) => { - all_records.extend(records); + all_records.extend(records.as_ref().iter().cloned()); if all_records.len() >= 4 { break; } @@ -190,7 +190,7 @@ async fn test_checkpoint_across_multiple_files() { .unwrap_or_else(|| panic!("expected batch on iteration {i}")) .unwrap_or_else(|e| panic!("error on iteration {i}: {e}")); k9::assert_equal!(batch.len(), 1); - all_records.push(batch[0].clone()); + all_records.push(batch.as_ref()[0].clone()); tailer.as_mut().close().await.unwrap(); } @@ -234,7 +234,7 @@ async fn test_close_advances_checkpoint_past_consumed_batch() { .await .expect("should yield a batch") .expect("batch should be Ok"); - k9::assert_equal!(first, vec![r#"{"n":1}"#.to_string()]); + k9::assert_equal!(first.as_ref(), &[r#"{"n":1}"#.to_string()]); tailer.as_mut().close().await.unwrap(); @@ -254,7 +254,7 @@ async fn test_close_advances_checkpoint_past_consumed_batch() { .await .expect("should yield a batch") .expect("batch should be Ok"); - k9::assert_equal!(second, vec![r#"{"n":2}"#.to_string()]); + k9::assert_equal!(second.as_ref(), &[r#"{"n":2}"#.to_string()]); tailer2.as_mut().close().await.unwrap(); } @@ -291,7 +291,7 @@ async fn test_drop_without_close_does_not_advance_checkpoint() { .await .expect("should yield a batch") .expect("batch should be Ok"); - k9::assert_equal!(first, vec![r#"{"n":1}"#.to_string()]); + k9::assert_equal!(first.as_ref(), &[r#"{"n":1}"#.to_string()]); // tailer is dropped here without calling close() } @@ -314,7 +314,7 @@ async fn test_drop_without_close_does_not_advance_checkpoint() { .await .expect("should yield a batch") .expect("batch should be Ok"); - k9::assert_equal!(second, vec![r#"{"n":1}"#.to_string()]); + k9::assert_equal!(second.as_ref(), &[r#"{"n":1}"#.to_string()]); tailer2.as_mut().close().await.unwrap(); } @@ -361,7 +361,7 @@ async fn test_tail_starts_from_latest_segment() { batch = tailer.next() => { match batch { Some(Ok(records)) => { - all_records.extend(records); + all_records.extend(records.as_ref().iter().cloned()); if all_records.len() >= 2 { break; } diff --git a/crates/tailer/src/main.rs b/crates/tailer/src/main.rs index 413e7359..3714152d 100644 --- a/crates/tailer/src/main.rs +++ b/crates/tailer/src/main.rs @@ -59,7 +59,7 @@ async fn main() -> anyhow::Result<()> { while let Some(result) = tailer.next().await { let batch = result?; - for line in &batch { + for line in batch.as_ref() { println!("{line}"); } }