blob: 7ae5a4c17c8b3d7c73f8f71e5cc7cc194cd96a31 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
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
|
