Active Job Continuation isolated steps (rails#55212)

Add an isolated option to steps. Defaults to false.

Isolated steps are always run in their own job execution. We do this by
interrupting the job before running a step if both:

1. It is not the first step we've run in this execution
2. Either this step or the previous one was marked as isolated.

The first to be run in an execution is always run inline so we don't
end up in an infinite loop of interruptions.

This allows you to execute a long running step separately which is
useful to ensure that progress is saved before it runs.

Author: djmb

activejob/lib/active_job/continuable.rb

@@ -33,7 +33,7 @@ def initialize(...)
     attr_accessor :continuation # :nodoc:
 
     # Start a new continuation step
-    def step(step_name, start: nil, &block)
+    def step(step_name, start: nil, isolated: false, &block)
       unless block_given?
         step_method = method(step_name)
 
@@ -46,7 +46,7 @@ def step(step_name, start: nil, &block)
         block = step_method.arity == 0 ? -> (_) { step_method.call } : step_method
       end
       checkpoint! if continuation.advanced?
-      continuation.step(step_name, start: start, &block)
+      continuation.step(step_name, start: start, isolated: isolated, &block)
     end
 
     def serialize # :nodoc:
@@ -60,7 +60,12 @@ def deserialize(job_data) # :nodoc:
     end
 
     def checkpoint! # :nodoc:
-      interrupt! if queue_adapter.stopping?
+      interrupt!(reason: :stopping) if queue_adapter.stopping?
+    end
+
+    def interrupt!(reason:) # :nodoc:
+      instrument :interrupt, reason: reason, **continuation.instrumentation
+      raise Continuation::Interrupt, "Interrupted #{continuation.description} (#{reason})"
     end
 
     private
@@ -91,11 +96,6 @@ def resume_job(exception) # :nodoc:
           raise Continuation::ResumeLimitError, "Job was resumed a maximum of #{max_resumptions} times"
         end
       end
-
-      def interrupt! # :nodoc:
-        instrument :interrupt, **continuation.instrumentation
-        raise Continuation::Interrupt, "Interrupted #{continuation.description}"
-      end
   end
 
   ActiveSupport.run_load_hooks(:active_job_continuable, Continuable)

activejob/lib/active_job/continuation.rb

@@ -150,6 +150,21 @@ module ActiveJob
   # - a list of the completed steps
   # - the current step and its cursor value (if one is in progress)
   #
+  # === Isolated Steps
+  #
+  # Steps run sequentially in a single job execution, unless the job is interrupted.
+  #
+  # You can specify that a step should always run in its own execution by passing the +isolated: true+ option.
+  #
+  # This is useful for long-running steps where it may not be possible to checkpoint within
+  # the job grace period - it ensures that progress is serialized back into the job data before
+  # the step starts.
+  #
+  #   step :quick_step1
+  #   step :slow_step, isolated: true
+  #   step :quick_step2
+  #   step :quick_step3
+  #
   # === Errors
   #
   # If a job raises an error and is not retried via Active Job, it will be passed back to the underlying
@@ -204,23 +219,24 @@ def initialize(job, serialized_progress) # :nodoc:
       @encountered = []
       @advanced = false
       @running_step = false
+      @isolating = false
     end
 
-    def step(name, start:, &block) # :nodoc:
+    def step(name, **options, &block) # :nodoc:
       validate_step!(name)
       encountered << name
 
       if completed?(name)
         skip_step(name)
       else
-        run_step(name, start: start, &block)
+        run_step(name, **options, &block)
       end
     end
 
     def to_h # :nodoc:
       {
         "completed" => completed.map(&:to_s),
-        "current" => current&.to_a
+        "current" => current&.to_a,
       }.compact
     end
 
@@ -255,6 +271,10 @@ def running_step?
         @running_step
       end
 
+      def isolating?
+        @isolating
+      end
+
       def completed?(name)
         completed.include?(name)
       end
@@ -267,7 +287,17 @@ def skip_step(name)
         instrument :step_skipped, step: name
       end
 
-      def run_step(name, start:, &block)
+      def run_step(name, start:, isolated:, &block)
+        @isolating ||= isolated
+
+        if isolating? && advanced?
+          job.interrupt!(reason: :isolating)
+        else
+          run_step_inline(name, start: start, &block)
+        end
+      end
+
+      def run_step_inline(name, start:, **options, &block)
         @running_step = true
         @current ||= new_step(name, start, resumed: false)
 

activejob/lib/active_job/log_subscriber.rb

@@ -141,7 +141,7 @@ def discard(event)
     def interrupt(event)
       job = event.payload[:job]
       info do
-        "Interrupted #{job.class} (Job ID: #{job.job_id}) #{event.payload[:description]}"
+        "Interrupted #{job.class} (Job ID: #{job.job_id}) #{event.payload[:description]} (#{event.payload[:reason]})"
       end
     end
     subscribe_log_level :interrupt, :info

activejob/test/cases/continuation_test.rb

@@ -258,7 +258,7 @@ def perform
       assert_no_match "Resuming", @logger.messages
       assert_match(/Step 'step_one' started/, @logger.messages)
       assert_match(/Step 'step_one' completed/, @logger.messages)
-      assert_match(/Interrupted ActiveJob::TestContinuation::LinearJob \(Job ID: [0-9a-f-]{36}\) after 'step_one'/, @logger.messages)
+      assert_match(/Interrupted ActiveJob::TestContinuation::LinearJob \(Job ID: [0-9a-f-]{36}\) after 'step_one' \(stopping\)/, @logger.messages)
     end
 
     perform_enqueued_jobs
@@ -278,7 +278,7 @@ def perform
       assert_no_match "Resuming", @logger.messages
       assert_match(/Step 'rename' started/, @logger.messages)
       assert_match(/Step 'rename' interrupted at cursor '433'/, @logger.messages)
-      assert_match(/Interrupted ActiveJob::TestContinuation::IteratingJob \(Job ID: [0-9a-f-]{36}\) at 'rename', cursor '433'/, @logger.messages)
+      assert_match(/Interrupted ActiveJob::TestContinuation::IteratingJob \(Job ID: [0-9a-f-]{36}\) at 'rename', cursor '433' \(stopping\)/, @logger.messages)
     end
 
     perform_enqueued_jobs
@@ -651,6 +651,73 @@ def perform(iterations)
     end
   end
 
+  class IsolatedStepsJob < ContinuableJob
+    cattr_accessor :items, default: []
+
+    def perform(*isolated)
+      step :step_one, isolated: isolated.include?(:step_one) do |step|
+        items << "step_one"
+      end
+      step :step_two, isolated: isolated.include?(:step_two) do |step|
+        items << "step_two"
+      end
+      step :step_three, isolated: isolated.include?(:step_three) do |step|
+        items << "step_three"
+      end
+      step :step_four, isolated: isolated.include?(:step_four) do |step|
+        items << "step_four"
+      end
+    end
+  end
+
+  test "runs isolated step separately" do
+    IsolatedStepsJob.items = []
+    IsolatedStepsJob.perform_later(:step_three)
+
+    assert_enqueued_jobs 1, only: IsolatedStepsJob do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one", "step_two" ], IsolatedStepsJob.items
+
+    assert_enqueued_jobs 1 do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one", "step_two", "step_three" ], IsolatedStepsJob.items
+
+    assert_enqueued_jobs 0 do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one", "step_two", "step_three", "step_four" ], IsolatedStepsJob.items
+    assert_match(/Interrupted ActiveJob::TestContinuation::IsolatedStepsJob \(Job ID: [0-9a-f-]{36}\) after 'step_two' \(isolating\)/, @logger.messages)
+    assert_match(/Interrupted ActiveJob::TestContinuation::IsolatedStepsJob \(Job ID: [0-9a-f-]{36}\) after 'step_three' \(isolating\)/, @logger.messages)
+  end
+
+  test "runs initial and final isolated steps separately" do
+    IsolatedStepsJob.items = []
+    IsolatedStepsJob.perform_later(:step_one, :step_four)
+
+    assert_enqueued_jobs 1, only: IsolatedStepsJob do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one" ], IsolatedStepsJob.items
+
+    assert_enqueued_jobs 1 do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one", "step_two", "step_three" ], IsolatedStepsJob.items
+
+    assert_enqueued_jobs 0 do
+      perform_enqueued_jobs
+    end
+
+    assert_equal [ "step_one", "step_two", "step_three", "step_four" ], IsolatedStepsJob.items
+  end
+
   private
     def capture_info_stdout(&block)
       ActiveJob::Base.logger.with(level: :info) do