ZJIT: Enable sample rate for side exit tracing (#14696)

Author: aidenfoxivey

doc/zjit.md

@@ -155,7 +155,7 @@ make -j
 
 ### Tracing side exits
 
-Through [Stackprof](https://github.com/tmm1/stackprof), detailed information about the methods that the JIT side-exits from can be displayed after some execution of a program.
+Through [Stackprof](https://github.com/tmm1/stackprof), detailed information about the methods that the JIT side-exits from can be displayed after some execution of a program. Optionally, you can use `--zjit-trace-exits-sample-rate=N` to sample every N-th occurrence. Enabling `--zjit-trace-exits-sample-rate=N` will automatically enable `--zjit-trace-exits`.
 
 ```bash
 ./miniruby --zjit-trace-exits script.rb

zjit.rb

@@ -128,6 +128,7 @@ def dump_exit_locations(filename)
 
     File.open(filename, "wb") do |file|
       Marshal.dump(RubyVM::ZJIT.exit_locations, file)
+      file.size
     end
   end
 
@@ -275,9 +276,9 @@ def print_stats
   def dump_locations # :nodoc:
     return unless trace_exit_locations_enabled?
 
-    filename = "zjit_exit_locations.dump"
-    dump_exit_locations(filename)
+    filename = "zjit_exits_#{Time.now.to_i}.dump"
+    n_bytes = dump_exit_locations(filename)
 
-    $stderr.puts("ZJIT exit locations dumped to `#{filename}`.")
+    $stderr.puts("#{n_bytes} bytes written to #{filename}.")
   end
 end

zjit/src/options.rs

@@ -72,6 +72,9 @@ pub struct Options {
     /// Trace and write side exit source maps to /tmp for stackprof.
     pub trace_side_exits: bool,
 
+    /// Frequency of tracing side exits.
+    pub trace_side_exits_sample_interval: usize,
+
     /// Dump code map to /tmp for performance profilers.
     pub perf: bool,
 
@@ -98,6 +101,7 @@ impl Default for Options {
             dump_lir: false,
             dump_disasm: false,
             trace_side_exits: false,
+            trace_side_exits_sample_interval: 0,
             perf: false,
             allowed_iseqs: None,
             log_compiled_iseqs: None,
@@ -120,7 +124,9 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
     ("--zjit-trace-exits",
-                     "Record Ruby source location when side-exiting.")
+                     "Record Ruby source location when side-exiting."),
+    ("--zjit-trace-exits-sample-rate",
+                     "Frequency at which to record side exits. Must be `usize`.")
 ];
 
 #[derive(Clone, Copy, Debug)]
@@ -245,6 +251,13 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             options.trace_side_exits = true;
         }
 
+        ("trace-exits-sample-rate", sample_interval) => {
+            // Even if `trace_side_exits` is already set, set it.
+            options.trace_side_exits = true;
+            // `sample_interval ` must provide a string that can be validly parsed to a `usize`.
+            options.trace_side_exits_sample_interval = sample_interval.parse::<usize>().ok()?;
+        }
+
         ("debug", "") => options.debug = true,
 
         ("disable-hir-opt", "") => options.disable_hir_opt = true,

zjit/src/state.rs

@@ -223,6 +223,16 @@ impl ZJITState {
     pub fn get_line_samples() -> Option<&'static mut Vec<i32>> {
         ZJITState::get_instance().exit_locations.as_mut().map(|el| &mut el.line_samples)
     }
+
+    /// Get number of skipped samples.
+    pub fn get_skipped_samples() -> Option<&'static mut usize> {
+        ZJITState::get_instance().exit_locations.as_mut().map(|el| &mut el.skipped_samples)
+    }
+
+    /// Get number of skipped samples.
+    pub fn set_skipped_samples(n: usize) -> Option<()> {
+        ZJITState::get_instance().exit_locations.as_mut().map(|el| el.skipped_samples = n)
+    }
 }
 
 /// Initialize ZJIT
@@ -354,6 +364,20 @@ pub extern "C" fn rb_zjit_record_exit_stack(exit_pc: *const VALUE) {
         return;
     }
 
+    // When `trace_side_exits_sample_interval` is zero, then the feature is disabled.
+    if get_option!(trace_side_exits_sample_interval) != 0 {
+        // If `trace_side_exits_sample_interval` is set, then can safely unwrap
+        // both `get_skipped_samples` and `set_skipped_samples`.
+        let skipped_samples = *ZJITState::get_skipped_samples().unwrap();
+        if skipped_samples < get_option!(trace_side_exits_sample_interval) {
+            // Skip sample and increment counter.
+            ZJITState::set_skipped_samples(skipped_samples + 1).unwrap();
+            return;
+        } else {
+            ZJITState::set_skipped_samples(0).unwrap();
+        }
+    }
+
     let (stack_length, frames_buffer, lines_buffer) = record_profiling_frames();
 
     // Can safely unwrap since `trace_side_exits` must be true at this point

zjit/src/stats.rs

@@ -501,6 +501,8 @@ pub struct SideExitLocations {
     pub raw_samples: Vec<VALUE>,
     /// Line numbers of the iseq caller.
     pub line_samples: Vec<i32>,
+    /// Skipped samples
+    pub skipped_samples: usize
 }
 
 /// Primitive called in zjit.rb