Active Job Continuations continued (rails#55174)

Follow up to rails#55127 and rails#55151

- Instrument steps in a block so that runtime is recorded
  (matching perform/perform_started)
- Allow job resumption configuration - max_resumptions, resume_options,
  resume_errors_after_advancing
- Add Active Record Railtie to prevent checkpoints in database
  transactions
- Checkpoint before each step except the first, rather than after
  each step.
- Error if order of completed steps changes when re-running

Author: djmb

activejob/lib/active_job/continuable.rb

@@ -11,15 +11,26 @@ module ActiveJob
   module Continuable
     extend ActiveSupport::Concern
 
-    CONTINUATION_KEY = "continuation"
-
     included do
-      retry_on Continuation::Interrupt, attempts: :unlimited
-      retry_on Continuation::AfterAdvancingError, attempts: :unlimited
+      class_attribute :max_resumptions, instance_writer: false
+      class_attribute :resume_options, instance_writer: false, default: { wait: 5.seconds }
+      class_attribute :resume_errors_after_advancing, instance_writer: false, default: true
 
       around_perform :continue
+
+      def initialize(...)
+        super(...)
+        self.resumptions = 0
+        self.continuation = Continuation.new(self, {})
+      end
     end
 
+    # The number of times the job has been resumed.
+    attr_accessor :resumptions
+
+    attr_accessor :continuation # :nodoc:
+
+    # Start a new continuation step
     def step(step_name, start: nil, &block)
       unless block_given?
         step_method = method(step_name)
@@ -32,25 +43,58 @@ 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)
     end
 
-    def serialize
-      super.merge(CONTINUATION_KEY => continuation.to_h)
+    def serialize # :nodoc:
+      super.merge("continuation" => continuation.to_h, "resumptions" => resumptions)
     end
 
-    def deserialize(job_data)
+    def deserialize(job_data) # :nodoc:
       super
-      @continuation = Continuation.new(self, job_data.fetch(CONTINUATION_KEY, {}))
+      self.continuation = Continuation.new(self, job_data.fetch("continuation", {}))
+      self.resumptions = job_data.fetch("resumptions", 0)
+    end
+
+    def checkpoint! # :nodoc:
+      interrupt! if queue_adapter.stopping?
     end
 
     private
-      def continuation
-        @continuation ||= Continuation.new(self, {})
+      def continue(&block)
+        if continuation.started?
+          self.resumptions += 1
+          instrument :resume, **continuation.instrumentation
+        end
+
+        block.call
+      rescue Continuation::Interrupt => e
+        resume_job(e)
+      rescue Continuation::Error
+        raise
+      rescue StandardError => e
+        if resume_errors_after_advancing? && continuation.advanced?
+          resume_job(exception: e)
+        else
+          raise
+        end
       end
 
-      def continue(&block)
-        continuation.continue(&block)
+      def resume_job(exception) # :nodoc:
+        executions_for(exception)
+        if max_resumptions.nil? || resumptions < max_resumptions
+          retry_job(**self.resume_options)
+        else
+          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)
 end

activejob/lib/active_job/continuation.rb

@@ -134,8 +134,8 @@ module ActiveJob
   # +queue_adapter.stopping?+. If it returns true, the job will raise an
   # ActiveJob::Continuation::Interrupt exception.
   #
-  # There is an automatic checkpoint at the end of each step. Within a step one is
-  # created when calling +set!+, +advance!+ or +checkpoint!+.
+  # There is an automatic checkpoint before the start of each step except for the first for
+  # each job execution. Within a step one is created when calling +set!+, +advance!+ or +checkpoint!+.
   #
   # Jobs are not automatically interrupted when the queue adapter is marked as stopping - they
   # will continue to run either until the next checkpoint, or when the process is stopped.
@@ -158,49 +158,57 @@ module ActiveJob
   # To mitigate this, the job will be automatically retried if it raises an error after it has made progress.
   # Making progress is defined as having completed a step or advanced the cursor within the current step.
   #
+  # === Configuration
+  #
+  # Continuable jobs have several configuration options:
+  # * :max_resumptions</tt> - The maximum number of times a job can be resumed. Defaults to +nil+ which means
+  #   unlimited resumptions.
+  # * <tt>:resume_options</tt> - Options to pass to +retry_job+ when resuming the job.
+  #   Defaults to +{ wait: 5.seconds }+.
+  #   See +ActiveJob::Exceptions#retry_job+ for available options.
+  # * <tt>:resume_errors_after_advancing</tt> - Whether to resume errors after advancing the continuation.
+  #   Defaults to +true+.
   class Continuation
     extend ActiveSupport::Autoload
 
     autoload :Step
+    autoload :Validation
 
     # Raised when a job is interrupted, allowing Active Job to requeue it.
     # This inherits from +Exception+ rather than +StandardError+, so it's not
     # caught by normal exception handling.
     class Interrupt < Exception; end
 
-    # Base error class for all Continuation errors.
+    # Base class for all Continuation errors.
     class Error < StandardError; end
 
     # Raised when a step is invalid.
     class InvalidStepError < Error; end
 
+    # Raised when there is an error with a checkpoint, such as open database transactions.
+    class CheckpointError < Error; end
+
     # Raised when attempting to advance a cursor that doesn't implement `succ`.
     class UnadvanceableCursorError < Error; end
 
-    # Raised when an error occurs after a job has made progress.
-    #
-    # The job will be automatically retried to ensure that the progress is serialized
-    # in the retried job.
-    class AfterAdvancingError < Error; end
+    # Raised when a job has reached its limit of the number of resumes.
+    # The limit is defined by the +max_resumes+ class attribute.
+    class ResumeLimitError < Error; end
 
-    def initialize(job, serialized_progress)
+    include Validation
+
+    def initialize(job, serialized_progress) # :nodoc:
       @job = job
       @completed = serialized_progress.fetch("completed", []).map(&:to_sym)
       @current = new_step(*serialized_progress["current"], resumed: true) if serialized_progress.key?("current")
-      @encountered_step_names = []
+      @encountered = []
       @advanced = false
       @running_step = false
     end
 
-    def continue(&block)
-      wrapping_errors_after_advancing do
-        instrument_job :resume if started?
-        block.call
-      end
-    end
-
-    def step(name, start:, &block)
+    def step(name, start:, &block) # :nodoc:
       validate_step!(name)
+      encountered << name
 
       if completed?(name)
         skip_step(name)
@@ -209,14 +217,14 @@ def step(name, start:, &block)
       end
     end
 
-    def to_h
+    def to_h # :nodoc:
       {
         "completed" => completed.map(&:to_s),
         "current" => current&.to_a
       }.compact
     end
 
-    def description
+    def description # :nodoc:
       if current
         current.description
       elsif completed.any?
@@ -226,36 +234,33 @@ def description
       end
     end
 
-    private
-      attr_reader :job, :encountered_step_names, :completed, :current
+    def started?
+      completed.any? || current.present?
+    end
 
-      def advanced?
-        @advanced
-      end
+    def advanced?
+      @advanced
+    end
+
+    def instrumentation
+      { description: description,
+        completed_steps: completed,
+        current_step: current }
+    end
+
+    private
+      attr_reader :job, :encountered, :completed, :current
 
       def running_step?
         @running_step
       end
 
-      def started?
-        completed.any? || current.present?
-      end
-
       def completed?(name)
         completed.include?(name)
       end
 
-      def validate_step!(name)
-        raise InvalidStepError, "Step '#{name}' must be a Symbol, found '#{name.class}'" unless name.is_a?(Symbol)
-        raise InvalidStepError, "Step '#{name}' has already been encountered" if encountered_step_names.include?(name)
-        raise InvalidStepError, "Step '#{name}' is nested inside step '#{current.name}'" if running_step?
-        raise InvalidStepError, "Step '#{name}' found, expected to resume from '#{current.name}'" if current && current.name != name && !completed?(name)
-
-        encountered_step_names << name
-      end
-
       def new_step(*args, **options)
-        Step.new(*args, **options) { checkpoint! }
+        Step.new(*args, job: job, **options)
       end
 
       def skip_step(name)
@@ -273,49 +278,24 @@ def run_step(name, start:, &block)
         @completed << current.name
         @current = nil
         @advanced = true
-
-        checkpoint!
       ensure
         @running_step = false
         @advanced ||= current&.advanced?
       end
 
-      def interrupt!
-        instrument_job :interrupt
-        raise Interrupt, "Interrupted #{description}"
-      end
-
-      def checkpoint!
-        interrupt! if job.queue_adapter.stopping?
-      end
+      def instrumenting_step(step, &block)
+        instrument :step, step: step, interrupted: false do |payload|
+          instrument :step_started, step: step
 
-      def wrapping_errors_after_advancing(&block)
-        block.call
-      rescue StandardError => e
-        if !e.is_a?(Error) && advanced?
-          raise AfterAdvancingError, "Advanced job failed with error: #{e.message}"
-        else
+          block.call
+        rescue Interrupt
+          payload[:interrupted] = true
           raise
         end
       end
 
-      def instrumenting_step(step, &block)
-        instrument (step.resumed? ? :step_resumed : :step_started), step: step
-
-        block.call
-
-        instrument :step_completed, step: step
-      rescue Interrupt
-        instrument :step_interrupted, step: step
-        raise
-      end
-
-      def instrument_job(event)
-        instrument event, description: description, completed_steps: completed, current_step: current
-      end
-
-      def instrument(event, payload = {})
-        job.instrument event, **payload
+      def instrument(...)
+        job.instrument(...)
       end
   end
 end

activejob/lib/active_job/continuation/step.rb

@@ -22,18 +22,18 @@ class Step
       # The cursor for the step.
       attr_reader :cursor
 
-      def initialize(name, cursor, resumed:, &checkpoint_callback)
+      def initialize(name, cursor, job:, resumed:)
         @name = name.to_sym
         @initial_cursor = cursor
         @cursor = cursor
         @resumed = resumed
-        @checkpoint_callback = checkpoint_callback
+        @job = job
       end
 
       # Check if the job should be interrupted, and if so raise an Interrupt exception.
       # The job will be requeued for retry.
       def checkpoint!
-        checkpoint_callback.call
+        job.checkpoint!
       end
 
       # Set the cursor and interrupt the job if necessary.
@@ -77,7 +77,7 @@ def description
       end
 
       private
-        attr_reader :checkpoint_callback, :initial_cursor
+        attr_reader :initial_cursor, :job
     end
   end
 end

activejob/lib/active_job/continuation/test_helper.rb

@@ -40,6 +40,8 @@ def interrupt_job_during_step(job, step, cursor: nil, &block)
 
       # Interrupt a job after a step.
       #
+      # Note that there's no checkpoint after the final step so it won't be interrupted.
+      #
       #  class MyJob < ApplicationJob
       #    include ActiveJob::Continuable
       #