From 9ab978b7183d92cf1451be5dce7c7cc72222a39a Mon Sep 17 00:00:00 2001
From: "Eileen M. Uchitelle" <eileencodes@users.noreply.github.com>
Date: Wed, 9 Nov 2022 17:09:16 -0500
Subject: [PATCH] Fix exit locations dump (#6703)

While I was working on my RubyConf talk for tracing yjit exit locations
I realized that there were exits from the dump code included in the
stats data. For example I saw 224 interp leave exits for a simple script
that should have had 1 or 2. I realized that the dump code needs to be
called _after_ the stats are generated, otherwise the dump code will be
counted in the stats exits.

I've added a `_dump_locations` method to the `at_exit` for stats
generation to ensure that it runs last. I've updated the documentation
to add a note that if you call `dump_exit_locations` directly, your
stats will include the dump code exits as well.
---
 yjit.rb | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/yjit.rb b/yjit.rb
index 5a00dcca575045..21f2eea4def39c 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -112,13 +112,25 @@ def self.exit_locations
   #
   # Usage:
   #
+  # If `--yjit-exit-locations` is passed, a file named
+  # "yjit_exit_locations.dump" will automatically be generated.
+  #
+  # If you want to collect traces manually, call `dump_exit_locations`
+  # directly.
+  #
+  # Note that calling this in a script will generate stats after the
+  # dump is created, so the stats data may include exits from the
+  # dump itself.
+  #
   # In a script call:
   #
-  #   RubyVM::YJIT.dump_exit_locations("my_file.dump")
+  #   at_exit do
+  #     RubyVM::YJIT.dump_exit_locations("my_file.dump")
+  #   end
   #
   # Then run the file with the following options:
   #
-  #   ruby --yjit --yjit-stats --yjit-trace-exits test.rb
+  #   ruby --yjit --yjit-trace-exits test.rb
   #
   # Once the code is done running, use Stackprof to read the dump file.
   # See Stackprof documentation for options.
@@ -196,12 +208,24 @@ def self.simulate_oom!
 
   # Avoid calling a method here to not interfere with compilation tests
   if Primitive.rb_yjit_stats_enabled_p
-    at_exit { _print_stats }
+    at_exit do
+      _print_stats
+      _dump_locations
+    end
   end
 
   class << self
     private
 
+    def _dump_locations
+      return unless trace_exit_locations_enabled?
+
+      filename = "yjit_exit_locations.dump"
+      dump_exit_locations(filename)
+
+      $stderr.puts("YJIT exit locations dumped to `#{filename}`.")
+    end
+
     # Format and print out counters
     def _print_stats
       stats = runtime_stats