gitlab-org--gitlab-foss/lib/gitlab/database/multi_threaded_migration.rb

53 lines
1.6 KiB
Ruby
Raw Normal View History

module Gitlab
module Database
module MultiThreadedMigration
MULTI_THREAD_AR_CONNECTION = :thread_local_ar_connection
# This overwrites the default connection method so that every thread can
# use a thread-local connection, while still supporting all of Rails'
# migration methods.
def connection
Thread.current[MULTI_THREAD_AR_CONNECTION] ||
ActiveRecord::Base.connection
end
# Starts a thread-pool for N threads, along with N threads each using a
# single connection. The provided block is yielded from inside each
# thread.
#
# Example:
#
# with_multiple_threads(4) do
# execute('SELECT ...')
# end
#
# thread_count - The number of threads to start.
#
# join - When set to true this method will join the threads, blocking the
# caller until all threads have finished running.
#
# Returns an Array containing the started threads.
def with_multiple_threads(thread_count, join: true)
pool = Gitlab::Database.create_connection_pool(thread_count)
threads = Array.new(thread_count) do
Thread.new do
pool.with_connection do |connection|
begin
Thread.current[MULTI_THREAD_AR_CONNECTION] = connection
yield
ensure
Thread.current[MULTI_THREAD_AR_CONNECTION] = nil
end
end
end
end
threads.each(&:join) if join
threads
end
end
end
end