class ActiveSupport::Testing::Parallelization

def self.after_fork_hook(&blk)
  @@after_fork_hooks << blk
end

def self.run_cleanup_hook(&blk)
  @@run_cleanup_hooks << blk
end

def <<(work)
  @queue << work
end

def add_setup_exception(result, setup_exception)
  result.failures.prepend Minitest::UnexpectedError.new(setup_exception)
end

def after_fork(worker)
  self.class.after_fork_hooks.each do |cb|
    cb.call(worker)
  end
end

def initialize(queue_size)
  @queue_size = queue_size
  @queue      = Server.new
  @pool       = []
  @url = DRb.start_service("drbunix:", @queue).uri
end

def run_cleanup(worker)
  self.class.run_cleanup_hooks.each do |cb|
    cb.call(worker)
  end
end

def shutdown
  @queue_size.times { @queue << nil }
  @pool.each { |pid| Process.waitpid pid }
  if @queue.length > 0
    raise "Queue not empty, but all workers have finished. This probably means that a worker crashed and #{@queue.length} tests were missed."
  end
end

def start
  @pool = @queue_size.times.map do |worker|
    fork do
      DRb.stop_service
      begin
        after_fork(worker)
      rescue => setup_exception; end
      queue = DRbObject.new_with_uri(@url)
      while job = queue.pop
        klass    = job[0]
        method   = job[1]
        reporter = job[2]
        result = klass.with_info_handler reporter do
          Minitest.run_one_method(klass, method)
        end
        add_setup_exception(result, setup_exception) if setup_exception
        begin
          queue.record(reporter, result)
        rescue DRb::DRbConnError
          result.failures.each do |failure|
            failure.exception = DRb::DRbRemoteError.new(failure.exception)
          end
          queue.record(reporter, result)
        end
      end
    ensure
      run_cleanup(worker)
    end
  end
end