Merge pull request #31900 from eileencodes/parallel-testing

Parallel testing

2 files changed, 188 insertions, 0 deletions

diff --git a/activesupport/lib/active_support/test_case.rb b/activesupport/lib/active_support/test_case.rb
index d1f7e6ea09..3588a5ac02 100644
--- a/activesupport/lib/active_support/test_case.rb
+++ b/activesupport/lib/active_support/test_case.rb
@@ -11,6 +11,7 @@ require "active_support/testing/isolation"
 require "active_support/testing/constant_lookup"
 require "active_support/testing/time_helpers"
 require "active_support/testing/file_fixtures"
+require "active_support/testing/parallelization"
 
 module ActiveSupport
   class TestCase < ::Minitest::Test
@@ -39,6 +40,91 @@ module ActiveSupport
       def test_order
         ActiveSupport.test_order ||= :random
       end
+
+      # Parallelizes the test suite.
+      #
+      # Takes a `workers` argument that controls how many times the process
+      # is forked. For each process a new database will be created suffixed
+      # with the worker number.
+      #
+      #   test-database-0
+      #   test-database-1
+      #
+      # If `ENV["PARALLEL_WORKERS"]` is set the workers argument will be ignored
+      # and the environment variable will be used instead. This is useful for CI
+      # environments, or other environments where you may need more workers than
+      # you do for local testing.
+      #
+      # If the number of workers is set to `1` or fewer, the tests will not be
+      # parallelized.
+      #
+      # The default parallelization method is to fork processes. If you'd like to
+      # use threads instead you can pass `with: :threads` to the `parallelize`
+      # method. Note the threaded parallelization does not create multiple
+      # database and will not work with system tests at this time.
+      #
+      #   parallelize(workers: 2, with: :threads)
+      #
+      # The threaded parallelization uses Minitest's parallel exector directly.
+      # The processes paralleliztion uses a Ruby Drb server.
+      def parallelize(workers: 2, with: :processes)
+        workers = ENV["PARALLEL_WORKERS"].to_i if ENV["PARALLEL_WORKERS"]
+
+        return if workers <= 1
+
+        executor = case with
+                   when :processes
+                     Testing::Parallelization.new(workers)
+                   when :threads
+                     Minitest::Parallel::Executor.new(workers)
+                   else
+                     raise ArgumentError, "#{with} is not a supported parallelization exectutor."
+        end
+
+        self.lock_threads = false if defined?(self.lock_threads) && with == :threads
+
+        Minitest.parallel_executor = executor
+
+        parallelize_me!
+      end
+
+      # Set up hook for parallel testing. This can be used if you have multiple
+      # databases or any behavior that needs to be run after the process is forked
+      # but before the tests run.
+      #
+      # Note: this feature is not available with the threaded parallelization.
+      #
+      # In your +test_helper.rb+ add the following:
+      #
+      #   class ActiveSupport::TestCase
+      #     parallelize_setup do
+      #       # create databases
+      #     end
+      #   end
+      def parallelize_setup(&block)
+        ActiveSupport::Testing::Parallelization.after_fork_hook do |worker|
+          yield worker
+        end
+      end
+
+      # Clean up hook for parallel testing. This can be used to drop databases
+      # if your app uses multiple write/read databases or other clean up before
+      # the tests finish. This runs before the forked process is closed.
+      #
+      # Note: this feature is not available with the threaded parallelization.
+      #
+      # In your +test_helper.rb+ add the following:
+      #
+      #   class ActiveSupport::TestCase
+      #     parallelize_teardown do
+      #       # drop databases
+      #     end
+      #   end
+      def parallelize_teardown(&block)
+        ActiveSupport::Testing::Parallelization.run_cleanup_hook do |worker|
+          yield worker
+        end
+      end
     end
 
     alias_method :method_name, :name
diff --git a/activesupport/lib/active_support/testing/parallelization.rb b/activesupport/lib/active_support/testing/parallelization.rb
new file mode 100644
index 0000000000..59c8486f41
--- /dev/null
+++ b/activesupport/lib/active_support/testing/parallelization.rb
@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require "drb"
+require "drb/unix"
+
+module ActiveSupport
+  module Testing
+    class Parallelization # :nodoc:
+      class Server
+        include DRb::DRbUndumped
+
+        def initialize
+          @queue = Queue.new
+        end
+
+        def record(reporter, result)
+          reporter.synchronize do
+            reporter.record(result)
+          end
+        end
+
+        def <<(o)
+          @queue << o
+        end
+
+        def pop; @queue.pop; end
+      end
+
+      @after_fork_hooks = []
+
+      def self.after_fork_hook(&blk)
+        @after_fork_hooks << blk
+      end
+
+      def self.after_fork_hooks
+        @after_fork_hooks
+      end
+
+      @run_cleanup_hooks = []
+
+      def self.run_cleanup_hook(&blk)
+        @run_cleanup_hooks << blk
+      end
+
+      def self.run_cleanup_hooks
+        @run_cleanup_hooks
+      end
+
+      def initialize(queue_size)
+        @queue_size = queue_size
+        @queue      = Server.new
+        @pool       = []
+
+        @url = DRb.start_service("drbunix:", @queue).uri
+      end
+
+      def after_fork(worker)
+        self.class.after_fork_hooks.each do |cb|
+          cb.call(worker)
+        end
+      end
+
+      def run_cleanup(worker)
+        self.class.run_cleanup_hooks.each do |cb|
+          cb.call(worker)
+        end
+      end
+
+      def start
+        @pool = @queue_size.times.map do |worker|
+          fork do
+            DRb.stop_service
+
+            after_fork(worker)
+
+            queue = DRbObject.new_with_uri(@url)
+
+            while job = queue.pop
+              klass    = job[0]
+              method   = job[1]
+              reporter = job[2]
+              result   = Minitest.run_one_method(klass, method)
+
+              queue.record(reporter, result)
+            end
+
+            run_cleanup(worker)
+          end
+        end
+      end
+
+      def <<(work)
+        @queue << work
+      end
+
+      def shutdown
+        @queue_size.times { @queue << nil }
+        @pool.each { |pid| Process.waitpid pid }
+      end
+    end
+  end
+end