Skip to main content

piw/bundle/
tail.rs

1//! Incremental NDJSON tailing. `trace.ndjson` and `session/entries.ndjson`
2//! are append-only, so a tailer only ever reads bytes past its offset. A
3//! partial trailing line (a writer mid-append) is buffered until the newline
4//! arrives. Truncation (which the format forbids) resets the tailer.
5
6use std::io::{Read, Seek, SeekFrom};
7use std::path::{Path, PathBuf};
8
9pub struct NdjsonTailer {
10    path: PathBuf,
11    /// When set, the path must canonicalize inside this directory on every
12    /// poll: the file name comes from an untrusted manifest and could be a
13    /// symlink out of the bundle.
14    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    /// A tailer that refuses to read once `path` resolves outside `base`.
32    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    /// Read complete new lines appended since the last poll and parse each
48    /// as `T`. Unparsable lines are omitted from values but retained as an
49    /// integrity flag.
50    pub fn poll<T: serde::de::DeserializeOwned>(&mut self) -> std::io::Result<Vec<T>> {
51        let path = match &self.base {
52            Some(base) => {
53                // A missing file canonicalizes to None too; treat both the
54                // not-yet-written and the escaping case as "nothing to read".
55                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            // Truncated (should never happen for append-only files): re-read.
70            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}