From 3f1c22273b1f5f8f934ed682e9c54f16e2598d97 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Wed, 1 May 2024 01:37:27 +0300
Subject: [PATCH] Do not run multiple background schema migrations on the same
 table at the same time

---
 CHANGELOG.md                                  |  2 +
 .../background_schema_migrations/scheduler.rb | 18 ++++++-
 .../scheduler_test.rb                         | 53 +++++++++++++++----
 3 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c49e1e0..373873f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## master (unreleased)
 
+- Do not run multiple background schema migrations on the same table at the same time
+
 ## 0.17.1 (2024-04-28)
 
 - Fix raising in development when using sharding and background index creation/removal was not enqueued
diff --git a/lib/online_migrations/background_schema_migrations/scheduler.rb b/lib/online_migrations/background_schema_migrations/scheduler.rb
index 4e39502..33c5a3f 100644
--- a/lib/online_migrations/background_schema_migrations/scheduler.rb
+++ b/lib/online_migrations/background_schema_migrations/scheduler.rb
@@ -3,7 +3,8 @@
 module OnlineMigrations
   module BackgroundSchemaMigrations
     # Class responsible for scheduling background schema migrations.
-    # It selects a single migration and runs it if there is no currently running migration.
+    # It selects a single migration and runs it if there is no currently
+    # running migration on the same table.
     #
     # Scheduler should be configured to run periodically, for example, via cron.
     # @example Run via whenever
@@ -19,12 +20,25 @@ def self.run
 
       # Runs Scheduler
       def run
-        migration = Migration.runnable.enqueued.queue_order.first || Migration.retriable.queue_order.first
+        migration = find_migration
         if migration
           runner = MigrationRunner.new(migration)
           runner.run
         end
       end
+
+      private
+        def find_migration
+          active_migrations = Migration.running.select(:table_name, :shard).to_a
+          runnable_migrations = Migration.runnable.enqueued.queue_order.to_a + Migration.retriable.queue_order.to_a
+
+          runnable_migrations.find do |runnable_migration|
+            active_migrations.none? do |active_migration|
+              active_migration.shard == runnable_migration.shard &&
+                active_migration.table_name == runnable_migration.table_name
+            end
+          end
+        end
     end
   end
 end
diff --git a/test/background_schema_migrations/scheduler_test.rb b/test/background_schema_migrations/scheduler_test.rb
index 6403420..20ad3f6 100644
--- a/test/background_schema_migrations/scheduler_test.rb
+++ b/test/background_schema_migrations/scheduler_test.rb
@@ -10,7 +10,12 @@ def teardown
     end
 
     def test_run
-      m = create_migration
+      m = create_migration(
+        name: "index_dogs_on_name",
+        table_name: "dogs",
+        definition: 'CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "index_dogs_on_name" ON "dogs" ("name")',
+        connection_class_name: "ShardRecord"
+      )
       child1, child2, child3 = m.children.to_a
 
       scheduler = OnlineMigrations::BackgroundSchemaMigrations::Scheduler.new
@@ -28,7 +33,12 @@ def test_run
     end
 
     def test_run_retries_failed_migrations
-      m = create_migration
+      m = create_migration(
+        name: "index_dogs_on_name",
+        table_name: "dogs",
+        definition: 'CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "index_dogs_on_name" ON "dogs" ("name")',
+        connection_class_name: "ShardRecord"
+      )
       child = m.children.first
 
       scheduler = OnlineMigrations::BackgroundSchemaMigrations::Scheduler.new
@@ -46,14 +56,39 @@ def test_run_retries_failed_migrations
       assert m.children.all?(&:succeeded?)
     end
 
+    def test_run_when_on_the_same_table_already_running
+      connection = ActiveRecord::Base.connection
+      connection.create_table(:users, force: true) do |t|
+        t.string :email
+        t.string :name
+      end
+
+      m1 = create_migration(
+        name: "index_users_on_email",
+        table_name: "users",
+        definition: 'CREATE UNIQUE INDEX CONCURRENTLY "index_users_on_email" ON "users" ("email")'
+      )
+      m1.update_column(:status, :running) # emulate running migration
+
+      _m2 = create_migration(
+        name: "index_users_on_name",
+        table_name: "users",
+        definition: 'CREATE INDEX CONCURRENTLY "index_users_on_name" ON "users" ("name")'
+      )
+
+      assert_equal 1, OnlineMigrations::BackgroundSchemaMigrations::Migration.running.count
+
+      scheduler = OnlineMigrations::BackgroundSchemaMigrations::Scheduler.new
+      scheduler.run
+
+      assert_equal 1, OnlineMigrations::BackgroundSchemaMigrations::Migration.running.count
+    ensure
+      connection.drop_table(:users)
+    end
+
     private
-      def create_migration
-        ActiveRecord::Base.connection.create_background_schema_migration(
-          "index_dogs_on_name",
-          "dogs",
-          definition: 'CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "index_dogs_on_name" ON "dogs" ("name")',
-          connection_class_name: "ShardRecord"
-        )
+      def create_migration(name:, table_name:, **options)
+        ActiveRecord::Base.connection.create_background_schema_migration(name, table_name, **options)
       end
   end
 end