From 9712669b2d1ac33784ee5225e9087d4a6259409c Mon Sep 17 00:00:00 2001
From: Randy Stauner <randy.stauner@shopify.com>
Date: Mon, 9 Dec 2024 11:39:03 -0700
Subject: [PATCH] Limit YJIT stats to the measured benchmark code

Reset stats at the beginning of the benchmark
and collect stats before loading any additional code.
---
 harness/harness-common.rb | 6 ++++--
 harness/harness.rb        | 2 ++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/harness/harness-common.rb b/harness/harness-common.rb
index 33e6315..dfab8a1 100644
--- a/harness/harness-common.rb
+++ b/harness/harness-common.rb
@@ -101,6 +101,9 @@ def return_results(warmup_iterations, bench_iterations)
     "bench" => bench_iterations,
   }
 
+  # Collect yjit stats before loading any additional code.
+  yjit_stats = RubyVM::YJIT.runtime_stats if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
+
   # Collect our own peak mem usage as soon as reasonable after finishing the last iteration.
   rss = get_rss
   yjit_bench_results["rss"] = rss
@@ -108,8 +111,7 @@ def return_results(warmup_iterations, bench_iterations)
     yjit_bench_results["maxrss"] = maxrss
   end
 
-  if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
-    yjit_stats = RubyVM::YJIT.runtime_stats
+  if yjit_stats
     yjit_bench_results["yjit_stats"] = yjit_stats
 
     formatted_stats = proc { |key| "%10s" % yjit_stats[key].to_s.reverse.scan(/\d{1,3}/).join(",").reverse }
diff --git a/harness/harness.rb b/harness/harness.rb
index 915292a..c32254d 100644
--- a/harness/harness.rb
+++ b/harness/harness.rb
@@ -38,6 +38,8 @@ def run_benchmark(_num_itrs_hint, &block)
   num_itrs = 0
   header = "itr:   time"
 
+  RubyVM::YJIT.reset_stats! if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
+
   # If $YJIT_BENCH_STATS is given, print the diff of these stats at each iteration.
   if ENV["YJIT_BENCH_STATS"]
     yjit_stats = ENV["YJIT_BENCH_STATS"].split(",").map { |key| [key.to_sym, nil] }.to_h