Merge pull request #38 from basecamp/concurrency-controls

Introduce concurrency controls

Author: rosa

Diff for: app/models/solid_queue/blocked_execution.rb

@@ -0,0 +1,60 @@
+module SolidQueue
+  class BlockedExecution < SolidQueue::Execution
+    assume_attributes_from_job :concurrency_key
+    before_create :set_expires_at
+
+    has_one :semaphore, foreign_key: :key, primary_key: :concurrency_key
+
+    scope :expired, -> { where(expires_at: ...Time.current) }
+
+    class << self
+      def unblock(count)
+        expired.distinct.limit(count).pluck(:concurrency_key).then do |concurrency_keys|
+          release_many releasable(concurrency_keys)
+        end
+      end
+
+      def release_many(concurrency_keys)
+        # We want to release exactly one blocked execution for each concurrency key, and we need to do it
+        # one by one, locking each record and acquiring the semaphore individually for each of them:
+        Array(concurrency_keys).each { |concurrency_key| release_one(concurrency_key) }
+      end
+
+      def release_one(concurrency_key)
+        ordered.where(concurrency_key: concurrency_key).limit(1).lock("FOR UPDATE SKIP LOCKED").each(&:release)
+      end
+
+      private
+        def releasable(concurrency_keys)
+          semaphores = Semaphore.where(key: concurrency_keys).select(:key, :value).index_by(&:key)
+
+          # Concurrency keys without semaphore + concurrency keys with open semaphore
+          (concurrency_keys - semaphores.keys) | semaphores.select { |key, semaphore| semaphore.value > 0 }.map(&:first)
+        end
+    end
+
+    def release
+      transaction do
+        if acquire_concurrency_lock
+          promote_to_ready
+          destroy!
+
+          SolidQueue.logger.info("[SolidQueue] Unblocked job #{job.id} under #{concurrency_key}")
+        end
+      end
+    end
+
+    private
+      def set_expires_at
+        self.expires_at = job.concurrency_duration.from_now
+      end
+
+      def acquire_concurrency_lock
+        Semaphore.wait(job)
+      end
+
+      def promote_to_ready
+        ReadyExecution.create!(ready_attributes)
+      end
+  end
+end

Diff for: app/models/solid_queue/claimed_execution.rb

@@ -31,6 +31,8 @@ def perform
     else
       failed_with(result.error)
     end
+  ensure
+    job.unblock_next_blocked_job
   end
 
   def release

Diff for: app/models/solid_queue/execution.rb

@@ -1,9 +1,17 @@
-class SolidQueue::Execution < SolidQueue::Record
-  include JobAttributes
+module SolidQueue
+  class Execution < SolidQueue::Record
+    include JobAttributes
 
-  self.abstract_class = true
+    self.abstract_class = true
 
-  belongs_to :job
+    scope :ordered, -> { order(priority: :asc, job_id: :asc) }
 
-  alias_method :discard, :destroy
+    belongs_to :job
+
+    alias_method :discard, :destroy
+
+    def ready_attributes
+      attributes.slice("job_id", "queue_name", "priority")
+    end
+  end
 end

Diff for: app/models/solid_queue/job.rb

@@ -18,7 +18,8 @@ def enqueue_active_job(active_job, scheduled_at: Time.current)
         priority: active_job.priority,
         scheduled_at: scheduled_at,
         class_name: active_job.class.name,
-        arguments: active_job.serialize
+        arguments: active_job.serialize,
+        concurrency_key: active_job.try(:concurrency_key)
     end
 
     def enqueue(**kwargs)

Diff for: app/models/solid_queue/job/concurrency_controls.rb

@@ -0,0 +1,48 @@
+module SolidQueue
+  class Job
+    module ConcurrencyControls
+      extend ActiveSupport::Concern
+
+      included do
+        has_one :blocked_execution, dependent: :destroy
+
+        delegate :concurrency_limit, :concurrency_duration, to: :job_class
+      end
+
+      def unblock_next_blocked_job
+        if release_concurrency_lock
+          release_next_blocked_job
+        end
+      end
+
+      private
+        def acquire_concurrency_lock
+          return true unless concurrency_limited?
+
+          Semaphore.wait(self)
+        end
+
+        def release_concurrency_lock
+          return false unless concurrency_limited?
+
+          Semaphore.signal(self)
+        end
+
+        def block
+          BlockedExecution.create_or_find_by!(job_id: id)
+        end
+
+        def release_next_blocked_job
+          BlockedExecution.release_one(concurrency_key)
+        end
+
+        def concurrency_limited?
+          concurrency_key.present? && concurrency_limit.to_i > 0
+        end
+
+        def job_class
+          @job_class ||= class_name.safe_constantize
+        end
+    end
+  end
+end

Diff for: app/models/solid_queue/job/executable.rb

@@ -1,64 +1,85 @@
 module SolidQueue
-  module Job::Executable
-    extend ActiveSupport::Concern
+  class Job
+    module Executable
+      extend ActiveSupport::Concern
 
-    included do
-      has_one :ready_execution, dependent: :destroy
-      has_one :claimed_execution, dependent: :destroy
-      has_one :failed_execution, dependent: :destroy
+      included do
+        include ConcurrencyControls
 
-      has_one :scheduled_execution, dependent: :destroy
+        has_one :ready_execution, dependent: :destroy
+        has_one :claimed_execution, dependent: :destroy
+        has_one :failed_execution, dependent: :destroy
 
-      after_create :prepare_for_execution
+        has_one :scheduled_execution, dependent: :destroy
 
-      scope :finished, -> { where.not(finished_at: nil) }
-    end
-
-    STATUSES = %w[ ready claimed failed scheduled ]
+        after_create :prepare_for_execution
 
-    STATUSES.each do |status|
-      define_method("#{status}?") { public_send("#{status}_execution").present? }
-    end
+        scope :finished, -> { where.not(finished_at: nil) }
+      end
 
-    def prepare_for_execution
-      if due?
-        ReadyExecution.create_or_find_by!(job_id: id)
-      else
-        ScheduledExecution.create_or_find_by!(job_id: id)
+      %w[ ready claimed failed scheduled ].each do |status|
+        define_method("#{status}?") { public_send("#{status}_execution").present? }
       end
-    end
 
-    def finished!
-      if delete_finished_jobs?
-        destroy!
-      else
-        touch(:finished_at)
+      def prepare_for_execution
+        if due? then dispatch
+        else
+          schedule
+        end
       end
-    end
 
-    def finished?
-      finished_at.present?
-    end
+      def finished!
+        if delete_finished_jobs?
+          destroy!
+        else
+          touch(:finished_at)
+        end
+      end
 
-    def failed_with(exception)
-      FailedExecution.create_or_find_by!(job_id: id, exception: exception)
-    end
+      def finished?
+        finished_at.present?
+      end
 
-    def discard
-      destroy unless claimed?
-    end
+      def failed_with(exception)
+        FailedExecution.create_or_find_by!(job_id: id, exception: exception)
+      end
 
-    def retry
-      failed_execution&.retry
-    end
+      def discard
+        destroy unless claimed?
+      end
 
-    private
-      def due?
-        scheduled_at.nil? || scheduled_at <= Time.current
+      def retry
+        failed_execution&.retry
       end
 
-      def delete_finished_jobs?
-        SolidQueue.delete_finished_jobs
+      def failed_with(exception)
+        FailedExecution.create_or_find_by!(job_id: id, exception: exception)
       end
+
+      private
+        def due?
+          scheduled_at.nil? || scheduled_at <= Time.current
+        end
+
+        def dispatch
+          if acquire_concurrency_lock then ready
+          else
+            block
+          end
+        end
+
+        def schedule
+          ScheduledExecution.create_or_find_by!(job_id: id)
+        end
+
+        def ready
+          ReadyExecution.create_or_find_by!(job_id: id)
+        end
+
+
+        def delete_finished_jobs?
+          SolidQueue.delete_finished_jobs
+        end
+    end
   end
 end

Diff for: app/models/solid_queue/ready_execution.rb

@@ -1,7 +1,6 @@
 module SolidQueue
   class ReadyExecution < Execution
     scope :queued_as, ->(queue_name) { where(queue_name: queue_name) }
-    scope :ordered, -> { order(priority: :asc, job_id: :asc) }
 
     assume_attributes_from_job
 

Diff for: app/models/solid_queue/scheduled_execution.rb

@@ -1,5 +1,5 @@
 class SolidQueue::ScheduledExecution < SolidQueue::Execution
-  scope :due, -> { where("scheduled_at <= ?", Time.current) }
+  scope :due, -> { where(scheduled_at: ..Time.current) }
   scope :ordered, -> { order(scheduled_at: :asc, priority: :asc) }
   scope :next_batch, ->(batch_size) { due.ordered.limit(batch_size) }
 
@@ -10,7 +10,7 @@ def prepare_batch(batch)
       prepared_at = Time.current
 
       rows = batch.map do |scheduled_execution|
-        scheduled_execution.execution_ready_attributes.merge(created_at: prepared_at)
+        scheduled_execution.ready_attributes.merge(created_at: prepared_at)
       end
 
       if rows.any?
@@ -23,8 +23,4 @@ def prepare_batch(batch)
       SolidQueue.logger.info("[SolidQueue] Prepared scheduled batch with #{rows.size} jobs at #{prepared_at}")
     end
   end
-
-  def execution_ready_attributes
-    attributes.slice("job_id", "queue_name", "priority")
-  end
 end

Diff for: app/models/solid_queue/semaphore.rb

@@ -0,0 +1,63 @@
+class SolidQueue::Semaphore < SolidQueue::Record
+  scope :available, -> { where("value > 0") }
+  scope :expired, -> { where(expires_at: ...Time.current) }
+
+  class << self
+    def wait(job)
+      Proxy.new(job, self).wait
+    end
+
+    def signal(job)
+      Proxy.new(job, self).signal
+    end
+  end
+
+  class Proxy
+    def initialize(job, proxied_class)
+      @job = job
+      @proxied_class = proxied_class
+    end
+
+    def wait
+      if semaphore = proxied_class.find_by(key: key)
+        semaphore.value > 0 && attempt_decrement
+      else
+        attempt_creation
+      end
+    end
+
+    def signal
+      attempt_increment
+    end
+
+    private
+      attr_reader :job, :proxied_class
+
+      def attempt_creation
+        proxied_class.create!(key: key, value: limit - 1, expires_at: expires_at)
+        true
+      rescue ActiveRecord::RecordNotUnique
+        attempt_decrement
+      end
+
+      def attempt_decrement
+        proxied_class.available.where(key: key).update_all([ "value = value - 1, expires_at = ?", expires_at ]) > 0
+      end
+
+      def attempt_increment
+        proxied_class.where(key: key, value: ...limit).update_all([ "value = value + 1, expires_at = ?", expires_at ]) > 0
+      end
+
+      def key
+        job.concurrency_key
+      end
+
+      def expires_at
+        job.concurrency_duration.from_now
+      end
+
+      def limit
+        job.concurrency_limit
+      end
+  end
+end