New performance timings (#3096)

Co-authored-by: Laurenz <laurmaedje@gmail.com>

1 files changed, 233 insertions, 0 deletions

diff --git a/crates/typst-timing/src/lib.rs b/crates/typst-timing/src/lib.rs
new file mode 100644
index 00000000..fe9abb51
--- /dev/null
+++ b/crates/typst-timing/src/lib.rs
@@ -0,0 +1,233 @@
+//! Performance timing for Typst.
+
+use std::hash::Hash;
+use std::io::Write;
+use std::thread::ThreadId;
+use std::time::{Duration, SystemTime};
+
+use parking_lot::Mutex;
+use serde::ser::SerializeSeq;
+use serde::{Serialize, Serializer};
+use typst_syntax::Span;
+
+/// Whether the timer is enabled. Defaults to `false`.
+///
+/// # Safety
+/// This is unsafe because it is a global variable that is not thread-safe.
+/// But at worst, if we have a race condition, we will just be missing some
+/// events. So it's not a big deal. And it avoids needing to do an atomic
+/// operation every time we want to check if the timer is enabled.
+static mut ENABLED: bool = false;
+
+/// The global event recorder.
+static RECORDER: Mutex<Recorder> = Mutex::new(Recorder::new());
+
+/// The recorder of events.
+struct Recorder {
+    /// The events that have been recorded.
+    events: Vec<Event>,
+    /// The discriminator of the next event.
+    discriminator: u64,
+}
+
+impl Recorder {
+    /// Create a new recorder.
+    const fn new() -> Self {
+        Self { events: Vec::new(), discriminator: 0 }
+    }
+}
+
+/// An event that has been recorded.
+#[derive(Clone, Copy, Eq, PartialEq, Hash)]
+struct Event {
+    /// Whether this is a start or end event.
+    kind: EventKind,
+    /// The start time of this event.
+    timestamp: SystemTime,
+    /// The discriminator of this event.
+    id: u64,
+    /// The name of this event.
+    name: &'static str,
+    /// The span of code that this event was recorded in.
+    span: Option<Span>,
+    /// The thread ID of this event.
+    thread_id: ThreadId,
+}
+
+/// Whether an event marks the start or end of a scope.
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+enum EventKind {
+    Start,
+    End,
+}
+
+/// Enable the timer.
+#[inline]
+pub fn enable() {
+    unsafe {
+        ENABLED = true;
+    }
+}
+
+/// Whether the timer is enabled.
+#[inline]
+pub fn is_enabled() -> bool {
+    unsafe { ENABLED }
+}
+
+/// Clears the recorded events.
+#[inline]
+pub fn clear() {
+    RECORDER.lock().events.clear();
+}
+
+/// A scope that records an event when it is dropped.
+pub struct TimingScope {
+    name: &'static str,
+    span: Option<Span>,
+    id: u64,
+    thread_id: ThreadId,
+}
+
+impl TimingScope {
+    /// Create a new scope if timing is enabled.
+    pub fn new(name: &'static str, span: Option<Span>) -> Option<Self> {
+        if !is_enabled() {
+            return None;
+        }
+
+        let timestamp = SystemTime::now();
+        let thread_id = std::thread::current().id();
+
+        let mut recorder = RECORDER.lock();
+        let id = recorder.discriminator;
+        recorder.discriminator += 1;
+        recorder.events.push(Event {
+            kind: EventKind::Start,
+            timestamp,
+            id,
+            name,
+            span,
+            thread_id,
+        });
+
+        Some(TimingScope { name, span, id, thread_id })
+    }
+}
+
+impl Drop for TimingScope {
+    fn drop(&mut self) {
+        let event = Event {
+            kind: EventKind::End,
+            timestamp: SystemTime::now(),
+            id: self.id,
+            name: self.name,
+            span: self.span,
+            thread_id: self.thread_id,
+        };
+
+        RECORDER.lock().events.push(event);
+    }
+}
+
+/// Creates a timing scope around an expression.
+///
+/// The output of the expression is returned.
+///
+/// The scope will be named `name` and will have the span `span`. The span is
+/// optional.
+///
+/// ## Example
+///
+/// ```rs
+/// // With a scope name and span.
+/// timed!(
+///     "my scope",
+///     span = Span::detached(),
+///     std::thread::sleep(std::time::Duration::from_secs(1)),
+/// );
+///
+/// // With a scope name and no span.
+/// timed!(
+///     "my scope",
+///     std::thread::sleep(std::time::Duration::from_secs(1)),
+/// );
+/// ```
+#[macro_export]
+macro_rules! timed {
+    ($name:expr, span = $span:expr, $body:expr $(,)?) => {{
+        let __scope = $crate::TimingScope::new($name, Some($span));
+        $body
+    }};
+    ($name:expr, $body:expr $(,)?) => {{
+        let __scope = $crate::TimingScope::new($name, None);
+        $body
+    }};
+}
+
+/// Export data as JSON for Chrome's tracing tool.
+///
+/// The `source` function is called for each span to get the source code
+/// location of the span. The first element of the tuple is the file path and
+/// the second element is the line number.
+pub fn export_json<W: Write>(
+    writer: W,
+    mut source: impl FnMut(Span) -> (String, u32),
+) -> Result<(), String> {
+    #[derive(Serialize)]
+    struct Entry {
+        name: &'static str,
+        cat: &'static str,
+        ph: &'static str,
+        ts: f64,
+        pid: u64,
+        tid: u64,
+        args: Option<Args>,
+    }
+
+    #[derive(Serialize)]
+    struct Args {
+        file: String,
+        line: u32,
+    }
+
+    let recorder = RECORDER.lock();
+    let run_start = recorder
+        .events
+        .first()
+        .map(|event| event.timestamp)
+        .unwrap_or_else(SystemTime::now);
+
+    let mut serializer = serde_json::Serializer::new(writer);
+    let mut seq = serializer
+        .serialize_seq(Some(recorder.events.len()))
+        .map_err(|e| format!("failed to serialize events: {e}"))?;
+
+    for event in recorder.events.iter() {
+        seq.serialize_element(&Entry {
+            name: event.name,
+            cat: "typst",
+            ph: match event.kind {
+                EventKind::Start => "B",
+                EventKind::End => "E",
+            },
+            ts: event
+                .timestamp
+                .duration_since(run_start)
+                .unwrap_or(Duration::ZERO)
+                .as_nanos() as f64
+                / 1_000.0,
+            pid: 1,
+            tid: unsafe {
+                // Safety: `thread_id` is a `ThreadId` which is a `u64`.
+                std::mem::transmute_copy(&event.thread_id)
+            },
+            args: event.span.map(&mut source).map(|(file, line)| Args { file, line }),
+        })
+        .map_err(|e| format!("failed to serialize event: {e}"))?;
+    }
+
+    seq.end().map_err(|e| format!("failed to serialize events: {e}"))?;
+
+    Ok(())
+}