fix(mm): close TOCTOU race in _restore_incomplete_installs#9142
Conversation
The restore path snapshotted active sources under the lock, then released the lock and iterated tmpdirs against the stale snapshot. A foreground import_model call landing during iteration could append a job the loop never saw, leading to a duplicate enqueue and a FileNotFoundError when the second download tried to rename the .downloading file the first had already moved. Move the membership check inside the lock immediately before the append so check-and-append is atomic. Fixes #9141. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
JPPhoto
left a comment
There was a problem hiding this comment.
This looks good when inspecting visually, but I'd also like to see a test added for this in tests/app/services/model_install/test_model_install.py, especially since it's a race condition. It may be a little challenging to get the test right, though:
- Create a tmp install dir with an active marker for a remote source.
- Monkeypatch
_resume_remote_download()so restore pauses at a controlled point. - Insert an active job for the same
sourceinto_install_jobswhile restore is between marker parsing and the new locked duplicate check. - Let restore continue.
- Assert restore did not append a duplicate job and did not enqueue/resume a second download.
That maps directly to the PR’s fix around invokeai/app/services/model_install/model_install_default.py:249.
The only slightly fiddly part is choosing the pause point. Because the fixed code does the duplicate check under _lock immediately before append, the test probably needs to monkeypatch either _guess_source() or _read_install_marker() to signal “restore has observed the marker”, then add the foreground job before restore reaches the locked block. That is manageable with two threading.Events, similar to the existing test_import_waits_for_startup_restore.
The test should not need real network downloads if _resume_remote_download() is monkeypatched to fail the test if called for the skipped duplicate.
Summary
This PR fixes an intermittent unit test failure in
test_model_install.py._restore_incomplete_installssnapshotted active sources under the lock, then iterated tmpdirs against the stale snapshot. A foregroundimport_modelcall landing during iteration could be missed, the same source enqueued twice, and the second download would fail withFileNotFoundErrorwhen trying to rename the.downloadingfile the first download had already moved._install_jobs.append(job)so check-and-append is atomic. Duplicate-tmpdir suppression within a single restore pass stays outside the lock since it doesn't race againstimport_model.Fixes #9141.
Test plan
pytest tests/app/services/model_install/test_model_install.py::test_simple_download— ran 10× consecutively, all pass (was flaky ~30–50% onmain)pytest tests/app/services/model_install/— 31/31 pass🤖 Generated with Claude Code