From 4c32f628e2fe8548895279c41a895f65ec2c7cc3 Mon Sep 17 00:00:00 2001
From: Adrianna Chang <adrianna.chang@shopify.com>
Date: Thu, 25 Feb 2021 09:41:43 -0500
Subject: [PATCH 1/2] Postpone reenqueuing the iteration job until after
 callbacks are done

This PR fixes a bug that occurs when a job gets interrupted and is
reenqueued faster than the original job takes to shut down. This results
in race conditions with the @run and results in errors due to invalid
status transitions (and notably results in a run showing up as interrupted
when it is actually running).

To solve this, we should delay reenqueuing the job until after all callbacks
have completed, which can be done by calling #reenqueue_iteration_job
in a prepended after_perform callback in TaskJob, and skipping the call
that JobIteration performs in #iterate_with_enumerator.
---
 app/jobs/maintenance_tasks/task_job.rb       |  8 ++++++++
 test/jobs/maintenance_tasks/task_job_test.rb | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/app/jobs/maintenance_tasks/task_job.rb b/app/jobs/maintenance_tasks/task_job.rb
index 7ad599711..b2d2a971c 100644
--- a/app/jobs/maintenance_tasks/task_job.rb
+++ b/app/jobs/maintenance_tasks/task_job.rb
@@ -97,8 +97,16 @@ def on_shutdown
       @ticker.persist
     end
 
+    def reenqueue_iteration_job(should_ignore: true)
+      super() unless should_ignore
+      @reenqueue_iteration_job = true
+    end
+
     def after_perform
       @run.save!
+      if defined?(@reenqueue_iteration_job) && @reenqueue_iteration_job
+        reenqueue_iteration_job(should_ignore: false)
+      end
     end
 
     def on_error(error)
diff --git a/test/jobs/maintenance_tasks/task_job_test.rb b/test/jobs/maintenance_tasks/task_job_test.rb
index f95081e52..714d8493c 100644
--- a/test/jobs/maintenance_tasks/task_job_test.rb
+++ b/test/jobs/maintenance_tasks/task_job_test.rb
@@ -164,6 +164,24 @@ class TaskJobTest < ActiveJob::TestCase
       assert_equal 'Task Maintenance::DeletedTask not found.', run.error_message
     end
 
+    test '.perform_now delays reenqueuing the job after interruption until all callbacks are finished' do
+      JobIteration.stubs(interruption_adapter: -> { true })
+
+      AnotherTaskJob = Class.new(TaskJob) do
+        after_perform { self.class.times_interrupted = times_interrupted }
+
+        class << self
+          attr_accessor :times_interrupted
+        end
+      end
+      AnotherTaskJob.perform_now(@run)
+
+      # The job should not yet have been reenqueued, so times_interrupted should
+      # be 0
+      assert_equal 0, AnotherTaskJob.times_interrupted
+      assert_enqueued_jobs 1
+    end
+
     test '.perform_now does not enqueue another job if Run errors' do
       run = Run.create!(task_name: 'Maintenance::ErrorTask')
 

From b91efac95da7cbfa19dafed33a06f5b266791784 Mon Sep 17 00:00:00 2001
From: Adrianna Chang <adrianna.chang@shopify.com>
Date: Fri, 26 Feb 2021 16:35:01 -0500
Subject: [PATCH 2/2] Ensure we are compatible with JobIteration::Iteration and
 that #reenqueue_iteration_job is defined upstream

---
 app/jobs/maintenance_tasks/task_job.rb | 11 +++++++++++
 lib/maintenance_tasks.rb               |  4 ++++
 2 files changed, 15 insertions(+)

diff --git a/app/jobs/maintenance_tasks/task_job.rb b/app/jobs/maintenance_tasks/task_job.rb
index b2d2a971c..ea8781683 100644
--- a/app/jobs/maintenance_tasks/task_job.rb
+++ b/app/jobs/maintenance_tasks/task_job.rb
@@ -97,6 +97,17 @@ def on_shutdown
       @ticker.persist
     end
 
+    # We are reopening a private part of Job Iteration's API here, so we should
+    # ensure the method is still defined upstream. This way, in the case where
+    # the method changes upstream, we catch it at load time instead of at
+    # runtime while calling `super`.
+    unless private_method_defined?(:reenqueue_iteration_job)
+      error_message = <<~HEREDOC
+        JobIteration::Iteration#reenqueue_iteration_job is expected to be
+        defined. Upgrading the maintenance_tasks gem should solve this problem.
+      HEREDOC
+      raise error_message
+    end
     def reenqueue_iteration_job(should_ignore: true)
       super() unless should_ignore
       @reenqueue_iteration_job = true
diff --git a/lib/maintenance_tasks.rb b/lib/maintenance_tasks.rb
index ad0964cd1..0804aa3a5 100644
--- a/lib/maintenance_tasks.rb
+++ b/lib/maintenance_tasks.rb
@@ -9,6 +9,10 @@
 require 'pagy'
 require 'pagy/extras/bulma'
 
+# Force the TaskJob class to load so we can verify upstream compatibility with
+# the JobIteration gem
+require_relative '../app/jobs/maintenance_tasks/task_job'
+
 # The engine's namespace module. It provides isolation between the host
 # application's code and the engine-specific code. Top-level engine constants
 # and variables are defined under this module.