fix(indexing): reissue a failed attempt's batches without racing its tasks

Merged
EV
15h ago
+547-148

Description

Fixes the race behind the flaky test_mock_connector_checkpoint_recovery (assert 101 == 102, and attempts stuck at completed < total with NoSuchKey).

When an attempt fails, its docprocessing tasks keep running. Each task increments the old attempt's completed_batches, then deletes its batch file. The resuming attempt's reissue_old_batches read these with no lock:

  1. Reused batch number (lost document). The next batch number was len(old_batches) + most_recent_attempt.completed_batches. That count came from an ORM object loaded earlier. A batch that finished after that load and was deleted before the listing was in neither count. So the next number was one too low and equal to a reissued batch's number. store_batch then called upsert_filerecord (ON CONFLICT DO UPDATE), which repointed the reissued batch's record at the new batch. The reissued batch's documents were lost.
  2. Rename vs delete of the same batch. change_file_id (rename, the object keeps its key) and delete_file (SELECT, delete object, DELETE row) took no lock. If the delete read the record, the rename committed, and then the delete removed the object, the renamed record pointed at a deleted object (NoSuchKey). The attempt never reached completed >= total. If the delete finished between the listing and the rename, change_file_id raised and crashed docfetching.

The fix:

  • reissue_old_batches moves the batches first, then reads the old attempt's completed_batches again with a column query (IndexingCoordination.get_completed_batches). A task counts before it deletes, so every old batch is now moved, counted, or both. "Both" only leaves an unused batch number.
  • The next batch number is max(fresh count + moved, highest moved batch + 1). The second term covers a case I saw locally: the attempt before the resume can be one that failed before it fetched anything (count 0), while the batches came from an earlier attempt.
  • The new attempt starts with completed_batches = next batch number - moved batches, so completion still fires at completed_batches >= total_batches. It is now set and committed before the reissued tasks go out. Before, the caller set it after the tasks were sent, so a fast task's increment could be overwritten.
  • change_file_id and delete_file lock the file_record row (SELECT ... FOR UPDATE) on all four backends. delete_file deletes the row and then the object while it holds the lock, so a rename that won leaves the object alone. change_file_id returns False if the record is gone. update_old_batches_to_new_index_attempt returns only the batches it moved, and only those are reissued and counted. The S3, GCS, and Azure stores share one helper, rename_filerecord.
  • store_batch refuses to overwrite an existing batch record. If a collision happens again, it raises instead of losing data. No other save_file caller changes.

The feature-map docs for the indexing pipeline and the file store describe the new behaviour.

How Has This Been Tested?

  • New external-dependency-unit test backend/tests/external_dependency_unit/indexing/test_reissue_old_batches_race.py (real Postgres and S3 file store). It runs the old task's increment and delete at each racy point:

    • a batch finishes after the old attempt is loaded,
    • a batch is deleted between the listing and the rename,
    • the rename runs while the delete holds the row (a thread, synchronized on pg_stat_activity lock waits),
    • the count comes from an attempt that fetched nothing,
    • store_batch refuses an overwrite.

    With the fix: 5/5 pass. On main (the same scenarios, adapted to main's reissue_old_batches signature), the first four fail with the real signatures: the reissued batch reads the new batch's document (the lost-document bug), FileRecordNotFoundError (the docfetching crash), and NoSuchKey (the stuck attempt). The overwrite check also fails, as expected.

  • The existing file store unit tests and the external-dependency-unit tests for the file store and indexing pass. I updated the unit tests that patched the old lookups.

  • test_mock_connector_checkpoint_recovery, run locally against a separate stack (the harness starts its own Celery workers):

    • with the fix: 20 runs, 19 pass. The one failure is a different, existing flake: the auto-retry attempt was still cancelling when the test called run-once. run-once skips a pair with an unfinished attempt, so no recovery attempt started (timeout at line 489). This happens before any reissue.
    • on unmodified main: 10 runs, 10 pass.
    • No NoSuchKey or "does not exist or was deleted" lines in the worker logs of either set. The race did not reproduce naturally on this machine, so the external-dependency-unit test is the deterministic proof.

Additional Options

  • [Optional] Please cherry-pick this PR to the latest release version.
  • [Optional] Override Linear Check

🤖 Generated with Claude Code

Summary by cubic

Fixes a race in reissuing a failed attempt's batches that could lose documents or leave the resuming attempt stuck. The reissue now moves the batches first, then reads completed_batches again from the database, so every old batch is moved, counted, or both. change_file_id and delete_file lock the FileRecord row, change_file_id returns False when the record is gone, and store_batch refuses to overwrite an existing batch. Also fixes a deadlock in the legacy copy test and tightens the reissue race test.

  • Adds an external-dependency test that exercises each racy point against real Postgres and S3; 5/5 pass with the fix, the equivalent scenarios fail on main.
  • Updates existing tests for the new locking helpers and reissue_old_batches signature.

Summary by CodeRabbit

  • Bug Fixes

    • Improved reliability of file renaming and deletion during concurrent operations.
    • Preserved file metadata when storage operations fail.
    • Prevented existing document batches from being overwritten.
    • Improved recovery of interrupted indexing jobs, including safer batch reassignment and numbering.
    • Missing batch metadata now surfaces an error instead of being silently skipped.

  • Documentation

    • Expanded documentation describing checkpoint recovery and batch resumption behavior.