turbopack_trace_utils/
raw_trace.rs

1use std::{
2    borrow::Cow, fmt::Write, marker::PhantomData, sync::atomic::AtomicU64, thread, time::Instant,
3};
4
5use tracing::{
6    Subscriber,
7    field::{Visit, display},
8    span,
9};
10use tracing_subscriber::{Layer, registry::LookupSpan};
11use turbo_tasks_malloc::TurboMalloc;
12
13use crate::{
14    flavor::WriteGuardFlavor,
15    trace_writer::TraceWriter,
16    tracing::{TraceRow, TraceValue},
17};
18
19pub struct RawTraceLayerOptions {}
20
21struct RawTraceLayerExtension {
22    id: u64,
23}
24
25fn get_id<S: Subscriber + for<'a> LookupSpan<'a>>(
26    ctx: tracing_subscriber::layer::Context<'_, S>,
27    id: &span::Id,
28) -> u64 {
29    ctx.span(id)
30        .unwrap()
31        .extensions()
32        .get::<RawTraceLayerExtension>()
33        .unwrap()
34        .id
35}
36
37/// A tracing layer that writes raw trace data to a writer. The data format is
38/// defined by [FullTraceRow].
39pub struct RawTraceLayer<S: Subscriber + for<'a> LookupSpan<'a>> {
40    trace_writer: TraceWriter,
41    start: Instant,
42    next_id: AtomicU64,
43    _phantom: PhantomData<fn(S)>,
44}
45
46impl<S: Subscriber + for<'a> LookupSpan<'a>> RawTraceLayer<S> {
47    pub fn new(trace_writer: TraceWriter) -> Self {
48        Self {
49            trace_writer,
50            start: Instant::now(),
51            next_id: AtomicU64::new(1),
52            _phantom: PhantomData,
53        }
54    }
55
56    fn write(&self, data: TraceRow<'_>) {
57        let start = TurboMalloc::allocation_counters();
58        let guard = self.trace_writer.start_write();
59        postcard::serialize_with_flavor(&data, WriteGuardFlavor { guard }).unwrap();
60        TurboMalloc::reset_allocation_counters(start);
61    }
62
63    fn report_allocations(&self, ts: u64, thread_id: u64) {
64        let allocation_counters = turbo_tasks_malloc::TurboMalloc::allocation_counters();
65        self.write(TraceRow::AllocationCounters {
66            ts,
67            thread_id,
68            allocations: allocation_counters.allocations as u64,
69            deallocations: allocation_counters.deallocations as u64,
70            allocation_count: allocation_counters.allocation_count as u64,
71            deallocation_count: allocation_counters.deallocation_count as u64,
72        });
73    }
74}
75
76impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for RawTraceLayer<S> {
77    fn on_new_span(
78        &self,
79        attrs: &span::Attributes<'_>,
80        id: &span::Id,
81        ctx: tracing_subscriber::layer::Context<'_, S>,
82    ) {
83        let ts = self.start.elapsed().as_micros() as u64;
84        let mut values = ValuesVisitor::new();
85        attrs.values().record(&mut values);
86        let external_id = self
87            .next_id
88            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
89        ctx.span(id)
90            .unwrap()
91            .extensions_mut()
92            .insert(RawTraceLayerExtension { id: external_id });
93        self.write(TraceRow::Start {
94            ts,
95            id: external_id,
96            parent: if attrs.is_contextual() {
97                ctx.current_span().id().map(|p| get_id(ctx, p))
98            } else {
99                attrs.parent().map(|p| get_id(ctx, p))
100            },
101            name: attrs.metadata().name().into(),
102            target: attrs.metadata().target().into(),
103            values: values.values,
104        });
105    }
106
107    fn on_close(&self, id: span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
108        let ts = self.start.elapsed().as_micros() as u64;
109        self.write(TraceRow::End {
110            ts,
111            id: get_id(ctx, &id),
112        });
113    }
114
115    fn on_enter(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
116        let ts = self.start.elapsed().as_micros() as u64;
117        let thread_id = thread::current().id().as_u64().into();
118        self.report_allocations(ts, thread_id);
119        self.write(TraceRow::Enter {
120            ts,
121            id: get_id(ctx, id),
122            thread_id,
123        });
124    }
125
126    fn on_exit(&self, id: &span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
127        let ts = self.start.elapsed().as_micros() as u64;
128        let thread_id = thread::current().id().as_u64().into();
129        self.report_allocations(ts, thread_id);
130        self.write(TraceRow::Exit {
131            ts,
132            id: get_id(ctx, id),
133            thread_id,
134        });
135    }
136
137    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
138        let ts = self.start.elapsed().as_micros() as u64;
139        let mut values = ValuesVisitor::new();
140        event.record(&mut values);
141        self.write(TraceRow::Event {
142            ts,
143            parent: if event.is_contextual() {
144                ctx.current_span().id().map(|p| get_id(ctx, p))
145            } else {
146                event.parent().map(|p| get_id(ctx, p))
147            },
148            values: values.values,
149        });
150    }
151
152    fn on_record(
153        &self,
154        id: &span::Id,
155        record: &span::Record<'_>,
156        ctx: tracing_subscriber::layer::Context<'_, S>,
157    ) {
158        let mut values = ValuesVisitor::new();
159        record.record(&mut values);
160        self.write(TraceRow::Record {
161            id: get_id(ctx, id),
162            values: values.values,
163        });
164    }
165}
166
167struct ValuesVisitor {
168    values: Vec<(Cow<'static, str>, TraceValue<'static>)>,
169}
170
171impl ValuesVisitor {
172    fn new() -> Self {
173        Self { values: Vec::new() }
174    }
175}
176
177impl Visit for ValuesVisitor {
178    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
179        let mut str = String::new();
180        let _ = write!(str, "{value:?}");
181        self.values
182            .push((field.name().into(), TraceValue::String(str.into())));
183    }
184
185    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
186        self.values
187            .push((field.name().into(), TraceValue::Float(value)));
188    }
189
190    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
191        self.values
192            .push((field.name().into(), TraceValue::Int(value)));
193    }
194
195    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
196        self.values
197            .push((field.name().into(), TraceValue::UInt(value)));
198    }
199
200    fn record_i128(&mut self, field: &tracing::field::Field, value: i128) {
201        self.record_debug(field, &value)
202    }
203
204    fn record_u128(&mut self, field: &tracing::field::Field, value: u128) {
205        self.record_debug(field, &value)
206    }
207
208    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
209        self.values
210            .push((field.name().into(), TraceValue::Bool(value)));
211    }
212
213    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
214        self.values.push((
215            field.name().into(),
216            TraceValue::String(value.to_string().into()),
217        ));
218    }
219
220    fn record_error(
221        &mut self,
222        field: &tracing::field::Field,
223        value: &(dyn std::error::Error + 'static),
224    ) {
225        self.record_debug(field, &display(value))
226    }
227}
turbopack_trace_utils/raw_trace.rs

turbopack_trace_utils/
raw_trace.rs
