diff --git a/src/kiro_crew/stt/decoder.py b/src/kiro_crew/stt/decoder.py index 22da003cca4c..110973b0887b 100644 --- a/src/kiro_crew/stt/decoder.py +++ b/src/kiro_crew/stt/decoder.py @@ -560,6 +560,19 @@ def _remove_retired() -> None: logger.debug("Could not remove retired decoder %s", name, exc_info=True) +async def _join_migration(task: asyncio.Task[Path | None]) -> Path | None: + """Wait for the migration *task* another caller started; ``None`` when it + runs on a different event loop, which this caller cannot await.""" + try: + same_loop = task.get_loop() is asyncio.get_running_loop() + except RuntimeError: + same_loop = False + if not same_loop: + return None + # Shielded: one caller giving up must not cancel the shared transfer. + return await asyncio.shield(task) + + async def migrate_retired_decoder() -> Path | None: """Fetch the current pin for an install whose store has only a retired one. @@ -569,19 +582,17 @@ async def migrate_retired_decoder() -> Path | None: by :meth:`DecoderStore.ensure` once the current pin is installed. """ global _migration_task - task = _migration_task - if task is not None: - try: - same_loop = task.get_loop() is asyncio.get_running_loop() - except RuntimeError: - same_loop = False - if not same_loop: - return None - # Shielded: one caller giving up must not cancel the shared transfer. - return await asyncio.shield(task) + if _migration_task is not None: + return await _join_migration(_migration_task) if bundle_carries_decoder() or artifact_for() is None: return None - if not await asyncio.to_thread(has_retired_decoder): + has_retired = await asyncio.to_thread(has_retired_decoder) + # Another caller can start the migration while this one awaits the probe, and + # its transfer removes the retired file, so a probe that ran after it finds + # none: join that transfer rather than report the decoder missing. + if _migration_task is not None: + return await _join_migration(_migration_task) + if not has_retired: return None # Started at most once per process: on a host that cannot reach the release, # every later transcode would otherwise wait out another failed fetch. The diff --git a/test/test_stt_decoder.py b/test/test_stt_decoder.py index a01c9cdeac4a..a44bd259e889 100644 --- a/test/test_stt_decoder.py +++ b/test/test_stt_decoder.py @@ -19,6 +19,7 @@ import hashlib import os import stat +import threading from pathlib import Path import pytest @@ -605,6 +606,54 @@ async def test_concurrent_callers_share_one_transfer(self, monkeypatch, tmp_path assert first == second == tmp_path / artifact.filename assert len(calls) == 1 + @pytest.mark.asyncio + async def test_a_caller_whose_probe_sees_the_migrated_store_joins_the_transfer( + self, monkeypatch, tmp_path + ): + """The second caller's retired-file probe runs only after the first caller's + migration installed the pin and removed the retired file, so it finds none. + A migration is under way all the same, and that caller still gets its path.""" + artifact = _artifact() + _pin(monkeypatch, tmp_path, artifact) + _allow_download(monkeypatch) + (tmp_path / decoder.RETIRED_FILENAMES[1]).write_bytes(b"old decoder") + calls: list[dict] = [] + _serve(monkeypatch, _PAYLOAD, calls) + + migrated = threading.Event() + probes = 0 + probes_lock = threading.Lock() + real_probe = decoder.has_retired_decoder + + def _probe() -> bool: + nonlocal probes + with probes_lock: + probes += 1 + later = probes > 1 + if later: + assert migrated.wait(timeout=10), "the first migration never finished" + return real_probe() + + shared_store = decoder.store() + real_ensure = shared_store.ensure + + async def _ensure() -> Path | None: + try: + return await real_ensure() + finally: + migrated.set() + + monkeypatch.setattr(decoder, "has_retired_decoder", _probe) + monkeypatch.setattr(shared_store, "ensure", _ensure) + + first, second = await asyncio.wait_for( + asyncio.gather(decoder.migrate_retired_decoder(), decoder.migrate_retired_decoder()), + timeout=30, + ) + + assert first == second == tmp_path / artifact.filename + assert len(calls) == 1 + @pytest.mark.asyncio async def test_no_retired_file_means_no_network(self, monkeypatch, tmp_path): _pin(monkeypatch, tmp_path, _artifact())