Add instrumentation (Part 1) (#761)

5 files changed, 192 insertions, 4 deletions

diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index 768c6053..7719874d 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -34,6 +34,12 @@ siphasher = "0.3"
 walkdir = "2"
 clap = { version = "4.2.1", features = ["derive", "env"] }
 open = "4.0.1"
+tracing = "0.1.37"
+tracing-subscriber = "0.3.16"
+tracing-flame = "0.2.0"
+tracing-error = "0.2"
+inferno = "0.11.15"
+tempfile = "3.5.0"
 
 [build-dependencies]
 clap = { version = "4.2.1", features = ["derive", "string"] }
diff --git a/cli/build.rs b/cli/build.rs
index 75e8492a..b6f7d6ee 100644
--- a/cli/build.rs
+++ b/cli/build.rs
@@ -24,9 +24,9 @@ pub fn typst_version() -> String {
     format!("{pkg} ({hash})")
 }
 
-mod args {
-    include!("src/args.rs");
-}
+#[path = "src/args.rs"]
+#[allow(dead_code)]
+mod args;
 
 fn main() {
     println!("cargo:rerun-if-env-changed=TYPST_VERSION");
diff --git a/cli/src/args.rs b/cli/src/args.rs
index 42d52140..699dbb97 100644
--- a/cli/src/args.rs
+++ b/cli/src/args.rs
@@ -17,6 +17,10 @@ pub struct CliArguments {
     /// The typst command to run
     #[command(subcommand)]
     pub command: Command,
+
+    /// Sets the level of verbosity: 0 = none, 1 = warning & error, 2 = info, 3 = debug, 4 = trace
+    #[clap(short, long, action = ArgAction::Count)]
+    pub verbosity: u8,
 }
 
 /// What to do.
@@ -35,6 +39,22 @@ pub enum Command {
     Fonts(FontsCommand),
 }
 
+impl Command {
+    /// Returns the compile command if this is a compile or watch command.
+    pub fn as_compile(&self) -> Option<&CompileCommand> {
+        match self {
+            Command::Compile(cmd) => Some(cmd),
+            Command::Watch(cmd) => Some(cmd),
+            Command::Fonts(_) => None,
+        }
+    }
+
+    /// Returns whether this is a watch command.
+    pub fn is_watch(&self) -> bool {
+        matches!(self, Command::Watch(_))
+    }
+}
+
 /// Compiles the input file into a PDF file
 #[derive(Debug, Clone, Parser)]
 pub struct CompileCommand {
@@ -47,6 +67,11 @@ pub struct CompileCommand {
     /// Opens the output file after compilation using the default PDF viewer
     #[arg(long = "open")]
     pub open: Option<Option<String>>,
+
+    /// Produces a flamegraph of the compilation process and saves it to the
+    /// given file or to `flamegraph.svg` in the current working directory.
+    #[arg(long = "flamegraph", value_name = "OUTPUT_SVG")]
+    pub flamegraph: Option<Option<PathBuf>>,
 }
 
 /// List all discovered fonts in system and custom font paths
diff --git a/cli/src/main.rs b/cli/src/main.rs
index 924ea9d7..476a82ee 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -1,4 +1,5 @@
 mod args;
+mod trace;
 
 use std::cell::{RefCell, RefMut};
 use std::collections::HashMap;
@@ -28,6 +29,7 @@ use typst::World;
 use walkdir::WalkDir;
 
 use crate::args::{CliArguments, Command, CompileCommand};
+use crate::trace::init_tracing;
 
 type CodespanResult<T> = Result<T, CodespanError>;
 type CodespanError = codespan_reporting::files::Error;
@@ -80,7 +82,7 @@ impl CompileSettings {
     /// Panics if the command is not a compile or watch command.
     pub fn with_arguments(args: CliArguments) -> Self {
         let watch = matches!(args.command, Command::Watch(_));
-        let CompileCommand { input, output, open } = match args.command {
+        let CompileCommand { input, output, open, .. } = match args.command {
             Command::Compile(command) => command,
             Command::Watch(command) => command,
             _ => unreachable!(),
@@ -118,6 +120,13 @@ impl FontsSettings {
 /// Entry point.
 fn main() {
     let arguments = CliArguments::parse();
+    let _guard = match init_tracing(&arguments) {
+        Ok(guard) => guard,
+        Err(err) => {
+            eprintln!("failed to initialize tracing, reason: {}", err);
+            return;
+        }
+    };
 
     let res = match &arguments.command {
         Command::Compile(_) | Command::Watch(_) => {
@@ -225,7 +234,10 @@ fn compile(mut command: CompileSettings) -> StrResult<()> {
 }
 
 /// Compile a single time.
+#[tracing::instrument(skip_all)]
 fn compile_once(world: &mut SystemWorld, command: &CompileSettings) -> StrResult<bool> {
+    tracing::info!("Starting compilation");
+
     status(command, Status::Compiling).unwrap();
 
     world.reset();
@@ -237,6 +249,8 @@ fn compile_once(world: &mut SystemWorld, command: &CompileSettings) -> StrResult
             let buffer = typst::export::pdf(&document);
             fs::write(&command.output, buffer).map_err(|_| "failed to write PDF file")?;
             status(command, Status::Success).unwrap();
+
+            tracing::info!("Compilation succeeded");
             Ok(false)
         }
 
@@ -245,12 +259,15 @@ fn compile_once(world: &mut SystemWorld, command: &CompileSettings) -> StrResult
             status(command, Status::Error).unwrap();
             print_diagnostics(world, *errors)
                 .map_err(|_| "failed to print diagnostics")?;
+
+            tracing::info!("Compilation failed");
             Ok(true)
         }
     }
 }
 
 /// Clear the terminal and render the status message.
+#[tracing::instrument(skip_all)]
 fn status(command: &CompileSettings, status: Status) -> io::Result<()> {
     if !command.watch {
         return Ok(());
@@ -431,6 +448,7 @@ impl World for SystemWorld {
         self.source(self.main)
     }
 
+    #[tracing::instrument(skip_all)]
     fn resolve(&self, path: &Path) -> FileResult<SourceId> {
         self.slot(path)?
             .source
@@ -469,6 +487,7 @@ impl World for SystemWorld {
 }
 
 impl SystemWorld {
+    #[tracing::instrument(skip_all)]
     fn slot(&self, path: &Path) -> FileResult<RefMut<PathSlot>> {
         let mut hashes = self.hashes.borrow_mut();
         let hash = match hashes.get(path).cloned() {
@@ -488,6 +507,7 @@ impl SystemWorld {
         }))
     }
 
+    #[tracing::instrument(skip_all)]
     fn insert(&self, path: &Path, text: String) -> SourceId {
         let id = SourceId::from_u16(self.sources.len() as u16);
         let source = Source::new(id, path, text);
@@ -520,6 +540,7 @@ impl SystemWorld {
                 .map_or(false, |hash| self.paths.borrow().contains_key(&hash))
     }
 
+    #[tracing::instrument(skip_all)]
     fn reset(&mut self) {
         self.sources.as_mut().clear();
         self.hashes.borrow_mut().clear();
@@ -542,6 +563,7 @@ impl PathHash {
 }
 
 /// Read a file.
+#[tracing::instrument(skip_all)]
 fn read(path: &Path) -> FileResult<Vec<u8>> {
     let f = |e| FileError::from_io(e, path);
     if fs::metadata(path).map_err(f)?.is_dir() {
diff --git a/cli/src/trace.rs b/cli/src/trace.rs
new file mode 100644
index 00000000..0b37c5fc
--- /dev/null
+++ b/cli/src/trace.rs
@@ -0,0 +1,135 @@
+use std::fs::File;
+use std::io::{BufReader, BufWriter, Error, ErrorKind, Seek, SeekFrom};
+use std::path::PathBuf;
+
+use inferno::flamegraph::Options;
+use tracing::metadata::LevelFilter;
+use tracing_error::ErrorLayer;
+use tracing_flame::{FlameLayer, FlushGuard};
+use tracing_subscriber::fmt;
+use tracing_subscriber::prelude::*;
+
+use crate::args::CliArguments;
+
+/// Will flush the flamegraph to disk when dropped.
+pub struct TracingGuard {
+    flush_guard: Option<FlushGuard<BufWriter<File>>>,
+    temp_file: File,
+    output_svg: PathBuf,
+}
+
+impl TracingGuard {
+    pub fn finish(&mut self) -> Result<(), Error> {
+        if self.flush_guard.is_none() {
+            return Ok(());
+        }
+
+        tracing::info!("Flushing tracing flamegraph...");
+
+        // At this point, we're done tracing, so we can drop the guard.
+        // This will flush the tracing output to disk.
+        // We can then read the file and generate the flamegraph.
+        drop(self.flush_guard.take());
+
+        // Reset the file pointer to the beginning.
+        self.temp_file.seek(SeekFrom::Start(0))?;
+
+        // Create the readers and writers.
+        let reader = BufReader::new(&mut self.temp_file);
+        let output = BufWriter::new(File::create(&self.output_svg)?);
+
+        // Create the options: default in flame chart mode
+        let mut options = Options::default();
+        options.flame_chart = true;
+
+        inferno::flamegraph::from_reader(&mut options, reader, output)
+            .map_err(|e| Error::new(ErrorKind::Other, e))?;
+
+        Ok(())
+    }
+}
+
+impl Drop for TracingGuard {
+    fn drop(&mut self) {
+        if !std::thread::panicking() {
+            if let Err(e) = self.finish() {
+                // Since we are finished, we cannot rely on tracing to log the
+                // error.
+                eprintln!("Failed to flush tracing flamegraph: {e}");
+            }
+        }
+    }
+}
+
+/// Initializes the tracing system and returns a guard that will flush the
+/// flamegraph to disk when dropped.
+pub fn init_tracing(args: &CliArguments) -> Result<Option<TracingGuard>, Error> {
+    let flamegraph = args.command.as_compile().and_then(|c| c.flamegraph.as_ref());
+
+    if flamegraph.is_some() && args.command.is_watch() {
+        return Err(Error::new(
+            ErrorKind::InvalidInput,
+            "cannot use --flamegraph with watch command",
+        ));
+    }
+
+    // Short circuit if we don't need to initialize flamegraph or debugging.
+    if flamegraph.is_none() && args.verbosity == 0 {
+        tracing_subscriber::fmt()
+            .without_time()
+            .with_max_level(level_filter(args))
+            .init();
+
+        return Ok(None);
+    }
+
+    // Build the FMT layer printing to the console.
+    let fmt_layer = fmt::Layer::default().without_time().with_filter(level_filter(args));
+
+    // Error layer for building backtraces
+    let error_layer = ErrorLayer::default();
+
+    // Build the registry.
+    let registry = tracing_subscriber::registry().with(fmt_layer).with(error_layer);
+
+    let Some(path) = flamegraph else {
+        registry.init();
+        return Ok(None);
+    };
+
+    // Create a temporary file to store the flamegraph data.
+    let temp_file = tempfile::tempfile()?;
+    let writer = BufWriter::new(temp_file.try_clone()?);
+
+    // Build the flamegraph layer.
+    let flame_layer = FlameLayer::new(writer)
+        .with_empty_samples(false)
+        .with_threads_collapsed(true)
+        .with_module_path(false)
+        .with_file_and_line(true);
+    let flush_guard = flame_layer.flush_on_drop();
+
+    // Build the subscriber.
+    registry.with(flame_layer).init();
+
+    tracing::warn!(
+        "Flamegraph is enabled, this can create a large temporary \
+         file and slow down the compilation process."
+    );
+
+    Ok(Some(TracingGuard {
+        flush_guard: Some(flush_guard),
+        temp_file,
+        output_svg: path.clone().unwrap_or_else(|| "flamegraph.svg".into()),
+    }))
+}
+
+/// Returns the log level filter for the given verbosity level.
+fn level_filter(args: &CliArguments) -> LevelFilter {
+    match args.verbosity {
+        0 => LevelFilter::WARN,
+        1 => LevelFilter::INFO,
+        2 => LevelFilter::DEBUG,
+        _ => LevelFilter::TRACE,
+    }
+}