fix: Don't start the reporter in build processes (#180)

* 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.

* Fix tests

* Require "set" for older Ruby versions

.standard.yml

@@ -1,2 +1,4 @@
+ruby_version: 3.0
+
 ignore:
   - "sample-apps/**/*"

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,

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"

judoscale-rails/Gemfile-rails-6-1

@@ -5,4 +5,5 @@ gemspec name: "judoscale-rails"
 gem "judoscale-ruby", path: "../judoscale-ruby"
 gem "railties", "~> 6.1"
 gem "minitest"
+gem "minitest-stub-const"
 gem "rake"

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

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

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

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,

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"

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
 

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)

judoscale-ruby/lib/judoscale/job_metrics_collector.rb

@@ -9,7 +9,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

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

judoscale-ruby/lib/judoscale/reporter.rb

@@ -20,23 +20,25 @@ def start!(config, adapters)
       @pid = Process.pid
 
       if !config.api_base_url
-        logger.debug "Reporter not started: JUDOSCALE_URL is not set"
+        logger.debug "Set api_base_url to enable metrics reporting"
         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)
 

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

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

judoscale-ruby/test/reporter_test.rb

@@ -119,22 +119,22 @@ def stub_reporter_loop
           Reporter.instance.start!(Config.instance, Judoscale.adapters)
         }
 
-        _(log_string).must_include "Reporter not started: JUDOSCALE_URL is not set"
+        _(log_string).must_include "Set api_base_url to enable metrics reporting"
       end
 
       it "does not run the reporter thread when there are no metrics collectors" do
         Reporter.instance.stub(:run_loop, ->(*) { raise "SHOULD NOT BE CALLED" }) {
           Reporter.instance.start!(Config.instance, Judoscale.adapters.select { |a| a.metrics_collector.nil? })
         }
 
-        _(log_string).must_include "Reporter not started: no metrics need to be collected in this process"
+        _(log_string).must_include "No metrics need to be collected"
       end
 
       it "logs when the reporter starts successfully" do
         stub_request(:post, "http://example.com/api/test-token/v3/reports")
         run_reporter_start_thread
 
-        _(log_string).must_include "Reporter starting, will report every 10 seconds or so. Adapters: [judoscale-ruby, test_web, test_job]"
+        _(log_string).must_include "Reporter starting, will report every ~10 seconds (adapters: judoscale-ruby, test_web, test_job)"
       end
 
       it "logs only enabled adapters" do
@@ -143,7 +143,7 @@ def stub_reporter_loop
         stub_request(:post, "http://example.com/api/test-token/v3/reports")
         run_reporter_start_thread
 
-        _(log_string).must_include "Reporter starting, will report every 10 seconds or so. Adapters: [judoscale-ruby, test_web]"
+        _(log_string).must_include "Reporter starting, will report every ~10 seconds (adapters: judoscale-ruby, test_web)"
       end
     end
 

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?

judoscale-sidekiq/Gemfile-sidekiq-5

@@ -5,4 +5,5 @@ gemspec name: "judoscale-sidekiq"
 gem "judoscale-ruby", path: "../judoscale-ruby"
 gem "sidekiq", "~> 5.0"
 gem "minitest"
+gem "minitest-stub-const"
 gem "rake"

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