Give concurrency slots back only to the semaphore they came from

UPGRADING.md

@@ -1,3 +1,10 @@
+# Upgrading to version 1.7.2
+Jobs with [concurrency controls](README.md#concurrency-controls) now remember which semaphore they got their slot from, in their serialized job data, and give the slot back only to that semaphore. Before, a job that ran past its `duration` gave its slot back to the semaphore a later job had created after the job's own one expired, so one more job than the limit allowed could run at the same time, and keep doing so until the jobs waiting for that key ran out.
+
+There's nothing to change or migrate. Jobs that got their slot before you upgrade have nothing recorded and give it back as before, so old and new versions can run side by side while you deploy.
+
+Solid Queue 2.0 will stop giving back slots that weren't recorded. Before upgrading to 2.0, make sure the concurrency-limited jobs that were ready or running when you upgraded to 1.7.2 have finished. Otherwise their slots won't be given back: a slot like that only comes back when its semaphore expires, and every job that takes or gives back a slot for the same key pushes that expiry out by another `duration`, so on a busy key with a limit above 1 it can stay lost for as long as the key stays busy.
+
 # Upgrading to version 1.7.x
 This version introduces support for grouping jobs into batches, which needs new tables. Fresh installs get them with the base schema; existing installations need to copy the migration that adds them and run it:
 

app/models/solid_queue/blocked_execution.rb

@@ -51,7 +51,7 @@ def release
             destroy!
 
             payload[:failed] = true
-          elsif acquire_concurrency_lock
+          elsif job.acquire_concurrency_lock
             promote_to_ready
             destroy!
 
@@ -63,12 +63,7 @@ def release
 
     private
       def set_expires_at
-        duration = job.job_class ? job.concurrency_duration : SolidQueue.default_concurrency_control_period
-        self.expires_at = duration.from_now
-      end
-
-      def acquire_concurrency_lock
-        Semaphore.wait(job)
+        self.expires_at = job.concurrency_duration.from_now
       end
 
       def promote_to_ready

app/models/solid_queue/job/concurrency_controls.rb

@@ -5,11 +5,12 @@ class Job
     module ConcurrencyControls
       extend ActiveSupport::Concern
 
+      # Stored in the job's serialized Active Job data, which Active Job ignores when running the job
+      CONCURRENCY_LOCK_KEY = "solid_queue_concurrency_lock"
+
       included do
         has_one :blocked_execution
 
-        delegate :concurrency_limit, :concurrency_duration, to: :job_class
-
         before_destroy :unblock_next_blocked_job, if: -> { concurrency_limited? && ready? }
       end
 
@@ -25,8 +26,39 @@ def unblock_next_blocked_job
         end
       end
 
+      def acquire_concurrency_lock
+        return true unless concurrency_limited?
+
+        if job_class.nil?
+          # A job whose class doesn't load skips its limits. Record that it holds no slot,
+          # so it doesn't give one back if it runs somewhere its class does load.
+          record_concurrency_lock(semaphore_id: nil)
+        elsif semaphore_id = Semaphore.wait(self)
+          record_concurrency_lock(semaphore_id:, limit: concurrency_limit)
+        end
+      end
+
       def concurrency_limited?
-        concurrency_key.present? && job_class.present?
+        concurrency_key.present?
+      end
+
+      # A job that kept its concurrency key after its class dropped limits_concurrency gets a limit of 1
+      def concurrency_limit
+        job_class.concurrency_limit || 1 if concurrency_limited?
+      end
+
+      def concurrency_duration
+        job_class ? job_class.concurrency_duration : SolidQueue.default_concurrency_control_period
+      end
+
+      # The semaphore the job got its concurrency slot from
+      def acquired_semaphore_id
+        concurrency_lock&.dig("semaphore_id")
+      end
+
+      # The concurrency limit the job had when it got its slot
+      def acquired_concurrency_limit
+        concurrency_lock&.dig("limit")
       end
 
       def blocked?
@@ -38,20 +70,35 @@ def job_class
       end
 
       private
-        def concurrency_on_conflict
-          job_class.concurrency_on_conflict.to_s.inquiry
+        def release_concurrency_lock
+          holds_concurrency_slot? && Semaphore.signal(self)
         end
 
-        def acquire_concurrency_lock
-          return true unless concurrency_limited?
+        def holds_concurrency_slot?
+          if concurrency_lock_recorded?
+            acquired_semaphore_id.present?
+          else
+            # Jobs dispatched before Solid Queue recorded their semaphore, which need their class to know their limit
+            concurrency_limited? && job_class.present?
+          end
+        end
 
-          Semaphore.wait(self)
+        def concurrency_lock_recorded?
+          arguments.key?(CONCURRENCY_LOCK_KEY)
         end
 
-        def release_concurrency_lock
-          return false unless concurrency_limited?
+        def release_next_blocked_job
+          BlockedExecution.release_one(concurrency_key)
+        end
+
+        def record_concurrency_lock(semaphore_id:, limit: nil)
+          arguments[CONCURRENCY_LOCK_KEY] = { "semaphore_id" => semaphore_id, "limit" => limit }
+          update_column :arguments, arguments
+          true
+        end
 
-          Semaphore.signal(self)
+        def concurrency_lock
+          arguments[CONCURRENCY_LOCK_KEY]
         end
 
         def handle_concurrency_conflict
@@ -62,12 +109,12 @@ def handle_concurrency_conflict
           end
         end
 
-        def block
-          BlockedExecution.create_or_find_by!(job: self)
+        def concurrency_on_conflict
+          job_class.concurrency_on_conflict.to_s.inquiry
         end
 
-        def release_next_blocked_job
-          BlockedExecution.release_one(concurrency_key)
+        def block
+          BlockedExecution.create_or_find_by!(job: self)
         end
 
         def execution

app/models/solid_queue/semaphore.rb

@@ -31,9 +31,10 @@ def initialize(job)
         @job = job
       end
 
+      # Returns the id of the semaphore the job got a slot on, or nil if there was none free
       def wait
         if semaphore = Semaphore.lock.find_by(key: key)
-          semaphore.value > 0 && attempt_decrement
+          semaphore.id if semaphore.value > 0 && attempt_decrement
         else
           attempt_creation
         end
@@ -46,24 +47,46 @@ def signal
       private
         attr_accessor :job
 
+        def attempt_decrement
+          Semaphore.available.where(key: key).update_all([ "value = value - 1, expires_at = ?", expires_at ]) > 0
+        end
+
         def attempt_creation
           if Semaphore.create_unique_by(key: key, value: limit - 1, expires_at: expires_at)
-            true
+            acquired_semaphore_id
           else
             check_limit_or_decrement
           end
         end
 
-        def check_limit_or_decrement
-          limit == 1 ? false : attempt_decrement
+        # A locking read, so it sees the latest committed rows rather than the transaction's
+        # snapshot, which on MySQL can still include a semaphore that expired and was deleted
+        def acquired_semaphore_id
+          Semaphore.lock.where(key: key).pick(:id)
         end
 
-        def attempt_decrement
-          Semaphore.available.where(key: key).update_all([ "value = value - 1, expires_at = ?", expires_at ]) > 0
+        def check_limit_or_decrement
+          acquired_semaphore_id if limit > 1 && attempt_decrement
         end
 
         def attempt_increment
-          Semaphore.where(key: key, value: ...limit).update_all([ "value = value + 1, expires_at = ?", expires_at ]) > 0
+          semaphore_to_release.where(value: ...limit_to_release).update_all([ "value = value + 1, expires_at = ?", expires_at ]) > 0
+        end
+
+        # A job only gives its slot back to the semaphore it got it from. If that one expired
+        # and was deleted, a semaphore for the same key now belongs to other jobs.
+        # Jobs dispatched before Solid Queue recorded their semaphore give it back by key.
+        def semaphore_to_release
+          if job.acquired_semaphore_id
+            Semaphore.where(id: job.acquired_semaphore_id)
+          else
+            Semaphore.where(key: key)
+          end
+        end
+
+        # Recorded when the job got its slot, so it's known even if the job's class no longer loads
+        def limit_to_release
+          job.acquired_concurrency_limit || limit
         end
 
         def key
@@ -75,7 +98,7 @@ def expires_at
         end
 
         def limit
-          job.concurrency_limit || 1
+          job.concurrency_limit
         end
     end
   end

test/integration/concurrency_controls_test.rb

@@ -116,7 +116,7 @@ class ConcurrencyControlsTest < ActiveSupport::TestCase
 
   test "rely on dispatcher to unblock blocked executions with an available semaphore" do
     # Simulate a scenario where we got an available semaphore and some stuck jobs
-    job = NonOverlappingUpdateResultJob.perform_later(@result, name: "A")
+    job = SolidQueue::Job.find_by!(active_job_id: NonOverlappingUpdateResultJob.perform_later(@result, name: "A").job_id)
 
     wait_for_jobs_to_finish_for(5.seconds)
     assert_no_unfinished_jobs
@@ -148,7 +148,7 @@ class ConcurrencyControlsTest < ActiveSupport::TestCase
 
   test "rely on dispatcher to unblock blocked executions with an expired semaphore" do
     # Simulate a scenario where we got an available semaphore and some stuck jobs
-    job = NonOverlappingUpdateResultJob.perform_later(@result, name: "A")
+    job = SolidQueue::Job.find_by!(active_job_id: NonOverlappingUpdateResultJob.perform_later(@result, name: "A").job_id)
     wait_for_jobs_to_finish_for(5.seconds)
     assert_no_unfinished_jobs
 

test/models/solid_queue/claimed_execution_test.rb

@@ -170,7 +170,7 @@ class SolidQueue::ClaimedExecutionTest < ActiveSupport::TestCase
   test "dispatch job with missing class and concurrency key skips concurrency controls" do
     job = create_job_with_missing_class(concurrency_key: "test_key")
 
-    assert_not job.concurrency_limited?
+    assert_nil job.job_class
 
     job.prepare_for_execution