Fix console detection and improve rake detection

judoscale-delayed_job/lib/judoscale/delayed_job.rb

@@ -6,8 +6,6 @@
 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/lib/judoscale/rails/railtie.rb

@@ -13,27 +13,36 @@ class Railtie < ::Rails::Railtie
       include Judoscale::Logger
 
       def in_rails_console?
-        defined?(::Rails::Console)
+        # This is gross, but we can't find a more reliable way to detect if we're in a Rails console.
+        caller.any? { |call| call.include?("console_command.rb") }
       end
 
-      def in_rake_task?
-        defined?(::Rake) && Rake.respond_to?(:application) && Rake.application.top_level_tasks.any?
+      def in_rake_task?(task_regex)
+        top_level_tasks = defined?(::Rake) && Rake.respond_to?(:application) && Rake.application.top_level_tasks || []
+        top_level_tasks.any? { |task| task_regex.match?(task) }
+      end
+
+      def judoscale_config
+        # Disambiguate from Judoscale::Rails::Config
+        ::Judoscale::Config.instance
       end
 
       initializer "Judoscale.logger" do |app|
-        ::Judoscale::Config.instance.logger = ::Rails.logger
+        judoscale_config.logger = ::Rails.logger
       end
 
       initializer "Judoscale.request_middleware" do |app|
-        if !in_rails_console? && !in_rake_task?
-          logger.debug "Preparing request middleware"
-          app.middleware.insert_before Rack::Runtime, RequestMiddleware
-        end
+        app.middleware.insert_before Rack::Runtime, RequestMiddleware
       end
 
       config.after_initialize do
-        # Don't suppress this in Rake tasks since some job adapters use Rake tasks to run jobs.
-        Reporter.start if !in_rails_console? && ::Judoscale::Config.instance.start_reporter_after_initialize
+        if in_rails_console?
+          logger.debug "No reporting since we're in a Rails console"
+        elsif in_rake_task?(/assets:|db:/)
+          logger.debug "No reporting since we're in a build process"
+        elsif judoscale_config.start_reporter_after_initialize
+          Reporter.start
+        end
       end
     end
   end

judoscale-rails/test/railtie_test.rb

@@ -13,13 +13,6 @@ 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,8 +7,6 @@
 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/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, :allow_rake_tasks
+      :max_request_size_bytes, :logger, :log_tag, :current_runtime_container
     attr_reader :log_level
 
     def initialize
@@ -77,7 +77,6 @@ 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/metric.rb

@@ -5,7 +5,7 @@ class Metric < Struct.new(:identifier, :value, :time, :queue_name)
     # No queue_name is assumed to be a web request metric
     # Metrics: qt = queue time (default), qd = queue depth, busy
     def initialize(identifier, value, time, queue_name = nil)
-      super identifier, value.to_i, time.utc, queue_name
+      super(identifier, value.to_i, time.utc, queue_name)
     end
   end
 end

judoscale-ruby/lib/judoscale/metrics_collector.rb

@@ -3,23 +3,11 @@
 module Judoscale
   class MetricsCollector
     def self.collect?(config)
-      in_rake_task = defined?(::Rake) && Rake.respond_to?(:application) && Rake.application.top_level_tasks.any?
-      in_generator = defined?(::Rails::Command::GenerateCommand)
-
-      !in_generator && (!in_rake_task || in_whitelisted_rake_tasks?(config.allow_rake_tasks))
+      true
     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/test/job_metrics_collector_test.rb

@@ -1,38 +1,12 @@
 # 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
-
-        Object.stub_const :Rake, RakeMock.new([]) 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(["foo"]) 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/support/config_helpers.rb

@@ -35,13 +35,13 @@ def swap_config(config_instance, options, example)
 
     options.each do |key, val|
       original_config[key] = config_instance.public_send(key)
-      config_instance.public_send "#{key}=", val
+      config_instance.public_send :"#{key}=", val
     end
 
     example.call
   ensure
     options.each do |key, val|
-      config_instance.public_send "#{key}=", original_config[key]
+      config_instance.public_send :"#{key}=", original_config[key]
     end
   end
 end

judoscale-ruby/test/web_metrics_collector_test.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "test_helper"
-require "rake_mock"
 require "minitest/stub_const"
 require "judoscale/web_metrics_collector"
 require "judoscale/config"
@@ -14,39 +13,6 @@ class GenerateCommand; end
 
 module Judoscale
   describe WebMetricsCollector 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
-
-        Object.stub_const :Rake, RakeMock.new([]) do
-          _(WebMetricsCollector.collect?(Config.instance)).must_equal true
-        end
-      end
-
-      it "returns false when running in a generator" do
-        Object.stub_const :Rails, RailsMock do
-          _(WebMetricsCollector.collect?(Config.instance)).must_equal false
-        end
-      end
-
-      it "returns false when running in a rake task" do
-        Object.stub_const :Rake, RakeMock.new(["foo"]) 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 }
 

sample-apps/rails-sample/Gemfile.lock

@@ -1,14 +1,14 @@
 PATH
   remote: ../../judoscale-rails
   specs:
-    judoscale-rails (1.5.0)
-      judoscale-ruby (= 1.5.0)
+    judoscale-rails (1.5.2)
+      judoscale-ruby (= 1.5.2)
       railties
 
 PATH
   remote: ../../judoscale-ruby
   specs:
-    judoscale-ruby (1.5.0)
+    judoscale-ruby (1.5.2)
 
 GEM
   remote: https://rubygems.org/