Add shard_selection_lambda

Author: alxckn

lib/active_job/queue_adapters/solid_queue_adapter.rb

@@ -17,23 +17,26 @@ def enqueue_after_transaction_commit?
       end
 
       def enqueue(active_job) # :nodoc:
-        select_shard { SolidQueue::Job.enqueue(active_job) }
+        select_shard(active_job:) { SolidQueue::Job.enqueue(active_job) }
       end
 
       def enqueue_at(active_job, timestamp) # :nodoc:
-        select_shard do
+        select_shard(active_job:) do
           SolidQueue::Job.enqueue(active_job, scheduled_at: Time.at(timestamp))
         end
       end
 
       def enqueue_all(active_jobs) # :nodoc:
-        select_shard { SolidQueue::Job.enqueue_all(active_jobs) }
+        select_shard(active_jobs:) { SolidQueue::Job.enqueue_all(active_jobs) }
       end
 
       private
 
-      def select_shard(&block)
-        shard = @db_shard || SolidQueue.primary_shard
+      def select_shard(active_job: nil, active_jobs: nil, &block)
+        shard =
+          SolidQueue.shard_selection_lambda&.call(active_job:, active_jobs:) ||
+            @db_shard ||
+              SolidQueue.primary_shard
 
         if shard
           ActiveRecord::Base.connected_to(shard: shard) { block.call }

lib/solid_queue.rb

@@ -41,7 +41,7 @@ module SolidQueue
   mattr_accessor :clear_finished_jobs_after, default: 1.day
   mattr_accessor :default_concurrency_control_period, default: 3.minutes
 
-  mattr_accessor :primary_shard, :active_shard
+  mattr_accessor :primary_shard, :active_shard, :shard_selection_lambda
 
   delegate :on_start, :on_stop, to: Supervisor
 

test/unit/multisharding_test.rb

@@ -16,6 +16,16 @@ class MultishardingTest < ActiveSupport::TestCase
     end
   end
 
+  test "shard_selection_lambda can override which shard is used to enqueue individual jobs" do
+    shard_selection_lambda = ->(active_job:, active_jobs:) { :queue_shard_two if active_job.arguments.first == "hey!" }
+
+    with_shard_selection_lambda(shard_selection_lambda) do
+      assert_difference -> { connected_to_shard_two { SolidQueue::Job.count } }, 1 do
+        AddToBufferJob.perform_later "hey!"
+      end
+    end
+  end
+
   test "jobs are enqueued for later in the right shard" do
     assert_difference -> { SolidQueue::ScheduledExecution.count }, 1 do
       assert_difference -> { connected_to_shard_two { SolidQueue::ScheduledExecution.count } }, 1 do
@@ -39,4 +49,31 @@ class MultishardingTest < ActiveSupport::TestCase
       end
     end
   end
+
+  test "shard_selection_lambda can override which shard is used to enqueue jobs in bulk" do
+    active_jobs = [
+      AddToBufferJob.new(2),
+      ShardTwoJob.new(6),
+      AddToBufferJob.new(3),
+      ShardTwoJob.new(7)
+    ]
+    shard_selection_lambda = ->(active_job:, active_jobs:) { :queue_shard_two if active_jobs.size == 2 }
+
+    with_shard_selection_lambda(shard_selection_lambda) do
+      assert_difference -> { SolidQueue::Job.count }, 0 do
+        assert_difference -> { connected_to_shard_two { SolidQueue::Job.count } }, 4 do
+          ActiveJob.perform_all_later(active_jobs)
+        end
+      end
+    end
+  end
+
+  private
+
+  def with_shard_selection_lambda(lambda, &block)
+    SolidQueue.shard_selection_lambda = lambda
+    block.call
+  ensure
+    SolidQueue.shard_selection_lambda = nil
+  end
 end