From 0da79561c32e41e7a51ba10622ff6d29cdd51dc4 Mon Sep 17 00:00:00 2001
From: Alan Guo Xiang Tan <gxtan1990@gmail.com>
Date: Fri, 30 Dec 2022 07:25:11 +0800
Subject: [PATCH] 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            |  2 +
 lib/tasks/profile.rake         | 24 ++++++++++
 script/bench.rb                | 85 ++++++++++++++++++++++++++--------
 script/profile_db_generator.rb | 75 +++++++++++++++++-------------
 4 files changed, 136 insertions(+), 50 deletions(-)
 create mode 100644 lib/tasks/profile.rake

diff --git a/lib/rate_limiter.rb b/lib/rate_limiter.rb
index 45f5fc52c92..1ec2d455add 100644
--- a/lib/rate_limiter.rb
+++ b/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
diff --git a/lib/tasks/profile.rake b/lib/tasks/profile.rake
new file mode 100644
index 00000000000..1b44ff3511c
--- /dev/null
+++ b/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
diff --git a/script/bench.rb b/script/bench.rb
index 86f6c6c73cf..e114cdd2cbf 100644
--- a/script/bench.rb
+++ b/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)
 
diff --git a/script/profile_db_generator.rb b/script/profile_db_generator.rb
index 9d4acf71fbf..c49ef6f6278 100644
--- a/script/profile_db_generator.rb
+++ b/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!