turbo_trace_server/reader/
nextjs.rs

1use std::{borrow::Cow, fmt::Display, sync::Arc};
2
3use rustc_hash::{FxHashMap, FxHashSet};
4use serde::Deserialize;
5
6use super::TraceFormat;
7use crate::{FxIndexMap, span::SpanIndex, store_container::StoreContainer, timestamp::Timestamp};
8
9pub struct NextJsFormat {
10    store: Arc<StoreContainer>,
11    id_mapping: FxHashMap<u64, SpanIndex>,
12    queued_children: FxHashMap<u64, Vec<SpanIndex>>,
13}
14
15impl NextJsFormat {
16    pub fn new(store: Arc<StoreContainer>) -> Self {
17        Self {
18            store,
19            id_mapping: FxHashMap::default(),
20            queued_children: FxHashMap::default(),
21        }
22    }
23}
24
25impl TraceFormat for NextJsFormat {
26    type Reused = ();
27
28    fn read(&mut self, mut buffer: &[u8], _reuse: &mut Self::Reused) -> anyhow::Result<usize> {
29        let mut bytes_read = 0;
30        let mut outdated_spans = FxHashSet::default();
31        loop {
32            let Some(line_end) = buffer.iter().position(|b| *b == b'\n') else {
33                break;
34            };
35            let line = &buffer[..line_end];
36            buffer = &buffer[line_end + 1..];
37            bytes_read += line.len() + 1;
38
39            let spans: Vec<NextJsSpan> = serde_json::from_slice(line)?;
40
41            let mut store = self.store.write();
42
43            for span in spans {
44                let NextJsSpan {
45                    name,
46                    duration,
47                    timestamp,
48                    id,
49                    parent_id,
50                    tags,
51                    start_time: _,
52                    trace_id: _,
53                } = span;
54                let timestamp = Timestamp::from_micros(timestamp);
55                let duration = Timestamp::from_micros(duration);
56                let (parent, queue_parent) = if let Some(parent) = parent_id {
57                    if let Some(parent) = self.id_mapping.get(&parent) {
58                        (Some(*parent), None)
59                    } else {
60                        (None, Some(parent))
61                    }
62                } else {
63                    (None, None)
64                };
65                let index = store.add_span(
66                    parent,
67                    timestamp,
68                    "nextjs".to_string(),
69                    name.into_owned(),
70                    tags.iter()
71                        .map(|(k, v)| {
72                            (
73                                k.to_string(),
74                                v.as_ref().map(|v| v.to_string()).unwrap_or_default(),
75                            )
76                        })
77                        .collect(),
78                    &mut outdated_spans,
79                );
80                self.id_mapping.insert(id, index);
81                if let Some(parent) = queue_parent {
82                    self.queued_children.entry(parent).or_default().push(index);
83                }
84                if let Some(children) = self.queued_children.remove(&id) {
85                    for child in children {
86                        store.set_parent(child, index, &mut outdated_spans);
87                    }
88                }
89                store.set_total_time(index, timestamp, duration, &mut outdated_spans);
90                store.complete_span(index);
91            }
92            store.invalidate_outdated_spans(&outdated_spans);
93            drop(store);
94        }
95        Ok(bytes_read)
96    }
97}
98
99#[derive(Debug, Deserialize)]
100#[serde(untagged)]
101enum TagValue<'a> {
102    String(Cow<'a, str>),
103    Number(f64),
104    Bool(bool),
105    Array(Vec<TagValue<'a>>),
106}
107
108impl Display for TagValue<'_> {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self {
111            TagValue::String(s) => write!(f, "{s}"),
112            TagValue::Number(n) => write!(f, "{n}"),
113            TagValue::Bool(b) => write!(f, "{b}"),
114            TagValue::Array(a) => {
115                write!(f, "[")?;
116                for (i, v) in a.iter().enumerate() {
117                    if i > 0 {
118                        write!(f, ", ")?;
119                    }
120                    write!(f, "{v}")?;
121                }
122                write!(f, "]")
123            }
124        }
125    }
126}
127
128#[derive(Debug, Deserialize)]
129#[serde(rename_all = "camelCase")]
130struct NextJsSpan<'a> {
131    name: Cow<'a, str>,
132    duration: u64,
133    timestamp: u64,
134    id: u64,
135    parent_id: Option<u64>,
136    tags: FxIndexMap<Cow<'a, str>, Option<TagValue<'a>>>,
137    #[allow(dead_code)]
138    start_time: u64,
139    #[allow(dead_code)]
140    trace_id: Cow<'a, str>,
141}