Merge pull request #93 from basecamp/multi-enqueue

Implement `perform_all_later` via `enqueue_all`

Author: rosa

app/models/solid_queue/blocked_execution.rb

@@ -2,7 +2,7 @@
 
 module SolidQueue
   class BlockedExecution < Execution
-    assume_attributes_from_job :concurrency_key
+    assumes_attributes_from_job :concurrency_key
     before_create :set_expires_at
 
     has_one :semaphore, foreign_key: :key, primary_key: :concurrency_key

app/models/solid_queue/execution.rb

@@ -11,5 +11,17 @@ class Execution < Record
     belongs_to :job
 
     alias_method :discard, :destroy
+
+    class << self
+      def create_all_from_jobs(jobs)
+        insert_all execution_data_from_jobs(jobs)
+      end
+
+      def execution_data_from_jobs(jobs)
+        jobs.collect do |job|
+          attributes_from_job(job).merge(job_id: job.id)
+        end
+      end
+    end
   end
 end

app/models/solid_queue/execution/job_attributes.rb

@@ -5,17 +5,24 @@ class Execution
     module JobAttributes
       extend ActiveSupport::Concern
 
-      ASSUMIBLE_ATTRIBUTES_FROM_JOB = %i[ queue_name priority ]
+      included do
+        class_attribute :assumible_attributes_from_job, instance_accessor: false, default: %i[ queue_name priority ]
+      end
 
       class_methods do
-        def assume_attributes_from_job(*attributes)
-          before_create -> { assume_attributes_from_job(ASSUMIBLE_ATTRIBUTES_FROM_JOB | attributes) }
+        def assumes_attributes_from_job(*attribute_names)
+          self.assumible_attributes_from_job |= attribute_names
+          before_create -> { assume_attributes_from_job }
+        end
+
+        def attributes_from_job(job)
+          job.attributes.symbolize_keys.slice(*assumible_attributes_from_job)
         end
       end
 
       private
-        def assume_attributes_from_job(attributes)
-          attributes.each do |attribute|
+        def assume_attributes_from_job
+          self.class.assumible_attributes_from_job.each do |attribute|
             send("#{attribute}=", job.send(attribute))
           end
         end

app/models/solid_queue/job.rb

@@ -1,38 +1,62 @@
 # frozen_string_literal: true
 
-class SolidQueue::Job < SolidQueue::Record
-  include Executable
+module SolidQueue
+  class Job < Record
+    include Executable
 
-  if Gem::Version.new(Rails.version) >= Gem::Version.new("7.1")
-    serialize :arguments, coder: JSON
-  else
-    serialize :arguments, JSON
-  end
-
-  DEFAULT_PRIORITY = 0
-  DEFAULT_QUEUE_NAME = "default"
-
-  class << self
-    def enqueue_active_job(active_job, scheduled_at: Time.current)
-      enqueue \
-        queue_name: active_job.queue_name,
-        active_job_id: active_job.job_id,
-        priority: active_job.priority,
-        scheduled_at: scheduled_at,
-        class_name: active_job.class.name,
-        arguments: active_job.serialize,
-        concurrency_key: active_job.try(:concurrency_key)
+    if Gem::Version.new(Rails.version) >= Gem::Version.new("7.1")
+      serialize :arguments, coder: JSON
+    else
+      serialize :arguments, JSON
     end
 
-    def enqueue(**kwargs)
-      create!(**kwargs.compact.with_defaults(defaults)).tap do
-        SolidQueue.logger.debug "[SolidQueue] Enqueued job #{kwargs}"
+    class << self
+      def enqueue_all(active_jobs)
+        active_jobs_by_job_id = active_jobs.index_by(&:job_id)
+
+        transaction do
+          jobs = create_all_from_active_jobs(active_jobs)
+          prepare_all_for_execution(jobs).tap do |enqueued_jobs|
+            enqueued_jobs.each do |enqueued_job|
+              active_jobs_by_job_id[enqueued_job.active_job_id].provider_job_id = enqueued_job.id
+            end
+          end
+        end
       end
-    end
 
-    private
-      def defaults
-        { queue_name: DEFAULT_QUEUE_NAME, priority: DEFAULT_PRIORITY }
+      def enqueue(active_job, scheduled_at: Time.current)
+        active_job.scheduled_at = scheduled_at
+
+        create_from_active_job(active_job).tap do |enqueued_job|
+          active_job.provider_job_id = enqueued_job.id
+        end
       end
+
+      private
+        DEFAULT_PRIORITY = 0
+        DEFAULT_QUEUE_NAME = "default"
+
+        def create_from_active_job(active_job)
+          create!(**attributes_from_active_job(active_job))
+        end
+
+        def create_all_from_active_jobs(active_jobs)
+          job_rows = active_jobs.map { |job| attributes_from_active_job(job) }
+          insert_all(job_rows)
+          where(active_job_id: active_jobs.map(&:job_id))
+        end
+
+        def attributes_from_active_job(active_job)
+          {
+            queue_name: active_job.queue_name || DEFAULT_QUEUE_NAME,
+            active_job_id: active_job.job_id,
+            priority: active_job.priority || DEFAULT_PRIORITY,
+            scheduled_at: active_job.scheduled_at,
+            class_name: active_job.class.name,
+            arguments: active_job.serialize,
+            concurrency_key: active_job.concurrency_key
+          }
+        end
+    end
   end
 end

app/models/solid_queue/job/executable.rb

@@ -6,21 +6,56 @@ module Executable
       extend ActiveSupport::Concern
 
       included do
-        include Clearable, ConcurrencyControls
+        include Clearable, ConcurrencyControls, Schedulable
 
         has_one :ready_execution, dependent: :destroy
         has_one :claimed_execution, dependent: :destroy
         has_one :failed_execution, dependent: :destroy
 
-        has_one :scheduled_execution, dependent: :destroy
-
         after_create :prepare_for_execution
 
         scope :finished, -> { where.not(finished_at: nil) }
-        scope :failed, -> { includes(:failed_execution).where.not(failed_execution: {id: nil}) }
+        scope :failed, -> { includes(:failed_execution).where.not(failed_execution: { id: nil }) }
       end
 
-      %w[ ready claimed failed scheduled ].each do |status|
+      class_methods do
+        def prepare_all_for_execution(jobs)
+          due, not_yet_due = jobs.partition(&:due?)
+          dispatch_all(due) + schedule_all(not_yet_due)
+        end
+
+        def dispatch_all(jobs)
+          with_concurrency_limits, without_concurrency_limits = jobs.partition(&:concurrency_limited?)
+
+          dispatch_all_at_once(without_concurrency_limits)
+          dispatch_all_one_by_one(with_concurrency_limits)
+
+          successfully_dispatched(jobs)
+        end
+
+        private
+          def dispatch_all_at_once(jobs)
+            ReadyExecution.create_all_from_jobs jobs
+          end
+
+          def dispatch_all_one_by_one(jobs)
+            jobs.each(&:dispatch)
+          end
+
+          def successfully_dispatched(jobs)
+            dispatched_and_ready(jobs) + dispatched_and_blocked(jobs)
+          end
+
+          def dispatched_and_ready(jobs)
+            where(id: ReadyExecution.where(job_id: jobs.map(&:id)).pluck(:job_id))
+          end
+
+          def dispatched_and_blocked(jobs)
+            where(id: BlockedExecution.where(job_id: jobs.map(&:id)).pluck(:job_id))
+          end
+      end
+
+      %w[ ready claimed failed ].each do |status|
         define_method("#{status}?") { public_send("#{status}_execution").present? }
       end
 
@@ -50,6 +85,10 @@ def finished?
         finished_at.present?
       end
 
+      def due?
+        scheduled_at.nil? || scheduled_at <= Time.current
+      end
+
       def discard
         destroy unless claimed?
       end
@@ -63,14 +102,6 @@ def failed_with(exception)
       end
 
       private
-        def due?
-          scheduled_at.nil? || scheduled_at <= Time.current
-        end
-
-        def schedule
-          ScheduledExecution.create_or_find_by!(job_id: id)
-        end
-
         def ready
           ReadyExecution.create_or_find_by!(job_id: id)
         end

app/models/solid_queue/job/schedulable.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module SolidQueue
+  class Job
+    module Schedulable
+      extend ActiveSupport::Concern
+
+      included do
+        has_one :scheduled_execution, dependent: :destroy
+
+        scope :scheduled, -> { where.not(finished_at: nil) }
+      end
+
+      class_methods do
+        def schedule_all(jobs)
+          schedule_all_at_once(jobs)
+          successfully_scheduled(jobs)
+        end
+
+        private
+          def schedule_all_at_once(jobs)
+            ScheduledExecution.create_all_from_jobs(jobs)
+          end
+
+          def successfully_scheduled(jobs)
+            where(id: ScheduledExecution.where(job_id: jobs.map(&:id)).pluck(:job_id))
+          end
+      end
+
+      def due?
+        scheduled_at.nil? || scheduled_at <= Time.current
+      end
+
+      private
+        def schedule
+          ScheduledExecution.create_or_find_by!(job_id: id)
+        end
+    end
+  end
+end

app/models/solid_queue/ready_execution.rb

@@ -4,7 +4,7 @@ module SolidQueue
   class ReadyExecution < Execution
     scope :queued_as, ->(queue_name) { where(queue_name: queue_name) }
 
-    assume_attributes_from_job
+    assumes_attributes_from_job
 
     class << self
       def claim(queue_list, limit, process_id)

app/models/solid_queue/scheduled_execution.rb

@@ -6,7 +6,7 @@ class ScheduledExecution < Execution
     scope :ordered, -> { order(scheduled_at: :asc, priority: :asc) }
     scope :next_batch, ->(batch_size) { due.ordered.limit(batch_size) }
 
-    assume_attributes_from_job :scheduled_at
+    assumes_attributes_from_job :scheduled_at
 
     class << self
       def dispatch_next_batch(batch_size)
@@ -22,44 +22,12 @@ def dispatch_next_batch(batch_size)
       private
         def dispatch_batch(job_ids)
           jobs = Job.where(id: job_ids)
-          with_concurrency_limits, without_concurrency_limits = jobs.partition(&:concurrency_limited?)
 
-          dispatch_at_once(without_concurrency_limits)
-          dispatch_one_by_one(with_concurrency_limits)
-
-          successfully_dispatched(job_ids).tap do |dispatched_job_ids|
+          Job.dispatch_all(jobs).map(&:id).tap do |dispatched_job_ids|
             where(job_id: dispatched_job_ids).delete_all
             SolidQueue.logger.info("[SolidQueue] Dispatched scheduled batch with #{dispatched_job_ids.size} jobs")
           end
         end
-
-        def dispatch_at_once(jobs)
-          ReadyExecution.insert_all ready_rows_from_batch(jobs)
-        end
-
-        def dispatch_one_by_one(jobs)
-          jobs.each(&:dispatch)
-        end
-
-        def ready_rows_from_batch(jobs)
-          prepared_at = Time.current
-
-          jobs.map do |job|
-            { job_id: job.id, queue_name: job.queue_name, priority: job.priority, created_at: prepared_at }
-          end
-        end
-
-        def successfully_dispatched(job_ids)
-          dispatched_and_ready(job_ids) + dispatched_and_blocked(job_ids)
-        end
-
-        def dispatched_and_ready(job_ids)
-          ReadyExecution.where(job_id: job_ids).pluck(:job_id)
-        end
-
-        def dispatched_and_blocked(job_ids)
-          BlockedExecution.where(job_id: job_ids).pluck(:job_id)
-        end
     end
   end
 end

lib/active_job/concurrency_controls.rb

@@ -36,6 +36,10 @@ def concurrency_key
       end
     end
 
+    def concurrency_limited?
+      concurrency_key.present?
+    end
+
     private
       def concurrency_group
         compute_concurrency_parameter(self.class.concurrency_group)

lib/active_job/queue_adapters/solid_queue_adapter.rb

@@ -9,15 +9,15 @@ module QueueAdapters
     #   Rails.application.config.active_job.queue_adapter = :solid_queue
     class SolidQueueAdapter
       def enqueue(active_job) # :nodoc:
-        SolidQueue::Job.enqueue_active_job(active_job).tap do |job|
-          active_job.provider_job_id = job.id
-        end
+        SolidQueue::Job.enqueue(active_job)
       end
 
       def enqueue_at(active_job, timestamp) # :nodoc:
-        SolidQueue::Job.enqueue_active_job(active_job, scheduled_at: Time.at(timestamp)).tap do |job|
-          active_job.provider_job_id = job.id
-        end
+        SolidQueue::Job.enqueue(active_job, scheduled_at: Time.at(timestamp))
+      end
+
+      def enqueue_all(active_jobs) # :nodoc:
+        SolidQueue::Job.enqueue_all(active_jobs)
       end
     end
   end