Commit 6a1101c7 authored by Alishan Ladhani's avatar Alishan Ladhani Committed by Alper Akgun

Refactor to improve readability

Move +1 into `actual_finish`
parent a8993d30
......@@ -49,6 +49,8 @@ module Gitlab
MAX_ALLOWED_LOOPS = 10_000
SLEEP_TIME_IN_SECONDS = 0.01 # 10 msec sleep
ALLOWED_MODES = [:itself, :distinct].freeze
FALLBACK_FINISH = 0
OFFSET_BY_ONE = 1
# Each query should take < 500ms https://gitlab.com/gitlab-org/gitlab/-/merge_requests/22705
DEFAULT_DISTINCT_BATCH_SIZE = 10_000
......@@ -65,7 +67,7 @@ module Gitlab
(@operation == :count && batch_size <= MIN_REQUIRED_BATCH_SIZE) ||
(@operation == :sum && batch_size < DEFAULT_SUM_BATCH_SIZE) ||
(finish - start) / batch_size >= MAX_ALLOWED_LOOPS ||
start > finish
start >= finish
end
def count(batch_size: nil, mode: :itself, start: nil, finish: nil)
......@@ -85,11 +87,13 @@ module Gitlab
results = nil
batch_start = start
while batch_start <= finish
batch_relation = build_relation_batch(batch_start, batch_start + batch_size, mode)
while batch_start < finish
batch_end = [batch_start + batch_size, finish].min
batch_relation = build_relation_batch(batch_start, batch_end, mode)
begin
results = merge_results(results, batch_relation.send(@operation, *@operation_args)) # rubocop:disable GitlabSecurity/PublicSend
batch_start += batch_size
batch_start = batch_end
rescue ActiveRecord::QueryCanceled => error
# retry with a safe batch size & warmer cache
if batch_size >= 2 * MIN_REQUIRED_BATCH_SIZE
......@@ -99,6 +103,7 @@ module Gitlab
return FALLBACK
end
end
sleep(SLEEP_TIME_IN_SECONDS)
end
......@@ -138,7 +143,7 @@ module Gitlab
end
def actual_finish(finish)
finish || @relation.unscope(:group, :having).maximum(@column) || 0
(finish || @relation.unscope(:group, :having).maximum(@column) || FALLBACK_FINISH) + OFFSET_BY_ONE
end
def check_mode!(mode)
......
......@@ -130,6 +130,16 @@ RSpec.describe Gitlab::Database::BatchCount do
expect(described_class.batch_count(model, start: model.minimum(:id), finish: model.maximum(:id))).to eq(5)
end
it 'stops counting when finish value is reached' do
stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
expect(described_class.batch_count(model,
start: model.minimum(:id),
finish: model.maximum(:id) - 1, # Do not count the last record
batch_size: model.count - 2 # Ensure there are multiple batches
)).to eq(model.count - 1)
end
it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE}" do
min_id = model.minimum(:id)
relation = instance_double(ActiveRecord::Relation)
......@@ -242,6 +252,19 @@ RSpec.describe Gitlab::Database::BatchCount do
expect(described_class.batch_distinct_count(model, column, start: model.minimum(column), finish: model.maximum(column))).to eq(2)
end
it 'stops counting when finish value is reached' do
# Create a new unique author that should not be counted
create(:issue)
stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
expect(described_class.batch_distinct_count(model, column,
start: User.minimum(:id),
finish: User.maximum(:id) - 1, # Do not count the newly created issue
batch_size: model.count - 2 # Ensure there are multiple batches
)).to eq(2)
end
it 'counts with User min and max as start and finish' do
expect(described_class.batch_distinct_count(model, column, start: User.minimum(:id), finish: User.maximum(:id))).to eq(2)
end
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment