53 lines
1.6 KiB
Ruby
53 lines
1.6 KiB
Ruby
|
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
|