Two queries too many: removing each_batch from GitLab's Maven snapshot cleanup - DEV Community

DEV Community

Cover image for Two queries too many: removing each_batch from GitLab's Maven snapshot cleanup
Priyanshu Jha
Priyanshu Jha

Posted on

Two queries too many: removing each_batch from GitLab's Maven snapshot cleanup

#gitlab #opensource #ruby #postgres

The final diff of gitlab-org/gitlab!270142 is tiny: a handful of lines in the service and a few more in its spec. But getting it merged took a database review with real production query plans, two spec improvements from a reviewer who tested them by breaking the code on purpose, and a Danger failure I caused myself.

This post covers what the issue was, why each_batch was the wrong tool here, how the specs got stronger in review, and what database review looks like when you're a community contributor without access to GitLab's production database tools.

The issue: a batch loop over something already batched

The issue was gitlab-org/gitlab#639257, opened by Tim Rizzi. It came from a suggestion by Daniyal Arshad in the review of an earlier MR (!266409), which added a Maven snapshot pass to the package cleanup policy.

Some background. Packages::Cleanup::ExecutePolicyService runs a project's package cleanup policy. For Maven snapshot packages, it works out which old snapshot builds to remove and marks their files as pending_destruction. A background job deletes them later.

The marking happened in slices of up to 200 ids. Each slice was handed to another service:

old_ids_by_name.values.flatten.each_slice(MARK_PACKAGE_FILES_FOR_DESTRUCTION_SERVICE_BATCH_SIZE) do |ids|
  result = ::Packages::MarkPackageFilesForDestructionService.new(::Packages::PackageFile.id_in(ids))
    .execute(batch_deadline: batch_deadline, batch_size: MARK_PACKAGE_FILES_FOR_DESTRUCTION_SERVICE_BATCH_SIZE)
  @counts[:marked_package_files_total_count] += result.payload[:marked_package_files_count]

  break :timeout unless result.success?
end
Enter fullscreen mode Exit fullscreen mode

And here's the core of that service:

@package_files.each_batch(of: min_batch_size) do |batched_package_files|
  if batch_deadline&.past?
    timeout = true
    break
  end

  updates_count += batched_package_files.update_all(status: :pending_destruction)
end
Enter fullscreen mode Exit fullscreen mode

each_batch is GitLab's standard way to walk a large relation safely. It finds a start id, then a boundary id, then yields a batch. That's exactly right when you don't know how big the relation is.

Here, though, the relation is PackageFile.id_in(ids) with at most 200 ids, and the batch size is 200. There is only ever one batch. So each slice paid for two lookups whose only job was to find out what we already knew.

Statements per slice, before and after

The plan

I commented on the issue with this plan:

My plan is to replace each_batch over the one-batch snapshot relation with each_slice plus Packages::PackageFile.id_in(ids).update_all(status: :pending_destruction), matching mark_packages_for_destruction, and to stub batch_deadline in the snapshot timeout spec.

mark_packages_for_destruction is a method in the same service that already slices ids and updates them directly. Matching an existing pattern in the same file is the easiest kind of change to review. Tim replied that the plan "sounds right", assigned me, and asked me to tag Tim and Fiona McCawley on the MR.

The change

The production change is the loop body, now one statement per slice:

old_ids_by_name.values.flatten.each_slice(MARK_PACKAGE_FILES_FOR_DESTRUCTION_SERVICE_BATCH_SIZE) do |ids|
  break :timeout if batch_deadline.past?

  @counts[:marked_package_files_total_count] +=
    ::Packages::PackageFile.id_in(ids).update_all(status: :pending_destruction)
end
Enter fullscreen mode Exit fullscreen mode

Two behaviors had to stay the same:

  • The deadline. The old service checked batch_deadline before each batch. The new loop checks it before each slice, so a long cleanup still stops on time and reports :timeout.
  • The count. update_all returns the number of updated rows, so marked_package_files_total_count still adds up the same way.

The spec needed a change too. The old timeout example stubbed MarkPackageFilesForDestructionService to return an error, and that service isn't called anymore.

CI lesson: Danger fails on long commit subjects

My first push failed the danger-review job, and it was entirely my fault. GitLab's Danger rules require commit subjects of 72 characters or fewer, and mine was longer. I reworded it to Use update_all for Maven snapshot cleanup slices and pushed again. The next pipeline passed.

The same Danger run also left a note on the update_all line:

When using any of these: … Bulk commands: bulk_insert, update_all, etc. … You must:

  1. Include the full database query in the merge request description (including scopes, pagination, and limits)
  2. Include the query execution plan
  3. Request a ~database review

Adding the SQL was easy: a primary-key UPDATE on packages_package_files with WHERE id IN (...) for up to 200 ids. The query plan was harder.

Database review without database access

GitLab's query plans come from Database Lab, which runs queries against a clone of production. It needs GitLab team access, and as a community contributor I don't have it. I put the full SQL in the description, explained that it's the same statement mark_packages_for_destruction already runs, and asked for a database review.

Fiona stepped in:

Database Lab needs GitLab team access, so I've captured the plans: the old and new SQL, each followed by its query plan, are now in the Database query section of the description.

Fiona went further than I expected and picked a real slice from production data: the first 200 ids this pass would mark for one snapshot package in gitlab-org/editor-extensions/gitlab-eclipse-plugin. That package has 601 installable files, and with keep_n_duplicated_package_files: 1 the pass marks 442 of them in slices of 200, 200 and 42.

Fiona's table, summarized:

Statement Plan Time (cold clone)
After UPDATE ... WHERE id IN (...) Index Scan on the primary key, 200 rows 968 ms
Before start SELECT ... LIMIT 1 Index Only Scan on the primary key 6 ms
boundary SELECT ... LIMIT 1 OFFSET 200 Index Only Scan on the primary key 92 ms
UPDATE ... AND id >= <start> Index Scan on the primary key, 200 rows 596 ms

The headline: the UPDATE plan is unchanged, and the two lookups disappear. The UPDATE times differ between runs mostly because of cold I/O on the clone. Fiona noted that a second run of the new UPDATE on a warmed clone took 760 ms with the same plan, and every UPDATE run touched about the same number of buffers.

This is the part of contributing that doesn't show up in the diff. The reviewer filled a gap I couldn't fill, and did it with real data instead of a guess.

The review that made the specs stronger

Fiona also left two spec suggestions. Both came with a "Why" section explaining how Fiona had tested them. I think this is the best kind of review feedback.

1. Guard the listing deadline check

The service checks the deadline in two places: while listing snapshot files, and before each update slice. My spec for "deadline already passed" asserted that no file changed and that timeout was true. Fiona found a hole:

I deleted break :timeout if batch_deadline.past? from the listing loop (execute_policy_service.rb:184) and ran this example locally: it still passed.

Without the listing check, the listing runs to the end, then the first slice's deadline check stops the marking, and no file changes. My assertions couldn't tell the two code paths apart. The suggested fix was one line on a counter that only grows after listing:

expect(execute.payload[:counts][:unique_package_id_and_file_name_total_count]).to eq(0)
Enter fullscreen mode Exit fullscreen mode

With this line, the example fails if the listing check is removed.

2. Don't rely on counting calls

My "timeout between slices" spec stubbed the deadline with a sequence: and_return(false, false, true). Fiona pointed out that the sequence is shared by every batch_deadline.past? call in execute, including other passes. So a new check anywhere earlier would shift it and break the test for the wrong reason.

The suggestion flips the deadline right after the first update, by wrapping PackageFile.id_in. That's the same approach an existing example uses for the generic rule:

allow(::Packages::PackageFile).to receive(:id_in).and_wrap_original do |method, ids|
  allow(service).to receive(:batch_deadline).and_return(1.second.ago)
  method.call(ids)
end
Enter fullscreen mode Exit fullscreen mode

With the batch size stubbed to 2, the example now asserts that exactly two files are marked before the timeout:

it 'stops after the current batch and reports a timeout' do
  expect { execute }.to change { installable_files.count }.by(-2)

  expect(execute.payload[:timeout]).to be(true)
  expect(execute.payload[:counts][:marked_package_files_total_count]).to eq(2)
end
Enter fullscreen mode Exit fullscreen mode

I applied both in one commit (028c6d2f), specs only, and credited Fiona in the commit message. Fiona then approved the backend side and asked Jean van der Walt for the database review.

From issue to merge

Database approval, and a follow-up idea

Jean reviewed the database side:

Database looks good: the UPDATE plan is unchanged, and each slice now skips the two each_batch lookups.

Jean also left a non-blocking note: the 200-row UPDATE takes about 600–960 ms on a cold clone, and a follow-up could try a smaller MARK_PACKAGE_FILES_FOR_DESTRUCTION_SERVICE_BATCH_SIZE. I agreed it's worth trying and kept it out of this MR. Changing the batch size changes the risk profile, and it deserves its own measurement.

Abdul Wadood gave the database maintainer approval ("~database LGTM! 🚀"), and Abdul merged the MR. Five pipelines in total: one failure (my commit subject) and four green.

What I took away

1. each_batch is for unknown sizes. If you've already sliced the ids, a direct update_all per slice is cheaper and easier to read. Look for an existing pattern in the same file before inventing one.

2. Keep deadline semantics when you refactor a loop. Check the deadline at the same granularity as before, and keep the counters adding up the same way.

3. Test that your test can fail. Fiona's method was simple: delete the line the spec claims to protect and see if the spec still passes. If it does, the spec isn't guarding anything.

4. Prefer event-based stubs over call-count sequences. and_return(false, false, true) breaks when someone adds an unrelated check. Flipping the stub when the interesting thing happens doesn't.

5. Community contributors can still do database work. Put the full SQL in the description, explain why it's safe, and ask. Reviewers with access can capture the plans. Just make their job easy.

6. Keep commit subjects within 72 characters. Danger will fail the pipeline if you don't.

Links

Thanks to Tim Rizzi for the issue, Daniyal Arshad for the original idea, Fiona McCawley for the plans and the spec reviews, Jean van der Walt for the database review, and Abdul Wadood for the final approval.

Top comments (0)