switchman-inst-jobs 3.2.10 → 4.0.0

data/lib/switchman_inst_jobs/jobs_migrator.rb

@@ -1,3 +1,7 @@
+# Just disabling all the rubocop metrics for this file for now,
+# as it is a direct port-in of existing code
+
+# rubocop:disable Metrics/BlockLength, Metrics/MethodLength, Metrics/AbcSize, Metrics/ClassLength
 require 'set'
 require 'parallel'
 
@@ -13,10 +17,10 @@ module SwitchmanInstJobs
         return yield if shards.empty?
 
         shard = shards.pop
-        current_shard = ::Switchman::Shard.current(:delayed_jobs)
-        shard.activate(:delayed_jobs) do
+        current_shard = ::Switchman::Shard.current(::Delayed::Backend::ActiveRecord::AbstractJob)
+        shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
           ::Delayed::Job.transaction do
-            current_shard.activate(:delayed_jobs) do
+            current_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
               transaction_on(shards, &block)
             end
           end
@@ -45,11 +49,11 @@ module SwitchmanInstJobs
         # rubocop:disable Style/CombinableLoops
         # We first migrate strands so that we can stop blocking strands before we migrate unstranded jobs
         source_shards.each do |s|
-          ::Switchman::Shard.lookup(s).activate(:delayed_jobs) { migrate_strands }
+          ::Switchman::Shard.lookup(s).activate(::Delayed::Backend::ActiveRecord::AbstractJob) { migrate_strands }
         end
 
         source_shards.each do |s|
-          ::Switchman::Shard.lookup(s).activate(:delayed_jobs) { migrate_everything }
+          ::Switchman::Shard.lookup(s).activate(::Delayed::Backend::ActiveRecord::AbstractJob) { migrate_everything }
         end
         ensure_unblock_stranded_for(shard_map.map(&:first))
         # rubocop:enable Style/CombinableLoops
@@ -89,9 +93,7 @@ module SwitchmanInstJobs
         migrate_everything
       end
 
-      def migrate_strands(batch_size: 1_000)
-        source_shard = ::Switchman::Shard.current(:delayed_jobs)
-
+      def migrate_strands
         # there are 4 scenarios to deal with here
         # 1) no running job, no jobs moved: do nothing
         # 2) running job, no jobs moved; create blocker with next_in_strand=false
@@ -100,89 +102,60 @@ module SwitchmanInstJobs
         #    those (= do nothing since it should already be false)
         # 4) no running job, jobs moved: set next_in_strand=true on the first of
         #    those (= do nothing since it should already be true)
-        handler = lambda { |scope, column, blocker_job_kwargs = {}, advisory_lock_cb = nil|
-          shard_map = build_shard_map(scope, source_shard)
-          shard_map.each do |(target_shard, source_shard_ids)|
-            shard_scope = scope.where(shard_id: source_shard_ids)
-
-            # 1) is taken care of because it should not show up here in strands
-            values = shard_scope.distinct.order(column).pluck(column)
-
-            target_shard.activate(:delayed_jobs) do
-              values.each do |value|
-                transaction_on([source_shard, target_shard]) do
-                  source_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
-                    advisory_lock_cb&.call(value)
-                  end
-
-                  value_scope = shard_scope.where(**{ column => value })
-                  # we want to copy all the jobs except the one that is still running.
-                  jobs_scope = value_scope.where(locked_by: nil)
-
-                  # 2) and part of 3) are taken care of here by creating a blocker
-                  # job with next_in_strand = false. as soon as the current
-                  # running job is finished it should set next_in_strand
-                  # We lock it to ensure that the jobs worker can't delete it until we are done moving the strand
-                  # Since we only unlock it on the new jobs queue *after* deleting from the original
-                  # the lock ensures the blocker always gets unlocked
-                  first = value_scope.where.not(locked_by: nil).next_in_strand_order.lock.first
-                  if first
-                    create_blocker_job(
-                      queue: first.queue,
-                      shard_id: first.shard_id,
-                      **{ column => value },
-                      **blocker_job_kwargs
-                    )
-
-                    # the rest of 3) is taken care of here
-                    # make sure that all the jobs moved over are NOT next in strand
-                    ::Delayed::Job.where(next_in_strand: true, locked_by: nil, **{ column => value }).
-                      update_all(next_in_strand: false)
-                  end
-
-                  # 4) is taken care of here, by leaving next_in_strand alone and
-                  # it should execute on the new shard
-                  batch_move_jobs(
-                    target_shard: target_shard,
-                    source_shard: source_shard,
-                    scope: jobs_scope,
-                    batch_size: batch_size
-                  ) do |job, new_job|
-                    # This ensures jobs enqueued on the old jobs shard run before jobs on the new jobs queue
-                    new_job.strand_order_override = job.strand_order_override - 1
-                  end
-                end
-              end
-            end
-          end
-        }
 
+        source_shard = ::Switchman::Shard.current(::Delayed::Backend::ActiveRecord::AbstractJob)
         strand_scope = ::Delayed::Job.shard(source_shard).where.not(strand: nil)
-        singleton_scope = ::Delayed::Job.shard(source_shard).where('strand IS NULL AND singleton IS NOT NULL')
-        all_scope = ::Delayed::Job.shard(source_shard).where('strand IS NOT NULL OR singleton IS NOT NULL')
-
-        singleton_blocker_additional_kwargs = {
-          locked_at: DateTime.now,
-          locked_by: ::Delayed::Backend::Base::ON_HOLD_BLOCKER
-        }
-
-        quoted_function_name = ::Delayed::Job.connection.quote_table_name('half_md5_as_bigint')
-        strand_advisory_lock_fn = lambda do |value|
-          ::Delayed::Job.connection.execute("SELECT pg_advisory_xact_lock(#{quoted_function_name}('#{value}'))")
-        end
-
-        singleton_advisory_lock_fn = lambda do |value|
-          ::Delayed::Job.connection.execute(
-            "SELECT pg_advisory_xact_lock(#{quoted_function_name}('singleton:#{value}'))"
-          )
-        end
+        shard_map = build_shard_map(strand_scope, source_shard)
+        shard_map.each do |(target_shard, source_shard_ids)|
+          shard_scope = strand_scope.where(shard_id: source_shard_ids)
+
+          # 1) is taken care of because it should not show up here in strands
+          strands = shard_scope.distinct.order(:strand).pluck(:strand)
+
+          target_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+            strands.each do |strand|
+              transaction_on([source_shard, target_shard]) do
+                this_strand_scope = shard_scope.where(strand: strand)
+                # we want to copy all the jobs except the one that is still running.
+                jobs_scope = this_strand_scope.where(locked_by: nil)
+
+                # 2) and part of 3) are taken care of here by creating a blocker
+                # job with next_in_strand = false. as soon as the current
+                # running job is finished it should set next_in_strand
+                # We lock it to ensure that the jobs worker can't delete it until we are done moving the strand
+                # Since we only unlock it on the new jobs queue *after* deleting from the original
+                # the lock ensures the blocker always gets unlocked
+                first = this_strand_scope.where.not(locked_by: nil).next_in_strand_order.lock.first
+                if first
+                  first_job = ::Delayed::Job.create!(strand: strand, next_in_strand: false)
+                  first_job.payload_object = ::Delayed::PerformableMethod.new(Kernel, :sleep, args: [0])
+                  first_job.queue = first.queue
+                  first_job.tag = 'Kernel.sleep'
+                  first_job.source = 'JobsMigrator::StrandBlocker'
+                  first_job.max_attempts = 1
+                  # If we ever have jobs left over from 9999 jobs moves of a single shard,
+                  # something has gone terribly wrong
+                  first_job.strand_order_override = -9999
+                  first_job.save!
+                  # the rest of 3) is taken care of here
+                  # make sure that all the jobs moved over are NOT next in strand
+                  ::Delayed::Job.where(next_in_strand: true, strand: strand, locked_by: nil).
+                    update_all(next_in_strand: false)
+                end
 
-        handler.call(strand_scope, :strand, {}, strand_advisory_lock_fn)
-        handler.call(singleton_scope, :singleton, singleton_blocker_additional_kwargs, singleton_advisory_lock_fn)
+                # 4) is taken care of here, by leaving next_in_strand alone and
+                # it should execute on the new shard
+                batch_move_jobs(
+                  target_shard: target_shard,
+                  source_shard: source_shard,
+                  scope: jobs_scope
+                ) do |job, new_job|
+                  # This ensures jobs enqueued on the old jobs shard run before jobs on the new jobs queue
+                  new_job.strand_order_override = job.strand_order_override - 1
+                end
+              end
+            end
 
-        shard_map = build_shard_map(all_scope, source_shard)
-        shard_map.each do |(target_shard, source_shard_ids)|
-          target_shard.activate(:delayed_jobs) do
             updated = ::Switchman::Shard.where(id: source_shard_ids, block_stranded: true).
               update_all(block_stranded: false)
             # If this is being manually re-run for some reason to clean something up, don't wait for nothing to happen
@@ -197,68 +170,41 @@ module SwitchmanInstJobs
         end
       end
 
-      def unblock_strands(target_shard, batch_size: 10_000)
-        block_stranded_ids = ::Switchman::Shard.where(block_stranded: true).pluck(:id)
-        query = lambda { |column, scope|
-          ::Delayed::Job.
-            where(id: ::Delayed::Job.select("DISTINCT ON (#{column}) id").
-              where(scope).
-              where.not(shard_id: block_stranded_ids).
-              where(
+      def unblock_strands(target_shard)
+        target_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+          loop do
+            # We only want to unlock stranded jobs where they don't belong to a blocked shard (if they *do* belong)
+            # to a blocked shard, they must be part of a concurrent jobs migration from a different source shard to
+            # this target shard, so we shouldn't unlock them yet.  We only ever unlock one job here to keep the
+            # logic cleaner; if the job is n-stranded, after the first one runs, the trigger will unlock larger
+            # batches
+            break if ::Delayed::Job.where(id: ::Delayed::Job.select('DISTINCT ON (strand) id').
+              where.not(strand: nil).
+              where.not(shard_id: ::Switchman::Shard.where(block_stranded: true).pluck(:id)).where(
                 ::Delayed::Job.select(1).from("#{::Delayed::Job.quoted_table_name} dj2").
                 where("dj2.next_in_strand = true OR dj2.source = 'JobsMigrator::StrandBlocker'").
-                where("dj2.#{column} = delayed_jobs.#{column}").arel.exists.not
-              ).
-              order(column, :strand_order_override, :id)).limit(batch_size)
-        }
-
-        target_shard.activate(:delayed_jobs) do
-          # We only want to unlock stranded jobs where they don't belong to a blocked shard (if they *do* belong)
-          # to a blocked shard, they must be part of a concurrent jobs migration from a different source shard to
-          # this target shard, so we shouldn't unlock them yet.  We only ever unlock one job here to keep the
-          # logic cleaner; if the job is n-stranded, after the first one runs, the trigger will unlock larger
-          # batches
-
-          loop do
-            break if query.call(:strand, 'strand IS NOT NULL').update_all(next_in_strand: true).zero?
-          end
-
-          loop do
-            break if query.call(:singleton,
-                                'strand IS NULL AND singleton IS NOT NULL').update_all(next_in_strand: true).zero?
+                where('dj2.strand = delayed_jobs.strand').arel.exists.not
+              ).order(:strand, :strand_order_override, :id)).limit(500).update_all(next_in_strand: true).zero?
           end
         end
       end
 
-      def migrate_everything(batch_size: 1_000)
-        source_shard = ::Switchman::Shard.current(:delayed_jobs)
-        scope = ::Delayed::Job.shard(source_shard).where('strand IS NULL')
+      def migrate_everything
+        source_shard = ::Switchman::Shard.current(::Delayed::Backend::ActiveRecord::AbstractJob)
+        scope = ::Delayed::Job.shard(source_shard).where(strand: nil)
 
         shard_map = build_shard_map(scope, source_shard)
         shard_map.each do |(target_shard, source_shard_ids)|
           batch_move_jobs(
             target_shard: target_shard,
             source_shard: source_shard,
-            scope: scope.where(shard_id: source_shard_ids).where(locked_by: nil),
-            batch_size: batch_size
+            scope: scope.where(shard_id: source_shard_ids).where(locked_by: nil)
           )
         end
       end
 
       private
 
-      def create_blocker_job(**kwargs)
-        first_job = ::Delayed::Job.create!(**kwargs, next_in_strand: false)
-        first_job.payload_object = ::Delayed::PerformableMethod.new(Kernel, :sleep, args: [0])
-        first_job.tag = 'Kernel.sleep'
-        first_job.source = 'JobsMigrator::StrandBlocker'
-        first_job.max_attempts = 1
-        # If we ever have jobs left over from 9999 jobs moves of a single shard,
-        # something has gone terribly wrong
-        first_job.strand_order_override = -9999
-        first_job.save!
-      end
-
       def build_shard_map(scope, source_shard)
         shard_ids = scope.distinct.pluck(:shard_id)
 
@@ -273,12 +219,12 @@ module SwitchmanInstJobs
         shard_map
       end
 
-      def batch_move_jobs(target_shard:, source_shard:, scope:, batch_size:)
+      def batch_move_jobs(target_shard:, source_shard:, scope:)
         while scope.exists?
           # Adapted from get_and_lock_next_available in delayed/backend/active_record.rb
-          target_jobs = scope.limit(batch_size).lock('FOR UPDATE SKIP LOCKED')
+          target_jobs = scope.limit(1000).lock('FOR UPDATE SKIP LOCKED')
 
-          query = source_shard.activate(:delayed_jobs) do
+          query = source_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
             "WITH limited_jobs AS (#{target_jobs.to_sql}) " \
             "UPDATE #{::Delayed::Job.quoted_table_name} " \
             "SET locked_by = #{::Delayed::Job.connection.quote(::Delayed::Backend::Base::ON_HOLD_LOCKED_BY)}, " \
@@ -287,7 +233,9 @@ module SwitchmanInstJobs
             "RETURNING #{::Delayed::Job.quoted_table_name}.*"
           end
 
-          jobs = source_shard.activate(:delayed_jobs) { ::Delayed::Job.find_by_sql(query) }
+          jobs = source_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
+            ::Delayed::Job.find_by_sql(query)
+          end
           new_jobs = jobs.map do |job|
             new_job = job.dup
             new_job.shard = target_shard
@@ -305,10 +253,10 @@ module SwitchmanInstJobs
             new_job
           end
           transaction_on([source_shard, target_shard]) do
-            target_shard.activate(:delayed_jobs) do
+            target_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
               bulk_insert_jobs(new_jobs)
             end
-            source_shard.activate(:delayed_jobs) do
+            source_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
               ::Delayed::Job.delete(jobs)
             end
           end
@@ -335,10 +283,7 @@ module SwitchmanInstJobs
         connection = ::Delayed::Job.connection
         quoted_keys = keys.map { |k| connection.quote_column_name(k) }.join(', ')
 
-        connection.execute 'DROP TABLE IF EXISTS delayed_jobs_bulk_copy'
-        connection.execute "CREATE TEMPORARY TABLE delayed_jobs_bulk_copy
-          (LIKE #{::Delayed::Job.quoted_table_name} INCLUDING DEFAULTS)"
-        connection.execute "COPY delayed_jobs_bulk_copy (#{quoted_keys}) FROM STDIN"
+        connection.execute "COPY #{::Delayed::Job.quoted_table_name} (#{quoted_keys}) FROM STDIN"
         records.map do |record|
           connection.raw_connection.put_copy_data("#{keys.map { |k| quote_text(record[k]) }.join("\t")}\n")
         end
@@ -350,9 +295,6 @@ module SwitchmanInstJobs
         rescue StandardError => e
           raise connection.send(:translate_exception, e, 'COPY FROM STDIN')
         end
-        connection.execute "INSERT INTO #{::Delayed::Job.quoted_table_name} (#{quoted_keys})
-          SELECT #{quoted_keys} FROM delayed_jobs_bulk_copy
-          ON CONFLICT (singleton) WHERE singleton IS NOT NULL AND locked_by IS NULL DO NOTHING"
         result.cmd_tuples
       end
 
@@ -370,3 +312,5 @@ module SwitchmanInstJobs
     end
   end
 end
+
+# rubocop:enable Metrics/BlockLength, Metrics/MethodLength, Metrics/AbcSize, Metrics/ClassLength

data/lib/switchman_inst_jobs/switchman/shard.rb

@@ -22,12 +22,12 @@ module SwitchmanInstJobs
       def hold_jobs!(wait: false)
         self.jobs_held = true
         save! if changed?
-        delayed_jobs_shard.activate(:delayed_jobs) do
+        delayed_jobs_shard.activate(::Delayed::Backend::ActiveRecord::AbstractJob) do
           lock_jobs_for_hold
         end
         return unless wait
 
-        delayed_jobs_shard.activate(:delayed_jobs) do
+        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?
@@ -47,7 +47,7 @@ module SwitchmanInstJobs
           Rails.logger.debug('Waiting for caches to clear')
           sleep(65)
         end
-        delayed_jobs_shard.activate(:delayed_jobs) do
+        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(
@@ -75,22 +75,14 @@ module SwitchmanInstJobs
           remove_instance_variable(:@delayed_jobs_shards) if instance_variable_defined?(:@delayed_jobs_shards)
         end
 
-        def current(category = :primary)
-          if category == :delayed_jobs
-            active_shards[category] || super(:primary).delayed_jobs_shard
-          else
-            super
-          end
-        end
-
         def activate!(categories)
           if !@skip_delayed_job_auto_activation &&
-             !categories[:delayed_jobs] &&
-             categories[:primary] &&
-             categories[:primary] != active_shards[:primary]
+             !categories[::Delayed::Backend::ActiveRecord::AbstractJob] &&
+             categories[::ActiveRecord::Base] &&
+             categories[::ActiveRecord::Base] != ::Switchman::Shard.current(::ActiveRecord::Base)
             skip_delayed_job_auto_activation do
-              categories[:delayed_jobs] =
-                categories[:primary].delayed_jobs_shard
+              categories[::Delayed::Backend::ActiveRecord::AbstractJob] =
+                categories[::ActiveRecord::Base].delayed_jobs_shard
             end
           end
           super
@@ -104,11 +96,6 @@ module SwitchmanInstJobs
           @skip_delayed_job_auto_activation = was
         end
 
-        def create
-          db = ::Switchman::DatabaseServer.server_for_new_shard
-          db.create_new_shard
-        end
-
         def periodic_clear_shard_cache
           # TODO: make this configurable
           @timed_cache ||= TimedCache.new(-> { 60.to_i.seconds.ago }) do

data/lib/switchman_inst_jobs/version.rb

@@ -1,3 +1,3 @@
 module SwitchmanInstJobs
-  VERSION = '3.2.10'.freeze
+  VERSION = '4.0.0'.freeze
 end

data/lib/switchman_inst_jobs.rb

@@ -5,6 +5,9 @@ module SwitchmanInstJobs
   cattr_accessor :delayed_jobs_shard_fallback
 
   def self.initialize_active_record
+    ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(
+      ActiveRecord::ConnectionAdapters::ConnectionPool
+    )
     ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(
       ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
     )
@@ -14,6 +17,10 @@ module SwitchmanInstJobs
     ::Delayed::Backend::ActiveRecord::Job.prepend(
       Delayed::Backend::Base
     )
+    ::Delayed::Backend::Redis::Job.prepend(
+      Delayed::Backend::Base
+    )
+    ::Delayed::Backend::Redis::Job.column :shard_id, :integer
     ::Delayed::Pool.prepend Delayed::Pool
     ::Delayed::Worker.prepend Delayed::Worker
     ::Delayed::Worker::HealthCheck.prepend Delayed::Worker::HealthCheck
@@ -32,6 +39,7 @@ module SwitchmanInstJobs
   end
 end
 
+require 'switchman_inst_jobs/active_record/connection_adapters/connection_pool'
 require 'switchman_inst_jobs/active_record/connection_adapters/postgresql_adapter'
 require 'switchman_inst_jobs/active_record/migration'
 require 'switchman_inst_jobs/delayed/settings'