From 91ccf91dcfa5664c0328cadb1645f89912629e55 Mon Sep 17 00:00:00 2001
From: Ed Robinson <edward-robinson@cookpad.com>
Date: Tue, 27 Feb 2024 14:00:22 +0000
Subject: [PATCH] Collect ratio_in_yjit from YJIT stats if avalible (#54)

* Collect ratio_in_yjit from YJIT stats if avalible

* This is only collected if --yjit-stats is set in RUBYOPT
* Since setting --yjit-stats has some overhead, we don't want to require it
* I added a bin that we can shell out to with various values of RUBYOPT
  to intergration test this stuff, rather than trying to enable /
disable yjit within the test suite process

* fmt

* Move test to run on all rubies
---
 bin/yjit_intergration_test  | 11 ++++++
 lib/promenade/yjit/stats.rb | 13 ++++++-
 spec/yjit_spec.rb           | 69 ++++++++++++++++++++++++++-----------
 3 files changed, 71 insertions(+), 22 deletions(-)
 create mode 100755 bin/yjit_intergration_test

diff --git a/bin/yjit_intergration_test b/bin/yjit_intergration_test
new file mode 100755
index 0000000..b753522
--- /dev/null
+++ b/bin/yjit_intergration_test
@@ -0,0 +1,11 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "promenade"
+require "promenade/yjit/stats"
+require "prometheus/client"
+require "prometheus/client/formats/text"
+
+Promenade.setup
+Promenade::YJIT::Stats.instrument
+puts Prometheus::Client::Formats::Text.marshal_multiprocess
diff --git a/lib/promenade/yjit/stats.rb b/lib/promenade/yjit/stats.rb
index d249f05..caea871 100644
--- a/lib/promenade/yjit/stats.rb
+++ b/lib/promenade/yjit/stats.rb
@@ -1,14 +1,25 @@
 module Promenade
   module YJIT
     class Stats
+      RUNTIME_STATS = %i(
+        code_region_size
+        ratio_in_yjit
+      ).freeze
+
       Promenade.gauge :ruby_yjit_code_region_size do
         doc "Ruby YJIT code size"
       end
 
+      Promenade.gauge :ruby_yjit_ratio_in_yjit do
+        doc "Shows the ratio of YJIT-executed instructions in %"
+      end
+
       def self.instrument
         return unless defined?(::RubyVM::YJIT) && ::RubyVM::YJIT.enabled?
 
-        Promenade.metric(:ruby_yjit_code_region_size).set({}, ::RubyVM::YJIT.runtime_stats[:code_region_size])
+        ::RubyVM::YJIT.runtime_stats.select { |stat, _| RUNTIME_STATS.include? stat }.each do |stat, value|
+          Promenade.metric(:"ruby_yjit_#{stat}").set({}, value)
+        end
       end
     end
   end
diff --git a/spec/yjit_spec.rb b/spec/yjit_spec.rb
index f1e29df..9147e66 100644
--- a/spec/yjit_spec.rb
+++ b/spec/yjit_spec.rb
@@ -1,31 +1,58 @@
 require "promenade/yjit/stats"
+require "open3"
 
 RSpec.describe Promenade::YJIT::Stats do
   describe "recording yjit stats" do
-    it "records code_region_size" do
-      # This method should not blow up in any case
+    it "doesn't explode" do
+      # This method should not blow up in any case, on any version of ruby
       expect { described_class.instrument }.not_to raise_error
 
-      if defined?(RubyVM::YJIT) && defined?(RubyVM::YJIT.enable)
-
-        # We want to test that this doesn't blow up when yjit is present but isn't enabled yet
-        # you need to run the testsuite with yjit disabled for this to work
-        expect(RubyVM::YJIT.enabled?).to be_falsey
-        expect { described_class.instrument }.not_to raise_error
-
-        # Then we enable yjit to test the instrumentation
-        RubyVM::YJIT.enable
-        described_class.instrument
-
-        expect(Promenade.metric(:ruby_yjit_code_region_size).get).to eq RubyVM::YJIT.runtime_stats[:code_region_size]
-      else
-        version = RUBY_VERSION.match(/(\d).(\d).\d/)
-        major = version[1].to_i
-        minor = version[2].to_i
-        if major >= 3 && minor >= 3
-          flunk "YJIT must be avalibe to test properly in ruby 3.3+"
-        end
+      metrics = run_yjit_metrics("")
+      expect(metrics).to be_empty
+    end
+
+    it "records yjit stats" do
+      version = RUBY_VERSION.match(/(\d).(\d).\d/)
+      major = version[1].to_i
+      minor = version[2].to_i
+      unless major >= 3 && minor >= 3
+        pending "YJIT metrics are only expected to work in ruby 3.3.0+"
       end
+
+      metrics = run_yjit_metrics("--yjit")
+      expect(metrics[:ruby_yjit_code_region_size]).to satisfy("be nonzero") { |n| n > 0 }
+      # ratio_in_yjit is only set when --yjit-stats is enabled
+      expect(metrics[:ruby_yjit_ratio_in_yjit]).to be_nil
+
+      metrics = run_yjit_metrics("--yjit --yjit-stats=quiet")
+      expect(metrics[:ruby_yjit_code_region_size]).to satisfy("be nonzero") { |n| n > 0 }
+      expect(metrics[:ruby_yjit_ratio_in_yjit]).to satisfy("be nonzero") { |n| n > 0 }
+    end
+  end
+
+  def run_yjit_metrics(rubyopt)
+    dir = Dir.mktmpdir
+    begin
+      output, status = Open3.capture2e({ "PROMETHEUS_MULTIPROC_DIR" => dir, "RUBYOPT" => rubyopt }, "bin/yjit_intergration_test")
+      expect(status).to eq 0
+      parse_metrics(output)
+    ensure
+      FileUtils.remove_entry dir
     end
   end
+
+  def parse_metrics(output)
+    output.lines.reject { |line| line.match("#") }.filter_map do |line|
+      match = line.match(/([a-z_]+)\{.+\} (\d+\.?\d*)/)
+      next unless match
+
+      [match[1].to_sym, parse_number(match[2])]
+    end.to_h
+  end
+
+  def parse_number(string)
+    Integer(string)
+  rescue StandardError
+    string.to_f
+  end
 end