Configure jobs to be enqueued on different shards

Author: alxckn

lib/active_job/queue_adapters/solid_queue_adapter.rb

@@ -8,20 +8,36 @@ module QueueAdapters
     #
     #   Rails.application.config.active_job.queue_adapter = :solid_queue
     class SolidQueueAdapter
+      def initialize(db_shard: nil)
+        @db_shard = db_shard
+      end
+
       def enqueue_after_transaction_commit?
         true
       end
 
       def enqueue(active_job) # :nodoc:
-        SolidQueue::Job.enqueue(active_job)
+        select_shard { SolidQueue::Job.enqueue(active_job) }
       end
 
       def enqueue_at(active_job, timestamp) # :nodoc:
-        SolidQueue::Job.enqueue(active_job, scheduled_at: Time.at(timestamp))
+        select_shard do
+          SolidQueue::Job.enqueue(active_job, scheduled_at: Time.at(timestamp))
+        end
       end
 
       def enqueue_all(active_jobs) # :nodoc:
-        SolidQueue::Job.enqueue_all(active_jobs)
+        select_shard { SolidQueue::Job.enqueue_all(active_jobs) }
+      end
+
+      private
+
+      def select_shard
+        if @db_shard
+          ActiveRecord::Base.connected_to(shard: @db_shard) { yield }
+        else
+          yield
+        end
       end
     end
   end

test/dummy/app/jobs/shard_two_job.rb

@@ -0,0 +1,8 @@
+class ShardTwoJob < ApplicationJob
+  self.queue_adapter = ActiveJob::QueueAdapters::SolidQueueAdapter.new(db_shard: :queue_shard_two)
+  queue_as :background
+
+  def perform(arg)
+    JobBuffer.add(arg)
+  end
+end

test/dummy/config/database.yml

@@ -48,6 +48,10 @@ development:
     <<: *default
     database: <%= database_name_from("development_queue") %>
     migrations_paths: db/queue_migrate
+  queue_shard_two:
+    <<: *default
+    database: <%= database_name_from("development_queue_shard_two") %>
+    migrations_paths: db/queue_migrate
 
 test:
   primary:
@@ -65,3 +69,7 @@ test:
     <<: *default
     database: <%= database_name_from("test_queue") %>
     migrations_paths: db/queue_migrate
+  queue_shard_two:
+    <<: *default
+    database: <%= database_name_from("test_queue_shard_two") %>
+    migrations_paths: db/queue_migrate

test/dummy/config/environments/test.rb

@@ -49,7 +49,12 @@
 
   # Replace the default in-process and non-durable queuing backend for Active Job.
   config.active_job.queue_adapter = :solid_queue
-  config.solid_queue.connects_to = { database: { writing: :queue } }
+  config.solid_queue.connects_to = {
+    shards: {
+      queue: { writing: :queue },
+      queue_shard_two:  { writing: :queue_shard_two }
+    }
+  }
 
   # Annotate rendered view with file names.
   # config.action_view.annotate_rendered_view_with_filenames = true

test/dummy/db/queue_shard_two_schema.rb

@@ -0,0 +1,141 @@
+# This file is auto-generated from the current state of the database. Instead
+# of editing this file, please use the migrations feature of Active Record to
+# incrementally modify your database, and then regenerate this schema definition.
+#
+# This file is the source Rails uses to define your schema when running `bin/rails
+# db:schema:load`. When creating a new database, `bin/rails db:schema:load` tends to
+# be faster and is potentially less error prone than running all of your
+# migrations from scratch. Old migrations may fail to apply correctly if those
+# migrations use external dependencies or application code.
+#
+# It's strongly recommended that you check this file into your version control system.
+
+ActiveRecord::Schema[7.1].define(version: 1) do
+  create_table "solid_queue_blocked_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.string "queue_name", null: false
+    t.integer "priority", default: 0, null: false
+    t.string "concurrency_key", null: false
+    t.datetime "expires_at", null: false
+    t.datetime "created_at", null: false
+    t.index ["concurrency_key", "priority", "job_id"], name: "index_solid_queue_blocked_executions_for_release"
+    t.index ["expires_at", "concurrency_key"], name: "index_solid_queue_blocked_executions_for_maintenance"
+    t.index ["job_id"], name: "index_solid_queue_blocked_executions_on_job_id", unique: true
+  end
+
+  create_table "solid_queue_claimed_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.bigint "process_id"
+    t.datetime "created_at", null: false
+    t.index ["job_id"], name: "index_solid_queue_claimed_executions_on_job_id", unique: true
+    t.index ["process_id", "job_id"], name: "index_solid_queue_claimed_executions_on_process_id_and_job_id"
+  end
+
+  create_table "solid_queue_failed_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.text "error"
+    t.datetime "created_at", null: false
+    t.index ["job_id"], name: "index_solid_queue_failed_executions_on_job_id", unique: true
+  end
+
+  create_table "solid_queue_jobs", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.string "queue_name", null: false
+    t.string "class_name", null: false
+    t.text "arguments"
+    t.integer "priority", default: 0, null: false
+    t.string "active_job_id"
+    t.datetime "scheduled_at"
+    t.datetime "finished_at"
+    t.string "concurrency_key"
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["active_job_id"], name: "index_solid_queue_jobs_on_active_job_id"
+    t.index ["class_name"], name: "index_solid_queue_jobs_on_class_name"
+    t.index ["finished_at"], name: "index_solid_queue_jobs_on_finished_at"
+    t.index ["queue_name", "finished_at"], name: "index_solid_queue_jobs_for_filtering"
+    t.index ["scheduled_at", "finished_at"], name: "index_solid_queue_jobs_for_alerting"
+  end
+
+  create_table "solid_queue_pauses", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.string "queue_name", null: false
+    t.datetime "created_at", null: false
+    t.index ["queue_name"], name: "index_solid_queue_pauses_on_queue_name", unique: true
+  end
+
+  create_table "solid_queue_processes", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.string "kind", null: false
+    t.datetime "last_heartbeat_at", null: false
+    t.bigint "supervisor_id"
+    t.integer "pid", null: false
+    t.string "hostname"
+    t.text "metadata"
+    t.datetime "created_at", null: false
+    t.string "name", null: false
+    t.index ["last_heartbeat_at"], name: "index_solid_queue_processes_on_last_heartbeat_at"
+    t.index ["name", "supervisor_id"], name: "index_solid_queue_processes_on_name_and_supervisor_id", unique: true
+    t.index ["supervisor_id"], name: "index_solid_queue_processes_on_supervisor_id"
+  end
+
+  create_table "solid_queue_ready_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.string "queue_name", null: false
+    t.integer "priority", default: 0, null: false
+    t.datetime "created_at", null: false
+    t.index ["job_id"], name: "index_solid_queue_ready_executions_on_job_id", unique: true
+    t.index ["priority", "job_id"], name: "index_solid_queue_poll_all"
+    t.index ["queue_name", "priority", "job_id"], name: "index_solid_queue_poll_by_queue"
+  end
+
+  create_table "solid_queue_recurring_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.string "task_key", null: false
+    t.datetime "run_at", null: false
+    t.datetime "created_at", null: false
+    t.index ["job_id"], name: "index_solid_queue_recurring_executions_on_job_id", unique: true
+    t.index ["task_key", "run_at"], name: "index_solid_queue_recurring_executions_on_task_key_and_run_at", unique: true
+  end
+
+  create_table "solid_queue_recurring_tasks", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.string "key", null: false
+    t.string "schedule", null: false
+    t.string "command", limit: 2048
+    t.string "class_name"
+    t.text "arguments"
+    t.string "queue_name"
+    t.integer "priority", default: 0
+    t.boolean "static", default: true, null: false
+    t.text "description"
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["key"], name: "index_solid_queue_recurring_tasks_on_key", unique: true
+    t.index ["static"], name: "index_solid_queue_recurring_tasks_on_static"
+  end
+
+  create_table "solid_queue_scheduled_executions", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.bigint "job_id", null: false
+    t.string "queue_name", null: false
+    t.integer "priority", default: 0, null: false
+    t.datetime "scheduled_at", null: false
+    t.datetime "created_at", null: false
+    t.index ["job_id"], name: "index_solid_queue_scheduled_executions_on_job_id", unique: true
+    t.index ["scheduled_at", "priority", "job_id"], name: "index_solid_queue_dispatch_all"
+  end
+
+  create_table "solid_queue_semaphores", charset: "utf8mb4", collation: "utf8mb4_0900_ai_ci", force: :cascade do |t|
+    t.string "key", null: false
+    t.integer "value", default: 1, null: false
+    t.datetime "expires_at", null: false
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["expires_at"], name: "index_solid_queue_semaphores_on_expires_at"
+    t.index ["key", "value"], name: "index_solid_queue_semaphores_on_key_and_value"
+    t.index ["key"], name: "index_solid_queue_semaphores_on_key", unique: true
+  end
+
+  add_foreign_key "solid_queue_blocked_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_claimed_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_failed_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_ready_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_recurring_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+  add_foreign_key "solid_queue_scheduled_executions", "solid_queue_jobs", column: "job_id", on_delete: :cascade
+end

test/unit/multisharding_test.rb

@@ -0,0 +1,51 @@
+require "test_helper"
+
+class MultishardingTest < ActiveSupport::TestCase
+  test "jobs are enqueued in the right shard" do
+    assert_difference -> { SolidQueue::Job.count }, 1 do
+      assert_difference -> do
+                          ActiveRecord::Base.connected_to(
+                            shard: :queue_shard_two
+                          ) { SolidQueue::Job.count }
+                        end,
+                        1 do
+        AddToBufferJob.perform_later "hey!"
+        ShardTwoJob.perform_later "coucou!"
+      end
+    end
+  end
+
+  test "jobs are enqueued for later in the right shard" do
+    assert_difference -> { SolidQueue::ScheduledExecution.count }, 1 do
+      assert_difference -> do
+                          ActiveRecord::Base.connected_to(
+                            shard: :queue_shard_two
+                          ) { SolidQueue::ScheduledExecution.count }
+                        end,
+                        1 do
+        AddToBufferJob.set(wait: 1).perform_later "hey!"
+        ShardTwoJob.set(wait: 1).perform_later "coucou!"
+      end
+    end
+  end
+
+  test "jobs are enqueued in bulk in the right shard" do
+    active_jobs = [
+      AddToBufferJob.new(2),
+      ShardTwoJob.new(6),
+      AddToBufferJob.new(3),
+      ShardTwoJob.new(7)
+    ]
+
+    assert_difference -> { SolidQueue::Job.count }, 2 do
+      assert_difference -> do
+                          ActiveRecord::Base.connected_to(
+                            shard: :queue_shard_two
+                          ) { SolidQueue::Job.count }
+                        end,
+                        2 do
+        ActiveJob.perform_all_later(active_jobs)
+      end
+    end
+  end
+end