RDoc Store Page Lookup Benchmark

benchmark.rb

#!/usr/bin/env ruby
# frozen_string_literal: true

require "fileutils"
require "json"
require "optparse"
require "rbconfig"

SCRIPT = File.expand_path(__FILE__)
TMP_ROOT = File.expand_path("tmp", __dir__)
DEFAULT_RDOC = File.expand_path("../rdoc", __dir__)
EXCLUDE = '(?:^|/)(?:app/assets|lib/templates)/|\.(?:json|rbs|xml|xsd)\z'

def install_variant(profile: false)
  require "rdoc"
  require "rdoc/store"

  implementation = if ENV.fetch("PAGE_LOOKUP_VARIANT") == "optimized"
    Module.new do
      def page(name)
        @page_index ||= @files_hash.each_value.each_with_object({}) do |file, index|
          index[file.page_name] ||= file
          index[file.base_name] ||= file
        end

        @page_index[name]
      end
    end
  else
    Module.new do
      def page(name)
        @files_hash.each_value.find do |file|
          file.page_name == name || file.base_name == name
        end
      end
    end
  end
  RDoc::Store.prepend(implementation)
  return unless profile

  result = {files: 0, calls: 0, hits: 0, misses: 0, elapsed_seconds: 0.0}
  profiler = Module.new do
    define_method(:page) do |name|
      started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
      found = super(name)
      elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at
      result[:files] = @files_hash.size
      result[:calls] += 1
      result[found ? :hits : :misses] += 1
      result[:elapsed_seconds] += elapsed
      found
    end
  end
  RDoc::Store.prepend(profiler)
  at_exit { File.write(ENV.fetch("PAGE_LOOKUP_PROFILE"), JSON.pretty_generate(result)) }
end

def synthetic_child
  require "rdoc/options"

  file_count = Integer(ARGV.fetch(0))
  lookup_count = Integer(ARGV.fetch(1))
  store = RDoc::Store.new(RDoc::Options.new)
  file_count.times { |index| store.add_file("lib/generated/file_#{index.to_s.rjust(5, "0")}.rb") }
  lookups = Array.new(lookup_count) { |index| "missing/reference_#{index.to_s.rjust(5, "0")}.md" }

  GC.start
  started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  misses = lookups.count { |name| !store.page(name) }
  elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at

  File.write(ENV.fetch("PAGE_LOOKUP_RESULT"), JSON.generate(
    files: store.all_files.size,
    lookups: lookup_count,
    misses:,
    elapsed_seconds: elapsed
  ))
end

def run!(*command, **options)
  return if system(*command, **options)

  abort "Command failed: #{command.join(" ")}"
end

def run_timed(name, root, env, command, chdir: nil)
  FileUtils.rm_rf(root)
  FileUtils.mkdir_p(root)
  timing = File.join(root, "time.txt")
  puts "Running #{name}..."
  run!(env, "/usr/bin/time", "--format=%e %M", "--output=#{timing}", *command, **(chdir ? {chdir:} : {}))
  elapsed, max_rss = File.read(timing).split
  [elapsed.to_f, max_rss.to_i]
end

def synthetic_variant(name, rdoc, files, lookups)
  root = File.join(TMP_ROOT, "synthetic", "#{files}-#{lookups}", name)
  result_path = File.join(root, "result.json")
  env = {
    "PAGE_LOOKUP_BENCHMARK" => "synthetic",
    "PAGE_LOOKUP_VARIANT" => name,
    "PAGE_LOOKUP_RESULT" => result_path
  }
  command = [RbConfig.ruby, "-I#{File.join(rdoc, "lib")}", SCRIPT, files.to_s, lookups.to_s]
  wall_seconds, max_rss_kb = run_timed(name, root, env, command)
  JSON.parse(File.read(result_path)).merge("wall_seconds" => wall_seconds, "max_rss_kb" => max_rss_kb)
end

def prepare_source(options)
  return options.fetch(:source) if options[:source]

  cache = File.join(TMP_ROOT, "source")
  source = File.join(cache, "#{options.fetch(:gem)}-#{options.fetch(:version)}")
  return source if File.directory?(source)

  FileUtils.mkdir_p(cache)
  run!("gem", "fetch", options.fetch(:gem), "--version", options.fetch(:version), chdir: cache)
  package = Dir[File.join(cache, "#{options.fetch(:gem)}-#{options.fetch(:version)}*.gem")].first
  abort "Downloaded gem package was not found" unless package

  run!("gem", "unpack", package, "--target", cache)
  source
end

def source_inputs(source, limit)
  return ["."] unless limit

  files = Dir.glob(File.join(source, "**", "*.rb")).sort.first(limit)
  abort "No Ruby files found under #{source}" if files.empty?

  files.map { |file| file.delete_prefix("#{source}/") }
end

def corpus_variant(name, rdoc, source, inputs)
  root = File.join(TMP_ROOT, "google", name)
  profile_path = File.join(root, "profile.json")
  env = {
    "PAGE_LOOKUP_BENCHMARK" => "profile",
    "PAGE_LOOKUP_VARIANT" => name,
    "PAGE_LOOKUP_PROFILE" => profile_path,
    "RUBYOPT" => "-r#{SCRIPT}"
  }
  command = [
    RbConfig.ruby, "-I#{File.join(rdoc, "lib")}", File.join(rdoc, "exe/rdoc"),
    "--format=aliki", "--dry-run", "--quiet", "--exclude=#{EXCLUDE}", *inputs
  ]
  wall_seconds, max_rss_kb = run_timed(name, root, env, command, chdir: source)
  JSON.parse(File.read(profile_path)).merge(
    "name" => name,
    "wall_seconds" => wall_seconds,
    "max_rss_kb" => max_rss_kb
  )
end

def mib(kb) = format("%.1f MiB", kb / 1024.0)

def benchmark
  options = {
    rdoc: DEFAULT_RDOC,
    files: 13_333,
    lookups: [1_000, 5_000, 10_000],
    gem: "google-api-client",
    version: "0.53.0",
    limit: 100
  }
  parser = OptionParser.new do |opts|
    opts.banner = "Usage: #{File.basename($PROGRAM_NAME)} [options]"
    opts.on("--google", "Run the Google API client benchmark") { options[:google] = true }
    opts.on("--files FILES", Integer, "Synthetic store files (default: 13333)") { |value| options[:files] = value }
    opts.on("--lookups LIST", "Synthetic unresolved lookup counts") do |value|
      options[:lookups] = value.split(",").map { |count| Integer(count) }
    end
    opts.on("--source PATH", "Use this corpus instead of downloading a gem") do |value|
      options[:google] = true
      options[:source] = File.expand_path(value)
    end
    opts.on("--limit FILES", Integer, "Process the first FILES corpus files") do |value|
      options[:google] = true
      options[:limit] = value
    end
    opts.on("--all", "Process the complete corpus") do
      options[:google] = true
      options[:limit] = nil
    end
    opts.on("--rdoc PATH", "RDoc worktree (default: ../rdoc)") { |value| options[:rdoc] = File.expand_path(value) }
    opts.on("-h", "--help", "Show this help") { puts opts; exit }
  end
  parser.parse!

  abort "/usr/bin/time is required" unless File.executable?("/usr/bin/time")
  rdoc = options.fetch(:rdoc)
  abort "RDoc worktree does not exist: #{rdoc}" unless File.file?(File.join(rdoc, "lib/rdoc.rb"))

  if options[:google]
    source = prepare_source(options)
    abort "Source directory does not exist: #{source}" unless File.directory?(source)
    inputs = source_inputs(source, options[:limit])
    baseline = corpus_variant("baseline", rdoc, source, inputs)
    optimized = corpus_variant("optimized", rdoc, source, inputs)
    comparable = baseline.values_at("files", "calls", "hits", "misses") == optimized.values_at("files", "calls", "hits", "misses")
    abort "Page lookup profiles differ" unless comparable

    puts
    puts "Input Ruby files: #{inputs == ["."] ? "all" : inputs.size}"
    puts "Parsed store files: #{baseline.fetch("files")}"
    puts "Page lookups: #{baseline.fetch("calls")} (#{baseline.fetch("hits")} hits, #{baseline.fetch("misses")} misses)"
    puts "| Variant | Store#page time | Wall time | Peak RSS |"
    puts "| --- | ---: | ---: | ---: |"
    [baseline, optimized].each do |result|
      puts "| #{result.fetch("name")} | #{format("%.3fs", result.fetch("elapsed_seconds"))} | #{format("%.2fs", result.fetch("wall_seconds"))} | #{mib(result.fetch("max_rss_kb"))} |"
    end
    puts
    puts "Store#page speedup: #{format("%.2fx", baseline.fetch("elapsed_seconds") / optimized.fetch("elapsed_seconds"))}"
    puts "Overall speedup: #{format("%.2fx", baseline.fetch("wall_seconds") / optimized.fetch("wall_seconds"))}"
    puts "Profiles equivalent: yes"
  else
    results = options.fetch(:lookups).map do |lookups|
      puts "Benchmarking #{options.fetch(:files)} files and #{lookups} unresolved lookups..."
      baseline = synthetic_variant("baseline", rdoc, options.fetch(:files), lookups)
      optimized = synthetic_variant("optimized", rdoc, options.fetch(:files), lookups)
      comparable = baseline.values_at("files", "lookups", "misses") == optimized.values_at("files", "lookups", "misses")
      abort "Results differ for #{lookups} lookups" unless comparable
      [lookups, baseline, optimized]
    end

    puts
    puts "| Files | Unresolved lookups | Baseline | Optimized | Speedup | Baseline peak | Optimized peak |"
    puts "| ---: | ---: | ---: | ---: | ---: | ---: | ---: |"
    results.each do |lookups, baseline, optimized|
      puts "| #{options.fetch(:files)} | #{lookups} | #{format("%.3fs", baseline.fetch("elapsed_seconds"))} | #{format("%.3fs", optimized.fetch("elapsed_seconds"))} | #{format("%.1fx", baseline.fetch("elapsed_seconds") / optimized.fetch("elapsed_seconds"))} | #{mib(baseline.fetch("max_rss_kb"))} | #{mib(optimized.fetch("max_rss_kb"))} |"
    end
    puts
    puts "Results equivalent: yes"
  end
end

case ENV["PAGE_LOOKUP_BENCHMARK"]
when "synthetic"
  install_variant
  synthetic_child
when "profile"
  install_variant(profile: true)
else
  benchmark
end