diff --git a/app/scripts/bulk_import.py b/app/scripts/bulk_import.py index 3c457bc..e2af9d6 100644 --- a/app/scripts/bulk_import.py +++ b/app/scripts/bulk_import.py @@ -87,25 +87,41 @@ async def _bulk_enqueue(entity_type: str, *, dry_run: bool, include_seasons: boo return queue = await get_arq_pool() - semaphore = asyncio.Semaphore(_CONCURRENCY) + # A fixed pool of _CONCURRENCY consumer tasks pulling from a plain + # asyncio.Queue -- NOT asyncio.gather() over one Task per id. gather() + # eagerly wraps every argument into a Task before awaiting any of them, + # so with a catalog in the millions that means millions of live Task + + # coroutine objects created up front regardless of how tightly a + # semaphore bounds actual concurrent execution -- a real bottleneck + # (and memory risk) discovered running this against the real 1.24M- + # movie export, where it looked hung for minutes just getting through + # that setup before any work could start. + id_queue: asyncio.Queue[str] = asyncio.Queue() + for external_id in ids: + id_queue.put_nowait(external_id) + enqueued = 0 lock = asyncio.Lock() - async def _worker(external_id: str) -> None: + async def _consumer() -> None: nonlocal enqueued - async with semaphore: + while True: + try: + external_id = id_queue.get_nowait() + except asyncio.QueueEmpty: + return await _enqueue_one(queue, entity_type, external_id, include_seasons=include_seasons) - async with lock: - enqueued += 1 - if enqueued % 5000 == 0: - logger.info( - "bulk_import_progress", - entity_type=entity_type, - enqueued=enqueued, - total=len(ids), - ) + async with lock: + enqueued += 1 + if enqueued % 5000 == 0: + logger.info( + "bulk_import_progress", + entity_type=entity_type, + enqueued=enqueued, + total=len(ids), + ) - await asyncio.gather(*(_worker(external_id) for external_id in ids)) + await asyncio.gather(*(_consumer() for _ in range(_CONCURRENCY))) logger.info("bulk_import_complete", entity_type=entity_type, enqueued=enqueued) diff --git a/tests/unit/test_bulk_import_script.py b/tests/unit/test_bulk_import_script.py new file mode 100644 index 0000000..5928458 --- /dev/null +++ b/tests/unit/test_bulk_import_script.py @@ -0,0 +1,104 @@ +"""Unit tests for app/scripts/bulk_import.py's enqueue logic. + +Exercises _bulk_enqueue's bounded worker-pool directly (network, arq, and +the DB are all monkeypatched out) -- the property under test is that every +id from the export gets enqueued exactly once, with no duplicates and none +skipped, which is the correctness risk of moving off asyncio.gather(*ids) +onto a fixed pool of consumers pulling from an asyncio.Queue (done after +the gather-based version was found, live, to effectively stall for +minutes on Task-creation overhead alone against a real 1.24M-id export).""" + +from __future__ import annotations + +import uuid + +import pytest + +from app.scripts import bulk_import + + +class _FakeQueue: + def __init__(self) -> None: + self.enqueued: list[tuple] = [] + + async def enqueue_job(self, task_name: str, *args): + self.enqueued.append((task_name, args)) + + +@pytest.mark.asyncio +async def test_bulk_enqueue_enqueues_every_id_exactly_once(monkeypatch): + ids = [str(n) for n in range(1, 501)] + fake_queue = _FakeQueue() + + async def fake_fetch(entity_type, *, client): + return ids + + async def fake_get_or_create_ingestion_job(job_type, entity_type, external_id): + return uuid.uuid4() + + async def fake_get_arq_pool(): + return fake_queue + + monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch) + monkeypatch.setattr( + bulk_import, "get_or_create_ingestion_job", fake_get_or_create_ingestion_job + ) + monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool) + + await bulk_import._bulk_enqueue("movie", dry_run=False, include_seasons=False) + + enqueued_ids = sorted(int(args[1]) for _, args in fake_queue.enqueued) + assert enqueued_ids == list(range(1, 501)) + assert len(fake_queue.enqueued) == 500 + assert all(task_name == "ingest_movie_task" for task_name, _ in fake_queue.enqueued) + + +@pytest.mark.asyncio +async def test_bulk_enqueue_dry_run_enqueues_nothing(monkeypatch): + ids = [str(n) for n in range(1, 51)] + fake_queue = _FakeQueue() + called = False + + async def fake_fetch(entity_type, *, client): + return ids + + async def fake_get_arq_pool(): + nonlocal called + called = True + return fake_queue + + monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch) + monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool) + + await bulk_import._bulk_enqueue("collection", dry_run=True, include_seasons=False) + + assert called is False + assert fake_queue.enqueued == [] + + +@pytest.mark.asyncio +async def test_bulk_enqueue_passes_include_seasons_flag_for_series(monkeypatch): + fake_queue = _FakeQueue() + + async def fake_fetch(entity_type, *, client): + return ["1396"] + + async def fake_get_or_create_ingestion_job(job_type, entity_type, external_id): + return uuid.uuid4() + + async def fake_get_arq_pool(): + return fake_queue + + monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch) + monkeypatch.setattr( + bulk_import, "get_or_create_ingestion_job", fake_get_or_create_ingestion_job + ) + monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool) + + await bulk_import._bulk_enqueue("series", dry_run=False, include_seasons=True) + + assert len(fake_queue.enqueued) == 1 + task_name, args = fake_queue.enqueued[0] + assert task_name == "ingest_series_task" + assert args[1] == "1396" + assert args[2] is True