debian-mirror-gitlab/lib/gitlab/database/migrations/test_batched_background_runner.rb
2023-05-27 22:25:52 +05:30

144 lines
5.4 KiB
Ruby

# frozen_string_literal: true
module Gitlab
module Database
module Migrations
class TestBatchedBackgroundRunner < BaseBackgroundRunner
include Gitlab::Database::DynamicModelHelpers
MIGRATION_DETAILS_FILE_NAME = 'details.json'
def initialize(result_dir:, connection:, from_id:)
super(result_dir: result_dir, connection: connection)
@connection = connection
@from_id = from_id
end
def jobs_by_migration_name
set_shared_model_connection do
Gitlab::Database::BackgroundMigration::BatchedMigration
.executable
.where('id > ?', from_id)
.to_h do |migration|
batching_strategy = migration.batch_class.new(connection: connection)
smallest_batch_start = migration.next_min_value
table_max_value = define_batchable_model(migration.table_name, connection: connection)
.maximum(migration.column_name)
largest_batch_start = [table_max_value - migration.batch_size, smallest_batch_start].max
# variance is the portion of the batch range that we shrink between variance * 0 and variance * 1
# to pick actual batches to sample.
variance = largest_batch_start - smallest_batch_start
batch_starts = uniform_fractions
.lazy # frac varies from 0 to 1, values in smallest_batch_start..largest_batch_start
.map { |frac| (variance * frac).to_i + smallest_batch_start }
# Track previously run batches so that we stop sampling if a new batch would intersect an older one
completed_batches = []
jobs_to_sample = batch_starts
# Stop sampling if a batch would intersect a previous batch
.take_while { |start| completed_batches.none? { |batch| batch.cover?(start) } }
.map do |batch_start|
# The current block is lazily evaluated as part of the jobs_to_sample enumerable
# so it executes after the enclosing using_connection block has already executed
# Therefore we need to re-associate with the explicit connection again
Gitlab::Database::SharedModel.using_connection(connection) do
next_bounds = batching_strategy.next_batch(
migration.table_name,
migration.column_name,
batch_min_value: batch_start,
batch_size: migration.batch_size,
job_class: migration.job_class,
job_arguments: migration.job_arguments
)
batch_min, batch_max = next_bounds
job = migration.create_batched_job!(batch_min, batch_max)
completed_batches << (batch_min..batch_max)
job
end
end
job_class_name = migration.job_class_name
export_migration_details(job_class_name, migration.slice(:interval, :total_tuple_count, :max_batch_size))
[job_class_name, jobs_to_sample]
end
end
end
def run_job(job)
set_shared_model_connection do
Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection).perform(job)
end
end
def uniform_fractions
Enumerator.new do |y|
# Generates equally distributed fractions between 0 and 1, with increasing detail as more are pulled from
# the enumerator.
# 0, 1 (special case)
# 1/2
# 1/4, 3/4
# 1/8, 3/8, 5/8, 7/8
# etc.
# The pattern here is at each outer loop, the denominator multiplies by 2, and at each inner loop,
# the numerator counts up all odd numbers 1 <= n < denominator.
y << 0
y << 1
# denominators are each increasing power of 2
denominators = (1..).lazy.map { |exponent| 2**exponent }
denominators.each do |denominator|
# Numerators at the current step are all odd numbers between 1 and the denominator
numerators = (1..denominator).step(2)
numerators.each do |numerator|
next_frac = numerator.fdiv(denominator)
y << next_frac
end
end
end
end
private
attr_reader :from_id
def set_shared_model_connection(&block)
Gitlab::Database::SharedModel.using_connection(connection, &block)
end
def job_meta(job)
set_shared_model_connection do
job.slice(:min_value, :max_value, :batch_size, :sub_batch_size, :pause_ms)
end
end
def export_migration_details(migration_name, attributes)
directory = result_dir.join(migration_name)
FileUtils.mkdir_p(directory) unless Dir.exist?(directory)
File.write(directory.join(MIGRATION_DETAILS_FILE_NAME), attributes.to_json)
end
def observers
::Gitlab::Database::Migrations::Observers.all_observers + [
::Gitlab::Database::Migrations::Observers::BatchDetails
]
end
end
end
end
end