Allow holding a relation worth of shards' jobs

.github/workflows/continuous_integration.yml

@@ -21,7 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         ruby: ['3.2', '3.3', '3.4']
-        lockfile: ['activerecord-7.1', 'activerecord-7.2', 'Gemfile.lock']
+        lockfile: ['Gemfile.activerecord-7.1.lock', 'Gemfile.activerecord-7.2.lock', 'Gemfile.lock']
 
     services:
       postgres:

Gemfile.activerecord-7.1.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    switchman-inst-jobs (4.3.2)
+    switchman-inst-jobs (4.3.3)
       inst-jobs (>= 2.4.9, < 4.0)
       parallel (>= 1.19)
       railties (>= 7.0, < 8.1)

Gemfile.activerecord-7.2.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    switchman-inst-jobs (4.3.2)
+    switchman-inst-jobs (4.3.3)
       inst-jobs (>= 2.4.9, < 4.0)
       parallel (>= 1.19)
       railties (>= 7.0, < 8.1)

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    switchman-inst-jobs (4.3.2)
+    switchman-inst-jobs (4.3.3)
       inst-jobs (>= 2.4.9, < 4.0)
       parallel (>= 1.19)
       railties (>= 7.0, < 8.1)

lib/switchman_inst_jobs/switchman/shard.rb

@@ -22,56 +22,82 @@ def delayed_jobs_shard
       # Adapted from hold/unhold methods in base delayed jobs base
       # Wait is required to be able to safely move jobs
       def hold_jobs!(wait: false)
-        self.jobs_held = true
-        save! if changed?
-        delayed_jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
-          lock_jobs_for_hold
-        end
-        return unless wait
-
-        delayed_jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
-          while ::Delayed::Job.where(shard_id: id)
-                              .where.not(locked_at: nil)
-                              .where.not(locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY).exists?
-            sleep 10
-            lock_jobs_for_hold
-          end
-        end
+        ::Switchman::Shard.where(id: self).hold_jobs!(wait:)
       end
 
       def unhold_jobs!
-        self.jobs_held = false
-        if changed?
-          save!
+        ::Switchman::Shard.where(id: self).unhold_jobs!
+      end
+
+      module ClassMethods
+        # Adapted from hold/unhold methods in base delayed jobs base
+        # Wait is required to be able to safely move jobs
+        def hold_jobs!(wait: false)
+          shards = all.to_a
+          wait_for_caches = false
+          shards.each do |shard|
+            shard.jobs_held = true
+            if shard.changed?
+              shard.save!
+              wait_for_caches = true if wait
+            end
+          end
+          shards_by_jobs_shard(shards).each do |jobs_shard, shard_ids|
+            jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+              lock_jobs_for_hold(shard_ids)
+            end
+          end
+          return unless wait
+
           # Wait a little over the 60 second in-process shard cache clearing
           # threshold to ensure that all new jobs are now being enqueued
-          # unlocked
+          # locked
           Rails.logger.debug("Waiting for caches to clear")
-          sleep(65)
-        end
-        delayed_jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
-          ::Delayed::Job.where(locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY, shard_id: id)
-                        .in_batches(of: 10_000)
-                        .update_all(
-                          locked_by: nil,
-                          locked_at: nil,
-                          attempts: 0,
-                          failed_at: nil
-                        )
+          sleep(65) if wait && wait_for_caches
+
+          shards_by_jobs_shard(shards).each do |jobs_shard, shard_ids|
+            jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+              while ::Delayed::Job.where(shard_id: shard_ids)
+                                  .where.not(locked_at: nil)
+                                  .where.not(locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY).exists?
+                sleep 10
+                lock_jobs_for_hold(shard_ids)
+              end
+            end
+          end
         end
-      end
-
-      private
 
-      def lock_jobs_for_hold
-        ::Delayed::Job.where(locked_at: nil, shard_id: id).in_batches(of: 10_000).update_all(
-          locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY,
-          locked_at: ::Delayed::Job.db_time_now,
-          attempts: ::Delayed::Backend::Base::ON_HOLD_COUNT
-        )
-      end
+        def unhold_jobs!
+          shards = all.to_a
+          waited = false
+          shards.each do |shard|
+            shard.jobs_held = false
+            next unless shard.changed?
+
+            shard.save!
+            next if waited
+
+            # Wait a little over the 60 second in-process shard cache clearing
+            # threshold to ensure that all new jobs are now being enqueued
+            # unlocked
+            Rails.logger.debug("Waiting for caches to clear")
+            sleep(65)
+            waited = true
+          end
+          shards_by_jobs_shard(shards).each do |jobs_shard, shard_ids|
+            jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+              ::Delayed::Job.where(locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY, shard_id: shard_ids)
+                            .in_batches(of: 10_000)
+                            .update_all(
+                              locked_by: nil,
+                              locked_at: nil,
+                              attempts: 0,
+                              failed_at: nil
+                            )
+            end
+          end
+        end
 
-      module ClassMethods
         def clear_cache
           super
           remove_instance_variable(:@delayed_jobs_shards) if instance_variable_defined?(:@delayed_jobs_shards)
@@ -130,6 +156,22 @@ def delayed_jobs_shards
 
           ::Switchman::Shard.merge(scope)
         end
+
+        private
+
+        # Group the given shards by the shard their jobs live on, returning a
+        # hash of delayed_jobs_shard => [shard_id, ...]
+        def shards_by_jobs_shard(shards)
+          shards.group_by(&:delayed_jobs_shard).transform_values { |group| group.map(&:id) }
+        end
+
+        def lock_jobs_for_hold(shard_ids)
+          ::Delayed::Job.where(locked_at: nil, shard_id: shard_ids).in_batches(of: 10_000).update_all(
+            locked_by: ::Delayed::Backend::Base::ON_HOLD_LOCKED_BY,
+            locked_at: ::Delayed::Job.db_time_now,
+            attempts: ::Delayed::Backend::Base::ON_HOLD_COUNT
+          )
+        end
       end
     end
   end

lib/switchman_inst_jobs/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module SwitchmanInstJobs
-  VERSION = "4.3.2"
+  VERSION = "4.3.3"
 end

spec/lib/switchman_inst_jobs/switchman/shard_spec.rb

@@ -81,4 +81,43 @@
       expect(job.reload.locked_by).to be_nil
     end
   end
+
+  describe ".hold_jobs!" do
+    it "locks existing jobs across a relation of shards" do
+      job1 = shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+      job2 = jobs_shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+
+      Switchman::Shard.where(id: [shard, jobs_shard]).hold_jobs!
+
+      expect(job1.reload.locked_by).to eq Delayed::Backend::Base::ON_HOLD_LOCKED_BY
+      expect(job2.reload.locked_by).to eq Delayed::Backend::Base::ON_HOLD_LOCKED_BY
+      expect(shard.reload.jobs_held).to be true
+      expect(jobs_shard.reload.jobs_held).to be true
+    end
+
+    it "only holds jobs for shards in the relation" do
+      held = shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+      other = jobs_shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+
+      Switchman::Shard.where(id: shard).hold_jobs!
+
+      expect(held.reload.locked_by).to eq Delayed::Backend::Base::ON_HOLD_LOCKED_BY
+      expect(other.reload.locked_by).to be_nil
+      expect(jobs_shard.reload.jobs_held).to be false
+    end
+  end
+
+  describe ".unhold_jobs!" do
+    it "unholds existing jobs across a relation of shards" do
+      job1 = shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+      job2 = jobs_shard.activate { Kernel.delay(ignore_transaction: true).sleep }
+      job1.hold!
+      job2.hold!
+
+      Switchman::Shard.where(id: [shard, jobs_shard]).unhold_jobs!
+
+      expect(job1.reload.locked_by).to be_nil
+      expect(job2.reload.locked_by).to be_nil
+    end
+  end
 end