From 1f2e1f5d2853813f630e6c1f235d2a4574df52f8 Mon Sep 17 00:00:00 2001
From: Adam McCrea <adam@adamlogic.com>
Date: Wed, 9 Aug 2023 12:51:13 -0400
Subject: [PATCH] Don't start the reporter in build processes

During build processes such as asset precompiling or database migrations, there's no need to report metrics. We've seen instances of these reports causing unexpected scaling.

There's no surefire way to know when we're in a build process, so we're checking for the presence of a Rake task. We need to allow for rake tasks that run jobs, though, and we do that through the `config.allow_rake_tasks` option. This is set automatically for DelayedJob and Resque.

We're also no longer including the middleware when running a Rails console. This wasn't causing problems, but it added noise to the console output.
---
 .../lib/judoscale/delayed_job.rb              |  2 ++
 judoscale-rails/Gemfile                       |  1 +
 judoscale-rails/Gemfile.lock                  |  2 ++
 .../lib/judoscale/rails/railtie.rb            | 21 +++++++++-----
 judoscale-rails/test/railtie_test.rb          |  8 ++++++
 judoscale-resque/lib/judoscale/resque.rb      |  2 ++
 judoscale-ruby/Gemfile                        |  1 +
 judoscale-ruby/Gemfile.lock                   |  2 ++
 judoscale-ruby/lib/judoscale/config.rb        |  3 +-
 .../lib/judoscale/job_metrics_collector.rb    |  3 +-
 .../lib/judoscale/metrics_collector.rb        | 13 ++++++++-
 judoscale-ruby/lib/judoscale/reporter.rb      |  8 ++++--
 .../test/job_metrics_collector_test.rb        | 23 +++++++++++++++
 judoscale-ruby/test/rake_mock.rb              | 12 ++++++++
 .../test/web_metrics_collector_test.rb        | 28 ++++++++++++++++++-
 sample-apps/delayed_job-sample/Gemfile.lock   | 11 ++++----
 sample-apps/resque-sample/Gemfile.lock        | 11 ++++----
 sample-apps/sidekiq-sample/Gemfile.lock       | 11 ++++----
 18 files changed, 132 insertions(+), 30 deletions(-)
 create mode 100644 judoscale-ruby/test/rake_mock.rb

diff --git a/judoscale-delayed_job/lib/judoscale/delayed_job.rb b/judoscale-delayed_job/lib/judoscale/delayed_job.rb
index aa3d2caf..85d55567 100644
--- a/judoscale-delayed_job/lib/judoscale/delayed_job.rb
+++ b/judoscale-delayed_job/lib/judoscale/delayed_job.rb
@@ -6,6 +6,8 @@
 require "judoscale/delayed_job/metrics_collector"
 require "delayed_job_active_record"
 
+Judoscale::Config.instance.allow_rake_tasks << /jobs:work/
+
 Judoscale.add_adapter :"judoscale-delayed_job",
   {
     adapter_version: Judoscale::DelayedJob::VERSION,
diff --git a/judoscale-rails/Gemfile b/judoscale-rails/Gemfile
index 499b8e99..98d4a18e 100644
--- a/judoscale-rails/Gemfile
+++ b/judoscale-rails/Gemfile
@@ -4,5 +4,6 @@ gemspec name: "judoscale-rails"
 
 gem "judoscale-ruby", path: "../judoscale-ruby"
 gem "minitest"
+gem "minitest-stub-const"
 gem "rake"
 gem "debug"
diff --git a/judoscale-rails/Gemfile.lock b/judoscale-rails/Gemfile.lock
index 70a68a30..9cf7fd0e 100644
--- a/judoscale-rails/Gemfile.lock
+++ b/judoscale-rails/Gemfile.lock
@@ -48,6 +48,7 @@ GEM
       nokogiri (>= 1.12.0)
     method_source (1.0.0)
     minitest (5.17.0)
+    minitest-stub-const (0.6)
     nokogiri (1.15.3-arm64-darwin)
       racc (~> 1.4)
     nokogiri (1.15.3-x86_64-darwin)
@@ -92,6 +93,7 @@ DEPENDENCIES
   judoscale-rails!
   judoscale-ruby!
   minitest
+  minitest-stub-const
   rake
 
 BUNDLED WITH
diff --git a/judoscale-rails/lib/judoscale/rails/railtie.rb b/judoscale-rails/lib/judoscale/rails/railtie.rb
index 1263308e..0d8c5fc8 100644
--- a/judoscale-rails/lib/judoscale/rails/railtie.rb
+++ b/judoscale-rails/lib/judoscale/rails/railtie.rb
@@ -12,21 +12,28 @@ module Rails
     class Railtie < ::Rails::Railtie
       include Judoscale::Logger
 
+      def in_rails_console?
+        defined?(::Rails::Console)
+      end
+
+      def in_rake_task?
+        defined?(::Rake) && Rake.respond_to?(:application)
+      end
+
       initializer "Judoscale.logger" do |app|
         Config.instance.logger = ::Rails.logger
       end
 
       initializer "Judoscale.request_middleware" do |app|
-        logger.debug "Preparing request middleware"
-        app.middleware.insert_before Rack::Runtime, RequestMiddleware
+        if !in_rails_console? && !in_rake_task?
+          logger.debug "Preparing request middleware"
+          app.middleware.insert_before Rack::Runtime, RequestMiddleware
+        end
       end
 
       config.after_initialize do
-        # Don't start the reporter in a Rails console.
-        # NOTE: This is untested because we initialize the Rails test app in test_helper.rb,
-        # so the reporter has already started before any of the tests run. You can manually
-        # test this by running `DYNO=web.1 rails c` in sample-apps/rails-sample.
-        Reporter.start unless defined?(::Rails::Console)
+        # Don't suppress this in Rake tasks since some job adapters use Rake tasks to run jobs.
+        Reporter.start unless in_rails_console?
       end
     end
   end
diff --git a/judoscale-rails/test/railtie_test.rb b/judoscale-rails/test/railtie_test.rb
index 26a8490b..566c65fd 100644
--- a/judoscale-rails/test/railtie_test.rb
+++ b/judoscale-rails/test/railtie_test.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "minitest/stub_const"
 
 module Judoscale
   describe Judoscale::Rails::Railtie do
@@ -12,6 +13,13 @@ module Judoscale
       _(::Rails.application.config.middleware).must_include Judoscale::RequestMiddleware
     end
 
+    # TODO: Fix this test. It fails because Rails initialization has already run in test_helper.
+    # it "skips the request middleware if running Rails console" do
+    #   ::Rails.stub_const :Console, Module.new do
+    #     _(::Rails.application.config.middleware).wont_include Judoscale::RequestMiddleware
+    #   end
+    # end
+
     it "boots up a reporter automatically after the app/config is initialized" do
       _(::Judoscale::Reporter.instance).must_be :started?
     end
diff --git a/judoscale-resque/lib/judoscale/resque.rb b/judoscale-resque/lib/judoscale/resque.rb
index 19a74f24..309d67e8 100644
--- a/judoscale-resque/lib/judoscale/resque.rb
+++ b/judoscale-resque/lib/judoscale/resque.rb
@@ -7,6 +7,8 @@
 require "resque"
 require "judoscale/resque/latency_extension"
 
+Judoscale::Config.instance.allow_rake_tasks << /resque:work/
+
 Judoscale.add_adapter :"judoscale-resque",
   {
     adapter_version: Judoscale::Resque::VERSION,
diff --git a/judoscale-ruby/Gemfile b/judoscale-ruby/Gemfile
index a7a673ef..a6826fb1 100644
--- a/judoscale-ruby/Gemfile
+++ b/judoscale-ruby/Gemfile
@@ -4,5 +4,6 @@ gemspec name: "judoscale-ruby"
 
 gem "rake", ">= 12.3.3"
 gem "minitest"
+gem "minitest-stub-const"
 gem "webmock"
 gem "debug"
diff --git a/judoscale-ruby/Gemfile.lock b/judoscale-ruby/Gemfile.lock
index 88a7a73e..6d1facf7 100644
--- a/judoscale-ruby/Gemfile.lock
+++ b/judoscale-ruby/Gemfile.lock
@@ -18,6 +18,7 @@ GEM
     irb (1.6.2)
       reline (>= 0.3.0)
     minitest (5.17.0)
+    minitest-stub-const (0.6)
     public_suffix (5.0.1)
     rake (13.0.6)
     reline (0.3.2)
@@ -39,6 +40,7 @@ DEPENDENCIES
   debug
   judoscale-ruby!
   minitest
+  minitest-stub-const
   rake (>= 12.3.3)
   webmock
 
diff --git a/judoscale-ruby/lib/judoscale/config.rb b/judoscale-ruby/lib/judoscale/config.rb
index 87f7ada9..727110a9 100644
--- a/judoscale-ruby/lib/judoscale/config.rb
+++ b/judoscale-ruby/lib/judoscale/config.rb
@@ -65,7 +65,7 @@ def self.expose_adapter_config(config_instance)
     end
 
     attr_accessor :api_base_url, :report_interval_seconds,
-      :max_request_size_bytes, :logger, :log_tag, :current_runtime_container
+      :max_request_size_bytes, :logger, :log_tag, :current_runtime_container, :allow_rake_tasks
     attr_reader :log_level
 
     def initialize
@@ -77,6 +77,7 @@ def reset
       @log_tag = "Judoscale"
       @max_request_size_bytes = 100_000 # ignore request payloads over 100k since they skew the queue times
       @report_interval_seconds = 10
+      @allow_rake_tasks = []
 
       self.log_level = ENV["JUDOSCALE_LOG_LEVEL"] || ENV["RAILS_AUTOSCALE_LOG_LEVEL"]
       @logger = ::Logger.new($stdout)
diff --git a/judoscale-ruby/lib/judoscale/job_metrics_collector.rb b/judoscale-ruby/lib/judoscale/job_metrics_collector.rb
index b7a92db0..87c614fb 100644
--- a/judoscale-ruby/lib/judoscale/job_metrics_collector.rb
+++ b/judoscale-ruby/lib/judoscale/job_metrics_collector.rb
@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "set"
 require "judoscale/metrics_collector"
 require "judoscale/logger"
 
@@ -9,7 +8,7 @@ class JobMetricsCollector < MetricsCollector
     include Judoscale::Logger
 
     def self.collect?(config)
-      !config.current_runtime_container.redundant_instance? && adapter_config.enabled
+      super && !config.current_runtime_container.redundant_instance? && adapter_config.enabled
     end
 
     def self.adapter_name
diff --git a/judoscale-ruby/lib/judoscale/metrics_collector.rb b/judoscale-ruby/lib/judoscale/metrics_collector.rb
index 326b02e3..de048453 100644
--- a/judoscale-ruby/lib/judoscale/metrics_collector.rb
+++ b/judoscale-ruby/lib/judoscale/metrics_collector.rb
@@ -3,11 +3,22 @@
 module Judoscale
   class MetricsCollector
     def self.collect?(config)
-      true
+      in_rake_task = defined?(::Rake) && Rake.respond_to?(:application)
+
+      !in_rake_task || in_whitelisted_rake_tasks?(config.allow_rake_tasks)
     end
 
     def collect
       []
     end
+
+    def self.in_whitelisted_rake_tasks?(allowed_rake_tasks)
+      # Get the tasks that were invoked from the command line.
+      tasks = Rake.application.top_level_tasks
+
+      allowed_rake_tasks.any? do |task_regex|
+        tasks.any? { |task| task =~ task_regex }
+      end
+    end
   end
 end
diff --git a/judoscale-ruby/lib/judoscale/reporter.rb b/judoscale-ruby/lib/judoscale/reporter.rb
index a78baf45..f82afcdf 100644
--- a/judoscale-ruby/lib/judoscale/reporter.rb
+++ b/judoscale-ruby/lib/judoscale/reporter.rb
@@ -24,19 +24,21 @@ def start!(config, adapters)
         return
       end
 
-      enabled_adapters = adapters.select { |adapter|
+      enabled_adapters, skipped_adapters = adapters.partition { |adapter|
+        # judoscale-ruby adapter does not have a metrics collector
         adapter.metrics_collector.nil? || adapter.metrics_collector.collect?(config)
       }
       metrics_collectors_classes = enabled_adapters.map(&:metrics_collector)
       metrics_collectors_classes.compact!
 
       if metrics_collectors_classes.empty?
-        logger.debug "Reporter not started: no metrics need to be collected in this process"
+        adapters_msg = skipped_adapters.map(&:identifier).join(", ")
+        logger.debug "No metrics need to be collected (adapters: #{adapters_msg})"
         return
       end
 
       adapters_msg = enabled_adapters.map(&:identifier).join(", ")
-      logger.info "Reporter starting, will report every #{config.report_interval_seconds} seconds or so. Adapters: [#{adapters_msg}]"
+      logger.info "Reporter starting, will report every ~#{config.report_interval_seconds} seconds (adapters: #{adapters_msg})"
 
       metrics_collectors = metrics_collectors_classes.map(&:new)
 
diff --git a/judoscale-ruby/test/job_metrics_collector_test.rb b/judoscale-ruby/test/job_metrics_collector_test.rb
index 913b3d66..855a2e3d 100644
--- a/judoscale-ruby/test/job_metrics_collector_test.rb
+++ b/judoscale-ruby/test/job_metrics_collector_test.rb
@@ -1,11 +1,34 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "rake_mock"
+require "minitest/stub_const"
 require "judoscale/job_metrics_collector"
 
 module Judoscale
   describe JobMetricsCollector do
     describe ".collect?" do
+      it "returns true when not running in a rake task" do
+        Object.stub_const :Rake, nil do
+          _(WebMetricsCollector.collect?(Config.instance)).must_equal true
+        end
+      end
+
+      it "returns false when running in a rake task" do
+        Object.stub_const :Rake, RakeMock.new do
+          _(WebMetricsCollector.collect?(Config.instance)).must_equal false
+        end
+      end
+
+      it "returns true when running in a whitelisted rake task" do
+        config = Config.instance
+        config.allow_rake_tasks << /foo/
+
+        Object.stub_const :Rake, RakeMock.new(["bar", "foo"]) do
+          _(WebMetricsCollector.collect?(config)).must_equal true
+        end
+      end
+
       it "collects only from the first container in the formation (if we know that), to avoid redundant collection from multiple containers when possible" do
         %w[
           web.1
diff --git a/judoscale-ruby/test/rake_mock.rb b/judoscale-ruby/test/rake_mock.rb
new file mode 100644
index 00000000..52ac4ebb
--- /dev/null
+++ b/judoscale-ruby/test/rake_mock.rb
@@ -0,0 +1,12 @@
+class RakeMock
+  def initialize(top_level_tasks = [])
+    @top_level_tasks = top_level_tasks
+  end
+
+  def application
+    Application.new(@top_level_tasks)
+  end
+
+  class Application < Struct.new(:top_level_tasks)
+  end
+end
diff --git a/judoscale-ruby/test/web_metrics_collector_test.rb b/judoscale-ruby/test/web_metrics_collector_test.rb
index ad8bd813..55504a58 100644
--- a/judoscale-ruby/test/web_metrics_collector_test.rb
+++ b/judoscale-ruby/test/web_metrics_collector_test.rb
@@ -1,13 +1,39 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "rake_mock"
+require "minitest/stub_const"
 require "judoscale/web_metrics_collector"
+require "judoscale/config"
 
 module Judoscale
   describe WebMetricsCollector do
-    let(:store) { MetricsStore.instance }
+    describe ".collect?" do
+      it "returns true when not running in a rake task" do
+        Object.stub_const :Rake, nil do
+          _(WebMetricsCollector.collect?(Config.instance)).must_equal true
+        end
+      end
+
+      it "returns false when running in a rake task" do
+        Object.stub_const :Rake, RakeMock.new do
+          _(WebMetricsCollector.collect?(Config.instance)).must_equal false
+        end
+      end
+
+      it "returns true when running in a whitelisted rake task" do
+        config = Config.instance
+        config.allow_rake_tasks << /foo/
+
+        Object.stub_const :Rake, RakeMock.new(["bar", "foo"]) do
+          _(WebMetricsCollector.collect?(config)).must_equal true
+        end
+      end
+    end
 
     describe "#collect" do
+      let(:store) { MetricsStore.instance }
+
       it "flushes the metrics previously collected from the store" do
         collector = WebMetricsCollector.new
         _(collector.collect).must_be :empty?
diff --git a/sample-apps/delayed_job-sample/Gemfile.lock b/sample-apps/delayed_job-sample/Gemfile.lock
index 4ba66291..f7c55443 100644
--- a/sample-apps/delayed_job-sample/Gemfile.lock
+++ b/sample-apps/delayed_job-sample/Gemfile.lock
@@ -1,21 +1,21 @@
 PATH
   remote: ../../judoscale-delayed_job
   specs:
-    judoscale-delayed_job (1.3.0)
+    judoscale-delayed_job (1.5.0)
       delayed_job_active_record (>= 4.0)
-      judoscale-ruby (= 1.3.0)
+      judoscale-ruby (= 1.5.0)
 
 PATH
   remote: ../../judoscale-rails
   specs:
-    judoscale-rails (1.3.0)
-      judoscale-ruby (= 1.3.0)
+    judoscale-rails (1.5.0)
+      judoscale-ruby (= 1.5.0)
       railties
 
 PATH
   remote: ../../judoscale-ruby
   specs:
-    judoscale-ruby (1.3.0)
+    judoscale-ruby (1.5.0)
 
 GEM
   remote: https://rubygems.org/
@@ -110,6 +110,7 @@ GEM
 PLATFORMS
   arm64-darwin-20
   arm64-darwin-21
+  arm64-darwin-22
   x86_64-darwin-21
   x86_64-linux
 
diff --git a/sample-apps/resque-sample/Gemfile.lock b/sample-apps/resque-sample/Gemfile.lock
index d6dfba57..150b38ca 100644
--- a/sample-apps/resque-sample/Gemfile.lock
+++ b/sample-apps/resque-sample/Gemfile.lock
@@ -1,21 +1,21 @@
 PATH
   remote: ../../judoscale-rails
   specs:
-    judoscale-rails (1.3.0)
-      judoscale-ruby (= 1.3.0)
+    judoscale-rails (1.5.0)
+      judoscale-ruby (= 1.5.0)
       railties
 
 PATH
   remote: ../../judoscale-resque
   specs:
-    judoscale-resque (1.3.0)
-      judoscale-ruby (= 1.3.0)
+    judoscale-resque (1.5.0)
+      judoscale-ruby (= 1.5.0)
       resque (>= 2.0)
 
 PATH
   remote: ../../judoscale-ruby
   specs:
-    judoscale-ruby (1.3.0)
+    judoscale-ruby (1.5.0)
 
 GEM
   remote: https://rubygems.org/
@@ -108,6 +108,7 @@ GEM
 PLATFORMS
   arm64-darwin-20
   arm64-darwin-21
+  arm64-darwin-22
   x86_64-darwin-21
   x86_64-linux
 
diff --git a/sample-apps/sidekiq-sample/Gemfile.lock b/sample-apps/sidekiq-sample/Gemfile.lock
index 3fc34d7f..64881ed4 100644
--- a/sample-apps/sidekiq-sample/Gemfile.lock
+++ b/sample-apps/sidekiq-sample/Gemfile.lock
@@ -1,20 +1,20 @@
 PATH
   remote: ../../judoscale-rails
   specs:
-    judoscale-rails (1.3.0)
-      judoscale-ruby (= 1.3.0)
+    judoscale-rails (1.5.0)
+      judoscale-ruby (= 1.5.0)
       railties
 
 PATH
   remote: ../../judoscale-ruby
   specs:
-    judoscale-ruby (1.3.0)
+    judoscale-ruby (1.5.0)
 
 PATH
   remote: ../../judoscale-sidekiq
   specs:
-    judoscale-sidekiq (1.3.0)
-      judoscale-ruby (= 1.3.0)
+    judoscale-sidekiq (1.5.0)
+      judoscale-ruby (= 1.5.0)
       sidekiq (>= 5.0)
 
 GEM
@@ -89,6 +89,7 @@ GEM
 PLATFORMS
   arm64-darwin-20
   arm64-darwin-21
+  arm64-darwin-22
   x86_64-darwin-21
   x86_64-linux