1use std::io::{Read, Seek, SeekFrom};
7use std::path::{Path, PathBuf};
8
9pub struct NdjsonTailer {
10 path: PathBuf,
11 base: Option<PathBuf>,
15 offset: u64,
16 partial: Vec<u8>,
17 malformed: bool,
18}
19
20impl NdjsonTailer {
21 pub fn new(path: &Path) -> Self {
22 Self {
23 path: path.to_path_buf(),
24 base: None,
25 offset: 0,
26 partial: Vec::new(),
27 malformed: false,
28 }
29 }
30
31 pub fn contained(path: &Path, base: &Path) -> Self {
33 Self {
34 base: Some(base.to_path_buf()),
35 ..Self::new(path)
36 }
37 }
38
39 pub fn malformed(&self) -> bool {
40 self.malformed
41 }
42
43 pub fn has_partial_line(&self) -> bool {
44 !self.partial.is_empty()
45 }
46
47 pub fn poll<T: serde::de::DeserializeOwned>(&mut self) -> std::io::Result<Vec<T>> {
51 let path = match &self.base {
52 Some(base) => {
53 match crate::bundle::reader::contained_path(base, &self.path) {
56 Some(path) => path,
57 None => return Ok(Vec::new()),
58 }
59 }
60 None => self.path.clone(),
61 };
62 let mut file = match std::fs::File::open(&path) {
63 Ok(file) => file,
64 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
65 Err(error) => return Err(error),
66 };
67 let len = file.metadata()?.len();
68 if len < self.offset {
69 self.offset = 0;
71 self.partial.clear();
72 }
73 if len == self.offset {
74 return Ok(Vec::new());
75 }
76 file.seek(SeekFrom::Start(self.offset))?;
77 let mut buffer = Vec::with_capacity((len - self.offset) as usize);
78 file.take(len - self.offset).read_to_end(&mut buffer)?;
79 self.offset = len;
80 self.partial.extend_from_slice(&buffer);
81
82 let mut records = Vec::new();
83 while let Some(newline) = self.partial.iter().position(|&byte| byte == b'\n') {
84 let line: Vec<u8> = self.partial.drain(..=newline).collect();
85 let line = &line[..line.len() - 1];
86 if line.iter().all(u8::is_ascii_whitespace) {
87 continue;
88 }
89 if let Ok(record) = serde_json::from_slice::<T>(line) {
90 records.push(record);
91 } else {
92 self.malformed = true;
93 }
94 }
95 Ok(records)
96 }
97}
98
99#[cfg(test)]
100mod tests {
101 use super::*;
102 use std::io::Write;
103
104 #[derive(serde::Deserialize, PartialEq, Debug)]
105 struct Row {
106 seq: u64,
107 }
108
109 #[test]
110 fn tails_appends_and_buffers_partial_lines() {
111 let dir = tempfile::tempdir().unwrap();
112 let path = dir.path().join("trace.ndjson");
113 let mut tailer = NdjsonTailer::new(&path);
114 assert_eq!(tailer.poll::<Row>().unwrap(), Vec::<Row>::new());
115
116 std::fs::write(&path, "{\"seq\":1}\n{\"seq\":2}\n{\"se").unwrap();
117 assert_eq!(
118 tailer.poll::<Row>().unwrap(),
119 vec![Row { seq: 1 }, Row { seq: 2 }]
120 );
121 assert!(tailer.has_partial_line());
122 assert!(!tailer.malformed());
123
124 let mut file = std::fs::OpenOptions::new()
125 .append(true)
126 .open(&path)
127 .unwrap();
128 file.write_all(b"q\":3}\n").unwrap();
129 drop(file);
130 assert_eq!(tailer.poll::<Row>().unwrap(), vec![Row { seq: 3 }]);
131 assert!(!tailer.has_partial_line());
132 assert_eq!(tailer.poll::<Row>().unwrap(), Vec::<Row>::new());
133
134 file = std::fs::OpenOptions::new()
135 .append(true)
136 .open(&path)
137 .unwrap();
138 file.write_all(b"not json\n").unwrap();
139 drop(file);
140 assert_eq!(tailer.poll::<Row>().unwrap(), Vec::<Row>::new());
141 assert!(tailer.malformed());
142 }
143}