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.
judoscale · Aug 9, 2023 · 1f2e1f5 · 1f2e1f5
1 parent 019a9cc
commit 1f2e1f5
Show file tree

Hide file tree

Showing 18 changed files with 132 additions and 30 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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