feat(tvix/eval/observer): Allow capturing timing of events r/7459

Add a new --trace-runtime-timing flag (probably a better bikeshed for
this) that enables capturing the time, relative to the last event, of
each event recorded with the tracing observer.

This probably isn't *super* useful yet, but I'd like to start here in
adding new profiling tools to the VM, specifically based on the runtime
observer

Change-Id: Id7f12077291c39bf3eef42ab6744bfba53687a65
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10713
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

2 files changed, 36 insertions, 0 deletions

diff --git a/tvix/cli/src/main.rs b/tvix/cli/src/main.rs
index 7683ad4c59f8..72f21ab5b185 100644
--- a/tvix/cli/src/main.rs
+++ b/tvix/cli/src/main.rs
@@ -31,6 +31,11 @@ struct Args {
     #[clap(long, env = "TVIX_TRACE_RUNTIME")]
     trace_runtime: bool,
 
+    /// Capture the time (relative to the start time of evaluation) of all events traced with
+    /// `--trace-runtime`
+    #[clap(long, env = "TVIX_TRACE_RUNTIME_TIMING", requires("trace_runtime"))]
+    trace_runtime_timing: bool,
+
     /// Only compile, but do not execute code. This will make Tvix act
     /// sort of like a linter.
     #[clap(long)]
@@ -113,6 +118,9 @@ fn interpret(code: &str, path: Option<PathBuf>, args: &Args, explain: bool) -> b
 
         let mut runtime_observer = TracingObserver::new(std::io::stderr());
         if args.trace_runtime {
+            if args.trace_runtime_timing {
+                runtime_observer.enable_timing()
+            }
             eval.runtime_observer = Some(&mut runtime_observer);
         }
 
diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index d805f0338d63..f5de399315c7 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -8,6 +8,7 @@
 /// what they are interested in observing.
 use std::io::Write;
 use std::rc::Rc;
+use std::time::Instant;
 use tabwriter::TabWriter;
 
 use crate::chunk::Chunk;
@@ -140,16 +141,31 @@ impl<W: Write> CompilerObserver for DisassemblingObserver<W> {
 /// An observer that collects a textual representation of an entire
 /// runtime execution.
 pub struct TracingObserver<W: Write> {
+    // If timing is enabled, contains the timestamp of the last-emitted trace event
+    last_event: Option<Instant>,
     writer: TabWriter<W>,
 }
 
 impl<W: Write> TracingObserver<W> {
     pub fn new(writer: W) -> Self {
         Self {
+            last_event: None,
             writer: TabWriter::new(writer),
         }
     }
 
+    /// Write the time of each runtime event, relative to when this method is called
+    pub fn enable_timing(&mut self) {
+        self.last_event = Some(Instant::now());
+    }
+
+    fn maybe_write_time(&mut self) {
+        if let Some(last_event) = &mut self.last_event {
+            let _ = write!(&mut self.writer, "+{}ns\t", last_event.elapsed().as_nanos());
+            *last_event = Instant::now();
+        }
+    }
+
     fn write_value(&mut self, val: &Value) {
         let _ = match val {
             // Potentially large types which we only want to print
@@ -191,6 +207,8 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
         lambda: &Rc<Lambda>,
         call_depth: usize,
     ) {
+        self.maybe_write_time();
+
         let _ = write!(&mut self.writer, "=== entering ");
 
         let _ = if arg_count == 0 {
@@ -212,17 +230,20 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
 
     /// Called when the runtime exits a call frame.
     fn observe_exit_call_frame(&mut self, frame_at: usize, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(&mut self.writer, "=== exiting frame {} ===\t ", frame_at);
         self.write_stack(stack);
     }
 
     fn observe_suspend_call_frame(&mut self, frame_at: usize, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(&mut self.writer, "=== suspending frame {} ===\t", frame_at);
 
         self.write_stack(stack);
     }
 
     fn observe_enter_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(
             &mut self.writer,
             "=== entering generator frame '{}' [{}] ===\t",
@@ -233,6 +254,7 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
     }
 
     fn observe_exit_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(
             &mut self.writer,
             "=== exiting generator '{}' [{}] ===\t",
@@ -243,6 +265,7 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
     }
 
     fn observe_suspend_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(
             &mut self.writer,
             "=== suspending generator '{}' [{}] ===\t",
@@ -253,6 +276,7 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
     }
 
     fn observe_generator_request(&mut self, name: &str, msg: &VMRequest) {
+        self.maybe_write_time();
         let _ = writeln!(
             &mut self.writer,
             "=== generator '{}' requested {} ===",
@@ -261,15 +285,18 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
     }
 
     fn observe_enter_builtin(&mut self, name: &'static str) {
+        self.maybe_write_time();
         let _ = writeln!(&mut self.writer, "=== entering builtin {} ===", name);
     }
 
     fn observe_exit_builtin(&mut self, name: &'static str, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(&mut self.writer, "=== exiting builtin {} ===\t", name);
         self.write_stack(stack);
     }
 
     fn observe_tail_call(&mut self, frame_at: usize, lambda: &Rc<Lambda>) {
+        self.maybe_write_time();
         let _ = writeln!(
             &mut self.writer,
             "=== tail-calling {:p} in frame[{}] ===",
@@ -278,6 +305,7 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
     }
 
     fn observe_execute_op(&mut self, ip: CodeIdx, op: &OpCode, stack: &[Value]) {
+        self.maybe_write_time();
         let _ = write!(&mut self.writer, "{:04} {:?}\t", ip.0, op);
         self.write_stack(stack);
     }