Module: Gitlab::Database::MultiThreadedMigration

Defined in:
lib/gitlab/database/multi_threaded_migration.rb

Constant Summary collapse

MULTI_THREAD_AR_CONNECTION =
:thread_local_ar_connection

Instance Method Summary collapse

Instance Method Details

#connectionObject

This overwrites the default connection method so that every thread can use a thread-local connection, while still supporting all of Rails' migration methods.


11
12
13
14
# File 'lib/gitlab/database/multi_threaded_migration.rb', line 11

def connection
  Thread.current[MULTI_THREAD_AR_CONNECTION] ||
    ActiveRecord::Base.connection
end

#with_multiple_threads(thread_count, join: true) ⇒ Object

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.


32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# File 'lib/gitlab/database/multi_threaded_migration.rb', line 32

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|
        Thread.current[MULTI_THREAD_AR_CONNECTION] = connection
        yield
      ensure
        Thread.current[MULTI_THREAD_AR_CONNECTION] = nil
      end
    end
  end

  threads.each(&:join) if join

  threads
end