Add ability to trace exit locations in yjit (#5970)

When running with `--yjit-stats` turned on, yjit can inform the user
what the most common exits are. While this is useful information it
doesn't tell you the source location of the code that exited or what the
code that exited looks like. This change intends to fix that.

To use the feature, run yjit with the `--yjit-trace-exits` option,
which will record the backtrace for every exit that occurs. This functionality
requires the stats feature to be turned on. Calling `--yjit-trace-exits`
will automatically set the `--yjit-stats` option.

Users must call `RubyVM::YJIT.dump_exit_locations(filename)` which will
Marshal dump the contents of `RubyVM::YJIT.exit_locations` into a file
based on the passed filename.

*Example usage:*

Given the following script, we write to a file called
`concat_array.dump` the results of `RubyVM::YJIT.exit_locations`.

```ruby
def concat_array
  ["t", "r", *x = "u", "e"].join
end

1000.times do
  concat_array
end

RubyVM::YJIT.dump_exit_locations("concat_array.dump")
```

When we run the file with this branch and the appropriate flags the
stacktrace will be recorded. Note Stackprof needs to be installed or you
need to point to the library directly.

```
./ruby --yjit --yjit-call-threshold=1 --yjit-trace-exits -I/Users/eileencodes/open_source/stackprof/lib test.rb
```

We can then read the dump file with Stackprof:

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump
```

Results will look similar to the following:

```
==================================
  Mode: ()
  Samples: 1817 (0.00% miss rate)
  GC: 0 (0.00%)
==================================
     TOTAL    (pct)     SAMPLES    (pct)     FRAME
      1001  (55.1%)        1001  (55.1%)     concatarray
       335  (18.4%)         335  (18.4%)     invokeblock
       178   (9.8%)         178   (9.8%)     send
       140   (7.7%)         140   (7.7%)     opt_getinlinecache
       ...etc...
```

Simply inspecting the `concatarray` method will give `SOURCE
UNAVAILABLE` because the source is insns.def.

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump --method concatarray
```

Result:

```
concatarray (nonexistent.def:1)
  samples:  1001 self (55.1%)  /   1001 total (55.1%)
  callers:
    1000  (   99.9%)  Object#concat_array
       1  (    0.1%)  Gem.suffixes
  callees (0 total):
  code:
        SOURCE UNAVAILABLE
```

However if we go deeper to the callee we can see the exact
source of the `concatarray` exit.

```
./ruby -I/Users/eileencodes/open_source/stackprof/lib/ /Users/eileencodes/open_source/stackprof/bin/stackprof --text concat_array.dump --method Object#concat_array
```

```
Object#concat_array (/Users/eileencodes/open_source/rust_ruby/test.rb:1)
  samples:     0 self (0.0%)  /   1000 total (55.0%)
  callers:
    1000  (  100.0%)  block in <main>
  callees (1000 total):
    1000  (  100.0%)  concatarray
  code:
                                  |     1  | def concat_array
 1000   (55.0%)                   |     2  |   ["t", "r", *x = "u", "e"].join
                                  |     3  | end
```

The `--walk` option is recommended for this feature as it make it
easier to traverse the tree of exits.

*Goals of this feature:*

This feature is meant to give more information when working on YJIT.
The idea is that if we know what code is exiting we can decide what
areas to prioritize when fixing exits. In some cases this means adding
prioritizing avoiding certain exits in yjit. In more complex cases it
might mean changing the Ruby code to be more performant when run with
yjit. Ultimately the more information we have about what code is exiting
AND why, the better we can make yjit.

*Known limitations:*

* Due to tracing exits, running this on large codebases like Rails
can be quite slow.
* On complex methods it can still be difficult to pinpoint the exact cause of
an exit.
* Stackprof is a requirement to to view the backtrace information from
the dump file.

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>

1 files changed, 100 insertions, 0 deletions

diff --git a/yjit.rb b/yjit.rb
index 7c5311d79a..cbaea7d003 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -18,11 +18,111 @@ module RubyVM::YJIT
     Primitive.rb_yjit_stats_enabled_p
   end
 
+  # Check if rb_yjit_trace_exit_locations_enabled_p is enabled.
+  def self.trace_exit_locations_enabled?
+    Primitive.rb_yjit_trace_exit_locations_enabled_p
+  end
+
   # Discard statistics collected for --yjit-stats.
   def self.reset_stats!
     Primitive.rb_yjit_reset_stats_bang
   end
 
+  # If --yjit-trace-exits is enabled parse the hashes from
+  # Primitive.rb_yjit_get_exit_locations into a format readable
+  # by Stackprof. This will allow us to find the exact location of a
+  # side exit in YJIT based on the instruction that is exiting.
+  def self.exit_locations
+    return unless trace_exit_locations_enabled?
+
+    results = Primitive.rb_yjit_get_exit_locations
+    raw_samples = results[:raw].dup
+    line_samples = results[:lines].dup
+    frames = results[:frames].dup
+    samples_count = 0
+
+    frames.each do |frame_id, frame|
+      frame[:samples] = 0
+      frame[:edges] = {}
+    end
+
+    # Loop through the instructions and set the frame hash with the data.
+    # We use nonexistent.def for the file name, otherwise insns.def will be displayed
+    # and that information isn't useful in this context.
+    RubyVM::INSTRUCTION_NAMES.each_with_index do |name, frame_id|
+      frame_hash = { samples: 0, total_samples: 0, edges: {}, name: name, file: "nonexistent.def", line: nil }
+      results[:frames][frame_id] = frame_hash
+      frames[frame_id] = frame_hash
+    end
+
+    # Loop through the raw_samples and build the hashes for StackProf.
+    # The loop is based off an example in the StackProf documentation and therefore
+    # this functionality can only work with that library.
+    while raw_samples.length > 0
+      stack_trace = raw_samples.shift(raw_samples.shift + 1)
+      lines = line_samples.shift(line_samples.shift + 1)
+      prev_frame_id = nil
+
+      stack_trace.each_with_index do |frame_id, idx|
+        if prev_frame_id
+          prev_frame = frames[prev_frame_id]
+          prev_frame[:edges][frame_id] ||= 0
+          prev_frame[:edges][frame_id] += 1
+        end
+
+        frame_info = frames[frame_id]
+        frame_info[:total_samples] ||= 0
+        frame_info[:total_samples] += 1
+
+        frame_info[:lines] ||= {}
+        frame_info[:lines][lines[idx]] ||= [0, 0]
+        frame_info[:lines][lines[idx]][0] += 1
+
+        prev_frame_id = frame_id
+      end
+
+      top_frame_id = stack_trace.last
+      top_frame_line = 1
+
+      frames[top_frame_id][:samples] += 1
+      frames[top_frame_id][:lines] ||= {}
+      frames[top_frame_id][:lines][top_frame_line] ||= [0, 0]
+      frames[top_frame_id][:lines][top_frame_line][1] += 1
+
+      samples_count += raw_samples.shift
+      line_samples.shift
+    end
+
+    results[:samples] = samples_count
+    # Set missed_samples and gc_samples to 0 as their values
+    # don't matter to us in this context.
+    results[:missed_samples] = 0
+    results[:gc_samples] = 0
+    results
+  end
+
+  # Marshal dumps exit locations to the given filename.
+  #
+  # Usage:
+  #
+  # In a script call:
+  #
+  #   RubyVM::YJIT.dump_exit_locations("my_file.dump")
+  #
+  # Then run the file with the following options:
+  #
+  #   ruby --yjit --yjit-stats --yjit-trace-exits test.rb
+  #
+  # Once the code is done running, use Stackprof to read the dump file.
+  # See Stackprof documentation for options.
+  def self.dump_exit_locations(filename)
+    unless trace_exit_locations_enabled?
+      raise ArgumentError, "--yjit-trace-exits must be enabled to use dump_exit_locations."
+    end
+
+    File.write(filename, Marshal.dump(RubyVM::YJIT.exit_locations))
+  end
+
   # Return a hash for statistics generated for the --yjit-stats command line option.
   # Return nil when option is not passed or unavailable.
   def self.runtime_stats