diff --git a/core/downloads/task_worker.py b/core/downloads/task_worker.py index 00086c8c..3850dd2d 100644 --- a/core/downloads/task_worker.py +++ b/core/downloads/task_worker.py @@ -276,7 +276,7 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke _t = download_tasks.get(task_id, {}) is_quarantine_retry = bool(_t.pop('_quarantine_retry', False)) if not is_quarantine_retry: - _t.pop('searched_sources', None) + _t.pop('searched_queries', None) if is_quarantine_retry and _try_cached_candidates(task_id, batch_id, track, deps): with tasks_lock: used_filename = download_tasks.get(task_id, {}).get('filename') @@ -373,16 +373,23 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke # source instead of re-fetching the same exhausted one (e.g. Soulseek # keeps returning fresh wrong peers — once its budget is gone, switch to # HiFi/Tidal/…). See monitor.requeue_quarantined_task_for_retry. - # On a quarantine retry we also exclude sources we've ALREADY searched: - # their candidates are walked via the cached-first path, so re-searching - # them is the wasteful repeat the cached-first design removes. The chain - # then falls through to a not-yet-searched source. Fresh / dead-connection - # runs cleared searched_sources above, so they search everything again. + # + # On a quarantine retry we do NOT exclude a source just because it was + # searched once: the first run only ran ONE query before starting a + # download, so the later queries (e.g. "artist + album") have never hit + # that source yet and may surface the correct upload. Instead we remember + # which QUERIES already ran (``searched_queries``) and skip re-running + # only those — their candidates are walked via the cached-first path + # above. The not-yet-searched queries still search the same source, so + # every query is exhausted per source before the chain switches sources. + # Fresh / dead-connection runs cleared searched_queries above, so they + # search everything again. with tasks_lock: _t = download_tasks.get(task_id, {}) _exhausted_sources = [str(s) for s in (_t.get('exhausted_download_sources') or ())] - if is_quarantine_retry: - _exhausted_sources.extend(str(s) for s in (_t.get('searched_sources') or ())) + _searched_queries = ( + set(_t.get('searched_queries') or ()) if is_quarantine_retry else set() + ) for query_index, query in enumerate(search_queries): # Cancellation check before each query with tasks_lock: @@ -395,6 +402,17 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke return download_tasks[task_id]['current_query_index'] = query_index + # Cached-first: a query already run last generation has its candidates + # sitting in cache (walked above) — re-searching it is the wasteful + # repeat the cached-first design removes. Skip it; the not-yet-run + # queries below still search this source. + if is_quarantine_retry and query in _searched_queries: + logger.debug( + f"[Modal Worker] Skipping already-searched query '{query}' " + f"(candidates served from cache) for task {task_id}" + ) + continue + logger.debug(f"[Modal Worker] Query {query_index + 1}/{len(search_queries)}: '{query}'") logger.debug(f"About to call soulseek search for task {task_id}") @@ -434,6 +452,16 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke if task_id not in download_tasks: logger.info(f"[Modal Worker] Task {task_id} was deleted after search returned") return + # Remember this query ran so a later quarantine retry skips + # re-searching it (its candidates are walked via cached-first). + # Recorded regardless of result count: re-running a query is + # deterministic, so a query that returned nothing won't return + # anything new next time either. + _sq = download_tasks[task_id].get('searched_queries') + if not isinstance(_sq, set): + _sq = set() + _sq.add(query) + download_tasks[task_id]['searched_queries'] = _sq if download_tasks[task_id]['status'] == 'cancelled': logger.warning(f"[Modal Worker] Task {task_id} cancelled after search returned - ignoring results") # Don't call _on_download_completed for cancelled tasks as it can stop monitoring @@ -456,19 +484,10 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke logger.warning(f"[Modal Worker] Task {task_id} cancelled before processing candidates") # Don't call _on_download_completed for cancelled tasks as it can stop monitoring return - # Store candidates for retry fallback (like GUI) + # Store candidates for retry fallback (like GUI). A + # later quarantine retry walks these via cached-first + # and skips re-searching this query (searched_queries). download_tasks[task_id]['cached_candidates'] = candidates - # Remember which sources produced candidates so a later - # quarantine retry walks them via cache instead of - # re-searching them (cached-first design). - _searched = download_tasks[task_id].get('searched_sources') - if not isinstance(_searched, set): - _searched = set() - for _c in candidates: - _u, _ = _cand_user_file(_c) - if _u: - _searched.add(_resolve_worker_source(_u)) - download_tasks[task_id]['searched_sources'] = _searched # Try to download with these candidates success = deps.attempt_download_with_candidates(task_id, candidates, track, batch_id) @@ -556,11 +575,6 @@ def download_track_worker(task_id: str, batch_id: Optional[str], deps: TaskWorke with tasks_lock: if task_id in download_tasks: download_tasks[task_id]['cached_candidates'] = fb_candidates - _searched = download_tasks[task_id].get('searched_sources') - if not isinstance(_searched, set): - _searched = set() - _searched.add(_resolve_worker_source(fallback_source)) - download_tasks[task_id]['searched_sources'] = _searched success = deps.attempt_download_with_candidates(task_id, fb_candidates, track, batch_id) if success: return diff --git a/tests/downloads/test_downloads_task_worker.py b/tests/downloads/test_downloads_task_worker.py index 84d9a4e8..39320940 100644 --- a/tests/downloads/test_downloads_task_worker.py +++ b/tests/downloads/test_downloads_task_worker.py @@ -570,14 +570,67 @@ def test_quarantine_retry_skips_cached_from_exhausted_source(): assert attempted == [['peerB']] # hifi candidate excluded -def test_quarantine_retry_searches_excluding_searched_source_when_cached_spent(): - # All cached used → Phase A finds nothing → search runs, but excludes the - # already-searched source so the chain falls through to a new source. +# A track_info that yields exactly the engine queries (no artist → no +# first-word legacy query; name equals a query → track-only query dedupes away), +# so the generated query set is deterministic for cached-first assertions. +_SOLO_TRACK = {'id': 'sp-1', 'name': 'Solo', 'artists': [], + 'album': 'A', 'duration_ms': 1000} + + +def test_quarantine_retry_searches_unsearched_query_without_excluding_source(): + # Cache spent + 'Solo' already searched, but 'q2' is NOT yet searched. The + # retry must search q2 against the SAME source (no source-level exclusion) so + # every query is exhausted per source before the chain switches sources + # (lazy multi-query retry). _seed_task( + track_info=dict(_SOLO_TRACK), _quarantine_retry=True, cached_candidates=[_Cand('peerA', 'f1.flac')], used_sources={'peerA_f1.flac'}, - searched_sources={'soulseek'}, + searched_queries={'Solo'}, + ) + sk = _FakeClient(results=[], mode='hybrid', + subclients={'hybrid_order': ['soulseek', 'hifi']}) + deps, _ = _build_deps( + soulseek=sk, + matching=_FakeMatchEngine(queries=['Solo', 'q2']), + ) + tw.download_track_worker('t1', 'b1', deps) + # Only q2 was searched — 'Solo' skipped because its candidates were cached. + assert [c[0] for c in sk.search_calls] == ['q2'] + # Soulseek NOT excluded: the not-yet-searched query still hits it. + assert all((not ex) or 'soulseek' not in ex for ex in sk.exclude_calls) + + +def test_quarantine_retry_skips_already_searched_query_no_research(): + # The only generated query is already searched and cache spent → it is NOT + # re-searched (its candidates live in cache). This is the wasteful repeat the + # cached-first design removes. + _seed_task( + track_info=dict(_SOLO_TRACK), + _quarantine_retry=True, + cached_candidates=[_Cand('peerA', 'f1.flac')], + used_sources={'peerA_f1.flac'}, + searched_queries={'Solo'}, + ) + sk = _FakeClient(results=[], mode='hybrid', + subclients={'hybrid_order': ['soulseek', 'hifi']}) + deps, _ = _build_deps( + soulseek=sk, + matching=_FakeMatchEngine(queries=['Solo']), + ) + tw.download_track_worker('t1', 'b1', deps) + assert sk.search_calls == [] # 'Solo' not re-searched + + +def test_quarantine_retry_still_excludes_budget_exhausted_source(): + # A source whose per-source budget is spent (exhaustive mode) stays excluded + # so the chain falls through to the next source. + _seed_task( + _quarantine_retry=True, + cached_candidates=[], + searched_queries=set(), + exhausted_download_sources={'soulseek'}, ) sk = _FakeClient(results=[], mode='hybrid', subclients={'hybrid_order': ['soulseek', 'hifi']}) @@ -586,11 +639,43 @@ def test_quarantine_retry_searches_excluding_searched_source_when_cached_spent() matching=_FakeMatchEngine(queries=['q1']), ) tw.download_track_worker('t1', 'b1', deps) - # Search ran (cached spent) and every call excluded the searched soulseek source. assert sk.search_calls assert all(ex and 'soulseek' in ex for ex in sk.exclude_calls) +def test_search_records_searched_queries(): + # Every query the worker actually runs is recorded so a later quarantine + # retry can skip re-searching it. + _seed_task(status='pending') + sk = _FakeClient(results=['r1']) + deps, _ = _build_deps( + soulseek=sk, + matching=_FakeMatchEngine(queries=['q1']), + get_valid_candidates=lambda r, t, q: [], # no candidates → loop completes + ) + tw.download_track_worker('t1', 'b1', deps) + assert 'q1' in download_tasks['t1'].get('searched_queries', set()) + + +def test_non_quarantine_run_resets_searched_queries(): + # A fresh / dead-connection retry starts a new search generation: the stale + # searched-query memory is cleared so every query can be searched again. + _seed_task( + track_info=dict(_SOLO_TRACK), + cached_candidates=[], + searched_queries={'stale-q'}, + ) + sk = _FakeClient(results=[]) + deps, _ = _build_deps( + soulseek=sk, + matching=_FakeMatchEngine(queries=['Solo']), + ) + tw.download_track_worker('t1', 'b1', deps) + sq = download_tasks['t1'].get('searched_queries', set()) + assert 'stale-q' not in sq # stale generation cleared + assert 'Solo' in sq # current generation recorded fresh + + def test_non_quarantine_run_ignores_cached_first_and_searches(): # A fresh (non-quarantine) run must NOT use cached-first — it searches. _seed_task( @@ -607,22 +692,6 @@ def test_non_quarantine_run_ignores_cached_first_and_searches(): assert sk.search_calls # searched, did not short-circuit on cache -def test_non_quarantine_run_resets_searched_sources(): - # A fresh / dead-connection retry starts a new search generation: the stale - # searched-source memory is cleared so sources can be searched again. - _seed_task( - cached_candidates=[], - searched_sources={'soulseek', 'hifi'}, - ) - sk = _FakeClient(results=[]) - deps, _ = _build_deps( - soulseek=sk, - matching=_FakeMatchEngine(queries=['q1']), - ) - tw.download_track_worker('t1', 'b1', deps) - assert download_tasks['t1'].get('searched_sources', set()) == set() - - # --------------------------------------------------------------------------- # Top-level exception path # ---------------------------------------------------------------------------