turbo_trace_server/reader/
mod.rs

1mod heaptrack;
2mod nextjs;
3mod turbopack;
4
5use std::{
6    any::Any,
7    env,
8    fs::File,
9    io::{self, BufReader, Read, Seek, SeekFrom, Write},
10    path::PathBuf,
11    sync::Arc,
12    thread::{self, JoinHandle},
13    time::{Duration, Instant},
14};
15
16use anyhow::Result;
17use flate2::bufread::GzDecoder;
18
19use crate::{
20    reader::{heaptrack::HeaptrackFormat, nextjs::NextJsFormat, turbopack::TurbopackFormat},
21    store_container::StoreContainer,
22};
23
24const MIN_INITIAL_REPORT_SIZE: u64 = 100 * 1024 * 1024;
25
26trait TraceFormat {
27    type Reused: Default;
28    /// Create the initial reused buffer. Override to pre-allocate capacity.
29    fn create_reused() -> Self::Reused {
30        Self::Reused::default()
31    }
32    fn read(&mut self, buffer: &[u8], reuse: &mut Self::Reused) -> Result<usize>;
33    fn stats(&self) -> String {
34        String::new()
35    }
36}
37
38type ErasedReused = Box<dyn Any>;
39
40struct ErasedTraceFormat(Box<dyn ObjectSafeTraceFormat>);
41
42trait ObjectSafeTraceFormat {
43    fn create_reused(&self) -> ErasedReused;
44    fn read(&mut self, buffer: &[u8], reuse: &mut ErasedReused) -> Result<usize>;
45    fn stats(&self) -> String;
46}
47
48impl<T: TraceFormat> ObjectSafeTraceFormat for T
49where
50    T::Reused: 'static,
51{
52    fn create_reused(&self) -> ErasedReused {
53        Box::new(T::create_reused())
54    }
55
56    fn read(&mut self, buffer: &[u8], reuse: &mut ErasedReused) -> Result<usize> {
57        let reuse = reuse.downcast_mut().expect("Type of reuse is invalid");
58        TraceFormat::read(self, buffer, reuse)
59    }
60
61    fn stats(&self) -> String {
62        TraceFormat::stats(self)
63    }
64}
65
66impl ObjectSafeTraceFormat for ErasedTraceFormat {
67    fn create_reused(&self) -> ErasedReused {
68        self.0.create_reused()
69    }
70
71    fn read(&mut self, buffer: &[u8], reuse: &mut ErasedReused) -> Result<usize> {
72        self.0.read(buffer, reuse)
73    }
74
75    fn stats(&self) -> String {
76        self.0.stats()
77    }
78}
79
80#[derive(Default)]
81enum TraceFile {
82    Raw(BufReader<File>),
83    Zstd(zstd::Decoder<'static, BufReader<File>>),
84    Gz(GzDecoder<BufReader<File>>),
85    #[default]
86    Unloaded,
87}
88
89impl TraceFile {
90    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
91        match self {
92            Self::Raw(file) => file.read(buffer),
93            Self::Zstd(decoder) => decoder.read(buffer),
94            Self::Gz(decoder) => decoder.read(buffer),
95            Self::Unloaded => unreachable!(),
96        }
97    }
98
99    fn stream_position(&mut self) -> io::Result<u64> {
100        match self {
101            Self::Raw(file) => file.stream_position(),
102            Self::Zstd(decoder) => decoder.get_mut().stream_position(),
103            Self::Gz(decoder) => decoder.get_mut().stream_position(),
104            Self::Unloaded => unreachable!(),
105        }
106    }
107
108    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
109        match self {
110            Self::Raw(file) => file.seek(pos),
111            Self::Zstd(decoder) => decoder.get_mut().seek(pos),
112            Self::Gz(decoder) => decoder.get_mut().seek(pos),
113            Self::Unloaded => unreachable!(),
114        }
115    }
116
117    fn size(&mut self) -> io::Result<u64> {
118        match self {
119            Self::Raw(file) => file.get_ref().metadata().map(|m| m.len()),
120            Self::Zstd(decoder) => decoder.get_mut().get_ref().metadata().map(|m| m.len()),
121            Self::Gz(decoder) => decoder.get_mut().get_ref().metadata().map(|m| m.len()),
122            Self::Unloaded => unreachable!(),
123        }
124    }
125}
126
127pub struct TraceReader {
128    store: Arc<StoreContainer>,
129    path: PathBuf,
130}
131
132impl TraceReader {
133    pub fn spawn(store: Arc<StoreContainer>, path: PathBuf) -> JoinHandle<()> {
134        let mut reader = Self { store, path };
135        std::thread::spawn(move || reader.run())
136    }
137
138    pub fn run(&mut self) {
139        let mut file_warning_printed = false;
140        loop {
141            let read_success = self.try_read();
142            if !file_warning_printed && !read_success {
143                println!("Unable to read trace file at {:?}, waiting...", self.path);
144                file_warning_printed = true;
145            }
146            thread::sleep(Duration::from_millis(500));
147        }
148    }
149
150    fn trace_file_from_file(&self, file: File) -> io::Result<TraceFile> {
151        let path = &self.path.to_string_lossy();
152        let mut file = BufReader::with_capacity(
153            // zstd max block size (1 << 17) + block header (3) + magic bytes (4)
154            (1 << 17) + 7,
155            file,
156        );
157        let magic_bytes = file.peek(4)?;
158        Ok(
159            if path.ends_with(".zst") || magic_bytes == [0x28, 0xb5, 0x2f, 0xfd] {
160                TraceFile::Zstd(zstd::Decoder::with_buffer(file)?)
161            } else if path.ends_with(".gz") || matches!(magic_bytes, [0x1f, 0x8b, _, _]) {
162                TraceFile::Gz(GzDecoder::new(file))
163            } else {
164                TraceFile::Raw(file)
165            },
166        )
167    }
168
169    fn try_read(&mut self) -> bool {
170        let Ok(mut file) = File::open(&self.path) else {
171            return false;
172        };
173        println!("Trace file opened");
174        let stop_at = env::var("STOP_AT")
175            .unwrap_or_default()
176            .parse()
177            .map_or(u64::MAX, |v: u64| v * 1024 * 1024);
178        if stop_at != u64::MAX {
179            println!("Will stop reading file at {} MB", stop_at / 1024 / 1024)
180        }
181
182        {
183            let mut store = self.store.write();
184            store.reset();
185        }
186
187        let mut format: Option<(ErasedTraceFormat, ErasedReused)> = None;
188
189        let mut current_read = 0;
190        let mut initial_read = file
191            .seek(SeekFrom::End(0))
192            .ok()
193            .map(|total| (total, Instant::now()));
194        if file.seek(SeekFrom::Start(0)).is_err() {
195            return false;
196        }
197        let mut file = match self.trace_file_from_file(file) {
198            Ok(f) => f,
199            Err(err) => {
200                println!("Error creating zstd decoder: {err}");
201                return false;
202            }
203        };
204
205        let mut buffer = Vec::new();
206        let mut index = 0;
207
208        let mut chunk = vec![0; 64 * 1024 * 1024];
209        loop {
210            match file.read(&mut chunk) {
211                Ok(bytes_read) => {
212                    if bytes_read == 0 {
213                        self.store.write().optimize();
214                        if let Some(value) = self.wait_for_more_data(
215                            &mut file,
216                            &mut initial_read,
217                            format.as_ref().map(|(f, _)| f),
218                        ) {
219                            return value;
220                        }
221                    } else {
222                        // If we have partially consumed some data, and we are at buffer capacity,
223                        // remove the consumed data to make more space.
224                        if index > 0 && buffer.len() + bytes_read > buffer.capacity() {
225                            buffer.splice(..index, std::iter::empty());
226                            index = 0;
227                        }
228                        buffer.extend_from_slice(&chunk[..bytes_read]);
229                        if format.is_none() && buffer.len() >= 8 {
230                            let erased_format = if buffer.starts_with(b"TRACEv0") {
231                                index = 7;
232                                ErasedTraceFormat(Box::new(TurbopackFormat::new(
233                                    self.store.clone(),
234                                )))
235                            } else if buffer.starts_with(b"[{\"name\"") {
236                                ErasedTraceFormat(Box::new(NextJsFormat::new(self.store.clone())))
237                            } else if buffer.starts_with(b"v ") {
238                                ErasedTraceFormat(Box::new(HeaptrackFormat::new(
239                                    self.store.clone(),
240                                )))
241                            } else {
242                                // Fallback to the format without magic bytes
243                                // TODO Remove this after a while and show an error instead
244                                ErasedTraceFormat(Box::new(TurbopackFormat::new(
245                                    self.store.clone(),
246                                )))
247                            };
248                            let reuse = erased_format.create_reused();
249                            format = Some((erased_format, reuse));
250                        }
251                        if let Some((format, reuse)) = &mut format {
252                            match format.read(&buffer[index..], reuse) {
253                                Ok(bytes_read) => {
254                                    index += bytes_read;
255                                }
256                                Err(err) => {
257                                    println!("Trace file error: {err}");
258                                    return true;
259                                }
260                            }
261                            if self.store.want_to_read() {
262                                thread::yield_now();
263                            }
264                            current_read += bytes_read as u64;
265                            if let Some((total, start)) = &mut initial_read {
266                                let pos = file.stream_position().unwrap_or(current_read);
267                                if pos > *total {
268                                    *total = file.size().unwrap_or(pos);
269                                }
270                                *total = (*total).max(pos);
271                                let total_bytes = *total;
272                                let percentage = pos * 100 / total_bytes;
273                                let read = pos / (1024 * 1024);
274                                let uncompressed = current_read / (1024 * 1024);
275                                let total = total_bytes / (1024 * 1024);
276                                let elapsed_ms = start.elapsed().as_millis() as u64;
277                                let stats = format.stats();
278                                let rate_mbs = read * 1000 / (elapsed_ms + 1);
279                                let mut line = format!(
280                                    "{percentage}% read ({read}/{total} MB, {rate_mbs} MB/s)"
281                                );
282                                // Estimate remaining time by linearly extrapolating the
283                                // elapsed time over the bytes still to be read.
284                                if pos > 0 && pos < total_bytes {
285                                    let eta_s = elapsed_ms * (total_bytes - pos) / pos / 1000;
286                                    line += &format!(", ETA {eta_s}s");
287                                }
288                                if uncompressed != read {
289                                    line += &format!(" ({uncompressed} MB uncompressed)");
290                                }
291                                if !stats.is_empty() {
292                                    line += &format!(" - {stats}");
293                                }
294
295                                // `\r` returns to the start of the line and `\x1b[2K` erases
296                                // it, so a shorter update doesn't leave behind characters from
297                                // a longer previous one.
298                                print!("\r\x1b[2K{line}");
299                                let _ = io::stdout().flush();
300                            }
301                            if current_read >= stop_at {
302                                println!(
303                                    "Stopped reading file as requested by STOP_AT env var. \
304                                     Waiting for new file..."
305                                );
306                                self.wait_for_new_file(&mut file);
307                                return true;
308                            }
309                        }
310                    }
311                }
312                Err(err) => {
313                    if err.kind() == io::ErrorKind::UnexpectedEof
314                        || err.kind() == io::ErrorKind::InvalidInput
315                    {
316                        self.store.write().optimize();
317                        if let Some(value) = self.wait_for_more_data(
318                            &mut file,
319                            &mut initial_read,
320                            format.as_ref().map(|(f, _)| f),
321                        ) {
322                            return value;
323                        }
324                    } else {
325                        // Error reading file, maybe it was removed
326                        println!("Error reading trace file: {err:?}");
327                        return true;
328                    }
329                }
330            }
331        }
332    }
333
334    fn wait_for_more_data(
335        &mut self,
336        file: &mut TraceFile,
337        initial_read: &mut Option<(u64, Instant)>,
338        format: Option<&ErasedTraceFormat>,
339    ) -> Option<bool> {
340        let Ok(pos) = file.stream_position() else {
341            return Some(true);
342        };
343        if let Some((total, start)) = initial_read.take() {
344            // Erase the in-place progress line (printed with a leading `\r` and
345            // no newline); it's no longer useful once the read is complete.
346            print!("\r\x1b[2K");
347            let stats = format.map(|format| format.stats()).unwrap_or_default();
348            if total > MIN_INITIAL_REPORT_SIZE {
349                let elapsed = (start.elapsed().as_millis() / 100) as f32 / 10.0;
350                print!(
351                    "Initial read completed ({} MB, {elapsed}s)",
352                    total / (1024 * 1024),
353                );
354                if !stats.is_empty() {
355                    print!(" - {stats}");
356                }
357                println!();
358            } else if !stats.is_empty() {
359                println!("{stats}");
360            }
361        }
362        loop {
363            // No more data to read, sleep for a while to wait for more data
364            thread::sleep(Duration::from_millis(100));
365            let Ok(mut real_file) = File::open(&self.path) else {
366                return Some(true);
367            };
368            let Ok(end) = real_file.seek(SeekFrom::End(0)) else {
369                return Some(true);
370            };
371            if end < pos {
372                // new file
373                return Some(true);
374            } else if end != pos {
375                // file has more data
376                return None;
377            }
378        }
379    }
380
381    fn wait_for_new_file(&self, file: &mut TraceFile) {
382        let Ok(pos) = file.stream_position() else {
383            return;
384        };
385        loop {
386            thread::sleep(Duration::from_millis(1000));
387            let Ok(end) = file.seek(SeekFrom::End(0)) else {
388                return;
389            };
390            if end < pos {
391                return;
392            }
393        }
394    }
395}
turbo_trace_server/reader/mod.rs

turbo_trace_server/reader/
mod.rs
