From ea6f76cf102f37416c2f123cf48c192901a6125e Mon Sep 17 00:00:00 2001
From: Randy Stauner <randy.stauner@shopify.com>
Date: Thu, 1 Aug 2024 15:56:22 -0700
Subject: [PATCH] Add harnesses to profile with stackprof and vernier

Auto install the gems outside of any gems the benchmark might require.
Save the file to the data dir and show a report when finished.
---
 harness-stackprof/harness.rb | 74 ++++++++++++++++++++++++++++++++++++
 harness-vernier/harness.rb   | 26 +++++++++++++
 harness/harness-extra.rb     | 63 ++++++++++++++++++++++++++++++
 3 files changed, 163 insertions(+)
 create mode 100644 harness-stackprof/harness.rb
 create mode 100644 harness-vernier/harness.rb
 create mode 100644 harness/harness-extra.rb

diff --git a/harness-stackprof/harness.rb b/harness-stackprof/harness.rb
new file mode 100644
index 0000000..0a2b963
--- /dev/null
+++ b/harness-stackprof/harness.rb
@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+# Usage:
+# STACKPROF_OPTS='mode:object' MIN_BENCH_TIME=0 MIN_BENCH_ITRS=1 ruby -v -I harness-stackprof benchmarks/.../benchmark.rb
+# STACKPROF_OPTS='interval:10,mode:cpu' MIN_BENCH_TIME=1 MIN_BENCH_ITRS=10 ruby -v -I harness-stackprof benchmarks/.../benchmark.rb
+
+require "yaml"
+
+require_relative "../harness/harness-common"
+require_relative "../harness/harness-extra"
+
+ensure_global_gem("stackprof")
+
+BOOLS = {"true" => true, "false" => false}
+def bool!(val)
+  case val
+  when TrueClass, FalseClass
+    val
+  else
+    BOOLS.fetch(val) { raise ArgumentError, "must be 'true' or 'false'" }
+  end
+end
+
+DEFAULTS = {
+  aggregate: true,
+  raw: true,
+}
+
+def parse_opts_string(str)
+  return {} unless str
+
+  str.split(/,/).map { |x| x.strip.split(/[=:]/, 2) }.to_h.transform_keys(&:to_sym)
+end
+
+def stackprof_opts
+  opts = DEFAULTS.merge(parse_opts_string(ENV['STACKPROF_OPTS']))
+
+  bool = method(:bool!)
+
+  {
+    aggregate: bool,
+    raw: bool,
+    mode: :to_sym,
+    interval: :to_i,
+  }.each do |key, method|
+    next unless opts.key?(key)
+
+    method = proc(&method) if method.is_a?(Symbol)
+    opts[key] = method.call(opts[key])
+  rescue => error
+    raise ArgumentError, "Option '#{key}' failed to convert: #{error}"
+  end
+
+  opts
+end
+
+def run_benchmark(n, &block)
+  require "stackprof"
+
+  opts = stackprof_opts
+  prefix = "stackprof"
+  prefix = "#{prefix}-#{opts[:mode]}" if opts[:mode]
+
+  out = output_file_path(prefix: prefix, ext: "dump")
+  StackProf.run(out: out, **opts) do
+    run_enough_to_profile(n, &block)
+  end
+
+  gem_exe("stackprof", "--text", out)
+  puts "Stackprof dump file:\n#{out}"
+
+  # Dummy results to satisfy ./run_benchmarks.rb
+  return_results([0], [1.0]) if ENV['RESULT_JSON_PATH']
+end
diff --git a/harness-vernier/harness.rb b/harness-vernier/harness.rb
new file mode 100644
index 0000000..d32c18c
--- /dev/null
+++ b/harness-vernier/harness.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+# Usage:
+# MIN_BENCH_TIME=1 MIN_BENCH_ITRS=1 ruby -v -I harness-vernier benchmarks/...
+# NO_VIEWER=1 MIN_BENCH_TIME=1 MIN_BENCH_ITRS=1 ruby -v -I harness-vernier benchmarks/...
+
+require_relative "../harness/harness-common"
+require_relative "../harness/harness-extra"
+
+ensure_global_gem("vernier")
+ensure_global_gem_exe("profile-viewer")
+
+def run_benchmark(n, &block)
+  require "vernier"
+
+  out = output_file_path(ext: "json")
+  Vernier.profile(out: out) do
+    run_enough_to_profile(n, &block)
+  end
+
+  puts "Vernier profile:\n#{out}"
+  gem_exe("profile-viewer", out) unless ENV['NO_VIEWER'] == '1'
+
+  # Dummy results to satisfy ./run_benchmarks.rb
+  return_results([0], [1.0])
+end
diff --git a/harness/harness-extra.rb b/harness/harness-extra.rb
new file mode 100644
index 0000000..42a70d3
--- /dev/null
+++ b/harness/harness-extra.rb
@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+def ensure_global_gem(name)
+  found = Gem.find_latest_files(name).first
+  unless found
+    Gem.install(name)
+    found = Gem.find_latest_files(name).first
+  end
+  warn "Adding to load path: #{File.dirname(found)}"
+  $LOAD_PATH << File.dirname(found)
+end
+
+def ensure_global_gem_exe(name, exe = name)
+  Gem.bin_path(name, exe)
+rescue Gem::GemNotFoundException
+  Gem.install(name)
+end
+
+def gem_exe(*args)
+  # Remove any bundler env from the benchmark, let the exe figure it out.
+  system({'RUBYOPT' => '', 'BUNDLER_SETUP' => nil}, *args)
+end
+
+def benchmark_name
+  $0.match(%r{([^/]+?)(?:(?:/benchmark)?\.rb)?$})[1]
+end
+
+def harness_name
+  $LOADED_FEATURES.reverse_each do |feat|
+    if m = feat.match(%r{/harness-([^/]+)/harness\.rb$})
+      return m[1]
+    end
+  end
+  raise "Unable to determine harness name"
+end
+
+# Share a single timestamp for everything from this execution.
+TIMESTAMP = Time.now.strftime('%F-%H%M%S')
+
+def output_file_path(prefix: harness_name, suffix: benchmark_name, ruby_info: ruby_version_info, timestamp: TIMESTAMP, ext: "bin")
+  File.expand_path("../data/#{prefix}-#{timestamp}-#{ruby_info}-#{suffix}.#{ext}", __dir__)
+end
+
+# Can we get the benchmark config name from somewhere?
+def ruby_version_info
+  "#{RUBY_ENGINE}-#{RUBY_ENGINE_VERSION}"
+end
+
+def get_time
+  Process.clock_gettime(Process::CLOCK_MONOTONIC)
+end
+
+MIN_BENCH_TIME = Integer(ENV.fetch('MIN_BENCH_TIME', 10))
+def run_enough_to_profile(n, &block)
+  start = get_time
+  loop do
+    # Allow MIN_BENCH_ITRS to override the argument.
+    n = ENV.fetch('MIN_BENCH_ITRS', n).to_i
+    n.times(&block)
+
+    break if (get_time - start) >= MIN_BENCH_TIME
+  end
+end