Infer interruption handler from a job's queue adapter

...and allow Iteration to be used with multiple job backends simultaneously

Removed test_mark_job_worker_as_interrupted since it was testing stubs

Co-authored-by: Justin Morris <[email protected]>

CHANGELOG.md

@@ -7,6 +7,7 @@
 - [310](https://github.com/Shopify/job-iteration/pull/310) - Support nested iteration
 - [338](https://github.com/Shopify/job-iteration/pull/338) - All logs are now `ActiveSupport::Notifications` events and logged using `ActiveSupport::LogSubscriber` to allow customization. Events now always include the `cursor_position` tag.
 - [341](https://github.com/Shopify/job-iteration/pull/341) - Add `JobIteration.default_retry_backoff`, which sets a default delay when jobs are re-enqueued after being interrupted. Defaults to `nil`, meaning no delay, which matches the current behaviour.
+- [367](https://github.com/Shopify/job-iteration/pull/367) - Interruption adapter is now inferred from the job's queue adapter instead of `JobIteration.interruption_adapter` (which has been removed). This allows using Iteration with multiple job backends simultaneously.
 
 ## v1.3.6 (Mar 9, 2022)
 

lib/job-iteration.rb

@@ -2,15 +2,12 @@
 
 require "active_job"
 require_relative "./job-iteration/version"
+require_relative "./job-iteration/integrations"
 require_relative "./job-iteration/enumerator_builder"
 require_relative "./job-iteration/iteration"
 require_relative "./job-iteration/log_subscriber"
 
 module JobIteration
-  IntegrationLoadError = Class.new(StandardError)
-
-  INTEGRATIONS = [:resque, :sidekiq]
-
   extend self
 
   attr_accessor :logger
@@ -45,10 +42,12 @@ module JobIteration
   # where the throttle backoff value will take precedence over this setting.
   attr_accessor :default_retry_backoff
 
-  # Used internally for hooking into job processing frameworks like Sidekiq and Resque.
-  attr_accessor :interruption_adapter
+  # Set if you want to use your own method of loading interruption handlers
+  # of the default way based on a job's queue adapter name.
 
-  self.interruption_adapter = -> { false }
+  attr_accessor :interruption_adapter_loader
+
+  self.interruption_adapter_loader = JobIteration::Integrations
 
   # Set if you want to use your own enumerator builder instead of default EnumeratorBuilder.
   # @example
@@ -61,29 +60,4 @@ module JobIteration
   attr_accessor :enumerator_builder
 
   self.enumerator_builder = JobIteration::EnumeratorBuilder
-
-  def load_integrations
-    loaded = nil
-    INTEGRATIONS.each do |integration|
-      load_integration(integration)
-      if loaded
-        raise IntegrationLoadError,
-          "#{loaded} integration has already been loaded, but #{integration} is also available. " \
-            "Iteration will only work with one integration."
-      end
-      loaded = integration
-    rescue LoadError
-    end
-  end
-
-  def load_integration(integration)
-    unless INTEGRATIONS.include?(integration)
-      raise IntegrationLoadError,
-        "#{integration} integration is not supported. Available integrations: #{INTEGRATIONS.join(", ")}"
-    end
-
-    require_relative "./job-iteration/integrations/#{integration}"
-  end
 end
-
-JobIteration.load_integrations unless ENV["ITERATION_DISABLE_AUTOCONFIGURE"]

lib/job-iteration/integrations.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module JobIteration
+  module Integrations # @private
+    LoadError = Class.new(StandardError)
+
+    autoload :Sidekiq, "job-iteration/integrations/sidekiq"
+    autoload :Resque, "job-iteration/integrations/resque"
+
+    extend self
+
+    def load(job)
+      camelized_name = job.class.queue_adapter_name.camelize
+      object = const_get(camelized_name)
+
+      if object.respond_to?(:call)
+        object
+      else
+        raise LoadError, "Integration for '#{camelized_name}' must respond to call"
+      end
+    rescue NameError
+      raise LoadError, "Could not find integration for '#{camelized_name}'"
+    end
+  end
+end

lib/job-iteration/integrations/resque.rb

@@ -4,21 +4,27 @@
 
 module JobIteration
   module Integrations
-    module ResqueIterationExtension # @private
-      def initialize(*) # @private
-        $resque_worker = self
-        super
+    module Resque
+      module IterationExtension
+        def initialize(*)
+          $resque_worker = self
+          super
+        end
       end
-    end
 
-    # @private
-    module ::Resque
-      class Worker
-        # The patch is required in order to call shutdown? on a Resque::Worker instance
-        prepend(ResqueIterationExtension)
+      # @private
+      module ::Resque
+        class Worker
+          # The patch is required in order to call shutdown? on a Resque::Worker instance
+          prepend(IterationExtension)
+        end
       end
-    end
 
-    JobIteration.interruption_adapter = -> { $resque_worker.try!(:shutdown?) }
+      extend self
+
+      def call
+        $resque_worker.try!(:shutdown?)
+      end
+    end
   end
 end

lib/job-iteration/integrations/sidekiq.rb

@@ -3,12 +3,16 @@
 require "sidekiq"
 
 module JobIteration
-  module Integrations # @private
-    JobIteration.interruption_adapter = -> do
-      if defined?(Sidekiq::CLI) && Sidekiq::CLI.instance
-        Sidekiq::CLI.instance.launcher.stopping?
-      else
-        false
+  module Integrations
+    module Sidekiq
+      extend self
+
+      def call
+        if defined?(::Sidekiq::CLI) && (instance = ::Sidekiq::CLI.instance)
+          instance.launcher.stopping?
+        else
+          false
+        end
       end
     end
   end

lib/job-iteration/iteration.rb

@@ -134,6 +134,10 @@ def retry_job(*, **)
 
     private
 
+    def interruption_adapter
+      @interruption_adapter ||= JobIteration.interruption_adapter_loader.load(self)
+    end
+
     def enumerator_builder
       JobIteration.enumerator_builder.new(self)
     end
@@ -293,7 +297,7 @@ def job_should_exit?
         return true
       end
 
-      JobIteration.interruption_adapter.call || (defined?(super) && super)
+      interruption_adapter.call || (defined?(super) && super)
     end
 
     def handle_completed(completed)

lib/job-iteration/test_helper.rb

@@ -23,7 +23,7 @@ def call
     #     MyJob.perform_now
     #   end
     def iterate_exact_times(n_times)
-      JobIteration.stubs(:interruption_adapter).returns(StoppingSupervisor.new(n_times.size))
+      JobIteration::Integrations.stubs(:load).returns(StoppingSupervisor.new(n_times.size))
     end
 
     # Stubs interruption adapter to interrupt the job after every sing iteration.
@@ -47,7 +47,7 @@ def mark_job_worker_as_interrupted
     def stub_shutdown_adapter_to_return(value)
       adapter = mock
       adapter.stubs(:call).returns(value)
-      JobIteration.stubs(:interruption_adapter).returns(adapter)
+      JobIteration::Integrations.stubs(:load).returns(adapter)
     end
   end
 end

test/integration/integrations_test.rb

@@ -1,55 +1,84 @@
 # frozen_string_literal: true
 
 require "test_helper"
-require "open3"
-
-class IntegrationsTest < ActiveSupport::TestCase
-  test "will prevent loading two integrations" do
-    with_env("ITERATION_DISABLE_AUTOCONFIGURE", nil) do
-      rubby = <<~RUBBY
-        require 'bundler/setup'
-        require 'job-iteration'
-      RUBBY
-      _stdout, stderr, status = run_ruby(rubby)
-
-      assert_equal false, status.success?
-      assert_match(/resque integration has already been loaded, but sidekiq is also available/, stderr)
+
+class IntegrationsTest < IterationUnitTest
+  class IterationJob < ActiveJob::Base
+    include JobIteration::Iteration
+
+    def build_enumerator(cursor:)
+      enumerator_builder.build_once_enumerator(cursor: cursor)
+    end
+
+    def each_iteration(*)
+    end
+  end
+
+  class ResqueJob < IterationJob
+    self.queue_adapter = :resque
+  end
+
+  class SidekiqJob < IterationJob
+    self.queue_adapter = :sidekiq
+  end
+
+  class StubOneAdapter
+    class << self
+      def enqueue(*); end
+      def enqueue_at(*); end
     end
   end
 
-  test "successfully loads one (resque) integration" do
-    with_env("ITERATION_DISABLE_AUTOCONFIGURE", nil) do
-      rubby = <<~RUBBY
-        require 'bundler/setup'
-        # Remove sidekiq, only resque will be left
-        $LOAD_PATH.delete_if { |p| p =~ /sidekiq/ }
-        require 'job-iteration'
-      RUBBY
-      _stdout, _stderr, status = run_ruby(rubby)
+  class StubOneJob < IterationJob
+    self.queue_adapter = StubOneAdapter
+  end
 
-      assert_equal true, status.success?
+  module ::ActiveJob
+    module QueueAdapters
+      class StubTwoAdapter
+        def enqueue(*); end
+        def enqueue_at(*); end
+      end
     end
   end
 
-  private
+  class StubTwoJob < IterationJob
+    self.queue_adapter = :stub_two
+  end
+
+  test "loads multiple integrations" do
+    resque_job = ResqueJob.new.serialize
+    ActiveJob::Base.execute(resque_job)
+
+    sidekiq_job = SidekiqJob.new.serialize
+    ActiveJob::Base.execute(sidekiq_job)
+  end
 
-  def run_ruby(body)
-    stdout, stderr, status = nil
-    Tempfile.open do |f|
-      f.write(body)
-      f.close
+  test "default interruption adapter loader raises for unknown Active Job queue adapter names" do
+    error1 = assert_raises(JobIteration::Integrations::LoadError) do
+      JobIteration::Integrations.load(StubOneJob.new)
+    end
+    assert_equal("Integration for 'Class' must respond to call", error1.message)
 
-      command = "ruby #{f.path}"
-      stdout, stderr, status = Open3.capture3(command)
+    error2 = assert_raises(JobIteration::Integrations::LoadError) do
+      JobIteration::Integrations.load(StubTwoJob.new)
     end
-    [stdout, stderr, status]
+    assert_equal("Could not find integration for 'StubTwo'", error2.message)
   end
 
-  def with_env(variable, value)
-    original = ENV[variable]
-    ENV[variable] = value
-    yield
+  test "override default interruption adapter loader" do
+    old_loader = JobIteration.interruption_adapter_loader
+    JobIteration.interruption_adapter_loader = Module.new do
+      module_function
+
+      def load(_job)
+        -> { false }
+      end
+    end
+
+    stub_job = StubOneJob.new.serialize
+    ActiveJob::Base.execute(stub_job)
   ensure
-    ENV[variable] = original
+    JobIteration.interruption_adapter_loader = old_loader
   end
 end

test/support/sidekiq/init.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "job-iteration"
-require "job-iteration/integrations/sidekiq"
 
 require "active_job"
 require "i18n"

test/test_helper.rb

@@ -3,8 +3,6 @@
 $LOAD_PATH.unshift(File.expand_path("../../lib", __FILE__))
 require "minitest/autorun"
 
-ENV["ITERATION_DISABLE_AUTOCONFIGURE"] = "true"
-
 require "job-iteration"
 require "job-iteration/test_helper"
 
@@ -39,6 +37,18 @@ def enqueue_at(job, timestamp)
   end
 end
 
+module JobIteration
+  module Integrations
+    module IterationTest
+      extend self
+
+      def call
+        false
+      end
+    end
+  end
+end
+
 ActiveJob::Base.queue_adapter = :iteration_test
 
 class Product < ActiveRecord::Base

test/unit/active_job_iteration_test.rb

@@ -724,16 +724,6 @@ def test_on_shutdown_called_before_reenqueue
       assert_jobs_in_queue(1)
     end
 
-    def test_mark_job_worker_as_interrupted
-      mark_job_worker_as_interrupted
-
-      assert_equal(true, JobIteration.interruption_adapter.call)
-
-      continue_iterating
-
-      assert_equal(false, JobIteration.interruption_adapter.call)
-    end
-
     def test_reenqueue_self
       iterate_exact_times(2.times)