DEV: Improve/Fix script/bench.rb (#19646)

1. Fix bug where we were not waiting for all unicorn workers to start up
before running benchmarks.

2. Fix a bug where headers were not used when benchmarking. Admin
benchmarks were basically running as anon user.

3. Disable rate limits when in profile env. We're pretty much going to
hit the rate limit every time as a normal user.

4. Benchmark against topic with a fixed posts count of 100. Previously profiling script was just randomly creating posts
and we would benchmark against a topic with a fixed posts count of 30.
Sometimes, the script fails because no topics with a posts count of 30
exists.

5. Benchmarks are not run against a normal user on top of anon and
admin.

6. Add script option to select tests that should be run.

lib/rate_limiter.rb

@@ -17,6 +17,8 @@ class RateLimiter
     @disabled = false
   end
 
+  disable if Rails.env.profile?
+
   # We don't observe rate limits in test mode
   def self.disabled?
     @disabled

lib/tasks/profile.rake

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+desc "generate a user api key for given user in profiling environment"
+task "user_api_key:create", [:username] => :environment do |task, args|
+  raise "user_api_key:create rake task is only meant for the profiling env" if ENV["RAILS_ENV"] != "profile"
+  raise "Supply a username for the key" if !args[:username]
+
+  user = User.find_by_username(args[:username])
+
+  raise "'#{args[:username]}' is not a valid username" if !user
+
+  application_name = 'perf test application'
+
+  user_api_key = UserApiKey.where(application_name: application_name).destroy_all
+
+  user_api_key = UserApiKey.create!(
+    application_name: application_name,
+    client_id: '1234',
+    scopes: ['read'].map { |name| UserApiKeyScope.new(name: name) },
+    user_id: user.id
+  )
+
+  puts user_api_key.key
+end

script/bench.rb

@@ -5,6 +5,8 @@ require "csv"
 require "yaml"
 require "optparse"
 require "fileutils"
+require "net/http"
+require "uri"
 
 @include_env = false
 @result_file = nil
@@ -52,6 +54,10 @@ opts = OptionParser.new do |o|
   o.on("-s", "--skip-bundle-assets", "Skip bundling assets") do
     @skip_asset_bundle = true
   end
+
+  o.on("-t", "--tests [STRING]", "List of tests to run. Example: '--tests topic,categories')") do |i|
+    @tests = i.split(",")
+  end
 end
 opts.parse!
 
@@ -182,17 +188,23 @@ end
 puts "Populating Profile DB"
 run("bundle exec ruby script/profile_db_generator.rb")
 
-puts "Getting api key"
-api_key = `bundle exec rake api_key:create_master[bench]`.split("\n")[-1]
+puts "Getting admin api key"
+admin_api_key = `bundle exec rake api_key:create_master[bench]`.split("\n")[-1]
+raise "Failed to obtain a user API key" if admin_api_key.to_s.empty?
 
-def bench(path, name)
+puts "Getting user api key"
+user_api_key = `bundle exec rake user_api_key:create[user1]`.split("\n")[-1]
+raise "Failed to obtain a user API key" if user_api_key.to_s.empty?
+
+def bench(path, name, headers)
   puts "Running apache bench warmup"
   add = ""
   add = "-c #{@concurrency} " if @concurrency > 1
-  `ab #{add} -n 20 -l "http://127.0.0.1:#{@port}#{path}"`
+  header_string = headers&.map { |k, v| "-H \"#{k}:#{v}\"" }&.join(" ")
+  `ab #{add} #{header_string} -n 20 -l "http://127.0.0.1:#{@port}#{path}"`
 
   puts "Benchmarking #{name} @ #{path}"
-  `ab #{add} -n #{@iterations} -l -e tmp/ab.csv "http://127.0.0.1:#{@port}#{path}"`
+  `ab #{add} #{header_string} -n #{@iterations} -l -e tmp/ab.csv "http://127.0.0.1:#{@port}#{path}"`
 
   percentiles = Hash[*[50, 75, 90, 99].zip([]).flatten]
   CSV.foreach("tmp/ab.csv") do |percent, time|
@@ -214,7 +226,17 @@ begin
       ENV['UNICORN_PORT'] = @port.to_s
       ENV['UNICORN_WORKERS'] = @unicorn_workers.to_s
       FileUtils.mkdir_p(File.join('tmp', 'pids'))
-      spawn("bundle exec unicorn -c config/unicorn.conf.rb")
+      unicorn_pid = spawn("bundle exec unicorn -c config/unicorn.conf.rb")
+
+      while (unicorn_master_pid = `ps aux | grep "unicorn master" | grep -v "grep" | awk '{print $2}'`.strip.to_i) == 0
+        sleep 1
+      end
+
+      while `ps -f --ppid #{unicorn_master_pid} | grep worker | awk '{ print $2 }'`.split("\n").map(&:to_i).size != @unicorn_workers.to_i
+        sleep 1
+      end
+
+      unicorn_pid
     else
       spawn("bundle exec puma -p #{@port} -e production")
     end
@@ -224,8 +246,15 @@ begin
   end
 
   puts "Starting benchmark..."
-  headers = { 'Api-Key' => api_key,
-              'Api-Username' => "admin1" }
+
+  admin_headers = {
+    'Api-Key' => admin_api_key,
+    'Api-Username' => "admin1"
+  }
+
+  user_headers = {
+    'User-Api-Key' => user_api_key
+  }
 
   # asset precompilation is a dog, wget to force it
   run "curl -s -o /dev/null http://127.0.0.1:#{@port}/"
@@ -237,20 +266,38 @@ begin
 
   topic_url = redirect_response.match(/^location: .+(\/t\/i-am-a-topic-used-for-perf-tests\/.+)$/i)[1].strip
 
-  tests = [
+  all_tests = [
     ["categories", "/categories"],
     ["home", "/"],
-    ["topic", topic_url]
-    # ["user", "/u/admin1/activity"],
+    ["topic", topic_url],
+    ["topic.json", "#{topic_url}.json"],
+    ["user activity", "/u/admin1/activity"],
   ]
 
-  tests.concat(tests.map { |k, url| ["#{k}_admin", "#{url}", headers] })
+  @tests ||= %w{categories home topic}
 
-  tests.each do |_, path, headers_for_path|
-    header_string = headers_for_path&.map { |k, v| "-H \"#{k}: #{v}\"" }&.join(" ")
+  tests_to_run = all_tests.select do |test_name, path|
+    @tests.include?(test_name)
+  end
 
-    if `curl -s -I "http://127.0.0.1:#{@port}#{path}" #{header_string}` !~ /200 OK/
-      raise "#{path} returned non 200 response code"
+  tests_to_run.concat(
+    tests_to_run.map { |k, url| ["#{k} user", "#{url}", user_headers] },
+    tests_to_run.map { |k, url| ["#{k} admin", "#{url}", admin_headers] }
+  )
+
+  tests_to_run.each do |test_name, path, headers_for_path|
+    uri = URI.parse("http://127.0.0.1:#{@port}#{path}")
+    http = Net::HTTP.new(uri.host, uri.port)
+    request = Net::HTTP::Get.new(uri.request_uri)
+
+    headers_for_path&.each do |key, value|
+      request[key] = value
+    end
+
+    response = http.request(request)
+
+    if response.code != "200"
+      raise "#{test_name} #{path} returned non 200 response code"
     end
   end
 
@@ -265,8 +312,8 @@ begin
 
   results = {}
   @best_of.times do
-    tests.each do |name, url|
-      results[name] = best_of(bench(url, name), results[name])
+    tests_to_run.each do |name, url, headers|
+      results[name] = best_of(bench(url, name, headers), results[name])
     end
   end
 
@@ -303,7 +350,7 @@ begin
   mem = get_mem(pid)
 
   results = results.merge("timings" => @timings,
-                          "ruby-version" => "#{RUBY_VERSION}-p#{RUBY_PATCHLEVEL}",
+                          "ruby-version" => "#{RUBY_DESCRIPTION}",
                           "rss_kb" => mem["rss_kb"],
                           "pss_kb" => mem["pss_kb"]).merge(facts)
 

script/profile_db_generator.rb

@@ -43,16 +43,20 @@ def sentence
   sentence
 end
 
-def create_admin(seq)
-  User.new.tap { |admin|
-    admin.email = "admin@localhost#{seq}.fake"
-    admin.username = "admin#{seq}"
-    admin.password = "password12345abc"
-    admin.save!
-    admin.grant_admin!
-    admin.change_trust_level!(TrustLevel[4])
-    admin.activate
-  }
+def create_user(seq, admin: false, username: nil)
+  User.new.tap do |user|
+    user.email = "user@localhost#{seq}.fake"
+    user.username = username || "user#{seq}"
+    user.password = "password12345abc"
+    user.save!
+
+    if admin
+      user.grant_admin!
+      user.change_trust_level!(TrustLevel[4])
+    end
+
+    user.activate
+  end
 end
 
 require File.expand_path(File.dirname(__FILE__) + "/../config/environment")
@@ -64,20 +68,9 @@ unless Rails.env == "profile"
   exit
 end
 
-def ensure_perf_test_topic_has_right_title!
-  title = "I am a topic used for perf tests"
-  # in case we have an old run and picked the wrong topic
-  Topic.where(title: title).update_all(title: "Test topic #{SecureRandom.hex}")
-  t = Topic.where(archetype: :regular, posts_count: 30).order(id: :desc).first
-  t.title = title
-  t.save!
-end
-
 # by default, Discourse has a "system" and `discobot` account
 if User.count > 2
   puts "Only run this script against an empty DB"
-
-  ensure_perf_test_topic_has_right_title!
   exit
 end
 
@@ -90,38 +83,58 @@ rescue LoadError
   unbundled_require 'gabbler'
 end
 
-puts "Creating 100 users"
-users = 100.times.map do |i|
+number_of_users = 100
+puts "Creating #{number_of_users} users"
+number_of_users.times.map do |i|
   putc "."
-  create_admin(i)
+  create_user(i)
 end
 
+puts
+puts "Creating 1 admin user"
+admin_user = create_user(number_of_users + 1, admin: true, username: "admin1")
+
+users = User.human_users.all
+
 puts
 puts "Creating 10 categories"
 categories = 10.times.map do |i|
   putc "."
-  Category.create(name: "category#{i}", text_color: "ffffff", color: "000000", user: users.first)
+  Category.create(name: "category#{i}", text_color: "ffffff", color: "000000", user: admin_user)
 end
 
 puts
 puts "Creating 100 topics"
-
 topic_ids = 100.times.map do
-  post = PostCreator.create(users.sample, raw: sentence, title: sentence[0..50].strip, category:  categories.sample.id, skip_validations: true)
-
+  post = PostCreator.create(admin_user, raw: sentence, title: sentence[0..50].strip, category:  categories.sample.id, skip_validations: true)
   putc "."
   post.topic_id
 end
 
 puts
-puts "creating 2000 replies"
+puts "Creating 2000 replies"
 2000.times do
   putc "."
   PostCreator.create(users.sample, raw: sentence, topic_id: topic_ids.sample, skip_validations: true)
 end
 
+puts
+puts "creating perf test topic"
+first_post = PostCreator.create(
+  users.sample,
+  raw: sentence,
+  title: "I am a topic used for perf tests",
+  category: categories.sample.id,
+  skip_validations: true
+)
+
+puts
+puts "Creating 100 replies for perf test topic"
+100.times do
+  putc "."
+  PostCreator.create(users.sample, raw: sentence, topic_id: first_post.topic_id, skip_validations: true)
+end
+
 # no sidekiq so update some stuff
 Category.update_stats
 Jobs::PeriodicalUpdates.new.execute(nil)
-
-ensure_perf_test_topic_has_right_title!