DEV: Resupport running multisite:migrate concurrently (#36853)

Previously, we were using multiple threads to run migrations for
multisite setups in parallel. However, it was discovered that
`db:migrate` is
not threadsafe and was not designed to be. As a result,
`multisite:migrate` was hardcoded to only support spinning up a single
thread.

This new attempt Uses the `parallel` gem with `in_processes` to run
migrations concurrently using multiple processes. Concurrency defaults
to 2 but can be configured via the
`DISCOURSE_MULTISITE_MIGRATE_CONCURRENCY` env. Example:
`DISCOURSE_MULTISITE_MIGRATE_CONCURRENCY=5 bin/rails multisite:migrate`

db/fixtures/600_themes.rb

@@ -6,7 +6,7 @@ SystemThemesManager.sync!
 
 # we can not guess what to do if customization already started, so skip it
 if !theme_exists
-  STDERR.puts "> Seeding theme and color schemes"
+  $stdout.puts "Seeding theme and color schemes"
 
   color_schemes = [
     { name: I18n.t("color_schemes.dark"), base_scheme_name: "Dark" },

lib/tasks/db.rake

@@ -86,33 +86,6 @@ end
 # this ensures we can run migrations concurrently to save huge amounts of time
 Rake::Task["multisite:migrate"].clear
 
-class StdOutDemux
-  def initialize(stdout)
-    @stdout = stdout
-    @data = {}
-  end
-
-  def write(data)
-    (@data[Thread.current] ||= +"") << data
-  end
-
-  def close
-    finish_chunk
-  end
-
-  def finish_chunk
-    data = @data[Thread.current]
-    if data
-      @stdout.write(data)
-      @data.delete Thread.current
-    end
-  end
-
-  def flush
-    # Do nothing
-  end
-end
-
 class SeedHelper
   def self.paths
     DiscoursePluginRegistry.seed_paths
@@ -128,106 +101,91 @@ class SeedHelper
   end
 end
 
+def execute_db_migration
+  ActiveRecord::Tasks::DatabaseTasks.migrate
+  Rake::Task["db:seed"].invoke if ENV["SKIP_SEED_FU"] != "1"
+
+  if !Discourse.skip_post_deployment_migrations? && ENV["SKIP_OPTIMIZE_ICONS"] != "1"
+    SiteIconManager.ensure_optimized!
+  end
+end
+
 task "multisite:migrate" => %w[
        db:load_config
        environment
        set_locale
        assets:precompile:asset_processor
      ] do |_, args|
-  raise "Multisite migrate is only supported in production" if ENV["RAILS_ENV"] != "production"
-
   DistributedMutex.synchronize(
     "db_migration",
     redis: Discourse.redis.without_namespace,
     validity: 1200,
   ) do
-    # TODO: Switch to processes for concurrent migrations because Rails migration
-    # is not thread safe by default.
-    concurrency = 1
+    concurrency = ENV["DISCOURSE_MULTISITE_MIGRATE_CONCURRENCY"]&.to_i || 2
 
-    puts "Multisite migrator is running using #{concurrency} threads"
-    puts
+    puts "Multisite migrator is running using #{concurrency} processes"
 
-    exceptions = Queue.new
+    databases = []
+    RailsMultisite::ConnectionManagement.each_connection { |db| databases << db }
 
-    if concurrency > 1
-      old_stdout = $stdout
-      $stdout = StdOutDemux.new($stdout)
-    end
+    puts "Running migrations and seeds for #{databases.join(", ")} database(s)"
 
-    SeedFu.quiet = true
+    results =
+      Parallel.map(
+        databases,
+        in_processes: concurrency,
+        isolation: true,
+        finish_in_order: true,
+        finish:
+          lambda do |db, _index, result|
+            result[:stdout]&.lines&.each { |line| puts "[#{db}] #{line}" }
+          end,
+      ) do |db|
+        Discourse.after_fork
 
-    def execute_concurrently(concurrency, exceptions)
-      queue = Queue.new
+        output_buffer = StringIO.new
+        original_stdout = $stdout
+        $stdout = output_buffer
 
-      RailsMultisite::ConnectionManagement.each_connection { |db| queue << db }
+        error = nil
 
-      concurrency.times { queue << :done }
-
-      (1..concurrency)
-        .map do
-          Thread.new do
-            while true
-              db = queue.pop
-              break if db == :done
-
-              RailsMultisite::ConnectionManagement.with_connection(db) do
-                begin
-                  yield(db) if block_given?
-                rescue => e
-                  exceptions << [db, e]
-                ensure
-                  begin
-                    $stdout.finish_chunk if concurrency > 1
-                  rescue => ex
-                    STDERR.puts ex.inspect
-                    STDERR.puts ex.backtrace
-                  end
-                end
-              end
-            end
+        begin
+          RailsMultisite::ConnectionManagement.with_connection(db) do
+            puts "Processing"
+            puts "-" * 40
+            execute_db_migration
           end
+        rescue => e
+          error = e
         end
-        .each(&:join)
-    end
 
-    def check_exceptions(exceptions)
-      if exceptions.length > 0
-        STDERR.puts
-        STDERR.puts "-" * 80
-        STDERR.puts "#{exceptions.length} migrations failed!"
-        while !exceptions.empty?
-          db, e = exceptions.pop
-          STDERR.puts
-          STDERR.puts "Failed to migrate #{db}"
-          STDERR.puts e.inspect
-          STDERR.puts e.backtrace
-          STDERR.puts
+        if error.nil?
+          puts "Completed"
+          puts "-" * 40
         end
-        exit 1
+
+        { db: db, stdout: output_buffer.string, error: error }
+      ensure
+        $stdout = original_stdout
       end
-    end
 
-    execute_concurrently(concurrency, exceptions) do |db|
-      puts "Migrating #{db}"
-      ActiveRecord::Tasks::DatabaseTasks.migrate
-    end
+    errors = results.select { |r| r[:error] }
 
-    check_exceptions(exceptions)
+    if errors.any?
+      $stderr.puts
+      $stderr.puts "-" * 80
+      $stderr.puts "#{errors.length} database(s) failed!"
 
-    SeedFu.seed(SeedHelper.paths, /001_refresh/)
-
-    execute_concurrently(concurrency, exceptions) do |db|
-      puts "Seeding #{db}"
-      SeedFu.seed(SeedHelper.paths, SeedHelper.filter)
-
-      if !Discourse.skip_post_deployment_migrations? && ENV["SKIP_OPTIMIZE_ICONS"] != "1"
-        SiteIconManager.ensure_optimized!
+      errors.each do |result|
+        $stderr.puts
+        $stderr.puts "Failed to process #{result[:db]}"
+        $stderr.puts result[:error].inspect
+        $stderr.puts result[:error].backtrace
+        $stderr.puts
       end
-    end
 
-    $stdout = old_stdout if concurrency > 1
-    check_exceptions(exceptions)
+      raise errors.first[:error]
+    end
 
     Rake::Task["db:_dump"].invoke
   end
@@ -264,15 +222,8 @@ task "db:migrate" => %w[
       end
     end
 
-    ActiveRecord::Tasks::DatabaseTasks.migrate
-
-    Rake::Task["db:seed"].invoke if ENV["SKIP_SEED_FU"] != "1"
-
+    execute_db_migration
     Rake::Task["db:schema:cache:dump"].invoke if Rails.env.development? && !ENV["RAILS_DB"]
-
-    if !Discourse.skip_post_deployment_migrations? && ENV["SKIP_OPTIMIZE_ICONS"] != "1"
-      SiteIconManager.ensure_optimized!
-    end
   end
 
   if !Discourse.is_parallel_test? && MultisiteTestHelpers.load_multisite?