From a5fde0502acc3c5a8068f7282faf0857e8606d66 Mon Sep 17 00:00:00 2001 From: Broque Thomas <26755000+Nezreka@users.noreply.github.com> Date: Tue, 5 May 2026 11:46:44 -0700 Subject: [PATCH] Engine state: nested-dict layout for O(source) iteration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per JohnBaumb's review: iter_records_for_source() walked every (source, id) tuple across the entire engine state to filter one source — O(total_records) instead of O(source_records). Fine in practice because total active downloads is usually small, but the shape was wrong. Switched the engine's _records storage from a single composite-key dict (Dict[Tuple[str, str], DownloadRecord]) to a nested dict (Dict[str, Dict[str, DownloadRecord]]). Per-source iteration now only touches that source's bucket. add/get/update/remove all adjusted to the nested layout. remove_record drops the empty source bucket so future iterations don't see stale source keys. Public surface unchanged. New test pins the empty-bucket-cleanup behavior. --- core/download_engine/engine.py | 46 ++++++++++++++++--------- tests/downloads/test_download_engine.py | 16 +++++++++ 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/core/download_engine/engine.py b/core/download_engine/engine.py index 59bd31c3..ed5df608 100644 --- a/core/download_engine/engine.py +++ b/core/download_engine/engine.py @@ -60,10 +60,12 @@ class DownloadEngine: def __init__(self) -> None: self.state_lock = threading.RLock() - # Composite key: (source_name, download_id) → record dict. - # RLock so a plugin's worker callback can re-enter while - # holding the lock for its own update. - self._records: Dict[Tuple[str, str], DownloadRecord] = {} + # Nested dict: source_name → {download_id → record}. Replaces + # the original single-dict composite-key layout so + # ``iter_records_for_source`` is O(source_records) instead of + # O(total_records). RLock so a plugin's worker callback can + # re-enter while holding the lock for its own update. + self._records: Dict[str, Dict[str, DownloadRecord]] = {} # Plugins that have registered with the engine. Source name # → plugin instance. self._plugins: Dict[str, Any] = {} @@ -162,16 +164,16 @@ class DownloadEngine: directly via their own dicts; Phase B2 routes them through here).""" with self.state_lock: - key = (source_name, download_id) - if key in self._records: + source_bucket = self._records.setdefault(source_name, {}) + if download_id in source_bucket: logger.warning("Replacing existing download record for %s/%s", source_name, download_id) - self._records[key] = dict(record) + source_bucket[download_id] = dict(record) def update_record(self, source_name: str, download_id: str, patch: DownloadRecord) -> None: """Apply a partial patch to an existing record. No-op if the record was already removed (e.g. cancelled mid-update).""" with self.state_lock: - existing = self._records.get((source_name, download_id)) + existing = self._records.get(source_name, {}).get(download_id) if existing is None: return existing.update(patch) @@ -191,7 +193,7 @@ class DownloadEngine: the check + write closes the window. """ with self.state_lock: - existing = self._records.get((source_name, download_id)) + existing = self._records.get(source_name, {}).get(download_id) if existing is None: return False if existing.get('state') in skip_if_state_in: @@ -203,25 +205,35 @@ class DownloadEngine: """Delete a record (cancellation cleanup). Returns the removed record or None if not found.""" with self.state_lock: - return self._records.pop((source_name, download_id), None) + source_bucket = self._records.get(source_name) + if not source_bucket: + return None + removed = source_bucket.pop(download_id, None) + # Drop the empty source bucket so iteration / membership + # checks don't see a stale source key. + if not source_bucket: + self._records.pop(source_name, None) + return removed def get_record(self, source_name: str, download_id: str) -> Optional[DownloadRecord]: """Return a SHALLOW COPY of the record. Caller mutations don't affect engine state — use ``update_record`` for that.""" with self.state_lock: - record = self._records.get((source_name, download_id)) + record = self._records.get(source_name, {}).get(download_id) return dict(record) if record is not None else None def iter_records_for_source(self, source_name: str) -> Iterator[DownloadRecord]: """Yield SHALLOW COPIES of every record owned by a source. Holds the lock briefly to snapshot, then yields outside the - lock so callers can spend arbitrary time on each record.""" + lock so callers can spend arbitrary time on each record. + + With the nested-dict layout this is O(source_records) — only + touches the bucket for the requested source, not every record + across every source. + """ with self.state_lock: - snapshot = [ - dict(record) - for (source, _), record in self._records.items() - if source == source_name - ] + source_bucket = self._records.get(source_name, {}) + snapshot = [dict(record) for record in source_bucket.values()] for record in snapshot: yield record diff --git a/tests/downloads/test_download_engine.py b/tests/downloads/test_download_engine.py index 853ce88e..aa597fa8 100644 --- a/tests/downloads/test_download_engine.py +++ b/tests/downloads/test_download_engine.py @@ -120,6 +120,22 @@ def test_remove_record_returns_none_when_missing(): assert engine.remove_record('qobuz', 'never-existed') is None +def test_remove_record_drops_empty_source_bucket(): + """Per JohnBaumb: nested layout makes per-source iteration + O(source_records). Removing the last record for a source must + also drop the empty source bucket so future iter_records_for_source + calls don't see a stale source key.""" + engine = DownloadEngine() + engine.add_record('youtube', 'yt-1', {'title': 'A'}) + engine.remove_record('youtube', 'yt-1') + + # Source bucket gone — internal state matches the contract. + assert 'youtube' not in engine._records + # Public surface still answers correctly. + assert list(engine.iter_records_for_source('youtube')) == [] + assert engine.get_record('youtube', 'yt-1') is None + + # --------------------------------------------------------------------------- # Iteration # ---------------------------------------------------------------------------